Source code for otx.types.image

# Copyright (C) 2023 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

"""OTX image type definition."""

from __future__ import annotations

from enum import Enum, IntEnum, auto

import numpy as np
from torchvision import tv_tensors


class ImageColorChannel(str, Enum):
    """ImageColorChannel definition."""

    RGB = "RGB"
    BGR = "BGR"


[docs] class ImageType(IntEnum): """Enum to indicate the image type in `ImageInfo` class.""" NUMPY = auto() TV_IMAGE = auto() NUMPY_LIST = auto() TV_IMAGE_LIST = auto()
[docs] @classmethod def get_image_type( cls, image: np.ndarray | tv_tensors.Image | list[np.ndarray] | list[tv_tensors.Image], ) -> ImageType: """Infer the image type from the given image object.""" if isinstance(image, np.ndarray): return ImageType.NUMPY if isinstance(image, tv_tensors.Image): return ImageType.TV_IMAGE if isinstance(image, list): image = next(iter(image)) if isinstance(image, np.ndarray): return ImageType.NUMPY_LIST if isinstance(image, tv_tensors.Image): return ImageType.TV_IMAGE_LIST raise TypeError(image)