docarray · JoanFM · Jan 3, 2023 · Dec 14, 2022 · Dec 14, 2022 · Dec 15, 2022
diff --git a/docarray/__init__.py b/docarray/__init__.py
@@ -2,6 +2,14 @@
 
 from docarray.array.array import DocumentArray
 from docarray.document.document import BaseDocument
-from docarray.predefined_document import Image, Mesh3D, PointCloud3D, Text
+from docarray.predefined_document import Audio, Image, Mesh3D, PointCloud3D, Text
 
-__all__ = ['BaseDocument', 'DocumentArray', 'Image', 'Text', 'Mesh3D', 'PointCloud3D']
+__all__ = [
+    'BaseDocument',
+    'DocumentArray',
+    'Image',
+    'Audio',
+    'Text',
+    'Mesh3D',
+    'PointCloud3D',
+]
diff --git a/docarray/computation/abstract_comp_backend.py b/docarray/computation/abstract_comp_backend.py
@@ -26,6 +26,11 @@ def stack(
         """
         ...
 
+    @staticmethod
+    @abstractmethod
+    def n_dim(array: 'TTensor') -> int:
+        ...
+
     class Retrieval(ABC, typing.Generic[TTensorRetrieval]):
         """
         Abstract class for retrieval and ranking functionalities

diff --git a/docarray/computation/numpy_backend.py b/docarray/computation/numpy_backend.py
@@ -40,6 +40,10 @@ def stack(
     ) -> 'np.ndarray':
         return np.stack(tensors, axis=dim)
 
+    @staticmethod
+    def n_dim(array: 'np.ndarray') -> int:
+        return array.ndim
+
     class Retrieval(AbstractComputationalBackend.Retrieval[np.ndarray]):
         """
         Abstract class for retrieval and ranking functionalities

diff --git a/docarray/computation/torch_backend.py b/docarray/computation/torch_backend.py
@@ -39,6 +39,10 @@ def stack(
     ) -> 'torch.Tensor':
         return torch.stack(tensors, dim=dim)
 
+    @staticmethod
+    def n_dim(array: 'torch.Tensor') -> int:
+        return array.ndim
+
     class Retrieval(AbstractComputationalBackend.Retrieval[torch.Tensor]):
         """
         Abstract class for retrieval and ranking functionalities

diff --git a/docarray/predefined_document/__init__.py b/docarray/predefined_document/__init__.py
@@ -1,6 +1,7 @@
+from docarray.predefined_document.audio import Audio
 from docarray.predefined_document.image import Image
 from docarray.predefined_document.mesh import Mesh3D
 from docarray.predefined_document.point_cloud import PointCloud3D
 from docarray.predefined_document.text import Text
 
-__all__ = ['Text', 'Image', 'Mesh3D', 'PointCloud3D']
+__all__ = ['Text', 'Image', 'Audio', 'Mesh3D', 'PointCloud3D']
diff --git a/docarray/predefined_document/audio.py b/docarray/predefined_document/audio.py
@@ -0,0 +1,77 @@
+from typing import Optional, TypeVar
+
+from docarray.document import BaseDocument
+from docarray.typing import AudioUrl, Embedding
+from docarray.typing.tensor.audio.audio_tensor import AudioTensor
+
+T = TypeVar('T', bound='Audio')
+
+
+class Audio(BaseDocument):
+    """
+    Document for handling audios.
+
+    The Audio Document can contain an AudioUrl (`Audio.url`), an AudioTensor
+    (`Audio.tensor`), and an Embedding (`Audio.embedding`).
+
+    EXAMPLE USAGE:
+
+    You can use this Document directly:
+
+    .. code-block:: python
+
+        from docarray import Audio
+
+        # use it directly
+        audio = Audio(
+            url='https://github.com/docarray/docarray/tree/feat-add-audio-v2/tests/toydata/hello.wav?raw=true'
+        )
+        audio.tensor = audio.url.load()
+        model = MyEmbeddingModel()
+        audio.embedding = model(audio.tensor)
+
+    You can extend this Document:
+
+    .. code-block:: python
+
+        from docarray import Audio, Text
+        from typing import Optional
+
+        # extend it
+        class MyAudio(Audio):
+            name: Optional[Text]
+
+
+        audio = MyAudio(
+            url='https://github.com/docarray/docarray/tree/feat-add-audio-v2/tests/toydata/hello.wav?raw=true'
+        )
+        audio.tensor = audio.url.load()
+        model = MyEmbeddingModel()
+        audio.embedding = model(audio.tensor)
+        audio.name = 'my first audio'
+
+
+    You can use this Document for composition:
+
+    .. code-block:: python
+
+        from docarray import Document, Audio, Text
+
+        # compose it
+        class MultiModalDoc(Document):
+            audio: Audio
+            text: Text
+
+
+        mmdoc = MultiModalDoc(
+            audio=Audio(
+                url='https://github.com/docarray/docarray/tree/feat-add-audio-v2/tests/toydata/hello.wav?raw=true'
+            ),
+            text=Text(text='hello world, how are you doing?'),
+        )
+        mmdoc.audio.tensor = mmdoc.audio.url.load()
+    """
+
+    url: Optional[AudioUrl]
+    tensor: Optional[AudioTensor]
+    embedding: Optional[Embedding]
diff --git a/docarray/proto/docarray.proto b/docarray/proto/docarray.proto
@@ -63,6 +63,11 @@ message NodeProto {
 
     string point_cloud_url = 13;
 
+    string audio_url = 14;
+
+    NdArrayProto audio_ndarray = 15;
+
+    NdArrayProto audio_torch_tensor = 16;
 
   }
 

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
@@ -1,13 +1,23 @@
 from docarray.typing.id import ID
+from docarray.typing.tensor.audio import AudioNdArray
 from docarray.typing.tensor.embedding.embedding import Embedding
 from docarray.typing.tensor.ndarray import NdArray
 from docarray.typing.tensor.tensor import AnyTensor
-from docarray.typing.url import AnyUrl, ImageUrl, Mesh3DUrl, PointCloud3DUrl, TextUrl
+from docarray.typing.url import (
+    AnyUrl,
+    AudioUrl,
+    ImageUrl,
+    Mesh3DUrl,
+    PointCloud3DUrl,
+    TextUrl,
+)
 
 __all__ = [
+    'AudioNdArray',
     'NdArray',
     'Embedding',
     'ImageUrl',
+    'AudioUrl',
     'TextUrl',
     'Mesh3DUrl',
     'PointCloud3DUrl',
@@ -22,5 +32,6 @@
     pass
 else:
     from docarray.typing.tensor import TorchEmbedding, TorchTensor  # noqa: F401
+    from docarray.typing.tensor.audio.audio_torch_tensor import AudioTorchTensor  # noqa
 
-    __all__.extend(['TorchEmbedding', 'TorchTensor'])
+    __all__.extend(['AudioTorchTensor', 'TorchEmbedding', 'TorchTensor'])
diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py
@@ -16,6 +16,7 @@
 class AbstractTensor(AbstractType, Generic[ShapeT], ABC):
 
     __parametrized_meta__ = type
+    _PROTO_FIELD_NAME: str
 
     @classmethod
     @abc.abstractmethod

diff --git a/docarray/typing/tensor/audio/__init__.py b/docarray/typing/tensor/audio/__init__.py
@@ -0,0 +1,12 @@
+from docarray.typing.tensor.audio.audio_ndarray import AudioNdArray
+
+__all__ = ['AudioNdArray']
+
+try:
+    import torch  # noqa: F401
+except ImportError:
+    pass
+else:
+    from docarray.typing.tensor.audio.audio_torch_tensor import AudioTorchTensor  # noqa
+
+    __all__.extend(['AudioTorchTensor'])
diff --git a/docarray/typing/tensor/audio/abstract_audio_tensor.py b/docarray/typing/tensor/audio/abstract_audio_tensor.py
@@ -0,0 +1,39 @@
+import wave
+from abc import ABC, abstractmethod
+from typing import BinaryIO, TypeVar, Union
+
+from docarray.typing.tensor.abstract_tensor import AbstractTensor
+
+T = TypeVar('T', bound='AbstractAudioTensor')
+
+
+class AbstractAudioTensor(AbstractTensor, ABC):
+    @abstractmethod
+    def to_audio_bytes(self):
+        """
+        Convert audio tensor to bytes.
+        """
+        ...
+
+    def save_to_wav_file(
+        self: 'T',
+        file_path: Union[str, BinaryIO],
+        sample_rate: int = 44100,
+        sample_width: int = 2,
+    ) -> None:
+        """
+        Save audio tensor to a .wav file. Mono/stereo is preserved.
+
+        :param file_path: path to a .wav file. If file is a string, open the file by
+            that name, otherwise treat it as a file-like object.
+        :param sample_rate: sampling frequency
+        :param sample_width: sample width in bytes
+        """
+        comp_backend = self.get_comp_backend()
+        n_channels = 2 if comp_backend.n_dim(array=self) > 1 else 1  # type: ignore
+
+        with wave.open(file_path, 'w') as f:
+            f.setnchannels(n_channels)
+            f.setsampwidth(sample_width)
+            f.setframerate(sample_rate)
+            f.writeframes(self.to_audio_bytes())
diff --git a/docarray/typing/tensor/audio/audio_ndarray.py b/docarray/typing/tensor/audio/audio_ndarray.py
@@ -0,0 +1,59 @@
+from typing import TypeVar
+
+from docarray.typing.tensor.audio.abstract_audio_tensor import AbstractAudioTensor
+from docarray.typing.tensor.ndarray import NdArray
+
+MAX_INT_16 = 2**15
+
+T = TypeVar('T', bound='AudioNdArray')
+
+
+class AudioNdArray(AbstractAudioTensor, NdArray):
+    """
+    Subclass of NdArray, to represent an audio tensor.
+    Adds audio-specific features to the tensor.
+
+
+    EXAMPLE USAGE
+
+    .. code-block:: python
+
+        from typing import Optional
+
+        from pydantic import parse_obj_as
+
+        from docarray import Document
+        from docarray.typing import AudioNdArray, AudioUrl
+        import numpy as np
+
+
+        class MyAudioDoc(Document):
+            title: str
+            audio_tensor: Optional[AudioNdArray]
+            url: Optional[AudioUrl]
+
+
+        # from tensor
+        doc_1 = MyAudioDoc(
+            title='my_first_audio_doc',
+            audio_tensor=np.random.rand(1000, 2),
+        )
+
+        doc_1.audio_tensor.save_to_wav_file(file_path='path/to/file_1.wav')
+
+        # from url
+        doc_2 = MyAudioDoc(
+            title='my_second_audio_doc',
+            url='https://www.kozco.com/tech/piano2.wav',
+        )
+
+        doc_2.audio_tensor = parse_obj_as(AudioNdArray, doc_2.url.load())
+        doc_2.audio_tensor.save_to_wav_file(file_path='path/to/file_2.wav')
+
+    """
+
+    _PROTO_FIELD_NAME = 'audio_ndarray'
+
+    def to_audio_bytes(self):
+        tensor = (self * MAX_INT_16).astype('<h')
+        return tensor.tobytes()
diff --git a/docarray/typing/tensor/audio/audio_tensor.py b/docarray/typing/tensor/audio/audio_tensor.py
@@ -0,0 +1,13 @@
+from typing import Union
+
+from docarray.typing.tensor.audio.audio_ndarray import AudioNdArray
+
+try:
+    import torch  # noqa: F401
+except ImportError:
+    AudioTensor = AudioNdArray
+
+else:
+    from docarray.typing.tensor.audio.audio_torch_tensor import AudioTorchTensor
+
+    AudioTensor = Union[AudioNdArray, AudioTorchTensor]  # type: ignore