feat: display video and add pydub to pyproject toml

anna-charlotte · anna-charlotte · commit 00e113092b5c · 2023-02-15T14:27:23.000+01:00
Signed-off-by: anna-charlotte &lt;charlotte.gerhaher@jina.ai&gt;
diff --git a/docarray/typing/tensor/audio/abstract_audio_tensor.py b/docarray/typing/tensor/audio/abstract_audio_tensor.py
@@ -3,6 +3,7 @@
 from typing import BinaryIO, TypeVar, Union
 
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
+from docarray.utils.misc import is_notebook
 
 T = TypeVar('T', bound='AbstractAudioTensor')
 
@@ -40,3 +41,35 @@ def save_to_wav_file(
             f.setsampwidth(sample_width)
             f.setframerate(sample_rate)
             f.writeframes(self.to_bytes())
+
+    def display(self, rate=44100):
+        audio_np = self.get_comp_backend().to_numpy(self)
+        if is_notebook():
+            from IPython.display import Audio, display
+
+            display(Audio(audio_np, rate=rate))
+        else:
+            # b = self.load()
+            # res = requests.get(self)
+            # print(f"type(res.text) = {type(res.text)}")
+            # print(f"type(res.content) = {type(res.content)}")
+            # sound = AudioSegment.from_file(BytesIO(res.content), "wav")
+            # sound = AudioSegment.from_file(self, format="wav")
+            import os
+            import tempfile
+
+            # sound = AudioSegment.from_file(self, format="wav")
+            # raise NotImplementedError
+            # from io import BytesIO
+            # import requests
+            from pydub import AudioSegment
+            from pydub.playback import play
+
+            tmp = tempfile.NamedTemporaryFile(delete=False)
+            try:
+                self.save_to_wav_file(tmp)
+                sound = AudioSegment.from_file(tmp, "wav")
+                play(sound)
+            finally:
+                tmp.close()
+                os.unlink(tmp.name)
diff --git a/docarray/typing/tensor/video/video_tensor_mixin.py b/docarray/typing/tensor/video/video_tensor_mixin.py
@@ -109,3 +109,9 @@ class MyDoc(BaseDocument):
 
             for packet in stream_video.encode(None):
                 container.mux(packet)
+
+    def display(self) -> None:
+        """
+        Display video data from tensor.
+        """
+        raise NotImplementedError
diff --git a/docarray/typing/url/video_url.py b/docarray/typing/url/video_url.py
@@ -8,6 +8,7 @@
 from docarray.typing.tensor.ndarray import NdArray
 from docarray.typing.tensor.video import VideoNdArray
 from docarray.typing.url.any_url import AnyUrl
+from docarray.utils.misc import is_notebook
 
 if TYPE_CHECKING:
     from pydantic import BaseConfig
@@ -130,3 +131,40 @@ class MyDoc(BaseDocument):
         indices = parse_obj_as(NdArray, keyframe_indices)
 
         return VideoLoadResult(video=video, audio=audio, key_frame_indices=indices)
+
+    def display(self):
+        """
+        Play video from url.
+        """
+        remote_url = True if self.startswith('http') else False
+
+        if is_notebook():
+            from IPython.display import Video, display
+
+            if remote_url:
+                display(Video(data=self))
+            else:
+                display(Video(filename=self))
+        else:
+            import pyglet
+
+            player = pyglet.media.Player()
+            pyglet.media.StreamingSource()
+            media = pyglet.media.load(self, streaming=True)
+            player.queue(media)
+
+            width = player.source.video_format.width
+            height = player.source.video_format.height
+            title = self.split('/')[-1]
+            window = pyglet.window.Window(width, height, title)
+
+            player.play()
+
+            @window.event
+            def on_draw():
+                window.clear()
+
+                if player.source and player.source.video_format:
+                    player.get_texture().blit(0, 0)
+
+            pyglet.app.run()
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -20,12 +20,14 @@ types-requests = ">=2.28.11.6"
 av = {version = ">=10.0.0", optional = true}
 fastapi = {version = ">=0.87.0", optional = true }
 rich = ">=13.1.0"
+pydub = {version = "^0.25.1", optional = true }
 
 [tool.poetry.extras]
 common = ["protobuf"]
 torch = ["torch"]
 image = ["pillow", "types-pillow"]
 video = ["av"]
+audio = ["pydub"]
 mesh = ["trimesh"]
 web = ["fastapi"]
 
