-
Notifications
You must be signed in to change notification settings - Fork 243
Expand file tree
/
Copy pathimage_tensor.py
More file actions
120 lines (92 loc) · 3.51 KB
/
Copy pathimage_tensor.py
File metadata and controls
120 lines (92 loc) · 3.51 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
from typing import TYPE_CHECKING, Any, Type, TypeVar, Union, cast
import numpy as np
from docarray.typing.tensor.image.abstract_image_tensor import AbstractImageTensor
from docarray.typing.tensor.image.image_ndarray import ImageNdArray
from docarray.typing.tensor.tensor import AnyTensor
from docarray.utils._internal.misc import (
is_jax_available,
is_tf_available,
is_torch_available,
)
jax_available = is_jax_available()
if jax_available:
import jax.numpy as jnp # type: ignore
from docarray.typing.tensor.image.image_jax_array import ImageJaxArray
from docarray.typing.tensor.jaxarray import JaxArray
torch_available = is_torch_available()
if torch_available:
import torch
from docarray.typing.tensor.image.image_torch_tensor import ImageTorchTensor
from docarray.typing.tensor.torch_tensor import TorchTensor
tf_available = is_tf_available()
if tf_available:
import tensorflow as tf # type: ignore
from docarray.typing.tensor.image.image_tensorflow_tensor import (
ImageTensorFlowTensor,
)
from docarray.typing.tensor.tensorflow_tensor import TensorFlowTensor
if TYPE_CHECKING:
from pydantic import BaseConfig
from pydantic.fields import ModelField
T = TypeVar("T", bound="ImageTensor")
class ImageTensor(AnyTensor, AbstractImageTensor):
"""
Represents an image tensor object that can be used with TensorFlow, PyTorch, and NumPy type.
---
'''python
from docarray import BaseDoc
from docarray.typing import ImageTensor
class MyImageDoc(BaseDoc):
image: ImageTensor
# Example usage with TensorFlow:
import tensorflow as tf
doc = MyImageDoc(image=tf.zeros((1000, 2)))
type(doc.image) # ImageTensorFlowTensor
# Example usage with PyTorch:
import torch
doc = MyImageDoc(image=torch.zeros((1000, 2)))
type(doc.image) # ImageTorchTensor
# Example usage with NumPy:
import numpy as np
doc = MyImageDoc(image=np.zeros((1000, 2)))
type(doc.image) # ImageNdArray
'''
---
Returns:
Union[ImageTorchTensor, ImageTensorFlowTensor, ImageNdArray]: The validated and converted image tensor.
Raises:
TypeError: If the input type is not one of [torch.Tensor, tensorflow.Tensor, numpy.ndarray].
"""
@classmethod
def __get_validators__(cls):
yield cls.validate
@classmethod
def validate(
cls: Type[T],
value: Union[T, np.ndarray, Any],
field: "ModelField",
config: "BaseConfig",
):
if torch_available:
if isinstance(value, TorchTensor):
return cast(ImageTorchTensor, value)
elif isinstance(value, torch.Tensor):
return ImageTorchTensor._docarray_from_native(value) # noqa
if tf_available:
if isinstance(value, TensorFlowTensor):
return cast(ImageTensorFlowTensor, value)
elif isinstance(value, tf.Tensor):
return ImageTensorFlowTensor._docarray_from_native(value) # noqa
if jax_available:
if isinstance(value, JaxArray):
return cast(ImageJaxArray, value)
elif isinstance(value, jnp.ndarray):
return ImageJaxArray._docarray_from_native(value) # noqa
try:
return ImageNdArray.validate(value, field, config)
except Exception: # noqa
pass
raise TypeError(
f"Expected one of [torch.Tensor, tensorflow.Tensor, numpy.ndarray] "
f"compatible type, got {type(value)}"
)