docarray · samsja · Jan 17, 2023 · Jan 3, 2023 · Jan 3, 2023 · Jan 4, 2023
diff --git a/docarray/computation/abstract_comp_backend.py b/docarray/computation/abstract_comp_backend.py
@@ -1,6 +1,9 @@
 import typing
 from abc import ABC, abstractmethod
-from typing import List, Optional, Tuple, TypeVar, Union, overload
+from typing import TYPE_CHECKING, List, Optional, Tuple, TypeVar, Union, overload
+
+if TYPE_CHECKING:
+    import numpy as np
 
 # In practice all of the below will be the same type
 TTensor = TypeVar('TTensor')
@@ -30,6 +33,17 @@ def stack(
     @staticmethod
     @abstractmethod
     def n_dim(array: 'TTensor') -> int:
+        """
+        Get the number of the array dimensions.
+        """
+        ...
+
+    @staticmethod
+    @abstractmethod
+    def to_numpy(array: 'TTensor') -> 'np.ndarray':
+        """
+        Convert array to np.ndarray.
+        """
         ...
 
     @staticmethod

diff --git a/docarray/computation/numpy_backend.py b/docarray/computation/numpy_backend.py
@@ -64,6 +64,10 @@ def to_device(
     def n_dim(array: 'np.ndarray') -> int:
         return array.ndim
 
+    @staticmethod
+    def to_numpy(array: 'np.ndarray') -> 'np.ndarray':
+        return array
+
     @staticmethod
     def empty(shape: Tuple[int, ...]) -> 'np.ndarray':
         return np.empty(shape)

diff --git a/docarray/computation/torch_backend.py b/docarray/computation/torch_backend.py
@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, overload
 
+import numpy as np
 import torch
 
 from docarray.computation.abstract_comp_backend import AbstractComputationalBackend
@@ -68,6 +69,10 @@ def empty(shape: Tuple[int, ...]) -> torch.Tensor:
     def n_dim(array: 'torch.Tensor') -> int:
         return array.ndim
 
+    @staticmethod
+    def to_numpy(array: 'torch.Tensor') -> 'np.ndarray':
+        return array.cpu().detach().numpy()
+
     @staticmethod
     def none_value() -> Any:
         """Provide a compatible value that represents None in torch."""

diff --git a/docarray/documents/__init__.py b/docarray/documents/__init__.py
@@ -3,5 +3,6 @@
 from docarray.documents.mesh import Mesh3D
 from docarray.documents.point_cloud import PointCloud3D
 from docarray.documents.text import Text
+from docarray.documents.video import Video
 
-__all__ = ['Text', 'Image', 'Audio', 'Mesh3D', 'PointCloud3D']
+__all__ = ['Text', 'Image', 'Audio', 'Mesh3D', 'PointCloud3D', 'Video']
diff --git a/docarray/documents/audio.py b/docarray/documents/audio.py
@@ -24,7 +24,7 @@ class Audio(BaseDocument):
 
         # use it directly
         audio = Audio(
-            url='https://github.com/docarray/docarray/tree/feat-add-audio-v2/tests/toydata/hello.wav?raw=true'
+            url='https://github.com/docarray/docarray/blob/feat-rewrite-v2/tests/toydata/hello.wav?raw=true'
         )
         audio.tensor = audio.url.load()
         model = MyEmbeddingModel()
@@ -43,12 +43,12 @@ class MyAudio(Audio):
 
 
         audio = MyAudio(
-            url='https://github.com/docarray/docarray/tree/feat-add-audio-v2/tests/toydata/hello.wav?raw=true'
+            url='https://github.com/docarray/docarray/blob/feat-rewrite-v2/tests/toydata/hello.wav?raw=true'
         )
         audio.tensor = audio.url.load()
         model = MyEmbeddingModel()
         audio.embedding = model(audio.tensor)
-        audio.name = 'my first audio'
+        audio.name = Text(text='my first audio')
 
 
     You can use this Document for composition:
@@ -66,7 +66,7 @@ class MultiModalDoc(Document):
 
         mmdoc = MultiModalDoc(
             audio=Audio(
-                url='https://github.com/docarray/docarray/tree/feat-add-audio-v2/tests/toydata/hello.wav?raw=true'
+                url='https://github.com/docarray/docarray/blob/feat-rewrite-v2/tests/toydata/hello.wav?raw=true'
             ),
             text=Text(text='hello world, how are you doing?'),
         )

diff --git a/docarray/documents/video.py b/docarray/documents/video.py
@@ -0,0 +1,85 @@
+from typing import Optional, TypeVar
+
+from docarray.base_document import BaseDocument
+from docarray.documents import Audio
+from docarray.typing import AnyEmbedding, AnyTensor
+from docarray.typing.tensor.video.video_tensor import VideoTensor
+from docarray.typing.url.video_url import VideoUrl
+
+T = TypeVar('T', bound='Video')
+
+
+class Video(BaseDocument):
+    """
+    Document for handling video.
+    The Video Document can contain a VideoUrl (`Video.url`), an Audio Document
+    (`Video.audio`), a VideoTensor (`Video.video_tensor`), an AnyTensor representing
+    the indices of the video's key frames (`Video.key_frame_indices`) and an
+    AnyEmbedding (`Video.embedding`).
+
+    EXAMPLE USAGE:
+
+    You can use this Document directly:
+
+    .. code-block:: python
+
+        from docarray.documents import Video
+
+        # use it directly
+        vid = Video(
+            url='https://github.com/docarray/docarray/tree/feat-add-video-v2/tests/toydata/mov_bbb.mp4?raw=true'
+        )
+        vid.audio.tensor, vid.video_tensor, vid.key_frame_indices = vid.url.load()
+        model = MyEmbeddingModel()
+        vid.embedding = model(vid.video_tensor)
+
+        You can extend this Document:
+
+    .. code-block:: python
+
+        from typing import Optional
+
+        from docarray.documents import Text, Video
+
+
+        # extend it
+        class MyVideo(Video):
+            name: Optional[Text]
+
+
+        video = MyVideo(
+            url='https://github.com/docarray/docarray/blob/feat-rewrite-v2/tests/toydata/mov_bbb.mp4?raw=true'
+        )
+        video.video_tensor = video.url.load_key_frames()
+        model = MyEmbeddingModel()
+        video.embedding = model(video.video_tensor)
+        video.name = Text(text='my first video')
+
+    You can use this Document for composition:
+
+    .. code-block:: python
+
+        from docarray import BaseDocument
+        from docarray.documents import Text, Video
+
+
+        # compose it
+        class MultiModalDoc(BaseDocument):
+            video: Video
+            text: Text
+
+
+        mmdoc = MultiModalDoc(
+            video=Video(
+                url='https://github.com/docarray/docarray/blob/feat-rewrite-v2/tests/toydata/mov_bbb.mp4?raw=true'
+            ),
+            text=Text(text='hello world, how are you doing?'),
+        )
+        mmdoc.video.video_tensor = mmdoc.video.url.load_key_frames()
+    """
+
+    url: Optional[VideoUrl]
+    audio: Optional[Audio] = Audio()
+    video_tensor: Optional[VideoTensor]
+    key_frame_indices: Optional[AnyTensor]
+    embedding: Optional[AnyEmbedding]
diff --git a/docarray/proto/docarray.proto b/docarray/proto/docarray.proto
@@ -69,6 +69,12 @@ message NodeProto {
 
     NdArrayProto audio_torch_tensor = 16;
 
+    string video_url = 17;
+
+    NdArrayProto video_ndarray = 18;
+
+    NdArrayProto video_torch_tensor = 19;
+
   }
 
 }

diff --git a/docarray/proto/pb2/docarray_pb2.py b/docarray/proto/pb2/docarray_pb2.py
diff --git a/docarray/typing/__init__.py b/docarray/typing/__init__.py
@@ -3,24 +3,28 @@
 from docarray.typing.tensor.embedding.embedding import AnyEmbedding
 from docarray.typing.tensor.ndarray import NdArray
 from docarray.typing.tensor.tensor import AnyTensor
+from docarray.typing.tensor.video import VideoNdArray
 from docarray.typing.url import (
     AnyUrl,
     AudioUrl,
     ImageUrl,
     Mesh3DUrl,
     PointCloud3DUrl,
     TextUrl,
+    VideoUrl,
 )
 
 __all__ = [
-    'AudioNdArray',
     'NdArray',
+    'AudioNdArray',
+    'VideoNdArray',
     'AnyEmbedding',
     'ImageUrl',
     'AudioUrl',
     'TextUrl',
     'Mesh3DUrl',
     'PointCloud3DUrl',
+    'VideoUrl',
     'AnyUrl',
     'ID',
     'AnyTensor',
@@ -33,5 +37,8 @@
 else:
     from docarray.typing.tensor import TorchEmbedding, TorchTensor  # noqa: F401
     from docarray.typing.tensor.audio.audio_torch_tensor import AudioTorchTensor  # noqa
+    from docarray.typing.tensor.video.video_torch_tensor import VideoTorchTensor  # noqa
 
-    __all__.extend(['AudioTorchTensor', 'TorchEmbedding', 'TorchTensor'])
+    __all__.extend(
+        ['AudioTorchTensor', 'TorchEmbedding', 'TorchTensor', 'VideoTorchTensor']
+    )
diff --git a/docarray/typing/tensor/video/__init__.py b/docarray/typing/tensor/video/__init__.py
@@ -0,0 +1,12 @@
+from docarray.typing.tensor.video.video_ndarray import VideoNdArray
+
+__all__ = ['VideoNdArray']
+
+try:
+    import torch  # noqa: F401
+except ImportError:
+    pass
+else:
+    from docarray.typing.tensor.video.video_torch_tensor import VideoTorchTensor  # noqa
+
+    __all__.extend(['VideoTorchTensor'])
diff --git a/docarray/typing/tensor/video/video_ndarray.py b/docarray/typing/tensor/video/video_ndarray.py
@@ -0,0 +1,34 @@
+from typing import TYPE_CHECKING, Any, List, Tuple, Type, TypeVar, Union
+
+import numpy as np
+
+from docarray.typing.tensor.ndarray import NdArray
+from docarray.typing.tensor.video.video_tensor_mixin import VideoTensorMixin
+
+T = TypeVar('T', bound='VideoNdArray')
+
+if TYPE_CHECKING:
+    from pydantic import BaseConfig
+    from pydantic.fields import ModelField
+
+
+class VideoNdArray(NdArray, VideoTensorMixin):
+    """
+    Subclass of NdArray, to represent a video tensor.
+    Adds video-specific features to the tensor.
+
+    EXAMPLE USAGE
+
+    """
+
+    _PROTO_FIELD_NAME = 'video_ndarray'
+
+    @classmethod
+    def validate(
+        cls: Type[T],
+        value: Union[T, np.ndarray, List[Any], Tuple[Any], Any],
+        field: 'ModelField',
+        config: 'BaseConfig',
+    ) -> T:
+        tensor = super().validate(value=value, field=field, config=config)
+        return cls.validate_shape(value=tensor)
@@ -0,0 +1,13 @@
+from typing import Union
+
+from docarray.typing.tensor.video.video_ndarray import VideoNdArray
+
+try:
+    import torch  # noqa: F401
+except ImportError:
+    VideoTensor = VideoNdArray
+
+else:
+    from docarray.typing.tensor.video.video_torch_tensor import VideoTorchTensor
+
+    VideoTensor = Union[VideoNdArray, VideoTorchTensor]  # type: ignore