diff --git a/depthai_nodes/node/parsers/utils/masks_utils.py b/depthai_nodes/node/parsers/utils/masks_utils.py index 506acd37..30b48b14 100644 --- a/depthai_nodes/node/parsers/utils/masks_utils.py +++ b/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) diff --git a/depthai_nodes/node/parsers/utils/rf_detr.py b/depthai_nodes/node/parsers/utils/rf_detr.py index 01e9ea29..546482ea 100644 --- a/depthai_nodes/node/parsers/utils/rf_detr.py +++ b/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