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

mblt_vision/utils/postprocess/depth_post.py

@@ -0,0 +1,106 @@
+"""Postprocessing for monocular depth-estimation models."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import torch
+import torch.nn.functional as functional
+
+from ..letterbox import RatioPad
+from ..types import ListTensorLike, TensorLike
+from ._letterbox import crop_letterbox, get_letterbox_input_shape, resolve_ratio_pads
+from .base import PostBase
+from .common import normalize_image_shapes
+
+
+class DepthPost(PostBase):
+    """Normalize depth outputs and undo letterbox padding when metadata is available."""
+
+    def __init__(self, pre_cfg: dict[str, Any], post_cfg: dict[str, Any]) -> None:
+        """Initialize depth restoration from the model letterbox configuration."""
+
+        super().__init__()
+        del post_cfg
+        self.input_shape = get_letterbox_input_shape(
+            pre_cfg, "Depth estimation", "Depth"
+        )
+
+    def __call__(
+        self,
+        x: TensorLike | ListTensorLike,
+        img0_shape: tuple[int, int] | Sequence[tuple[int, int]] | None = None,
+        ratio_pad: RatioPad | Sequence[RatioPad | None] | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | list[torch.Tensor]:
+        """Return normalized depth maps, optionally restored to original image sizes."""
+
+        if kwargs:
+            raise TypeError(
+                f"Unexpected depth postprocess kwargs: {', '.join(sorted(kwargs))}"
+            )
+        depth = self._normalize_output(x)
+        if img0_shape is None:
+            return depth
+
+        shapes = normalize_image_shapes(img0_shape, depth.shape[0])
+        pads = resolve_ratio_pads(ratio_pad, depth.shape[0], shapes, self.input_shape)
+        restored = [
+            self._restore(depth[index], shapes[index], pads[index])
+            for index in range(depth.shape[0])
+        ]
+        return restored[0] if len(restored) == 1 else restored
+
+    def _normalize_output(self, x: TensorLike | ListTensorLike) -> torch.Tensor:
+        """Validate dense output and normalize it to input-sized ``[B, H, W]`` tensors."""
+
+        if isinstance(x, (list, tuple)):
+            if len(x) != 1:
+                raise ValueError(
+                    f"Depth estimation expects one output tensor, received {len(x)}."
+                )
+            x = x[0]
+        depth = torch.as_tensor(x)
+        if depth.ndim == 4:
+            if depth.shape[1] == 1:
+                depth = depth[:, 0]
+            elif depth.shape[-1] == 1:
+                depth = depth[..., 0]
+            else:
+                raise ValueError(
+                    f"Depth estimation expects [B, 1, H, W] or [B, H, W, 1], got {tuple(depth.shape)}."
+                )
+        elif depth.ndim == 3 and depth.shape[-1] == 1:
+            depth = depth[..., 0].unsqueeze(0)
+        elif depth.ndim != 3:
+            raise ValueError(
+                "Depth estimation expects [B, H, W], [B, 1, H, W], [H, W, 1], or [B, H, W, 1], "
+                f"got {tuple(depth.shape)}."
+            )
+        depth = depth.to(device=self.device, dtype=torch.float32)
+        if not bool(torch.isfinite(depth).all()):
+            raise ValueError("Depth estimation output must contain only finite values.")
+        if tuple(depth.shape[-2:]) == self.input_shape:
+            return depth
+
+        quarter_shape = tuple(dimension // 4 for dimension in self.input_shape)
+        if tuple(depth.shape[-2:]) == quarter_shape:
+            return functional.interpolate(
+                depth[:, None], scale_factor=4.0, mode="bilinear", align_corners=False
+            )[:, 0]
+
+        raise ValueError(
+            f"Depth estimation output spatial shape must be {self.input_shape} or quarter-resolution {quarter_shape}, "
+            f"got {tuple(depth.shape[-2:])}."
+        )
+
+    def _restore(
+        self, depth: torch.Tensor, shape: tuple[int, int], ratio_pad: RatioPad
+    ) -> torch.Tensor:
+        """Crop padded depth pixels and bilinearly resize to an original image shape."""
+
+        cropped = crop_letterbox(depth, shape, ratio_pad, self.input_shape, "Depth")
+        return functional.interpolate(
+            cropped[None, None], size=shape, mode="bilinear", align_corners=False
+        )[0, 0]

mblt_vision/utils/postprocess/semantic_seg_post.py

@@ -0,0 +1,192 @@
+"""Postprocessing for semantic-segmentation models."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import torch
+import torch.nn.functional as functional
+
+from ..letterbox import RatioPad
+from ..types import ListTensorLike, TensorLike
+from ._letterbox import crop_letterbox, get_letterbox_input_shape, resolve_ratio_pads
+from .base import PostBase
+from .common import normalize_image_shapes
+
+
+class SemanticSegPost(PostBase):
+    """Convert semantic logits to class maps and undo letterbox padding."""
+
+    NC_BY_DATASET: dict[str, int] = {
+        "ade20k": 150,
+        "cityscapes": 19,
+    }
+
+    def __init__(self, pre_cfg: dict[str, Any], post_cfg: dict[str, Any]) -> None:
+        """Initialize semantic output handling for a configured dataset taxonomy."""
+
+        super().__init__()
+        self.input_shape = get_letterbox_input_shape(
+            pre_cfg, "Semantic segmentation", "Semantic"
+        )
+        dataset = post_cfg.get("dataset")
+        if not isinstance(dataset, str):
+            raise ValueError(
+                "Semantic segmentation requires a string dataset in post_cfg."
+            )
+        self.dataset = dataset.lower()
+        dataset_nc = self.NC_BY_DATASET.get(self.dataset)
+        configured_nc = post_cfg.get("nc")
+        if configured_nc is None:
+            if dataset_nc is None:
+                raise ValueError(
+                    f"Semantic segmentation requires nc for unknown dataset '{self.dataset}'."
+                )
+            self.nc = dataset_nc
+        else:
+            self.nc = int(configured_nc)
+        if dataset_nc is not None and self.nc != dataset_nc:
+            raise ValueError(
+                f"nc={configured_nc} conflicts with semantic dataset '{self.dataset}', which requires nc={dataset_nc}."
+            )
+
+    def __call__(
+        self,
+        x: TensorLike | ListTensorLike,
+        img0_shape: tuple[int, int] | Sequence[tuple[int, int]] | None = None,
+        ratio_pad: RatioPad | Sequence[RatioPad | None] | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | list[torch.Tensor]:
+        """Return class maps, optionally restored to original image sizes."""
+
+        if kwargs:
+            raise TypeError(
+                f"Unexpected semantic postprocess kwargs: {', '.join(sorted(kwargs))}"
+            )
+        output, is_logits = self._normalize_output(x)
+        if img0_shape is None:
+            return self._to_input_space(output, is_logits)
+
+        shapes = normalize_image_shapes(img0_shape, output.shape[0])
+        pads = resolve_ratio_pads(ratio_pad, output.shape[0], shapes, self.input_shape)
+        restored = [
+            self._restore(output[index], is_logits, shapes[index], pads[index])
+            for index in range(output.shape[0])
+        ]
+        return restored[0] if len(restored) == 1 else restored
+
+    def _normalize_output(
+        self, x: TensorLike | ListTensorLike
+    ) -> tuple[torch.Tensor, bool]:
+        """Validate logits in NCHW/NHWC layout or a baked class-map tensor."""
+
+        if isinstance(x, (list, tuple)):
+            if len(x) != 1:
+                raise ValueError(
+                    f"Semantic segmentation expects one output tensor, received {len(x)}."
+                )
+            x = x[0]
+        output = torch.as_tensor(x, device=self.device)
+        if output.ndim == 4:
+            if output.shape[1] == self.nc:
+                return self._validate_logits(output), True
+            if output.shape[-1] == self.nc:
+                return self._validate_logits(output.permute(0, 3, 1, 2)), True
+            raise ValueError(
+                f"Semantic segmentation for '{self.dataset}' expects [B, {self.nc}, H, W] or "
+                f"[B, H, W, {self.nc}] logits, got {tuple(output.shape)}."
+            )
+        if output.ndim == 3:
+            if output.shape[-1] == self.nc and output.is_floating_point():
+                return self._validate_logits(output.permute(2, 0, 1).unsqueeze(0)), True
+            if tuple(output.shape[:2]) == self.input_shape:
+                raise ValueError(
+                    f"Semantic segmentation for '{self.dataset}' expects [H, W, {self.nc}] MXQ logits, "
+                    f"got {tuple(output.shape)}."
+                )
+            if output.is_complex():
+                raise ValueError("Semantic class-map values must be finite integers.")
+            if output.is_floating_point():
+                if not bool(torch.isfinite(output).all()):
+                    raise ValueError("Semantic class-map values must be finite.")
+                if not bool(torch.eq(output, output.trunc()).all()):
+                    raise ValueError(
+                        "Semantic class-map values must be integer-valued."
+                    )
+            if output.numel() and (
+                int(output.min()) < 0 or int(output.max()) >= self.nc
+            ):
+                raise ValueError(
+                    f"Semantic class-map values must be in [0, {self.nc - 1}]."
+                )
+            return output.to(dtype=torch.int64), False
+        raise ValueError(
+            f"Semantic segmentation expects [B, C, H, W] or [B, H, W, C] logits, or [B, H, W] class maps, "
+            f"got {tuple(output.shape)}."
+        )
+
+    @staticmethod
+    def _validate_logits(output: torch.Tensor) -> torch.Tensor:
+        """Convert semantic logits to float while rejecting invalid artifact output."""
+
+        logits = output.to(dtype=torch.float32)
+        if not bool(torch.isfinite(logits).all()):
+            raise ValueError("Semantic logits must contain only finite values.")
+        return logits
+
+    def _to_input_space(self, output: torch.Tensor, is_logits: bool) -> torch.Tensor:
+        """Restore model output resolution to configured input space."""
+
+        if is_logits:
+            if tuple(output.shape[-2:]) != self.input_shape:
+                output = functional.interpolate(
+                    output, size=self.input_shape, mode="bilinear", align_corners=False
+                )
+            return output.argmax(dim=1).to(dtype=torch.int64)
+        if tuple(output.shape[-2:]) != self.input_shape:
+            output = functional.interpolate(
+                output[:, None].float(), size=self.input_shape, mode="nearest"
+            )[:, 0]
+        return output.to(dtype=torch.int64)
+
+    def _restore(
+        self,
+        output: torch.Tensor,
+        is_logits: bool,
+        shape: tuple[int, int],
+        ratio_pad: RatioPad,
+    ) -> torch.Tensor:
+        """Undo letterboxing, preserving logits until after bilinear restoration."""
+
+        if is_logits:
+            if tuple(output.shape[-2:]) != self.input_shape:
+                output = functional.interpolate(
+                    output[None],
+                    size=self.input_shape,
+                    mode="bilinear",
+                    align_corners=False,
+                )[0]
+            channels = [
+                crop_letterbox(channel, shape, ratio_pad, self.input_shape, "Semantic")
+                for channel in output
+            ]
+            cropped_logits = torch.stack(channels)
+            restored_logits = functional.interpolate(
+                cropped_logits[None],
+                size=shape,
+                mode="bilinear",
+                align_corners=False,
+            )[0]
+            return restored_logits.argmax(dim=0).to(dtype=torch.int64)
+
+        if tuple(output.shape[-2:]) != self.input_shape:
+            output = functional.interpolate(
+                output[None, None].float(),
+                size=self.input_shape,
+                mode="nearest",
+            )[0, 0].to(dtype=torch.int64)
+        cropped = crop_letterbox(output, shape, ratio_pad, self.input_shape, "Semantic")
+        return functional.interpolate(
+            cropped[None, None].float(), size=shape, mode="nearest"
+        )[0, 0].to(torch.int64)