src/postprocessor.py

"""
postprocessor.py (Optimized)
----------------------------
Vectorized Non-Maximum Suppression (NMS) utilizing NumPy matrix operations.
Eliminates heavy Python O(N^2) loops for extreme execution speed.
"""

import numpy as np
from typing import List, Dict


class PostProcessor:
    def __init__(
        self,
        iou_threshold: float = 0.3,
        min_confidence: float = 0.4,
        min_box_area: int = 16,
    ):
        self.iou_threshold = iou_threshold
        self.min_confidence = min_confidence
        self.min_box_area = min_box_area

    def _filter(self, detections: List[Dict]) -> List[Dict]:
        filtered = []
        for d in detections:
            if d["confidence"] < self.min_confidence:
                continue
            x, y, w, h = d["bbox"]
            if w * h < self.min_box_area or w <= 0 or h <= 0:
                continue
            filtered.append(d)
        return filtered

    def nms(self, detections: List[Dict]) -> List[Dict]:
        """Thuật toán Vectorized NMS xử lý hàng loạt box trùng lặp."""
        candidates = self._filter(detections)
        if not candidates:
            return []

        bboxes = np.array([c["bbox"] for c in candidates], dtype=np.float32)
        scores = np.array([c["confidence"] for c in candidates], dtype=np.float32)

        x1 = bboxes[:, 0]
        y1 = bboxes[:, 1]
        w  = bboxes[:, 2]
        h  = bboxes[:, 3]
        x2 = x1 + w
        y2 = y1 + h

        areas = w * h
        order = scores.argsort()[::-1]

        keep_indices = []
        while order.size > 0:
            i = order[0]
            keep_indices.append(i)

            if order.size == 1:
                break

            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])

            inter_w = np.maximum(0.0, xx2 - xx1)
            inter_h = np.maximum(0.0, yy2 - yy1)
            intersection = inter_w * inter_h

            union = areas[i] + areas[order[1:]] - intersection
            iou = np.where(union > 0, intersection / union, 0.0)

            inds = np.where(iou <= self.iou_threshold)[0]
            order = order[inds + 1]

        return [candidates[idx] for idx in keep_indices]

    def clamp_boxes(
        self, detections: List[Dict], img_h: int, img_w: int
    ) -> List[Dict]:
        clamped = []
        for d in detections:
            x, y, w, h = d["bbox"]
            cx = max(0, min(x, img_w - 1))
            cy = max(0, min(y, img_h - 1))
            cw = max(1, min(w, img_w - cx))
            ch = max(1, min(h, img_h - cy))
            clamped.append({**d, "bbox": (cx, cy, cw, ch)})
        return clamped

    def process(
        self,
        detections: List[Dict],
        img_h: int,
        img_w: int,
    ) -> List[Dict]:
        deduplicated = self.nms(detections)
        clamped = self.clamp_boxes(deduplicated, img_h, img_w)
        return clamped