Refactor segmentation handling to use RotatedSegmentationResult and update related tests

Author: eugene123tw

model_api/python/model_api/models/__init__.py

@@ -20,8 +20,7 @@
     ImageResultWithSoftPrediction,
     InstanceSegmentationResult,
     PredictedMask,
-    SegmentedObject,
-    SegmentedObjectWithRects,
+    RotatedSegmentationResult,
     VisualPromptingResult,
     ZSLVisualPromptingResult,
 )
@@ -94,8 +93,7 @@
     "classification_models",
     "detection_models",
     "segmentation_models",
-    "SegmentedObject",
-    "SegmentedObjectWithRects",
+    "RotatedSegmentationResult",
     "add_rotated_rects",
     "get_contours",
 ]

model_api/python/model_api/models/instance_segmentation.py

@@ -9,7 +9,7 @@
 from model_api.adapters.inference_adapter import InferenceAdapter
 
 from .image_model import ImageModel
-from .result_types import InstanceSegmentationResult, SegmentedObject
+from .result_types import InstanceSegmentationResult
 from .types import BooleanValue, ListValue, NumericalValue, StringValue
 from .utils import load_labels
 
@@ -176,27 +176,31 @@ def postprocess(self, outputs: dict, meta: dict) -> InstanceSegmentationResult:
             out=boxes,
         )
 
-        objects = []
         has_feature_vector_name = _feature_vector_name in self.outputs
         if has_feature_vector_name:
             if not self.labels:
                 self.raise_error("Can't get number of classes because labels are empty")
             saliency_maps: list = [[] for _ in range(len(self.labels))]
         else:
             saliency_maps = []
-        for box, confidence, cls, raw_mask in zip(boxes, scores, labels, masks):
-            x1, y1, x2, y2 = box
-            if (x2 - x1) * (y2 - y1) < 1 or (confidence <= self.confidence_threshold and not has_feature_vector_name):
-                continue
 
-            # Skip if label index is out of bounds
-            if self.labels and cls >= len(self.labels):
-                continue
+        # Apply confidence threshold, bounding box area filter and label index filter.
+        keep = (scores > self.confidence_threshold) & ((boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) > 1)
+
+        if self.labels:
+            keep &= labels < len(self.labels)
+
+        boxes = boxes[keep].astype(np.int32)
+        scores = scores[keep]
+        labels = labels[keep]
+        masks = masks[keep]
 
-            # Get label string
-            str_label = self.labels[cls] if self.labels else f"#{cls}"
+        resized_masks, label_names = [], []
+        for box, label_idx, raw_mask in zip(boxes, labels, masks, strict=True):
+            if self.labels:
+                label_names.append(self.labels[label_idx])
 
-            raw_cls_mask = raw_mask[cls, ...] if self.is_segmentoly else raw_mask
+            raw_cls_mask = raw_mask[label_idx, ...] if self.is_segmentoly else raw_mask
             if self.postprocess_semantic_masks or has_feature_vector_name:
                 resized_mask = _segm_postprocess(
                     box,
@@ -205,27 +209,21 @@ def postprocess(self, outputs: dict, meta: dict) -> InstanceSegmentationResult:
                 )
             else:
                 resized_mask = raw_cls_mask
-            if confidence > self.confidence_threshold:
-                output_mask = resized_mask if self.postprocess_semantic_masks else raw_cls_mask
-                xmin, ymin, xmax, ymax = box.astype(int)
-                objects.append(
-                    SegmentedObject(
-                        xmin,
-                        ymin,
-                        xmax,
-                        ymax,
-                        score=confidence,
-                        id=cls,
-                        str_label=str_label,
-                        mask=output_mask,
-                    ),
-                )
-            if has_feature_vector_name and confidence > self.confidence_threshold:
-                saliency_maps[cls - 1].append(resized_mask)
+
+            output_mask = resized_mask if self.postprocess_semantic_masks else raw_cls_mask
+            resized_masks.append(output_mask)
+            if has_feature_vector_name:
+                saliency_maps[label_idx - 1].append(resized_mask)
+
+        _masks = np.stack(resized_masks) if len(resized_masks) > 0 else np.empty((0, 16, 16), dtype=np.uint8)
         return InstanceSegmentationResult(
-            objects,
-            _average_and_normalize(saliency_maps),
-            outputs.get(_feature_vector_name, np.ndarray(0)),
+            bboxes=boxes,
+            labels=labels,
+            scores=scores,
+            masks=_masks,
+            label_names=label_names if label_names else None,
+            saliency_map=_average_and_normalize(saliency_maps),
+            feature_vector=outputs.get(_feature_vector_name, np.ndarray(0)),
         )
 
 

model_api/python/model_api/models/result_types/__init__.py

@@ -16,8 +16,7 @@
     Contour,
     ImageResultWithSoftPrediction,
     InstanceSegmentationResult,
-    SegmentedObject,
-    SegmentedObjectWithRects,
+    RotatedSegmentationResult,
 )
 from .visual_prompting import PredictedMask, VisualPromptingResult, ZSLVisualPromptingResult
 
@@ -29,12 +28,11 @@
     "DetectionResult",
     "DetectedKeypoints",
     "MultipleOutputParser",
-    "SegmentedObject",
-    "SegmentedObjectWithRects",
     "SingleOutputParser",
     "ImageResultWithSoftPrediction",
     "InstanceSegmentationResult",
     "PredictedMask",
     "VisualPromptingResult",
     "ZSLVisualPromptingResult",
+    "RotatedSegmentationResult",
 ]

model_api/python/model_api/models/result_types/detection.py

@@ -53,16 +53,22 @@ def __init__(
         self._feature_vector = feature_vector
 
     def __len__(self) -> int:
-        return len(self._bboxes)
+        return len(self.bboxes)
 
     def __str__(self) -> str:
-        return (
-            f"Num of boxes: {self._bboxes.shape}, "
-            f"Num of labels: {len(self._labels)}, "
-            f"Num of scores: {len(self._scores)}, "
-            f"Saliency Map: {array_shape_to_str(self._saliency_map)}, "
-            f"Feature Vec: {array_shape_to_str(self._feature_vector)}"
-        )
+        repr_str = ""
+        for box, score, label, name in zip(
+            self.bboxes,
+            self.scores,
+            self.labels,
+            self.label_names,
+            strict=True,
+        ):
+            x1, y1, x2, y2 = box
+            repr_str += f"{x1}, {y1}, {x2}, {y2}, {label} ({name}): {score:.3f}; "
+
+        repr_str += f"{array_shape_to_str(self.saliency_map)}; {array_shape_to_str(self.feature_vector)}"
+        return repr_str
 
     def get_obj_sizes(self) -> np.ndarray:
         """Get object sizes.
@@ -117,11 +123,11 @@ def label_names(self, value):
         self._label_names = value
 
     @property
-    def saliency_map(self) -> np.ndarray:
+    def saliency_map(self):
         return self._saliency_map
 
     @saliency_map.setter
-    def saliency_map(self, value):
+    def saliency_map(self, value: np.ndarray):
         if not isinstance(value, np.ndarray):
             msg = "Saliency map must be numpy array."
             raise ValueError(msg)
@@ -168,8 +174,12 @@ def __call__(self, outputs) -> DetectionResult:
             labels.append(label)
         bboxes = np.array(bboxes)
         scores = np.array(scores)
-        labels = np.array(labels)
-        return DetectionResult(bboxes, scores, labels)
+        labels = np.array(labels).astype(np.int32)
+        return DetectionResult(
+            bboxes=bboxes,
+            labels=labels,
+            scores=scores,
+        )
 
 
 class MultipleOutputParser:

model_api/python/model_api/models/result_types/segmentation.py

@@ -10,81 +10,143 @@
 import cv2
 import numpy as np
 
+from .detection import DetectionResult
 from .utils import array_shape_to_str
 
 if TYPE_CHECKING:
     from cv2.typing import RotatedRect
 
 
-class SegmentedObject:
-    def __init__(
-        self,
-        xmin: int,
-        ymin: int,
-        xmax: int,
-        ymax: int,
-        score: float,
-        id: int,
-        str_label: str,
-        mask: np.ndarray,
-    ) -> None:
-        self.xmin = xmin
-        self.ymin = ymin
-        self.xmax = xmax
-        self.ymax = ymax
-        self.score = score
-        self.id = id
-        self.str_label = str_label
-        self.mask = mask
-
-    def __str__(self):
-        return (
-            f"{self.xmin}, {self.ymin}, {self.xmax}, {self.ymax}, {self.id} ({self.str_label}): {self.score:.3f}"
-            f", {(self.mask > 0.5).sum()}"
-        )
-
+class InstanceSegmentationResult(DetectionResult):
+    """Instance segmentation result type.
 
-class SegmentedObjectWithRects(SegmentedObject):
-    def __init__(self, segmented_object: SegmentedObject, rotated_rect: RotatedRect) -> None:
-        super().__init__(
-            segmented_object.xmin,
-            segmented_object.ymin,
-            segmented_object.xmax,
-            segmented_object.ymax,
-            segmented_object.score,
-            segmented_object.id,
-            segmented_object.str_label,
-            segmented_object.mask,
-        )
-        self.rotated_rect = rotated_rect
+    Args:
+        bboxes (np.ndarray): bounding boxes in dim (N, 4) where N is the number of boxes.
+        labels (np.ndarray): labels for each bounding box in dim (N,).
+        masks (np.ndarray): masks for each bounding box in dim (N, H, W).
+        scores (np.ndarray | None, optional): confidence scores for each bounding box in dim (N,). Defaults to None.
+        label_names (list[str] | None, optional): class names for each label. Defaults to None.
+        saliency_map (list[np.ndarray] | None, optional): saliency maps for XAI. Defaults to None.
+        feature_vector (np.ndarray | None, optional): feature vector for XAI. Defaults to None.
+    """
 
-    def __str__(self):
-        res = super().__str__()
-        rect = self.rotated_rect
-        res += f", RotatedRect: {rect[0][0]:.3f} {rect[0][1]:.3f} {rect[1][0]:.3f} {rect[1][1]:.3f} {rect[2]:.3f}"
-        return res
+    def __init__(
+        self,
+        bboxes: np.ndarray,
+        labels: np.ndarray,
+        masks: np.ndarray,
+        scores: np.ndarray | None = None,
+        label_names: list[str] | None = None,
+        saliency_map: list[np.ndarray] | None = None,
+        feature_vector: np.ndarray | None = None,
+    ):
+        super().__init__(bboxes, labels, scores, label_names, saliency_map, feature_vector)
+        self._masks = masks
+
+    def __str__(self) -> str:
+        repr_str = ""
+        for box, score, label, name, mask in zip(
+            self.bboxes,
+            self.scores,
+            self.labels,
+            self.label_names,
+            self.masks,
+            strict=True,
+        ):
+            x1, y1, x2, y2 = box
+            repr_str += f"{x1}, {y1}, {x2}, {y2}, {label} ({name}): {score:.3f}, {(mask > 0.5).sum()}; "
 
+        filled = 0
+        for cls_map in self.saliency_map:
+            if cls_map.size:
+                filled += 1
+        prefix = f"{repr_str}" if len(repr_str) else ""
+        return prefix + f"{filled}; {array_shape_to_str(self.feature_vector)}"
+
+    @property
+    def masks(self) -> np.ndarray:
+        return self._masks
+
+    @masks.setter
+    def masks(self, value):
+        if not isinstance(value, np.ndarray):
+            msg = "Masks must be numpy array."
+            raise ValueError(msg)
+        self._masks = value
+
+    @property
+    def saliency_map(self):
+        return self._saliency_map
+
+    @saliency_map.setter
+    def saliency_map(self, value: list[np.ndarray]):
+        if not isinstance(value, list):
+            msg = "Saliency maps must be list."
+            raise ValueError(msg)
+        self._saliency_map = value
+
+
+class RotatedSegmentationResult(InstanceSegmentationResult):
+    """Rotated instance segmentation result type.
+
+    Args:
+        bboxes (np.ndarray): bounding boxes in dim (N, 4) where N is the number of boxes.
+        labels (np.ndarray): labels for each bounding box in dim (N,).
+        masks (np.ndarray): masks for each bounding box in dim (N, H, W).
+        rotated_rects (list[RotatedRect]): rotated rectangles for each bounding box.
+        scores (np.ndarray | None, optional): confidence scores for each bounding box in dim (N,). Defaults to None.
+        label_names (list[str] | None, optional): class names for each label. Defaults to None.
+        saliency_map (list[np.ndarray] | None, optional): saliency maps for XAI. Defaults to None.
+        feature_vector (np.ndarray | None, optional): feature vector for XAI. Defaults to None.
+    """
 
-class InstanceSegmentationResult:
     def __init__(
         self,
-        segmentedObjects: list[SegmentedObject | SegmentedObjectWithRects],
-        saliency_map: list[np.ndarray],
-        feature_vector: np.ndarray,
+        bboxes: np.ndarray,
+        labels: np.ndarray,
+        masks: np.ndarray,
+        rotated_rects: list[RotatedRect],
+        scores: np.ndarray | None = None,
+        label_names: list[str] | None = None,
+        saliency_map: list[np.ndarray] | None = None,
+        feature_vector: np.ndarray | None = None,
     ):
-        self.segmentedObjects = segmentedObjects
-        # Contain per class saliency_maps and "feature_vector" model output if feature_vector exists
-        self.saliency_map = saliency_map
-        self.feature_vector = feature_vector
+        super().__init__(bboxes, labels, masks, scores, label_names, saliency_map, feature_vector)
+        self.rotated_rects = rotated_rects
+
+    def __str__(self) -> str:
+        repr_str = ""
+        for box, score, label, name, mask, rotated_rect in zip(
+            self.bboxes,
+            self.scores,
+            self.labels,
+            self.label_names,
+            self.masks,
+            self.rotated_rects,
+            strict=True,
+        ):
+            x1, y1, x2, y2 = box
+            (cx, cy), (w, h), angle = rotated_rect
+            repr_str += f"{x1}, {y1}, {x2}, {y2}, {label} ({name}): {score:.3f}, {(mask > 0.5).sum()},"
+            repr_str += f" RotatedRect: {cx:.3f} {cy:.3f} {w:.3f} {h:.3f} {angle:.3f}; "
 
-    def __str__(self):
-        obj_str = "; ".join(str(obj) for obj in self.segmentedObjects)
         filled = 0
         for cls_map in self.saliency_map:
             if cls_map.size:
                 filled += 1
-        prefix = f"{obj_str}; " if len(obj_str) else ""
-        return prefix + f"{filled}; [{','.join(str(i) for i in self.feature_vector.shape)}]"
+        prefix = f"{repr_str}" if len(repr_str) else ""
+        return prefix + f"{filled}; {array_shape_to_str(self.feature_vector)}"
+
+    @property
+    def rotated_rects(self) -> list[RotatedRect]:
+        return self._rotated_rects
+
+    @rotated_rects.setter
+    def rotated_rects(self, value):
+        if not isinstance(value, list):
+            msg = "RotatedRects must be list."
+            raise ValueError(msg)
+        self._rotated_rects = value
 
 
 class Contour: