From 7432f54bd0e7ae84249cd300edbda6c1907498a2 Mon Sep 17 00:00:00 2001 From: anna-charlotte Date: Wed, 5 Apr 2023 19:14:25 +0200 Subject: [PATCH 1/5] fix: return video bytes in video tensors to bytes method Signed-off-by: anna-charlotte --- docarray/typing/bytes/video_bytes.py | 8 ++++---- docarray/typing/tensor/video/video_tensor_mixin.py | 9 +++++---- tests/units/typing/tensor/test_video_tensor.py | 4 +++- 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/docarray/typing/bytes/video_bytes.py b/docarray/typing/bytes/video_bytes.py index 1f4e22e4fd1..557f8fbbf74 100644 --- a/docarray/typing/bytes/video_bytes.py +++ b/docarray/typing/bytes/video_bytes.py @@ -7,21 +7,21 @@ from docarray.typing.abstract_type import AbstractType from docarray.typing.proto_register import _register_proto -from docarray.typing.tensor import AudioNdArray, NdArray, VideoNdArray from docarray.utils._internal.misc import import_library if TYPE_CHECKING: from pydantic.fields import BaseConfig, ModelField from docarray.proto import NodeProto + from docarray.typing.tensor import AudioNdArray, NdArray, VideoNdArray T = TypeVar('T', bound='VideoBytes') class VideoLoadResult(NamedTuple): - video: VideoNdArray - audio: AudioNdArray - key_frame_indices: NdArray + video: 'VideoNdArray' + audio: 'AudioNdArray' + key_frame_indices: 'NdArray' @_register_proto(proto_type_name='video_bytes') diff --git a/docarray/typing/tensor/video/video_tensor_mixin.py b/docarray/typing/tensor/video/video_tensor_mixin.py index 4ce6a1b7260..a302e4d2275 100644 --- a/docarray/typing/tensor/video/video_tensor_mixin.py +++ b/docarray/typing/tensor/video/video_tensor_mixin.py @@ -5,6 +5,7 @@ import numpy as np +from docarray.typing.bytes.video_bytes import VideoBytes from docarray.typing.tensor.abstract_tensor import AbstractTensor from docarray.typing.tensor.audio.audio_tensor import AudioTensor from docarray.utils._internal.misc import import_library, is_notebook @@ -130,9 +131,9 @@ def to_bytes( audio_frame_rate: int = 48000, audio_codec: str = 'aac', audio_format: str = 'fltp', - ) -> bytes: + ) -> VideoBytes: """ - Convert video tensor to bytes. + Convert video tensor to VideoBytes. :param audio_tensor: AudioTensor containing the video's soundtrack. :param video_frame_rate: video frames per second. @@ -142,7 +143,7 @@ def to_bytes( :param audio_format: the name of one of the audio formats supported by PyAV, such as 'flt', 'fltp', 's16' or 's16p'. - :return: bytes + :return: a VideoBytes object """ bytes = BytesIO() self.save( @@ -154,7 +155,7 @@ def to_bytes( audio_codec=audio_codec, audio_format=audio_format, ) - return bytes.getvalue() + return VideoBytes(bytes.getvalue()) def display(self, audio: Optional[AudioTensor] = None) -> None: """ diff --git a/tests/units/typing/tensor/test_video_tensor.py b/tests/units/typing/tensor/test_video_tensor.py index 23ae8115e64..551b52986e5 100644 --- a/tests/units/typing/tensor/test_video_tensor.py +++ b/tests/units/typing/tensor/test_video_tensor.py @@ -9,6 +9,7 @@ from docarray.typing import ( AudioNdArray, AudioTorchTensor, + VideoBytes, VideoNdArray, VideoTorchTensor, ) @@ -137,9 +138,10 @@ def test_save_video_tensor_to_file(video_tensor, tmpdir): parse_obj_as(VideoNdArray, np.zeros((1, 224, 224, 3))), ], ) -def test_save_video_tensor_to_bytes(video_tensor, tmpdir): +def test_save_video_tensor_to_bytes(video_tensor): b = video_tensor.to_bytes() isinstance(b, bytes) + isinstance(b, VideoBytes) @pytest.mark.tensorflow From c5cf9f278cd0364e809ed98ecbb5c671e38a4fb1 Mon Sep 17 00:00:00 2001 From: anna-charlotte Date: Thu, 6 Apr 2023 08:30:27 +0200 Subject: [PATCH 2/5] fix: return audio bytes Signed-off-by: anna-charlotte --- .../typing/tensor/audio/abstract_audio_tensor.py | 9 ++++++--- docarray/typing/tensor/audio/audio_ndarray.py | 4 ++-- .../typing/tensor/audio/audio_tensorflow_tensor.py | 4 ++-- docarray/typing/tensor/audio/audio_torch_tensor.py | 4 ++-- .../typing/tensor/image/abstract_image_tensor.py | 14 ++++++++++---- tests/units/typing/tensor/test_audio_tensor.py | 14 ++++++++++++++ 6 files changed, 36 insertions(+), 13 deletions(-) diff --git a/docarray/typing/tensor/audio/abstract_audio_tensor.py b/docarray/typing/tensor/audio/abstract_audio_tensor.py index 05a455e2c6c..1c2a6e6f2f2 100644 --- a/docarray/typing/tensor/audio/abstract_audio_tensor.py +++ b/docarray/typing/tensor/audio/abstract_audio_tensor.py @@ -5,19 +5,22 @@ from docarray.typing.tensor.abstract_tensor import AbstractTensor from docarray.utils._internal.misc import import_library, is_notebook +if TYPE_CHECKING: + from docarray.typing.bytes.audio_bytes import AudioBytes + T = TypeVar('T', bound='AbstractAudioTensor') MAX_INT_16 = 2**15 class AbstractAudioTensor(AbstractTensor, ABC): - def to_bytes(self): + def to_bytes(self) -> AudioBytes: """ - Convert audio tensor to bytes. + Convert audio tensor to AudioBytes. """ tensor = self.get_comp_backend().to_numpy(self) tensor = (tensor * MAX_INT_16).astype(' bytes: + def to_bytes(self, format: str = 'PNG') -> 'ImageBytes': """ - Convert image tensor to bytes. + Convert image tensor to ImageBytes. :param format: the image format use to store the image, can be 'PNG' , 'JPG' ... - :return: bytes + :return: an ImageBytes object """ PIL = import_library('PIL', raise_error=True) # noqa: F841 from PIL import Image as PILImage @@ -31,7 +35,9 @@ def to_bytes(self, format: str = 'PNG') -> bytes: pil_image.save(buffer, format=format) img_byte_arr = buffer.getvalue() - return img_byte_arr + from docarray.typing.bytes.image_bytes import ImageBytes + + return ImageBytes(img_byte_arr) def save(self, file_path: str) -> None: """ diff --git a/tests/units/typing/tensor/test_audio_tensor.py b/tests/units/typing/tensor/test_audio_tensor.py index 4a145f3c4b7..4f5b3f92c13 100644 --- a/tests/units/typing/tensor/test_audio_tensor.py +++ b/tests/units/typing/tensor/test_audio_tensor.py @@ -6,6 +6,7 @@ from pydantic import parse_obj_as from docarray import BaseDoc +from docarray.typing.bytes.audio_bytes import AudioBytes from docarray.typing.tensor.audio.audio_ndarray import AudioNdArray from docarray.typing.tensor.audio.audio_torch_tensor import AudioTorchTensor from docarray.utils._internal.misc import is_tf_available @@ -120,3 +121,16 @@ def test_save_audio_tensorflow_tensor_to_wav_file(tmpdir): audio_tensor = parse_obj_as(AudioTensorFlowTensor, tf.zeros((1000, 2))) audio_tensor.save(tmp_file) assert os.path.isfile(tmp_file) + + +@pytest.mark.parametrize( + 'audio_tensor', + [ + parse_obj_as(AudioTorchTensor, torch.zeros(1000, 2)), + parse_obj_as(AudioNdArray, np.zeros((1000, 2))), + ], +) +def test_save_audio_tensor_to_bytes(audio_tensor): + b = audio_tensor.to_bytes() + isinstance(b, bytes) + isinstance(b, AudioBytes) From 2138043ef69d2648cd0d5a0192c4227c3b75f12a Mon Sep 17 00:00:00 2001 From: anna-charlotte Date: Thu, 6 Apr 2023 08:31:08 +0200 Subject: [PATCH 3/5] fix: return image bytes Signed-off-by: anna-charlotte --- docarray/typing/bytes/image_bytes.py | 2 +- docarray/typing/tensor/image/image_ndarray.py | 4 ++-- .../tensor/image/image_tensorflow_tensor.py | 4 ++-- .../typing/tensor/image/image_torch_tensor.py | 4 ++-- tests/units/typing/tensor/test_image_tensor.py | 15 ++++++++++++++- 5 files changed, 21 insertions(+), 8 deletions(-) diff --git a/docarray/typing/bytes/image_bytes.py b/docarray/typing/bytes/image_bytes.py index b3b1261f6b2..fe13f99e39c 100644 --- a/docarray/typing/bytes/image_bytes.py +++ b/docarray/typing/bytes/image_bytes.py @@ -7,7 +7,7 @@ from docarray.typing.abstract_type import AbstractType from docarray.typing.proto_register import _register_proto -from docarray.typing.tensor.image import ImageNdArray +from docarray.typing.tensor.image.image_ndarray import ImageNdArray from docarray.utils._internal.misc import import_library if TYPE_CHECKING: diff --git a/docarray/typing/tensor/image/image_ndarray.py b/docarray/typing/tensor/image/image_ndarray.py index 4ac022b300b..3366b3a7c28 100644 --- a/docarray/typing/tensor/image/image_ndarray.py +++ b/docarray/typing/tensor/image/image_ndarray.py @@ -20,14 +20,14 @@ class ImageNdArray(AbstractImageTensor, NdArray): from typing import Optional from docarray import BaseDoc - from docarray.typing import ImageNdArray, ImageUrl + from docarray.typing import ImageBytes, ImageNdArray, ImageUrl class MyImageDoc(BaseDoc): title: str tensor: Optional[ImageNdArray] url: Optional[ImageUrl] - bytes: Optional[bytes] + bytes: Optional[ImageBytes] # from url diff --git a/docarray/typing/tensor/image/image_tensorflow_tensor.py b/docarray/typing/tensor/image/image_tensorflow_tensor.py index 95aba9d5337..ed89868239a 100644 --- a/docarray/typing/tensor/image/image_tensorflow_tensor.py +++ b/docarray/typing/tensor/image/image_tensorflow_tensor.py @@ -24,14 +24,14 @@ class ImageTensorFlowTensor( from typing import Optional from docarray import BaseDoc - from docarray.typing import ImageTensorFlowTensor, ImageUrl + from docarray.typing import ImageBytes, ImageTensorFlowTensor, ImageUrl class MyImageDoc(BaseDoc): title: str tensor: Optional[ImageTensorFlowTensor] url: Optional[ImageUrl] - bytes: Optional[bytes] + bytes: Optional[ImageBytes] doc = MyImageDoc( diff --git a/docarray/typing/tensor/image/image_torch_tensor.py b/docarray/typing/tensor/image/image_torch_tensor.py index bc2a8dbc7db..55f22cf5310 100644 --- a/docarray/typing/tensor/image/image_torch_tensor.py +++ b/docarray/typing/tensor/image/image_torch_tensor.py @@ -22,14 +22,14 @@ class ImageTorchTensor(AbstractImageTensor, TorchTensor, metaclass=metaTorchAndN from typing import Optional from docarray import BaseDoc - from docarray.typing import ImageTorchTensor, ImageUrl + from docarray.typing import ImageBytes, ImageTorchTensor, ImageUrl class MyImageDoc(BaseDoc): title: str tensor: Optional[ImageTorchTensor] url: Optional[ImageUrl] - bytes: Optional[bytes] + bytes: Optional[ImageBytes] doc = MyImageDoc( diff --git a/tests/units/typing/tensor/test_image_tensor.py b/tests/units/typing/tensor/test_image_tensor.py index 73c7b538d45..22c248b3cca 100644 --- a/tests/units/typing/tensor/test_image_tensor.py +++ b/tests/units/typing/tensor/test_image_tensor.py @@ -5,7 +5,7 @@ import torch from pydantic import parse_obj_as -from docarray.typing import ImageNdArray, ImageTorchTensor +from docarray.typing import ImageBytes, ImageNdArray, ImageTorchTensor from docarray.utils._internal.misc import is_tf_available tf_available = is_tf_available() @@ -35,3 +35,16 @@ def test_save_image_tensorflow_tensor_to_file(tmpdir): image_tensor = parse_obj_as(ImageTensorFlowTensor, tf.zeros((224, 224, 3))) image_tensor.save(tmp_file) assert os.path.isfile(tmp_file) + + +@pytest.mark.parametrize( + 'image_tensor', + [ + parse_obj_as(ImageTorchTensor, torch.zeros(224, 224, 3)), + parse_obj_as(ImageNdArray, np.zeros((224, 224, 3))), + ], +) +def test_save_image_tensor_to_bytes(image_tensor): + b = image_tensor.to_bytes() + isinstance(b, bytes) + isinstance(b, ImageBytes) From 68ad410060d8668a551e1548d7ef6310877a7b27 Mon Sep 17 00:00:00 2001 From: anna-charlotte Date: Thu, 6 Apr 2023 08:42:19 +0200 Subject: [PATCH 4/5] fix: audio bytes import Signed-off-by: anna-charlotte --- docarray/typing/tensor/audio/abstract_audio_tensor.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docarray/typing/tensor/audio/abstract_audio_tensor.py b/docarray/typing/tensor/audio/abstract_audio_tensor.py index 1c2a6e6f2f2..0d1636dc164 100644 --- a/docarray/typing/tensor/audio/abstract_audio_tensor.py +++ b/docarray/typing/tensor/audio/abstract_audio_tensor.py @@ -14,10 +14,12 @@ class AbstractAudioTensor(AbstractTensor, ABC): - def to_bytes(self) -> AudioBytes: + def to_bytes(self) -> 'AudioBytes': """ Convert audio tensor to AudioBytes. """ + from docarray.typing.bytes.audio_bytes import AudioBytes + tensor = self.get_comp_backend().to_numpy(self) tensor = (tensor * MAX_INT_16).astype(' Date: Thu, 6 Apr 2023 09:13:39 +0200 Subject: [PATCH 5/5] fix: circular import Signed-off-by: anna-charlotte --- docarray/typing/bytes/video_bytes.py | 8 ++++---- docarray/typing/tensor/video/video_tensor_mixin.py | 8 ++++++-- tests/integrations/predefined_document/test_image.py | 1 + 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/docarray/typing/bytes/video_bytes.py b/docarray/typing/bytes/video_bytes.py index 557f8fbbf74..1f4e22e4fd1 100644 --- a/docarray/typing/bytes/video_bytes.py +++ b/docarray/typing/bytes/video_bytes.py @@ -7,21 +7,21 @@ from docarray.typing.abstract_type import AbstractType from docarray.typing.proto_register import _register_proto +from docarray.typing.tensor import AudioNdArray, NdArray, VideoNdArray from docarray.utils._internal.misc import import_library if TYPE_CHECKING: from pydantic.fields import BaseConfig, ModelField from docarray.proto import NodeProto - from docarray.typing.tensor import AudioNdArray, NdArray, VideoNdArray T = TypeVar('T', bound='VideoBytes') class VideoLoadResult(NamedTuple): - video: 'VideoNdArray' - audio: 'AudioNdArray' - key_frame_indices: 'NdArray' + video: VideoNdArray + audio: AudioNdArray + key_frame_indices: NdArray @_register_proto(proto_type_name='video_bytes') diff --git a/docarray/typing/tensor/video/video_tensor_mixin.py b/docarray/typing/tensor/video/video_tensor_mixin.py index a302e4d2275..d2ed61eacee 100644 --- a/docarray/typing/tensor/video/video_tensor_mixin.py +++ b/docarray/typing/tensor/video/video_tensor_mixin.py @@ -5,11 +5,13 @@ import numpy as np -from docarray.typing.bytes.video_bytes import VideoBytes from docarray.typing.tensor.abstract_tensor import AbstractTensor from docarray.typing.tensor.audio.audio_tensor import AudioTensor from docarray.utils._internal.misc import import_library, is_notebook +if TYPE_CHECKING: + from docarray.typing.bytes.video_bytes import VideoBytes + T = TypeVar('T', bound='VideoTensorMixin') @@ -131,7 +133,7 @@ def to_bytes( audio_frame_rate: int = 48000, audio_codec: str = 'aac', audio_format: str = 'fltp', - ) -> VideoBytes: + ) -> 'VideoBytes': """ Convert video tensor to VideoBytes. @@ -145,6 +147,8 @@ def to_bytes( :return: a VideoBytes object """ + from docarray.typing.bytes.video_bytes import VideoBytes + bytes = BytesIO() self.save( file_path=bytes, diff --git a/tests/integrations/predefined_document/test_image.py b/tests/integrations/predefined_document/test_image.py index b9fa418e29c..77b1b37bf6c 100644 --- a/tests/integrations/predefined_document/test_image.py +++ b/tests/integrations/predefined_document/test_image.py @@ -84,4 +84,5 @@ def test_byte_from_tensor(): img.bytes_ = img.tensor.to_bytes() assert isinstance(img.bytes_, bytes) + assert isinstance(img.bytes_, ImageBytes) assert len(img.bytes_) > 0