sukhitashvili · September 10, 2023 13:59 · Sep 10, 2023 · Sep 10, 2023
diff --git a/extract_bboxes_from_binary_mask.py b/extract_bboxes_from_binary_mask.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, List
 
 import numpy as np
 from skimage.measure import label, regionprops, find_contours
@@ -78,7 +78,27 @@ def extract_bboxes(mask: np.ndarray,
         return None
 
     new_raw_bboxes = np.array(bboxes).reshape(-1, 4)
-    new_raw_bboxes[:, :4] = rescale_bounding_boxes(original_shape=(mask_w, mask_h), new_shape=(new_w, new_h),
+    new_raw_bboxes = rescale_bounding_boxes(original_shape=(mask_w, mask_h), new_shape=(new_w, new_h),
                                                    bounding_boxes=new_raw_bboxes[:, :4])
     new_raw_bboxes = new_raw_bboxes.tolist()
-    return new_raw_bboxes
+    return new_raw_bboxes
+
+
+def get_overlap(pred1: np.array, pred2: np.array) -> float:
+    x1 = np.max(pred1[0], pred2[0])
+    y1 = np.max(pred1[1], pred2[1])
+    x2 = np.min(pred1[2], pred2[2])
+    y2 = np.min(pred1[3], pred2[3])
+    inter = (x2 - x1).clamp(0) * (y2 - y1).clamp(0)  # clamp to zero if they don't intersect
+    return inter
+
+
+def if_overlaps(box: list, boxes: List[list], min_distance: int) -> bool:
+    box = np.array(box)
+    for b in boxes:
+        x1, y1, x2, y2 = b
+        b = x1 - min_distance, y1 - min_distance, x2 + min_distance, y2 + min_distance
+        b = np.array(b)
+        if get_overlap(box, b) > 0:
+            return True
+    return False
diff --git a/extract_bboxes_from_binary_mask.py b/extract_bboxes_from_binary_mask.py
@@ -0,0 +1,84 @@
+from typing import Optional
+
+import numpy as np
+from skimage.measure import label, regionprops, find_contours
+
+
+def mask_to_border(mask):
+    h, w = mask.shape
+    border = np.zeros((h, w))
+    contours = find_contours(mask, 128)
+    for contour in contours:
+        for c in contour:
+            x = int(c[0])
+            y = int(c[1])
+            border[x][y] = 255
+
+    return border
+
+
+def mask_to_bbox(mask):
+    """ Mask to bounding boxes """
+    bboxes = []
+    mask = mask_to_border(mask)
+    lbl = label(mask)
+    props = regionprops(lbl)
+    for prop in props:
+        x1 = prop.bbox[1]
+        y1 = prop.bbox[0]
+        x2 = prop.bbox[3]
+        y2 = prop.bbox[2]
+        bboxes.append([x1, y1, x2, y2])
+
+    return bboxes
+
+
+def rescale_bounding_boxes(
+        original_shape: tuple, new_shape: tuple, bounding_boxes: np.ndarray
+) -> np.ndarray:
+    """
+    Rescales bounding box coords according to new image size
+    :param original_shape: (W, H)
+    :param new_shape: (W, H)
+    :param bounding_boxes: 2D numpy array [[x1, y1, x2, y2], ...]
+    :return: scaled bbox coords 2D numpy array
+    """
+    original_w, original_h = original_shape
+    new_w, new_h = new_shape
+    bounding_boxes = bounding_boxes.astype(np.float64)
+    scale_h, scale_w = new_h / original_h, new_w / original_w
+    bounding_boxes[:, 0] *= scale_w
+    bounding_boxes[:, 1] *= scale_h
+    bounding_boxes[:, 2] *= scale_w
+    bounding_boxes[:, 3] *= scale_h
+    bounding_boxes = np.clip(bounding_boxes, a_min=0, a_max=None)
+    bounding_boxes = bounding_boxes.astype(np.uint32)
+    return bounding_boxes
+
+
+def extract_bboxes(mask: np.ndarray,
+                   output_size: list,
+                   thr_val: float) -> Optional[list]:
+    """
+    Extracts bounding boxes from mask
+    Args:
+        mask: [H, W] mask
+        output_size:
+        thr_val:
+
+    Returns:
+        list of bounding boxes [[x1, y1, x2, y2], ...]
+    """
+    mask_h, mask_w = mask.shape
+    new_h, new_w = output_size
+
+    binary_mask = np.where(mask >= thr_val, 255, 0).astype(np.uint8)
+    bboxes = mask_to_bbox(mask=binary_mask)
+    if len(bboxes) == 0:  # if no bbox was found, return dummy bbox with overall score
+        return None
+
+    new_raw_bboxes = np.array(bboxes).reshape(-1, 4)
+    new_raw_bboxes[:, :4] = rescale_bounding_boxes(original_shape=(mask_w, mask_h), new_shape=(new_w, new_h),
+                                                   bounding_boxes=new_raw_bboxes[:, :4])
+    new_raw_bboxes = new_raw_bboxes.tolist()
+    return new_raw_bboxes