mblt-vision-python 0.0.0__py3-none-any.whl

mblt_vision/utils/preprocess/normalize.py

@@ -0,0 +1,94 @@
+"""Normalization operation for image preprocessing."""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from PIL import Image
+
+from ..types import TensorLike
+from .base import PreOps
+
+STYLE_PARAMS = {
+    "torch": ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    "tf": ([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+    "openai": (
+        [0.48145466, 0.4578275, 0.40821073],
+        [0.26862954, 0.26130258, 0.27577711],
+    ),
+    "cv": ([0.0, 0.0, 0.0], [1.0, 1.0, 1.0]),
+}
+STYLE_LIST = list(STYLE_PARAMS.keys())
+
+
+class Normalize(PreOps):
+    """Normalization layer to scale and shift image data.
+
+    Attributes:
+        style: Data source style (e.g., 'torch', 'tf', 'openai', 'cv').
+        mean: Array of mean values for normalization.
+        std: Array of standard deviation values for normalization.
+    """
+
+    def __init__(self, style: str) -> None:
+        """Initializes the Normalize layer with a specific style.
+
+        Args:
+            style: The preprocessing style to use. Must be one of STYLE_LIST.
+        """
+        super().__init__()
+
+        if not isinstance(style, str):
+            raise TypeError(
+                f"Normalize style must be a string, got {type(style).__name__}."
+            )
+        if style.lower() not in STYLE_LIST:
+            raise ValueError(
+                f"Unsupported Normalize style {style!r}; expected one of {STYLE_LIST}."
+            )
+
+        self.style = style.lower()
+        mean, std = STYLE_PARAMS[self.style]
+        self.mean = np.array(mean)
+        self.std = np.array(std)
+
+    def __call__(self, x: TensorLike | Image.Image) -> np.ndarray:
+        """Applies normalization to the input image or tensor.
+
+        Args:
+            x (TensorLike | Image.Image): Input data as a torch.Tensor, PIL Image, or numpy-like array.
+
+        Returns:
+            np.ndarray: The normalized image as a float32 numpy array.
+        """
+        if isinstance(x, torch.Tensor):
+            x = x.detach().cpu().numpy()
+        elif isinstance(x, Image.Image):
+            x = np.array(x)
+        elif not isinstance(x, np.ndarray):
+            raise TypeError(
+                f"Normalize expects a NumPy array, tensor, or PIL image, got {type(x).__name__}."
+            )
+        if x.ndim != 3:
+            raise ValueError(
+                f"Normalize expects a three-dimensional image, got shape {x.shape}."
+            )
+        x = x.astype(np.float32) / 255.0
+        channels_first = x.shape[0] == len(self.mean)
+        channels_last = x.shape[-1] == len(self.mean)
+        if channels_first and channels_last:
+            raise ValueError(
+                f"Normalize cannot infer channel order from ambiguous shape {x.shape}."
+            )
+        if channels_last:
+            mean, std = self.mean, self.std
+        elif channels_first:
+            mean = self.mean[:, None, None]
+            std = self.std[:, None, None]
+        else:
+            raise ValueError(
+                f"Normalize expects HWC or CHW data with {len(self.mean)} channels, "
+                f"got shape {x.shape}."
+            )
+        x = (x - mean) / std
+        return x.astype(np.float32)

mblt_vision/utils/preprocess/order.py

@@ -0,0 +1,74 @@
+"""
+Channel order preprocessing.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+from ..types import TensorLike
+from .base import PreOps
+
+
+class SetOrder(PreOps):
+    """Sets the channel order of the image to either HWC or CHW format."""
+
+    def __init__(self, shape: str = "HWC") -> None:
+        """Initializes the SetOrder operation.
+
+        Args:
+            shape (str, optional): Target channel order, either "HWC" or "CHW".
+                Defaults to "HWC".
+        """
+        super().__init__()
+        if not isinstance(shape, str):
+            raise TypeError(
+                f"SetOrder shape must be a string, got {type(shape).__name__}."
+            )
+        if shape.lower() not in {"hwc", "chw"}:
+            raise ValueError(
+                f"Unsupported channel order {shape!r}; expected 'HWC' or 'CHW'."
+            )
+        self.shape = shape.lower()
+
+    def __call__(self, x: TensorLike) -> TensorLike:
+        """Reorders the dimensions of the input image.
+
+        Args:
+            x (TensorLike): Input image of shape (3, H, W) or (H, W, 3).
+
+        Returns:
+            TensorLike: Image with the specified channel order.
+        """
+        if not isinstance(x, (np.ndarray, torch.Tensor)):
+            raise TypeError(
+                f"SetOrder expects a NumPy array or tensor, got {type(x).__name__}."
+            )
+        if x.ndim != 3:
+            raise ValueError(
+                f"SetOrder expects a three-dimensional color image, got shape {x.shape}."
+            )
+        channels_first = x.shape[0] == 3
+        channels_last = x.shape[-1] == 3
+        if channels_first and channels_last:
+            raise ValueError(
+                f"SetOrder cannot infer channel order from ambiguous shape {x.shape}."
+            )
+        if channels_first:
+            cdim = 0
+        elif channels_last:
+            cdim = 2
+        else:
+            raise ValueError(
+                f"SetOrder expects HWC or CHW data with three channels, got shape {x.shape}."
+            )
+        if cdim == 0 and self.shape == "hwc":
+            if isinstance(x, torch.Tensor):
+                return torch.permute(x, (1, 2, 0))
+            return np.transpose(x, (1, 2, 0))
+        elif cdim == 2 and self.shape == "chw":
+            if isinstance(x, torch.Tensor):
+                return torch.permute(x, (2, 0, 1))
+            return np.transpose(x, (2, 0, 1))
+        return x

mblt_vision/utils/preprocess/reader.py

@@ -0,0 +1,94 @@
+"""
+Image reader preprocessing.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+
+from ..types import TensorLike
+from ._validation import normalize_uint8_rgb_array
+from .base import PreOps
+
+
+class Reader(PreOps):
+    """
+    Reader for loading images from file paths or converting existing objects.
+    Supports "pil" and "numpy" reading styles.
+
+    For ``style="pil"``, arrays must be ``uint8`` RGB values or finite floating-point
+    RGB values. Floating-point arrays in ``[0, 1]`` are treated as normalized RGB and
+    scaled to ``[0, 255]``; other floating-point arrays must already be in ``[0, 255]``.
+    """
+
+    def __init__(self, style: str) -> None:
+        """Initializes the Reader operation.
+
+        Args:
+            style (str): Reading style, either "pil" or "numpy".
+        """
+        super().__init__()
+        if not isinstance(style, str):
+            raise TypeError(
+                f"Reader style must be a string, got {type(style).__name__}."
+            )
+        if style.lower() not in {"pil", "numpy"}:
+            raise ValueError(
+                f"Unsupported Reader style {style!r}; expected 'pil' or 'numpy'."
+            )
+        self.style = style.lower()
+
+    def __call__(
+        self, x: str | Path | TensorLike | Image.Image
+    ) -> np.ndarray | Image.Image:
+        """Reads/converts the input into an image object.
+
+        Args:
+            x (str | Path | TensorLike | Image.Image): Input image path or image object.
+
+        Returns:
+            np.ndarray | Image.Image: Read image in the specified style.
+        """
+        if self.style == "numpy":
+            if isinstance(x, np.ndarray):
+                return x
+            elif isinstance(x, torch.Tensor):
+                return x.detach().cpu().numpy()
+            elif isinstance(x, (str, Path)):
+                image = cv2.imread(str(x))
+                if image is None:
+                    raise FileNotFoundError(f"Image not found: {x}")
+                return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            elif isinstance(x, Image.Image):
+                return np.array(x)
+            else:
+                raise TypeError(
+                    f"Reader(style='numpy') does not support input type {type(x).__name__}."
+                )
+        elif self.style == "pil":
+            if isinstance(x, np.ndarray):
+                return Image.fromarray(
+                    normalize_uint8_rgb_array(x, operation="Reader(style='pil')")
+                )
+            elif isinstance(x, torch.Tensor):
+                x = x.detach().cpu().numpy()
+                return Image.fromarray(
+                    normalize_uint8_rgb_array(x, operation="Reader(style='pil')")
+                )
+            elif isinstance(x, (str, Path)):
+                return Image.open(x).convert("RGB")
+            elif isinstance(x, Image.Image):
+                return x
+            else:
+                raise TypeError(
+                    f"Reader(style='pil') does not support input type {type(x).__name__}."
+                )
+        else:
+            raise RuntimeError(
+                f"Reader has an invalid validated style: {self.style!r}."
+            )

mblt_vision/utils/preprocess/resize.py

@@ -0,0 +1,202 @@
+"""
+Image resizing preprocessing.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+from ..types import TensorLike
+from ._validation import normalize_image_size
+from .base import PreOps
+
+PIL_INTERP_CODES = {
+    "nearest": Image.Resampling.NEAREST,
+    "bilinear": Image.Resampling.BILINEAR,
+    "bicubic": Image.Resampling.BICUBIC,
+    "box": Image.Resampling.BOX,
+    "hamming": Image.Resampling.HAMMING,
+    "lanczos": Image.Resampling.LANCZOS,
+}
+TORCH_INTERPOLATION_MODES = frozenset({"nearest", "bilinear", "bicubic"})
+
+
+class Resize(PreOps):
+    """Resizes the image to a specified size using various interpolation modes.
+
+    Supports PyTorch tensors in CHW or BCHW format, HWC NumPy arrays, and PIL images.
+    """
+
+    def __init__(
+        self,
+        size: int | list[int],
+        interpolation: str,
+    ) -> None:
+        """
+        Initialize the Resize operation.
+        Args:
+            size (int | list[int]): Target size. If int, the shorter edge is resized to this size
+                maintaining aspect ratio. If [h, w], it is resized to exactly this size.
+            interpolation (str): Interpolation mode (e.g., "bilinear", "bicubic", "nearest").
+        """
+        # Note that this behaves different for npy image and PIL image
+        super().__init__()
+        self.size = (
+            size
+            if isinstance(size, int) and not isinstance(size, bool)
+            else normalize_image_size(size)
+        )
+        if isinstance(self.size, int) and self.size <= 0:
+            raise ValueError(f"size must be positive, got {self.size}.")
+        if interpolation not in PIL_INTERP_CODES:
+            raise ValueError(
+                f"Unsupported resize interpolation {interpolation!r}; expected one of {sorted(PIL_INTERP_CODES)}."
+            )
+        self.interpolation = interpolation
+
+    def __call__(
+        self, x: TensorLike | Image.Image
+    ) -> np.ndarray | torch.Tensor | Image.Image:
+        """Resizes the input image.
+
+        Args:
+            x (TensorLike | Image.Image): Image to be resized.
+
+        Returns:
+            np.ndarray | torch.Tensor | Image.Image: Resized image in the same format as input.
+
+        Raises:
+            TypeError: If input type is not supported.
+            ValueError: If an input has an invalid layout or uses a PIL-only interpolation mode.
+        """
+        if isinstance(x, np.ndarray):
+            self._validate_tensor_interpolation()
+            if x.ndim != 3:
+                raise ValueError(
+                    f"Expected an HWC NumPy array, but got x.shape={x.shape}."
+                )
+            img_h, img_w = x.shape[:2]
+            new_h, new_w = self._compute_resized_output_size(img_h, img_w)
+            if [img_h, img_w] == [new_h, new_w]:
+                return x
+
+            tensor_x = torch.from_numpy(x).to(self.device)
+            tensor_x = tensor_x.permute(2, 0, 1)
+            tensor_x, need_cast, need_squeeze, out_dtype = self._cast_squeeze_in(
+                tensor_x, [torch.float32, torch.float64]
+            )
+            tensor_x = F.interpolate(
+                tensor_x,
+                size=(new_h, new_w),
+                mode=self.interpolation,
+                align_corners=(
+                    False if self.interpolation in ["bilinear", "bicubic"] else None
+                ),
+                antialias=self.interpolation in ["bilinear", "bicubic"],
+            )
+            tensor_x = self._cast_squeeze_out(
+                tensor_x, need_cast, need_squeeze, out_dtype
+            )
+            return tensor_x.permute(1, 2, 0).cpu().numpy()
+        elif isinstance(x, torch.Tensor):
+            self._validate_tensor_interpolation()
+            tensor_x = x.to(self.device)
+        elif isinstance(x, Image.Image):
+            img_w, img_h = x.size
+            new_h, new_w = self._compute_resized_output_size(img_h, img_w)
+            if [img_h, img_w] == [new_h, new_w]:
+                return x
+            return x.resize(
+                size=(new_w, new_h),
+                resample=PIL_INTERP_CODES[self.interpolation],
+            )
+        else:
+            raise TypeError(f"Got unexpected type for x={type(x)}.")
+
+        if tensor_x.ndim not in (3, 4):
+            raise ValueError(
+                f"Expected a CHW or BCHW tensor, but got x.shape={tensor_x.shape}."
+            )
+        img_h, img_w = tensor_x.shape[-2:]
+        new_h, new_w = self._compute_resized_output_size(img_h, img_w)
+        if [img_h, img_w] == [new_h, new_w]:
+            return tensor_x
+        tensor_x, need_cast, need_squeeze, out_dtype = self._cast_squeeze_in(
+            tensor_x, [torch.float32, torch.float64]
+        )
+        tensor_x = F.interpolate(
+            tensor_x,
+            size=(new_h, new_w),
+            mode=self.interpolation,
+            align_corners=(
+                False if self.interpolation in ["bilinear", "bicubic"] else None
+            ),
+            antialias=self.interpolation in ["bilinear", "bicubic"],
+        )
+        tensor_x = self._cast_squeeze_out(tensor_x, need_cast, need_squeeze, out_dtype)
+        return tensor_x.to(self.device)
+
+    def _validate_tensor_interpolation(self) -> None:
+        """Reject interpolation modes unsupported by PyTorch tensors."""
+
+        if self.interpolation not in TORCH_INTERPOLATION_MODES:
+            raise ValueError(
+                f"Resize interpolation {self.interpolation!r} is supported only for PIL images; "
+                f"NumPy arrays and tensors require one of {sorted(TORCH_INTERPOLATION_MODES)}."
+            )
+
+    def _compute_resized_output_size(self, img_h: int, img_w: int) -> list[int]:
+        if isinstance(self.size, int):
+            # to match the shortest side to self.size with the same ratio
+            if img_w <= img_h:
+                new_w = self.size
+                new_h = int(self.size * img_h / img_w)
+            else:
+                new_h = self.size
+                new_w = int(self.size * img_w / img_h)
+        elif isinstance(self.size, list):
+            new_h, new_w = self.size
+        else:
+            raise RuntimeError(f"Resize has an invalid validated size: {self.size!r}.")
+        return [new_h, new_w]
+
+    def _cast_squeeze_in(
+        self, img: torch.Tensor, req_dtypes: list[torch.dtype]
+    ) -> tuple[torch.Tensor, bool, bool, torch.dtype]:
+        need_squeeze = False
+        # make image NCHW
+        if img.ndim < 4:
+            img = img.unsqueeze(dim=0)
+            need_squeeze = True
+        out_dtype = img.dtype
+        need_cast = False
+        if out_dtype not in req_dtypes:
+            need_cast = True
+            req_dtype = req_dtypes[0]
+            img = img.to(req_dtype)
+        return img, need_cast, need_squeeze, out_dtype
+
+    def _cast_squeeze_out(
+        self,
+        img: torch.Tensor,
+        need_cast: bool,
+        need_squeeze: bool,
+        out_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        if need_squeeze:
+            img = img.squeeze(dim=0)
+        if need_cast:
+            if out_dtype in (
+                torch.uint8,
+                torch.int8,
+                torch.int16,
+                torch.int32,
+                torch.int64,
+            ):
+                # it is better to round before cast
+                img = torch.round(img)
+            img = img.to(out_dtype)
+        return img

mblt_vision/utils/preprocess/yolo_pre.py

@@ -0,0 +1,7 @@
+"""Compatibility name for the extracted letterbox operation."""
+
+from .letterbox import LetterBox
+
+YoloPre = LetterBox
+
+__all__ = ["YoloPre"]