Revert crop order changes in

depthai_nodes/node/parsers/utils/masks_utils.py

@@ -12,7 +12,9 @@ def probability_to_logit_threshold(probability: float) -> float:
     return float(np.log(probability / (1.0 - probability)))
 
 
-def crop_mask(mask: np.ndarray, bbox: np.ndarray) -> np.ndarray:
+def crop_mask(
+    mask: np.ndarray, bbox: np.ndarray, fill_value: float | int = 0
+) -> np.ndarray:
     """It takes a mask and a bounding box, and returns a mask that is cropped to the
     bounding box.
 
@@ -21,6 +23,8 @@ def crop_mask(mask: np.ndarray, bbox: np.ndarray) -> np.ndarray:
     @param bbox: A numpy array of bbox coordinates in (x_center, y_center, width,
         height) format
     @type bbox: np.ndarray
+    @param fill_value: Value assigned to pixels outside the bounding box.
+    @type fill_value: float | int
     @return: A mask that is cropped to the bounding box
     @rtype: np.ndarray
     """
@@ -33,7 +37,8 @@ def crop_mask(mask: np.ndarray, bbox: np.ndarray) -> np.ndarray:
     r = np.arange(w).reshape(1, w)
     c = np.arange(h).reshape(h, 1)
 
-    return mask * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
+    inside_bbox = (r >= x1) * (r < x2) * (c >= y1) * (c < y2)
+    return np.where(inside_bbox, mask, fill_value)
 
 
 def process_single_mask(
@@ -59,19 +64,28 @@ def process_single_mask(
     @return: Processed binary mask resized to `output_shape`.
     @rtype: np.ndarray
     """
+    _, mask_h, mask_w = protos.shape  # CHW
+    scaled_bbox = bbox * np.array([mask_w, mask_h, mask_w, mask_h])
+
     mask_logits = np.sum(protos * mask_coeff[..., np.newaxis, np.newaxis], axis=0)
+    logit_threshold = probability_to_logit_threshold(mask_conf)
+    # OpenCV interpolation with infinite fill values produces NaNs at crop edges.
+    if mask_conf <= 0.0:
+        logit_threshold = np.finfo(mask_logits.dtype).min
+    if mask_conf >= 1.0:
+        logit_threshold = np.finfo(mask_logits.dtype).max
+
+    mask_logits = crop_mask(
+        mask_logits,
+        scaled_bbox,
+        fill_value=logit_threshold,
+    )
     mask_logits = cv2.resize(
         mask_logits,
         (output_shape[1], output_shape[0]),
         interpolation=cv2.INTER_LINEAR,
     )
-    logit_threshold = probability_to_logit_threshold(mask_conf)
-    mask = (mask_logits > logit_threshold).astype(np.uint8)
-
-    scaled_bbox = bbox * np.array(
-        [output_shape[1], output_shape[0], output_shape[1], output_shape[0]]
-    )
-    return crop_mask(mask, scaled_bbox)
+    return (mask_logits > logit_threshold).astype(np.uint8)
 
 
 def get_segmentation_outputs(
@@ -101,7 +115,6 @@ def get_segmentation_outputs(
 def process_single_mask_rfdetr(
     mask_logits: np.ndarray,
     mask_conf: float,
-    bbox: np.ndarray,
     input_shape: tuple[int, int],
 ) -> np.ndarray:
     """Process a single RF-DETR instance segmentation mask.
@@ -110,9 +123,6 @@ def process_single_mask_rfdetr(
     @type mask_logits: np.ndarray
     @param mask_conf: Mask confidence threshold.
     @type mask_conf: float
-    @param bbox: A numpy array of bbox coordinates in (x_center, y_center, width,
-        height) normalized format.
-    @type bbox: np.ndarray
     @param input_shape: Target output mask shape as (height, width).
     @type input_shape: tuple[int, int]
     @return: Processed mask resized to the model input shape.
@@ -129,9 +139,4 @@ def process_single_mask_rfdetr(
         interpolation=cv2.INTER_LINEAR,
     )
     logit_threshold = probability_to_logit_threshold(mask_conf)
-    mask = (resized_mask_logits > logit_threshold).astype(np.uint8)
-
-    scaled_bbox = bbox * np.array(
-        [input_shape[1], input_shape[0], input_shape[1], input_shape[0]]
-    )
-    return crop_mask(mask, scaled_bbox)
+    return (resized_mask_logits > logit_threshold).astype(np.uint8)

depthai_nodes/node/parsers/utils/rf_detr.py

@@ -66,11 +66,10 @@ def compute_rfdetr_detections(
             )
 
         final_mask = np.full(input_shape, 255, dtype=np.uint8)
-        for i, (mask_logits, bbox) in enumerate(zip(masks, boxes_cxcywh)):
+        for i, mask_logits in enumerate(masks):
             resized_mask = process_single_mask_rfdetr(
                 mask_logits=mask_logits,
                 mask_conf=mask_conf,
-                bbox=bbox,
                 input_shape=input_shape,
             )
             foreground = resized_mask > 0