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

mblt_vision/utils/postprocess/yolo_anchor_post.py

@@ -0,0 +1,469 @@
+"""
+YOLO anchor-based postprocessing.
+"""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import torch
+
+from .base import YOLODetectionPostBase
+from .common import YOLOSegPostMixin, non_max_suppression
+
+
+class YOLOAnchorDetectionPost(YOLODetectionPostBase):
+    """Postprocessing for YOLO models with anchors."""
+
+    def __init__(
+        self, pre_cfg: dict[str, Any], post_cfg: dict[str, Any], **kwargs: Any
+    ) -> None:
+        """Initialize anchor-based YOLO detection postprocessing.
+
+        Args:
+            pre_cfg (dict): Preprocessing configuration.
+            post_cfg (dict): Postprocessing configuration.
+            **kwargs: Optional runtime overrides for postprocess behavior.
+        """
+        super().__init__(pre_cfg, post_cfg, **kwargs)
+        self.no = self.nc + 5 + self.n_extra
+        self.grid: torch.Tensor
+        self.anchor_grid: torch.Tensor
+        self.make_anchor_grid()
+
+    def non_e2e(self, x: list[torch.Tensor]) -> torch.Tensor | list[torch.Tensor]:
+        """Return the export-style output tensor for anchor-based YOLO models."""
+        if len(x) == 1:
+            converted = self.conversion(x)
+            if isinstance(converted, torch.Tensor):
+                return self._converted_to_batch_output(converted)
+            det_out, proto_out = converted
+            return [self._converted_to_batch_output(det_out), proto_out]
+
+        rearranged = self.rearrange(x)
+        if isinstance(rearranged, tuple):
+            det_out, proto_out = rearranged
+            return [self.decode_batch(det_out), proto_out.permute(0, 3, 1, 2)]
+        return self.decode_batch(rearranged)
+
+    def _converted_to_batch_output(self, x: torch.Tensor) -> torch.Tensor:
+        """Normalize converted outputs to the export-style batched layout."""
+        while x.ndim == 4 and 1 in (x.shape[0], x.shape[1]):
+            if x.shape[0] == 1:
+                x = x.squeeze(0)
+            elif x.shape[1] == 1:
+                x = x.squeeze(1)
+        if x.ndim != 3:
+            raise ValueError(
+                f"Expected 3D converted tensor, got shape {tuple(x.shape)}."
+            )
+        if x.shape[-1] == self.no:
+            return x
+        if x.shape[1] == self.no:
+            return x.transpose(1, 2)
+        raise ValueError(
+            f"Unsupported converted tensor shape {tuple(x.shape)} for non-e2e output."
+        )
+
+    def decode_batch(self, x: torch.Tensor) -> torch.Tensor:
+        """Decode every anchor without filtering and preserve batch shape."""
+        batch_size = x.shape[0]
+        grid = self.grid.unsqueeze(0).expand(batch_size, -1, -1)
+        anchor_grid = self.anchor_grid.unsqueeze(0).expand(batch_size, -1, -1)
+        stride = self.stride_as_tensor().unsqueeze(0).expand(batch_size, -1, -1)
+
+        decoded = x.clone()
+        decoded[..., :2] = (
+            decoded[..., :2].sigmoid().mul(2.0).add(grid).add(-0.5).mul(stride)
+        )
+        decoded[..., 2:4] = (
+            decoded[..., 2:4].sigmoid().mul(2.0).pow(2.0).mul(anchor_grid)
+        )
+        conf = decoded[..., 4:5].sigmoid()
+        decoded[..., 4:5] = conf
+        decoded[..., 5 : 5 + self.nc] = decoded[..., 5 : 5 + self.nc].sigmoid()
+        if self.task == "instance_segmentation" and self.n_extra > 0:
+            decoded[..., 5 + self.nc :] = decoded[..., 5 + self.nc :] * conf
+        return decoded
+
+    def rearrange(
+        self, x: list[torch.Tensor]
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        """Rearranges raw model output tensors into a concatenated decode input.
+
+        Args:
+            x (list[torch.Tensor]): Raw output tensors from the model detection heads.
+
+        Returns:
+            torch.Tensor | tuple[torch.Tensor, torch.Tensor]: Concatenated tensor in
+                ``(batch, anchors, no)`` format, optionally paired with prototype masks in
+                segmentation subclasses.
+        """
+        if len(x) != self.nl:
+            raise ValueError(f"Expected {self.nl} detection heads, got {len(x)}.")
+        y = []
+        for i in range(self.nl):
+            tmp = x[i]
+            if tmp.shape[3] == self.no * self.na:
+                y.append(
+                    tmp.permute(0, 3, 1, 2)
+                )  # (b, 80, 80, 255) -> (b, 255, 80, 80)
+            else:
+                raise NotImplementedError(
+                    f"Got unsupported shape for input: {tmp.shape}."
+                )
+        # sort by image size descending
+        y = sorted(y, key=lambda x: x.numel(), reverse=True)
+        return torch.cat(
+            [
+                xi.reshape(xi.shape[0], self.na, self.no, xi.shape[-2], xi.shape[-1])
+                .permute(0, 1, 3, 4, 2)
+                .reshape(xi.shape[0], -1, self.no)
+                for xi in y
+            ],
+            dim=1,
+        )
+
+    def decode(self, x: torch.Tensor) -> list[torch.Tensor]:
+        """Decodes model outputs into box coordinates and class scores.
+
+        Applies sigmoid to predictions and transforms boxes from anchor-relative
+        to image-relative coordinates.
+
+        Args:
+            x (torch.Tensor): Concatenated output tensor from `rearrange`.
+
+        Returns:
+            list[torch.Tensor]: Per-image decoded detections after confidence filtering.
+        """
+        return [self.process_box_cls(box_cls) for box_cls in x]
+
+    def process_box_cls(self, x: torch.Tensor) -> torch.Tensor:
+        """Processes a single image's detection tensor.
+
+        Args:
+            x: Raw detections for one image.
+
+        Returns:
+            Decoded boxes, confidence, and scores.
+        """
+        ic = x[:, 4] > self.inv_conf_thres  # candidates
+        box_cls = x[ic]  # (n, 85)
+        if box_cls.numel() == 0:
+            return box_cls.new_zeros((0, 5 + self.nc + self.n_extra))
+
+        grid = self.grid[ic, :]  # (n, 2)
+        anchor_grid = self.anchor_grid[ic, :]  # (n, 2)
+        stride = self.stride_as_tensor()[ic, :]  # (n, 2)
+
+        # Advanced indexing above materializes ``box_cls``, so in-place decode avoids a second output allocation.
+        box_cls[:, :2] = (
+            box_cls[:, :2].sigmoid_().mul_(2.0).add_(grid).add_(-0.5).mul_(stride)
+        )
+        box_cls[:, 2:4] = (
+            box_cls[:, 2:4].sigmoid_().mul_(2.0).pow_(2.0).mul_(anchor_grid)
+        )
+        conf = box_cls[:, 4:5].sigmoid_()
+        box_cls[:, 5 : 5 + self.nc].sigmoid_()
+        if self.task == "instance_segmentation" and self.n_extra > 0:
+            box_cls[:, 5 + self.nc :] *= conf
+        return box_cls
+
+    def filter_conversion(self, x: torch.Tensor) -> list[torch.Tensor]:
+        """Filters out low-confidence detections from a single concatenated output tensor.
+
+        Args:
+            x (torch.Tensor): Concatenated output tensor from the model.
+
+        Returns:
+            list[torch.Tensor]: Filtered detections for each image in the batch.
+        """
+        x_list = torch.split(
+            self._converted_to_batch_output(x), 1, dim=0
+        )  # [(1, 25200, 85), ...]
+
+        def process_conversion(x: torch.Tensor) -> torch.Tensor:
+            x = x.squeeze(0)  # (25200, 85)
+            ic = x[:, 4] > self.conf_thres  # candidates
+            x = x[ic]  # (n, 85)
+            if len(x) == 0:
+                return x.new_zeros((0, self.no))
+            return x
+
+        return [process_conversion(xi) for xi in x_list]
+
+    def _nms_single(
+        self,
+        xi: torch.Tensor,
+        max_det: int,
+        max_nms: int,
+        max_wh: int,
+        *,
+        multi_label: bool,
+    ) -> torch.Tensor:
+        """Apply anchor-based NMS to a single decoded image tensor."""
+        mi = 5 + self.nc  # mask index
+        if xi.numel() == 0:
+            return xi.new_zeros((0, 6 + self.n_extra))
+
+        scores = xi[:, 5:mi] * xi[:, 4:5]
+        if multi_label:
+            match_index = (scores > self.conf_thres).nonzero(as_tuple=False)
+            if match_index.numel() == 0:
+                return xi.new_zeros((0, 6 + self.n_extra))
+            i, j = match_index[:, 0], match_index[:, 1]
+            rows = xi[i]
+            row_scores = scores[i, j]
+        else:
+            row_scores, j = scores.max(dim=1)
+            keep = row_scores > self.conf_thres
+            if not bool(keep.any()):
+                return xi.new_zeros((0, 6 + self.n_extra))
+            rows, row_scores, j = xi[keep], row_scores[keep], j[keep]
+        boxes_xywh = rows[:, :4]
+        out = torch.empty(
+            (rows.shape[0], 6 + self.n_extra), dtype=rows.dtype, device=rows.device
+        )
+        out[:, 0] = boxes_xywh[:, 0] - boxes_xywh[:, 2] / 2
+        out[:, 1] = boxes_xywh[:, 1] - boxes_xywh[:, 3] / 2
+        out[:, 2] = boxes_xywh[:, 0] + boxes_xywh[:, 2] / 2
+        out[:, 3] = boxes_xywh[:, 1] + boxes_xywh[:, 3] / 2
+        out[:, 4] = row_scores
+        out[:, 5] = j.to(rows.dtype)
+        if self.n_extra > 0:
+            out[:, 6:] = rows[:, mi:]
+        out = out[out[:, 4].argsort(descending=True)[:max_nms]]
+        c = out[:, 5:6] * max_wh
+        boxes, score = out[:, :4] + c, out[:, 4]
+        i_idx = non_max_suppression(boxes, score, self.iou_thres, max_det)
+        return out[i_idx]
+
+    def nms(
+        self,
+        x: 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 (NMS) on the decoded detections.
+        Args:
+            x (list[torch.Tensor]): Decoded detections for each image.
+            max_det (int, optional): Maximum number of detections to keep. Defaults to 300.
+            max_nms (int, optional): Maximum number of candidates to consider for NMS.
+                Defaults to 30000.
+            max_wh (int, optional): Maximum box width/height for offset calculation.
+                Defaults to 7680.
+        Returns:
+            list[torch.Tensor]: Post-NMS detections for each image.
+        """
+        if isinstance(x, list):
+            return [
+                self._nms_single(
+                    xi,
+                    max_det=max_det,
+                    max_nms=max_nms,
+                    max_wh=max_wh,
+                    multi_label=multi_label,
+                )
+                for xi in x
+            ]
+        return [
+            self._nms_single(
+                xi,
+                max_det=max_det,
+                max_nms=max_nms,
+                max_wh=max_wh,
+                multi_label=multi_label,
+            )
+            for xi in x
+        ]
+
+    def nms_multilabel(
+        self, x: torch.Tensor | list[torch.Tensor]
+    ) -> list[torch.Tensor]:
+        """Perform Ultralytics-compatible multi-label NMS for validation."""
+
+        return self.nms(x, multi_label=True)
+
+    def make_anchor_grid(self) -> None:
+        """
+        Pre-calculate the anchor grid for decoding.
+        """
+        grid_parts: list[torch.Tensor] = []
+        anchor_grid_parts: list[torch.Tensor] = []
+        stride_parts: list[torch.Tensor] = []
+        strides = [2 ** (3 + i) for i in range(self.nl)]
+        if self.nl == 2:
+            strides = [strd * 2 for strd in strides]
+        out_sizes = [
+            [self.imh // strd, self.imw // strd] for strd in strides
+        ]  # (80, 80), (40, 40), (20, 20)
+        for anchr, (ny, nx), strd in zip(self.anchors_as_list(), out_sizes, strides):
+            yv, xv = torch.meshgrid(
+                torch.arange(ny, dtype=torch.float32, device=self.device),
+                torch.arange(nx, dtype=torch.float32, device=self.device),
+                indexing="ij",
+            )
+            grid = torch.stack((xv, yv), 2).expand(self.na, ny, nx, 2)
+            grid_parts.append(grid)
+            anchr_tensor = torch.broadcast_to(
+                torch.tensor(anchr).reshape(self.na, 1, 1, 2),
+                (self.na, ny, nx, 2),
+            )
+            anchor_grid_parts.append(anchr_tensor)
+            stride_parts.append(strd * torch.ones(self.na, ny, nx, 2))
+        self.grid = torch.cat([grd.reshape(-1, 2) for grd in grid_parts], dim=0)
+        self.anchor_grid = torch.cat(
+            [anc.reshape(-1, 2) for anc in anchor_grid_parts], dim=0
+        )
+        self.stride = torch.cat([strd.reshape(-1, 2) for strd in stride_parts], dim=0)
+
+    def chop(self, npu_out: torch.Tensor, idx: int = 0) -> tuple[torch.Tensor, ...]:
+        """Splits the detection tensor into individual components (xy, wh, conf, scores, extra).
+
+        Args:
+            npu_out (torch.Tensor): Raw detection tensor from one detection head.
+            idx (int, optional): Detection head index. Defaults to 0.
+
+        Returns:
+            tuple: (xy, wh, conf, scores, extra).
+        """
+        xy, wh, conf, scores, extra = torch.split(
+            npu_out, [2, 2, 1, self.nc, self.n_extra], dim=-1
+        )
+        return xy, wh, conf, scores, extra
+
+
+class YOLOAnchorSegPost(YOLOSegPostMixin, YOLOAnchorDetectionPost):
+    """Postprocessing for YOLO segmentation models with anchors."""
+
+    def non_e2e(self, x: list[torch.Tensor]) -> torch.Tensor | list[torch.Tensor]:
+        """Return the export-style output tensor for anchor-based YOLO segmentation models.
+
+        Args:
+            x: Checked raw model outputs.
+
+        Returns:
+            A detection tensor, or detections paired with prototype masks.
+        """
+        if any(xi.ndim <= 4 and self.no in xi.shape[1:] for xi in x):
+            converted, proto_outs = self.conversion(x)
+            return [
+                self._converted_to_batch_output(converted),
+                self._proto_to_nchw(proto_outs),
+            ]
+        return super().non_e2e(x)
+
+    def _proto_to_nchw(self, proto: torch.Tensor) -> torch.Tensor:
+        """Convert prototype tensors to ``(B, C, H, W)`` if needed.
+
+        Args:
+            proto: Prototype tensor from a model runtime.
+
+        Returns:
+            Prototype tensor in channel-first batch layout.
+        """
+        if proto.ndim == 4 and proto.shape[1] == self.n_extra:
+            return proto
+        if proto.ndim == 4 and proto.shape[-1] == self.n_extra:
+            return proto.permute(0, 3, 1, 2)
+        raise ValueError(
+            f"Unsupported proto tensor shape {tuple(proto.shape)} for non-e2e output."
+        )
+
+    def _pre_process(self, x: list[torch.Tensor]) -> tuple[Any, torch.Tensor | None]:
+        """Preprocesses intermediate inputs into (boxes, proto) format.
+
+        Args:
+            x (list[torch.Tensor]): Raw model output tensors.
+
+        Returns:
+            tuple: (decoded_detections, prototype_masks).
+        """
+        if any(xi.ndim <= 4 and self.no in xi.shape[1:] for xi in x):
+            converted, proto_outs = cast(
+                tuple[torch.Tensor, torch.Tensor], self.conversion(x)
+            )
+            return self.filter_conversion(converted), proto_outs
+        rearranged, proto_outs = self.rearrange(x)
+        return self.decode(rearranged), proto_outs
+
+    def conversion(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+        """Converts raw model output tensors into detections and prototypes.
+
+        Args:
+            x (list[torch.Tensor]): List of raw output tensors.
+
+        Returns:
+            tuple: (detections, prototypes)
+        """
+        det_out: torch.Tensor | None = None
+        proto_out: torch.Tensor | None = None
+        for xi in x:
+            if xi.ndim <= 4 and self.no in xi.shape[1:]:
+                det_out = xi
+            elif xi.ndim == 4 and self.n_extra in xi.shape[1:]:
+                proto_out = xi
+        if det_out is None or proto_out is None:
+            shapes = ", ".join(str(tuple(xi.shape)) for xi in x)
+            raise NotImplementedError(
+                f"Input shapes not supported for anchor segmentation: {shapes}."
+            )
+        return det_out, proto_out
+
+    def rearrange(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+        """Rearranges model output tensors for segmentation tasks.
+
+        Args:
+            x (list[torch.Tensor]): Raw output tensors from detection and prototype heads.
+
+        Returns:
+            tuple[torch.Tensor, torch.Tensor]: Concatenated detections and prototype masks.
+        """
+        proto: torch.Tensor | None = None
+        for i, xi in enumerate(x):
+            if self.n_extra == xi.shape[-1]:
+                proto = x.pop(i)
+                break
+        if proto is None:
+            raise ValueError("Proto output is missing.")
+        y = []
+        for xi in x:
+            if xi.shape[-1] == self.no * self.nl:
+                y.append(xi.permute(0, 3, 1, 2))
+            else:
+                raise ValueError(f"Wrong shape of input: {xi.shape}")
+        # sort by image size descending
+        y = sorted(y, key=lambda x: x.numel(), reverse=True)
+        return (
+            torch.cat(
+                [
+                    xi.reshape(
+                        xi.shape[0], self.na, self.no, xi.shape[-2], xi.shape[-1]
+                    )
+                    .permute(0, 1, 3, 4, 2)
+                    .reshape(xi.shape[0], -1, self.no)
+                    for xi in y
+                ],
+                dim=1,
+            ),
+            proto,
+        )
+
+    def chop(self, npu_out: torch.Tensor, idx: int = 0) -> tuple[torch.Tensor, ...]:
+        """Splits the detection tensor for segmentation tasks.
+
+        Args:
+            npu_out (torch.Tensor): Raw detection tensor.
+            idx (int, optional): Detection head index. Defaults to 0.
+
+        Returns:
+            tuple: (xy, wh, conf, scores, masks).
+        """
+        xy, wh, conf, scores, masks = torch.split(
+            npu_out, [2, 2, 1, self.nc, self.n_extra], dim=-1
+        )
+        masks = masks * conf.sigmoid()
+        return xy, wh, conf, scores, masks