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

mblt_vision/utils/postprocess/build_post.py

@@ -0,0 +1,120 @@
+"""
+Postprocessing builder.
+"""
+
+from __future__ import annotations
+
+from ..._tasks import normalize_vision_task
+from .base import PostBase
+from .cls_post import ClsPost
+from .depth_post import DepthPost
+from .semantic_seg_post import SemanticSegPost
+from .yolo_anchor_post import YOLOAnchorDetectionPost, YOLOAnchorSegPost
+from .yolo_anchorless_post import (
+    YOLOAnchorlessDetectionPost,
+    YOLOAnchorlessOBBPost,
+    YOLOAnchorlessPosePost,
+    YOLOAnchorlessSegPost,
+)
+from .yolo_dflfree_post import (
+    YOLODFLFreeDetectionPost,
+    YOLODFLFreeOBBPost,
+    YOLODFLFreePosePost,
+    YOLODFLFreeSegPost,
+)
+from .yolo_nmsfree_post import YOLONMSFreeDetectionPost
+
+
+def build_postprocess(
+    pre_cfg: dict,
+    post_cfg: dict,
+    **kwargs: object,
+) -> PostBase:
+    """Builds a postprocessing object based on the model configuration.
+
+    Args:
+        pre_cfg (dict): Preprocessing configuration from the model info.
+        post_cfg (dict): Postprocessing configuration from the model info.
+            Must contain "task" and relevant flags for the specific task.
+        **kwargs: Optional runtime overrides passed to the postprocessor.
+
+    Returns:
+        PostBase: An instance of a postprocessing class tailored for the task.
+
+    Raises:
+        NotImplementedError: If the specified task is not supported.
+    """
+    task = normalize_vision_task(post_cfg["task"])
+    if task == "image_classification":
+        return ClsPost(pre_cfg, post_cfg)
+    if task == "depth_estimation":
+        return DepthPost(pre_cfg, post_cfg)
+    if task == "semantic_segmentation":
+        return SemanticSegPost(pre_cfg, post_cfg)
+    if task in {"object_detection", "face_detection"}:
+        if post_cfg.get("anchors", False):
+            return YOLOAnchorDetectionPost(
+                pre_cfg,
+                post_cfg,
+                **kwargs,
+            )
+        if post_cfg.get("dflfree", False):  # nms free is only available for detection
+            return YOLODFLFreeDetectionPost(
+                pre_cfg,
+                post_cfg,
+                **kwargs,
+            )
+        if post_cfg.get("nmsfree", False):
+            return YOLONMSFreeDetectionPost(
+                pre_cfg,
+                post_cfg,
+                **kwargs,
+            )
+        return YOLOAnchorlessDetectionPost(
+            pre_cfg,
+            post_cfg,
+            **kwargs,
+        )
+    if task == "instance_segmentation":
+        if post_cfg.get("anchors", False):
+            return YOLOAnchorSegPost(
+                pre_cfg,
+                post_cfg,
+                **kwargs,
+            )
+        if post_cfg.get("dflfree", False):
+            return YOLODFLFreeSegPost(
+                pre_cfg,
+                post_cfg,
+                **kwargs,
+            )
+        return YOLOAnchorlessSegPost(
+            pre_cfg,
+            post_cfg,
+            **kwargs,
+        )
+    if task == "pose_estimation":
+        if post_cfg.get("dflfree", False):
+            return YOLODFLFreePosePost(
+                pre_cfg,
+                post_cfg,
+                **kwargs,
+            )
+        return YOLOAnchorlessPosePost(
+            pre_cfg,
+            post_cfg,
+            **kwargs,
+        )
+    if task == "obb":
+        if post_cfg.get("dflfree", False):
+            return YOLODFLFreeOBBPost(
+                pre_cfg,
+                post_cfg,
+                **kwargs,
+            )
+        return YOLOAnchorlessOBBPost(
+            pre_cfg,
+            post_cfg,
+            **kwargs,
+        )
+    raise NotImplementedError(f"Task {post_cfg['task']} is not implemented yet")

mblt_vision/utils/postprocess/cls_post.py

@@ -0,0 +1,117 @@
+"""
+Classification postprocessing.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import numpy as np
+import torch
+
+from ...datasets import get_dataset_class_names, get_dataset_config
+from ..types import ListTensorLike, TensorLike
+from .base import PostBase
+
+
+class ClsPost(PostBase):
+    """Post-processing for image classification models.
+
+    Typically applies softmax to logits and ensures correct output shape.
+    """
+
+    def __init__(self, pre_cfg: dict, post_cfg: dict) -> None:
+        """Initializes the classification post-processing.
+
+        Args:
+            pre_cfg (dict): Preprocessing configuration.
+            post_cfg (dict): Postprocessing configuration.
+        """
+        super().__init__()
+        self.softmax = post_cfg.get("softmax", False)
+        self.dataset: str | None = None
+        self.num_classes: int | None = None
+        dataset = post_cfg.get("dataset")
+        if dataset is not None:
+            if not isinstance(dataset, str) or not dataset:
+                raise ValueError(
+                    "Classification postprocessing requires post_cfg.dataset to be a non-empty string."
+                )
+            dataset = dataset.lower()
+            dataset_config = get_dataset_config(dataset)
+            if "image_classification" not in dataset_config["tasks"]:
+                raise ValueError(
+                    f"Dataset '{dataset}' does not support image classification."
+                )
+            self.dataset = dataset
+            self.num_classes = len(get_dataset_class_names(dataset))
+
+    def __call__(self, x: TensorLike | ListTensorLike) -> torch.Tensor:
+        """Executes classification post-processing.
+
+        Typically applies softmax to convert logits to probabilities.
+
+        Args:
+            x (TensorLike | ListTensorLike): Raw model outputs.
+                Expected to be pre-softmax logits.
+
+        Returns:
+            torch.Tensor: Softmax probabilities of shape (N, C).
+        """
+        if isinstance(x, Sequence):
+            if len(x) != 1:
+                raise ValueError(
+                    f"Classification postprocessing expects one output tensor, got {len(x)}."
+                )
+            x = x[0]
+        if isinstance(x, np.ndarray):
+            x = torch.from_numpy(x).to(self.device)
+        elif isinstance(x, torch.Tensor):
+            x = x.to(self.device)
+        else:
+            raise TypeError(f"Got unexpected type for x={type(x)}.")
+        if x.ndim == 2:
+            x = x.unsqueeze(-1).unsqueeze(-1)
+        elif x.ndim == 3:
+            if self.num_classes is None:
+                raise ValueError(
+                    "Classification 3D outputs require a configured class count to disambiguate layout."
+                )
+            if x.shape[1] == self.num_classes and x.shape[-1] == 1:
+                x = x.unsqueeze(-1)
+            elif x.shape[0] == self.num_classes:
+                x = x.unsqueeze(0)
+            else:
+                raise ValueError(
+                    f"Unsupported 3D classification output shape {tuple(x.shape)} for {self.num_classes} classes."
+                )
+        if x.ndim != 4:
+            raise ValueError(
+                f"Classification output must be convertible to NCHW, got shape {tuple(x.shape)}."
+            )
+        x = x.flatten(
+            1
+        )  # Classification heads may retain singleton spatial dimensions.
+        if self.num_classes is not None and x.shape[1] != self.num_classes:
+            raise ValueError(
+                f"Classification output has {x.shape[1]} classes, but dataset "
+                f"'{self.dataset}' requires {self.num_classes}."
+            )
+        if not torch.isfinite(x).all():
+            raise ValueError("Classification output scores must all be finite.")
+        if self.softmax:
+            if not bool(((x >= 0) & (x <= 1)).all()):
+                raise ValueError(
+                    "Classification probability outputs must be in [0, 1] when post_cfg.softmax is true."
+                )
+            if not torch.allclose(
+                x.sum(dim=-1),
+                torch.ones(x.shape[0], dtype=x.dtype, device=x.device),
+                rtol=1e-4,
+                atol=1e-4,
+            ):
+                raise ValueError(
+                    "Classification probability outputs must sum to 1 per sample when post_cfg.softmax is true."
+                )
+            return x
+        return x.softmax(dim=-1)