ImageTensor

@classmethod
def __docarray_validate_getitem__(cls, item: Any) -> Tuple[int]:
    """This method validates the input to `AbstractTensor.__class_getitem__`.

    It is called at "class creation time",
    i.e. when a class is created with syntax of the form AnyTensor[shape].

    The default implementation tries to cast any `item` to a tuple of ints.
    A subclass can override this method to implement custom validation logic.

    The output of this is eventually passed to
    [`AbstractTensor.__docarray_validate_shape__`]
    [docarray.typing.tensor.abstract_tensor.AbstractTensor.__docarray_validate_shape__]
    as its `shape` argument.

    Raises `ValueError` if the input `item` does not pass validation.

    :param item: The item to validate, passed to `__class_getitem__` (`Tensor[item]`).
    :return: The validated item == the target shape of this tensor.
    """
    if isinstance(item, int):
        item = (item,)
    try:
        item = tuple(item)
    except TypeError:
        raise TypeError(f'{item} is not a valid tensor shape.')
    return item

@classmethod
def __docarray_validate_shape__(cls, t: T, shape: Tuple[Union[int, str], ...]) -> T:
    """Every tensor has to implement this method in order to
    enable syntax of the form AnyTensor[shape].
    It is called when a tensor is assigned to a field of this type.
    i.e. when a tensor is passed to a Document field of type AnyTensor[shape].

    The intended behaviour is as follows:

    - If the shape of `t` is equal to `shape`, return `t`.
    - If the shape of `t` is not equal to `shape`,
        but can be reshaped to `shape`, return `t` reshaped to `shape`.
    - If the shape of `t` is not equal to `shape`
        and cannot be reshaped to `shape`, raise a ValueError.

    :param t: The tensor to validate.
    :param shape: The shape to validate against.
    :return: The validated tensor.
    """
    comp_be = t.get_comp_backend()
    tshape = comp_be.shape(t)
    if tshape == shape:
        return t
    elif any(isinstance(dim, str) or dim == Ellipsis for dim in shape):
        ellipsis_occurrences = [
            pos for pos, dim in enumerate(shape) if dim == Ellipsis
        ]
        if ellipsis_occurrences:
            if len(ellipsis_occurrences) > 1:
                raise ValueError(
                    f'Cannot use Ellipsis (...) more than once for the shape {shape}'
                )
            ellipsis_pos = ellipsis_occurrences[0]
            # Calculate how many dimensions to add. Should be at least 1.
            dimensions_needed = max(len(tshape) - len(shape) + 1, 1)
            shape = (
                shape[:ellipsis_pos]
                + tuple(
                    f'__dim_var_{index}__' for index in range(dimensions_needed)
                )
                + shape[ellipsis_pos + 1 :]
            )

        if len(tshape) != len(shape):
            raise ValueError(
                f'Tensor shape mismatch. Expected {shape}, got {tshape}'
            )
        known_dims: Dict[str, int] = {}
        for tdim, dim in zip(tshape, shape):
            if isinstance(dim, int) and tdim != dim:
                raise ValueError(
                    f'Tensor shape mismatch. Expected {shape}, got {tshape}'
                )
            elif isinstance(dim, str):
                if dim in known_dims and known_dims[dim] != tdim:
                    raise ValueError(
                        f'Tensor shape mismatch. Expected {shape}, got {tshape}'
                    )
                else:
                    known_dims[dim] = tdim
        else:
            return t
    else:
        shape = cast(Tuple[int], shape)
        warnings.warn(
            f'Tensor shape mismatch. Reshaping tensor '
            f'of shape {tshape} to shape {shape}'
        )
        try:
            value = cls._docarray_from_native(comp_be.reshape(t, shape))
            return cast(T, value)
        except RuntimeError:
            raise ValueError(
                f'Cannot reshape tensor of shape {tshape} to shape {shape}'
            )

@abc.abstractmethod
def __getitem__(self: T, item) -> T:
    """Get a slice of this tensor."""
    ...

@abc.abstractmethod
def __iter__(self):
    """Iterate over the elements of this tensor."""
    ...

@abc.abstractmethod
def __setitem__(self, index, value):
    """Set a slice of this tensor."""
    ...

def display(self) -> None:
    """
    Display image data from tensor in notebook.
    """
    if is_notebook():
        PIL = import_library('PIL', raise_error=True)  # noqa: F841
        from PIL import Image as PILImage

        np_array = self.get_comp_backend().to_numpy(self)
        img = PILImage.fromarray(np_array)

        from IPython.display import display

        display(img)
    else:
        warnings.warn('Display of image is only possible in a notebook.')

@staticmethod
@abc.abstractmethod
def get_comp_backend() -> AbstractComputationalBackend:
    """The computational backend compatible with this tensor type."""
    ...

def save(self, file_path: str) -> None:
    """
    Save image tensor to an image file.

    :param file_path: path to an image file. If file is a string, open the file by
        that name, otherwise treat it as a file-like object.
    """
    PIL = import_library('PIL', raise_error=True)  # noqa: F841
    from PIL import Image as PILImage

    comp_backend = self.get_comp_backend()
    np_img = comp_backend.to_numpy(self).astype(np.uint8)

    pil_img = PILImage.fromarray(np_img)
    pil_img.save(file_path)

def to_bytes(self, format: str = 'PNG') -> 'ImageBytes':
    """
    Convert image tensor to [`ImageBytes`][docarray.typing.ImageBytes].

    :param format: the image format use to store the image, can be 'PNG' , 'JPG' ...
    :return: an ImageBytes object
    """
    PIL = import_library('PIL', raise_error=True)  # noqa: F841
    from PIL import Image as PILImage

    if format == 'jpg':
        format = 'jpeg'  # unify it to ISO standard

    tensor = self.get_comp_backend().to_numpy(self)

    mode = 'RGB' if tensor.ndim == 3 else 'L'
    pil_image = PILImage.fromarray(tensor, mode=mode)

    with io.BytesIO() as buffer:
        pil_image.save(buffer, format=format)
        img_byte_arr = buffer.getvalue()

    from docarray.typing.bytes.image_bytes import ImageBytes

    return ImageBytes(img_byte_arr)

@abc.abstractmethod
def to_protobuf(self) -> 'NdArrayProto':
    """Convert DocList into a Protobuf message"""
    ...

def unwrap(self):
    """Return the native tensor object that this DocList tensor wraps."""