xinference 1.10.0__py3-none-any.whl → 1.11.0__py3-none-any.whl

xinference/thirdparty/audiotools/core/playback.py

@@ -0,0 +1,252 @@
+"""
+These are utilities that allow one to embed an AudioSignal
+as a playable object in a Jupyter notebook, or to play audio from
+the terminal, etc.
+"""  # fmt: skip
+import base64
+import io
+import random
+import string
+import subprocess
+from tempfile import NamedTemporaryFile
+
+import importlib_resources as pkg_resources
+
+from . import templates
+from .util import _close_temp_files
+from .util import format_figure
+
+headers = pkg_resources.files(templates).joinpath("headers.html").read_text()
+widget = pkg_resources.files(templates).joinpath("widget.html").read_text()
+
+DEFAULT_EXTENSION = ".wav"
+
+
+def _check_imports():  # pragma: no cover
+    try:
+        import ffmpy
+    except:
+        ffmpy = False
+
+    try:
+        import IPython
+    except:
+        raise ImportError("IPython must be installed in order to use this function!")
+    return ffmpy, IPython
+
+
+class PlayMixin:
+    def embed(self, ext: str = None, display: bool = True, return_html: bool = False):
+        """Embeds audio as a playable audio embed in a notebook, or HTML
+        document, etc.
+
+        Parameters
+        ----------
+        ext : str, optional
+            Extension to use when saving the audio, by default ".wav"
+        display : bool, optional
+            This controls whether or not to display the audio when called. This
+            is used when the embed is the last line in a Jupyter cell, to prevent
+            the audio from being embedded twice, by default True
+        return_html : bool, optional
+            Whether to return the data wrapped in an HTML audio element, by default False
+
+        Returns
+        -------
+        str
+            Either the element for display, or the HTML string of it.
+        """
+        if ext is None:
+            ext = DEFAULT_EXTENSION
+        ext = f".{ext}" if not ext.startswith(".") else ext
+        ffmpy, IPython = _check_imports()
+        sr = self.sample_rate
+        tmpfiles = []
+
+        with _close_temp_files(tmpfiles):
+            tmp_wav = NamedTemporaryFile(mode="w+", suffix=".wav", delete=False)
+            tmpfiles.append(tmp_wav)
+            self.write(tmp_wav.name)
+            if ext != ".wav" and ffmpy:
+                tmp_converted = NamedTemporaryFile(mode="w+", suffix=ext, delete=False)
+                tmpfiles.append(tmp_wav)
+                ff = ffmpy.FFmpeg(
+                    inputs={tmp_wav.name: None},
+                    outputs={
+                        tmp_converted.name: "-write_xing 0 -codec:a libmp3lame -b:a 128k -y -hide_banner -loglevel error"
+                    },
+                )
+                ff.run()
+            else:
+                tmp_converted = tmp_wav
+
+            audio_element = IPython.display.Audio(data=tmp_converted.name, rate=sr)
+            if display:
+                IPython.display.display(audio_element)
+
+        if return_html:
+            audio_element = (
+                f"<audio "
+                f"  controls "
+                f"  src='{audio_element.src_attr()}'> "
+                f"</audio> "
+            )
+        return audio_element
+
+    def widget(
+        self,
+        title: str = None,
+        ext: str = ".wav",
+        add_headers: bool = True,
+        player_width: str = "100%",
+        margin: str = "10px",
+        plot_fn: str = "specshow",
+        return_html: bool = False,
+        **kwargs,
+    ):
+        """Creates a playable widget with spectrogram. Inspired (heavily) by
+        https://sjvasquez.github.io/blog/melnet/.
+
+        Parameters
+        ----------
+        title : str, optional
+            Title of plot, placed in upper right of top-most axis.
+        ext : str, optional
+            Extension for embedding, by default ".mp3"
+        add_headers : bool, optional
+            Whether or not to add headers (use for first embed, False for later embeds), by default True
+        player_width : str, optional
+            Width of the player, as a string in a CSS rule, by default "100%"
+        margin : str, optional
+            Margin on all sides of player, by default "10px"
+        plot_fn : function, optional
+            Plotting function to use (by default self.specshow).
+        return_html : bool, optional
+            Whether to return the data wrapped in an HTML audio element, by default False
+        kwargs : dict, optional
+            Keyword arguments to plot_fn (by default self.specshow).
+
+        Returns
+        -------
+        HTML
+            HTML object.
+        """
+        import matplotlib.pyplot as plt
+
+        def _save_fig_to_tag():
+            buffer = io.BytesIO()
+
+            plt.savefig(buffer, bbox_inches="tight", pad_inches=0)
+            plt.close()
+
+            buffer.seek(0)
+            data_uri = base64.b64encode(buffer.read()).decode("ascii")
+            tag = "data:image/png;base64,{0}".format(data_uri)
+
+            return tag
+
+        _, IPython = _check_imports()
+
+        header_html = ""
+
+        if add_headers:
+            header_html = headers.replace("PLAYER_WIDTH", str(player_width))
+            header_html = header_html.replace("MARGIN", str(margin))
+            IPython.display.display(IPython.display.HTML(header_html))
+
+        widget_html = widget
+        if isinstance(plot_fn, str):
+            plot_fn = getattr(self, plot_fn)
+            kwargs["title"] = title
+        plot_fn(**kwargs)
+
+        fig = plt.gcf()
+        pixels = fig.get_size_inches() * fig.dpi
+
+        tag = _save_fig_to_tag()
+
+        # Make the source image for the levels
+        self.specshow()
+        format_figure((12, 1.5))
+        levels_tag = _save_fig_to_tag()
+
+        player_id = "".join(random.choice(string.ascii_uppercase) for _ in range(10))
+
+        audio_elem = self.embed(ext=ext, display=False)
+        widget_html = widget_html.replace("AUDIO_SRC", audio_elem.src_attr())
+        widget_html = widget_html.replace("IMAGE_SRC", tag)
+        widget_html = widget_html.replace("LEVELS_SRC", levels_tag)
+        widget_html = widget_html.replace("PLAYER_ID", player_id)
+
+        # Calculate width/height of figure based on figure size.
+        widget_html = widget_html.replace("PADDING_AMOUNT", f"{int(pixels[1])}px")
+        widget_html = widget_html.replace("MAX_WIDTH", f"{int(pixels[0])}px")
+
+        IPython.display.display(IPython.display.HTML(widget_html))
+
+        if return_html:
+            html = header_html if add_headers else ""
+            html += widget_html
+            return html
+
+    def play(self):
+        """
+        Plays an audio signal if ffplay from the ffmpeg suite of tools is installed.
+        Otherwise, will fail. The audio signal is written to a temporary file
+        and then played with ffplay.
+        """
+        tmpfiles = []
+        with _close_temp_files(tmpfiles):
+            tmp_wav = NamedTemporaryFile(suffix=".wav", delete=False)
+            tmpfiles.append(tmp_wav)
+            self.write(tmp_wav.name)
+            print(self)
+            subprocess.call(
+                [
+                    "ffplay",
+                    "-nodisp",
+                    "-autoexit",
+                    "-hide_banner",
+                    "-loglevel",
+                    "error",
+                    tmp_wav.name,
+                ]
+            )
+        return self
+
+
+if __name__ == "__main__":  # pragma: no cover
+    from audiotools import AudioSignal
+
+    signal = AudioSignal(
+        "tests/audio/spk/f10_script4_produced.mp3", offset=5, duration=5
+    )
+
+    wave_html = signal.widget(
+        "Waveform",
+        plot_fn="waveplot",
+        return_html=True,
+    )
+
+    spec_html = signal.widget("Spectrogram", return_html=True, add_headers=False)
+
+    combined_html = signal.widget(
+        "Waveform + spectrogram",
+        plot_fn="wavespec",
+        return_html=True,
+        add_headers=False,
+    )
+
+    signal.low_pass(8000)
+    lowpass_html = signal.widget(
+        "Lowpassed audio",
+        plot_fn="wavespec",
+        return_html=True,
+        add_headers=False,
+    )
+
+    with open("/tmp/index.html", "w") as f:
+        f.write(wave_html)
+        f.write(spec_html)
+        f.write(combined_html)
+        f.write(lowpass_html)