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

mblt_vision/utils/postprocess/yolo_nmsfree_post.py

@@ -0,0 +1,176 @@
+"""
+YOLO NMS-free postprocessing.
+"""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import torch
+
+from .common import dist2bbox, dual_topk
+from .yolo_anchorless_post import YOLOAnchorlessDetectionPost, _AnchorlessNMSInput
+
+
+class YOLONMSFreeDetectionPost(YOLOAnchorlessDetectionPost):
+    """Postprocessing for YOLO NMS-free models."""
+
+    max_det = 300
+
+    def non_e2e(self, x: list[torch.Tensor]) -> torch.Tensor:
+        """Return the export-style output tensor for NMS-free YOLO models."""
+        if len(x) == 2:
+            converted = cast(torch.Tensor, self.conversion(x))
+            return self._stack_topk_outputs(self.filter_conversion(converted))
+
+        rearranged = cast(torch.Tensor, self.rearrange(x))
+        return self.decode_batch(rearranged)
+
+    def _stack_topk_outputs(self, outputs: list[torch.Tensor]) -> torch.Tensor:
+        """Pad or trim per-image detections to a fixed batch tensor."""
+        padded_outputs = []
+        for output in outputs:
+            output = output[: self.max_det]
+            if output.shape[0] < self.max_det:
+                pad = torch.zeros(
+                    (self.max_det - output.shape[0], 6),
+                    dtype=output.dtype,
+                    device=output.device,
+                )
+                output = torch.cat([output, pad], dim=0)
+            padded_outputs.append(output)
+        return torch.stack(padded_outputs, dim=0)
+
+    def decode_batch(self, x: torch.Tensor) -> torch.Tensor:
+        """Decode every anchor, then apply batched top-k selection for export-style output."""
+        box, scores = torch.split(x, [self.reg_max * 4, self.nc], dim=1)
+        anchors = self.anchors_as_tensor().unsqueeze(0)
+        stride = self.stride_as_tensor().unsqueeze(0)
+        dbox = dist2bbox(self.dfl(box), anchors, xywh=False, dim=1) * stride
+        decoded = torch.cat([dbox, scores], dim=1).transpose(1, 2)
+        return self._stack_topk_outputs(
+            [
+                dual_topk(
+                    image,
+                    self.nc,
+                    self.n_extra,
+                    max_det=self.max_det,
+                    conf_thres=self.conf_thres,
+                    score_is_logits=True,
+                )
+                for image in decoded
+            ]
+        )
+
+    def _pre_process(self, x: list[torch.Tensor]) -> tuple[Any, torch.Tensor | None]:
+        """Preprocesses inputs for NMS-free models.
+
+        Args:
+            x (list[torch.Tensor]): Raw model outputs.
+
+        Returns:
+            tuple: (processed_detections, None).
+        """
+        if len(x) == 2:
+            converted = cast(torch.Tensor, self.conversion(x))
+            return self.filter_conversion(converted), None
+        rearranged = cast(torch.Tensor, self.rearrange(x))
+        return self.decode(rearranged), None
+
+    def conversion(
+        self, x: list[torch.Tensor]
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        """Convert input tensors.
+        Args:
+            x (list[torch.Tensor]): Input tensors.
+        Returns:
+            torch.Tensor: Converted tensor.
+        """
+        # sort by element number
+        x = sorted(x, key=lambda x: x.size(), reverse=self.nc < 4)
+        return torch.cat(x, dim=-1).squeeze(1)  # [b, 8400, 84]
+
+    def filter_conversion(self, x: torch.Tensor) -> list[torch.Tensor]:
+        """Filters out low-confidence detections from a single output tensor.
+
+        Args:
+            x (torch.Tensor): Model output tensor.
+
+        Returns:
+            list[torch.Tensor]: Decoded and filtered outputs for each image.
+        """
+        x_list = torch.split(x, 1, dim=0)  # [(1, 8400, 84), (1, 8400, 84), ...]
+
+        return [
+            dual_topk(xi.squeeze(0), self.nc, self.n_extra, conf_thres=self.conf_thres)
+            for xi in x_list
+        ]
+
+    def process_box_cls(self, box_cls: torch.Tensor) -> torch.Tensor:
+        """Processes detection results for a single image.
+
+        Args:
+            box_cls: Raw detections for one image.
+
+        Returns:
+            Decoded and top-k filtered detections.
+        """
+        ic = torch.amax(box_cls[-self.nc :, :], dim=0) > self.inv_conf_thres
+        box_cls = box_cls[:, ic]  # (144, *)
+        if box_cls.numel() == 0:
+            return box_cls.new_zeros((0, 6))
+        anchors = self.anchors_as_tensor()
+        stride = self.stride_as_tensor()
+        box, scores = torch.split(
+            box_cls[None], [self.reg_max * 4, self.nc], dim=1
+        )  # (1, 64, *), (1, 80, *)
+        dbox = (
+            dist2bbox(
+                self.dfl(box),
+                anchors[:, ic],
+                xywh=False,
+                dim=1,
+            )
+            * stride[:, ic]
+        )
+        pre_topk = (
+            torch.cat([dbox, scores], dim=1).squeeze(0).transpose(0, 1)
+        )  # (*, 84)
+        return dual_topk(
+            pre_topk,
+            self.nc,
+            self.n_extra,
+            conf_thres=self.conf_thres,
+            score_is_logits=True,
+        )
+
+    def nms(
+        self,
+        x: _AnchorlessNMSInput | torch.Tensor | list[torch.Tensor],
+        max_det: int = 300,
+        max_nms: int = 30000,
+        max_wh: int = 7680,
+        multi_label: bool = False,
+    ) -> list[torch.Tensor]:
+        """Perform Non-Maximum Suppression (no-op for NMS-free models).
+
+        Args:
+            x: Decoded detections, optionally with source-layout provenance.
+            max_det (int, optional): Maximum number of detections to keep. Defaults to 300.
+            max_nms (int, optional): Maximum candidates for NMS. Defaults to 30000.
+            max_wh (int, optional): Maximum box width/height. Defaults to 7680.
+            multi_label: Ignored because NMS-free outputs already select one
+                class per candidate.
+
+        Returns:
+            list[torch.Tensor]: Per-image detections with padded zero rows removed.
+        """
+        del max_det, max_nms, max_wh, multi_label
+        if isinstance(x, _AnchorlessNMSInput):
+            x = x.detections
+        if isinstance(x, list):
+            return x
+        return [xi[xi[:, 4] > 0] for xi in x]
+
+
+YOLONMSFreePost = YOLONMSFreeDetectionPost

mblt_vision/utils/preprocess/__init__.py

@@ -0,0 +1,27 @@
+"""
+Preprocessing utilities for vision models.
+"""
+
+from .base import PreBase, PreOps
+from .build_pre import build_preprocess
+from .center_crop import CenterCrop
+from .letterbox import LetterBox, letterbox_semantic_mask
+from .normalize import Normalize
+from .order import SetOrder
+from .reader import Reader
+from .resize import Resize
+from .yolo_pre import YoloPre
+
+__all__ = [
+    "CenterCrop",
+    "LetterBox",
+    "Normalize",
+    "PreBase",
+    "PreOps",
+    "Reader",
+    "Resize",
+    "SetOrder",
+    "YoloPre",
+    "build_preprocess",
+    "letterbox_semantic_mask",
+]

mblt_vision/utils/preprocess/_validation.py

@@ -0,0 +1,66 @@
+"""Validation helpers shared by image preprocessing operations."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import numpy as np
+
+
+def normalize_uint8_rgb_array(image: np.ndarray, *, operation: str) -> np.ndarray:
+    """Return byte RGB data after validating or scaling a floating-point image.
+
+    ``[0, 1]`` floating-point input is treated as normalized RGB and scaled to
+    ``[0, 255]``. Other floating-point input must already lie in ``[0, 255]``.
+    """
+
+    if image.dtype == np.uint8:
+        return image
+    if not np.issubdtype(image.dtype, np.floating):
+        raise TypeError(
+            f"{operation} accepts uint8 arrays or floating-point arrays with RGB "
+            f"values in [0, 1] or [0, 255]; got {image.dtype}."
+        )
+    if not np.isfinite(image).all():
+        raise ValueError(
+            f"{operation} requires floating-point image arrays to contain only "
+            "finite RGB values."
+        )
+
+    min_value = float(image.min())
+    max_value = float(image.max())
+    if min_value < 0.0 or max_value > 255.0:
+        raise ValueError(
+            f"{operation} accepts floating-point RGB values only in [0, 1] or "
+            f"[0, 255]; got range [{min_value}, {max_value}]."
+        )
+    if max_value <= 1.0:
+        image = image * 255.0
+    return np.rint(image).astype(np.uint8)
+
+
+def normalize_image_size(size: int | Sequence[int], *, name: str = "size") -> list[int]:
+    """Normalize a positive scalar or two-dimensional image size to ``[height, width]``."""
+
+    if isinstance(size, bool):
+        raise TypeError(
+            f"{name} must be a positive integer or a two-item integer sequence, got bool."
+        )
+    if isinstance(size, int):
+        if size <= 0:
+            raise ValueError(f"{name} must be positive, got {size}.")
+        return [size, size]
+    if isinstance(size, Sequence) and not isinstance(size, (str, bytes)):
+        if len(size) != 2:
+            raise ValueError(f"{name} must contain exactly two items, got {size!r}.")
+        if not all(
+            isinstance(value, int) and not isinstance(value, bool) for value in size
+        ):
+            raise TypeError(f"{name} items must be integers, got {size!r}.")
+        normalized = [int(size[0]), int(size[1])]
+        if any(value <= 0 for value in normalized):
+            raise ValueError(f"{name} items must be positive, got {size!r}.")
+        return normalized
+    raise TypeError(
+        f"{name} must be a positive integer or a two-item integer sequence, got {type(size).__name__}."
+    )

mblt_vision/utils/preprocess/base.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+import torch
+
+
+def _spatial_shape(value: Any) -> tuple[int, int] | None:
+    """Return the height and width of image-like preprocessing input."""
+
+    shape = getattr(value, "shape", None)
+    if shape is None or len(shape) < 2:
+        return None
+    if len(shape) == 2:
+        return int(shape[0]), int(shape[1])
+    if len(shape) == 3 and int(shape[-1]) in {1, 3, 4}:
+        return int(shape[0]), int(shape[1])
+    return int(shape[-2]), int(shape[-1])
+
+
+class PreOps(ABC):
+    """Abstract base class for individual preprocessing operations.
+
+    Attributes:
+            device: The torch device where tensors should be placed.
+    """
+
+    def __init__(self) -> None:
+        """Initializes the preprocessing operation."""
+        super().__init__()
+        self.device = torch.device("cpu")
+
+    @abstractmethod
+    def __call__(
+        self,
+        x: Any,
+    ) -> Any:
+        """Executes the preprocess operation.
+
+        Args:
+                x: Input data to be processed.
+
+        Returns:
+                Processed data.
+        """
+
+    def to(
+        self,
+        device: str | torch.device,
+    ) -> None:
+        """Move the operation to the specified device.
+
+        Args:
+                device: Device to move the operation to.
+        """
+        if isinstance(device, str):
+            self.device = torch.device(device)
+        elif isinstance(device, torch.device):
+            self.device = device
+        else:
+            raise TypeError(f"Got unexpected type for device={type(device)}.")
+        for name, value in self.__dict__.items():
+            if isinstance(value, torch.Tensor):
+                setattr(self, name, value.to(self.device))
+
+
+class PreBase:
+    """Base class for orchestrating a series of preprocessing operations.
+
+    Attributes:
+            Ops: List of ordered PreOps instances to be applied.
+            device: The torch device being used.
+    """
+
+    def __init__(
+        self,
+        Ops: list[PreOps],
+    ) -> None:
+        """Initializes the PreBase class with a list of operations.
+
+        Args:
+                Ops: List of ordered PreOps instances to be applied.
+        """
+        self.Ops = Ops
+        self._check_ops()
+        self.device = torch.device("cpu")
+
+    def _check_ops(self) -> None:
+        """Check if the operations are valid."""
+        for op in self.Ops:
+            if not isinstance(op, PreOps):
+                raise TypeError(f"Got unsupported type={type(op)}.")
+
+    def __call__(
+        self,
+        x: Any,
+    ) -> Any:
+        """Applies the sequence of preprocessing operations to the input.
+
+        Args:
+                x: Initial input data.
+
+        Returns:
+                Fully processed data.
+        """
+        for op in self.Ops:
+            x = op(x)
+        return x
+
+    def with_metadata(
+        self,
+        x: Any,
+    ) -> tuple[Any, dict[str, Any]]:
+        """Apply preprocessing and return metadata produced by preprocessing operations.
+
+        Args:
+                x: Initial input data.
+
+        Returns:
+                A tuple of the processed data and collected metadata.
+        """
+        metadata: dict[str, Any] = {}
+        img0_shape = _spatial_shape(x)
+        if img0_shape is not None:
+            metadata["img0_shape"] = img0_shape
+        for op in self.Ops:
+            x = op(x)
+            if "img0_shape" not in metadata:
+                img0_shape = _spatial_shape(x)
+                if img0_shape is not None:
+                    metadata["img0_shape"] = img0_shape
+            ratio_pad = getattr(op, "ratio_pad", None)
+            if ratio_pad is not None:
+                metadata["ratio_pad"] = ratio_pad
+        return x, metadata
+
+    def to(
+        self,
+        device: str | torch.device,
+    ) -> None:
+        """Move the operations to the specified device.
+
+        Args:
+                device: Device to move the operations to.
+        """
+        if isinstance(device, str):
+            self.device = torch.device(device)
+        elif isinstance(device, torch.device):
+            self.device = device
+        else:
+            raise TypeError(f"Got unexpected type for device={type(device)}.")
+        for name, value in self.__dict__.items():
+            if isinstance(value, torch.Tensor):
+                setattr(self, name, value.to(self.device))
+        for op in self.Ops:
+            op.to(self.device)

mblt_vision/utils/preprocess/build_pre.py

@@ -0,0 +1,48 @@
+"""
+Preprocessing builder.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from .base import PreBase
+from .center_crop import CenterCrop
+from .letterbox import LetterBox
+from .normalize import Normalize
+from .order import SetOrder
+from .reader import Reader
+from .resize import Resize
+
+
+def build_preprocess(
+    pre_cfg: Mapping[str, Mapping[str, Any]],
+) -> PreBase:
+    """Builds a preprocessing pipeline based on the configuration.
+
+    Args:
+            pre_cfg: Preprocessing configuration mapping operations to attributes.
+
+    Returns:
+            An orchestrator for the sequence of preprocessing steps.
+    """
+    res = []
+    for pre_type, pre_attr in pre_cfg.items():
+        pre_type_lower = pre_type.lower()
+        if pre_type_lower == Reader.__name__.lower():
+            res.append(Reader(**pre_attr))
+        elif pre_type_lower == Resize.__name__.lower():
+            res.append(Resize(**pre_attr))
+        elif pre_type_lower == CenterCrop.__name__.lower():
+            res.append(CenterCrop(**pre_attr))
+        elif pre_type_lower == SetOrder.__name__.lower():
+            res.append(SetOrder(**pre_attr))
+        elif pre_type_lower == LetterBox.__name__.lower():
+            res.append(LetterBox(**pre_attr))
+        elif pre_type_lower == Normalize.__name__.lower():
+            res.append(Normalize(**pre_attr))
+        else:
+            raise ValueError(f"Got unsupported pre_type={pre_type}.")
+
+    return PreBase(res)

mblt_vision/utils/preprocess/center_crop.py

@@ -0,0 +1,75 @@
+"""
+Center crop preprocessing.
+"""
+
+from __future__ import annotations
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+
+from ..types import TensorLike
+from ._validation import normalize_image_size
+from .base import PreOps
+
+
+class CenterCrop(PreOps):
+    """
+    Center crop the image to a specified size.
+    """
+
+    def __init__(self, size: int | list[int] | tuple[int, int]) -> None:
+        """Initializes the CenterCrop operation.
+
+        Args:
+            size (int | list[int]): Target size [h, w]. If int, size is [size, size].
+        """
+        super().__init__()
+        self.size = normalize_image_size(size)
+
+    def __call__(self, x: TensorLike | Image.Image) -> np.ndarray:
+        """Applies center crop to the image.
+
+        Args:
+            x (np.ndarray | torch.Tensor | Image.Image): Input image.
+
+        Returns:
+            np.ndarray: Center-cropped image in HWC format.
+        """
+        if isinstance(x, torch.Tensor):
+            image = x.detach().cpu().numpy()
+        elif isinstance(x, Image.Image):
+            image = np.array(x)
+        elif isinstance(x, np.ndarray):
+            image = x
+        else:
+            raise TypeError(
+                f"CenterCrop expects a NumPy array, tensor, or PIL image, got {type(x).__name__}."
+            )
+        if image.ndim != 3:
+            raise ValueError(
+                f"CenterCrop expects a three-dimensional image, got shape {image.shape}."
+            )
+        H, W = image.shape[:2]
+        if (self.size[0] == H) and (self.size[1] == W):
+            return image
+        elif (self.size[1] > W) or (self.size[0] > H):
+            image = cv2.copyMakeBorder(
+                image,
+                (self.size[0] - H) // 2 if self.size[0] > H else 0,
+                (self.size[0] - H + 1) // 2 if self.size[0] > H else 0,
+                (self.size[1] - W) // 2 if self.size[1] > W else 0,
+                (self.size[1] - W + 1) // 2 if self.size[1] > W else 0,
+                cv2.BORDER_CONSTANT,
+                value=(0.0,),
+            )
+            H, W = image.shape[:2]
+        crop_top = round((H - self.size[0]) / 2.0)
+        crop_left = round((W - self.size[1]) / 2.0)
+        image = image[
+            crop_top : crop_top + self.size[0],
+            crop_left : crop_left + self.size[1],
+            :,
+        ]
+        return image.astype(np.uint8)

mblt_vision/utils/preprocess/letterbox.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import cv2
+import numpy as np
+import torch
+
+from ..letterbox import LetterBoxGeometry, RatioPad
+from ..types import TensorLike
+from ._validation import normalize_image_size, normalize_uint8_rgb_array
+from .base import PreOps
+
+
+def _apply_letterbox(
+    image: np.ndarray,
+    img_size: list[int],
+    interpolation: int,
+    padding_value: int | tuple[int, int, int],
+) -> tuple[np.ndarray, RatioPad]:
+    """Resize and pad an array while preserving its aspect ratio.
+
+    Args:
+        image: Image or two-dimensional semantic mask.
+        img_size: Target size as ``[height, width]``.
+        interpolation: OpenCV interpolation mode.
+        padding_value: Constant border value.
+
+    Returns:
+        The letterboxed array and its resize/padding metadata.
+    """
+
+    input_shape = (int(img_size[0]), int(img_size[1]))
+    original_shape = (int(image.shape[0]), int(image.shape[1]))
+    geometry = LetterBoxGeometry.from_shapes(input_shape, original_shape)
+    resized_height, resized_width = geometry.resized_shape
+    if image.shape[:2] != geometry.resized_shape:
+        image = cv2.resize(
+            image, (resized_width, resized_height), interpolation=interpolation
+        )
+    top, bottom, left, right = geometry.borders
+    image = cv2.copyMakeBorder(
+        image,
+        top,
+        bottom,
+        left,
+        right,
+        cv2.BORDER_CONSTANT,
+        value=padding_value,
+    )
+    return image, geometry.ratio_pad
+
+
+def letterbox_semantic_mask(
+    mask: np.ndarray,
+    img_size: list[int],
+    ignore_label: int = 255,
+) -> tuple[np.ndarray, RatioPad]:
+    """Letterbox a semantic mask without interpolating class IDs.
+
+    Args:
+        mask: Two-dimensional semantic class map.
+        img_size: Target size as ``[height, width]``.
+        ignore_label: Class value used for padded pixels.
+
+    Returns:
+        The letterboxed mask and its resize/padding metadata.
+
+    Raises:
+        ValueError: If the mask is not two-dimensional.
+    """
+
+    if mask.ndim != 2:
+        raise ValueError(
+            f"Semantic masks must be two-dimensional, got shape {mask.shape}."
+        )
+    return _apply_letterbox(mask, img_size, cv2.INTER_NEAREST, ignore_label)
+
+
+class LetterBox(PreOps):
+    """Preprocessing for YOLO models, implementing letterbox resizing.
+
+    Resizes the image while maintaining aspect ratio, adding padding to meet
+    target dimensions. Floating-point RGB inputs in ``[0, 1]`` are scaled to
+    byte RGB; other floating-point values must be finite and in ``[0, 255]``.
+    Based on Ultralytics implementation.
+
+    Ref: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/data/augment.py#L1535
+    """
+
+    def __init__(self, img_size: list[int]) -> None:
+        """Initializes LetterBox with target image size.
+
+        Args:
+            img_size (list[int]): Target image size [h, w].
+        """
+        super().__init__()
+        self.img_size = normalize_image_size(img_size, name="img_size")
+        self.ratio_pad: tuple[tuple[float, float], tuple[float, float]] | None = None
+
+    def __call__(self, x: TensorLike) -> torch.Tensor:
+        """Executes YOLO preprocessing (letterbox resizing).
+
+        Args:
+            x (TensorLike): Input image.
+
+        Returns:
+            torch.Tensor: Preprocessed image in HWC format on the selected device.
+        """
+        if isinstance(x, torch.Tensor):
+            x = x.detach().cpu().numpy()
+        elif not isinstance(x, np.ndarray):
+            raise TypeError(
+                f"LetterBox expects a NumPy array or tensor, got {type(x).__name__}."
+            )
+        if x.ndim != 3:
+            raise ValueError(f"LetterBox expects an HWC image, got shape {x.shape}.")
+        x = normalize_uint8_rgb_array(x, operation="LetterBox")
+        img, self.ratio_pad = _apply_letterbox(
+            x,
+            self.img_size,
+            cv2.INTER_LINEAR,
+            (114, 114, 114),
+        )
+        return torch.from_numpy(img).to(self.device).byte()