File size: 3,421 Bytes
39fac4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import numpy as np
import cv2
def compute_iou_matrix(boxes):
'''
Given a set of bboxes (in [x1, y1, x2, y2]), output an IOU matrix,
while ignoring pairs where one box fully contains the other.
'''
N = boxes.shape[0]
iou_matrix = np.zeros((N, N), dtype=np.float32)
def is_contained(box_a, box_b):
return (
box_a[0] <= box_b[0] and box_a[1] <= box_b[1] and
box_a[2] >= box_b[2] and box_a[3] >= box_b[3]
)
def is_almost_contained(inner, outer, epsilon=2):
x1_i, y1_i, w_i, h_i = inner
x1_o, y1_o, w_o, h_o = outer
x2_i, y2_i = x1_i + w_i, y1_i + h_i
x2_o, y2_o = x1_o + w_o, y1_o + h_o
return (
x1_i >= x1_o - epsilon and
y1_i >= y1_o - epsilon and
x2_i <= x2_o + epsilon and
y2_i <= y2_o + epsilon
)
for i in range(N):
x1_i, y1_i, x2_i, y2_i = boxes[i]
area_i = (x2_i - x1_i) * (y2_i - y1_i)
for j in range(i + 1, N):
x1_j, y1_j, x2_j, y2_j = boxes[j]
area_j = (x2_j - x1_j) * (y2_j - y1_j)
box_i = boxes[i]
box_j = boxes[j]
# Skip if one box fully contains the other
if is_almost_contained(box_i, box_j) or is_almost_contained(box_j, box_i):
iou = 0.0
else:
inter_x1 = max(x1_i, x1_j)
inter_y1 = max(y1_i, y1_j)
inter_x2 = min(x2_i, x2_j)
inter_y2 = min(y2_i, y2_j)
inter_w = max(0, inter_x2 - inter_x1)
inter_h = max(0, inter_y2 - inter_y1)
inter_area = inter_w * inter_h
union_area = area_i + area_j - inter_area
iou = inter_area / union_area if union_area > 0 else 0.0
iou_matrix[i, j] = iou
iou_matrix[j, i] = iou # symmetric
return iou_matrix
def draw_bboxes(image, bbox_xyxy, color=(0, 255, 0), thickness=2):
'''
given an image and set of bboxes, output a bbox annotated image
'''
image_copy = image.copy()
for box in bbox_xyxy:
x1, y1, x2, y2 = map(int, box)
cv2.rectangle(image_copy, (x1, y1), (x2, y2), color, thickness)
return image_copy
def add_black_border(patch, border_width=1):
"""
Set the outermost border of the patch to white.
Args:
patch (np.ndarray): An image patch of shape (H, W, 3) for RGB.
border_width (int): Width of the white border (default: 1 pixel).
Returns:
np.ndarray: The patch with a black border.
"""
patch[:border_width, :, 0:3] = 255 # Top border
patch[-border_width:, :, 0:3] = 255 # Bottom border
patch[:, :border_width, 0:3] = 255 # Left border
patch[:, -border_width:, 0:3] = 255 # Right border
return patch
def mask_to_bbox_xywh(mask):
"""
Obtain bbox (xywh) from the mask.
Args:
mask (np.ndarray)
Returns:
np.ndarray: a bbox of shape (1, 4), in the format of [x, y, w, h]。
if empty, return shape (0, 4)。
"""
ys, xs = np.where(mask)
if len(xs) == 0 or len(ys) == 0:
return np.zeros((0, 4), dtype=np.int32)
x_min, x_max = xs.min(), xs.max()
y_min, y_max = ys.min(), ys.max()
bbox = [x_min, y_min, x_max - x_min, y_max - y_min]
return np.array(bbox, dtype=np.int32)
|