86 lines
3 KiB
Python
86 lines
3 KiB
Python
import numpy as np
|
|
|
|
|
|
def create_box(y_pos: float, x_pos: float, height: float, width: float) -> tuple[float, float, float, float]:
|
|
y_min, x_min, y_max, x_max = check_rectangle(
|
|
y_pos - (height / 2), x_pos - (width / 2), y_pos + (height / 2), x_pos + (width / 2))
|
|
return (y_min + y_max) / 2, (x_min + x_max) / 2, y_max - y_min, x_max - x_min
|
|
|
|
|
|
def check_rectangle(y_min: float, x_min: float, y_max: float, x_max: float) -> tuple[float, float, float, float]:
|
|
if y_min < 0:
|
|
y_min = 0
|
|
if x_min < 0:
|
|
x_min = 0
|
|
if y_min > 1:
|
|
y_min = 1
|
|
if x_min > 1:
|
|
x_min = 1
|
|
if y_max < 0:
|
|
y_max = 0
|
|
if x_max < 0:
|
|
x_max = 0
|
|
if y_max >= 1:
|
|
y_max = 1
|
|
if x_max >= 1:
|
|
x_max = 1
|
|
return y_min, x_min, y_max, x_max
|
|
|
|
|
|
def get_boxes(predictions: np.ndarray, anchors: np.ndarray, class_index: int) -> np.ndarray:
|
|
boxes = np.zeros(anchors.shape)
|
|
boxes[:, 0] = (predictions[:, 0] * anchors[:, 2]) + anchors[:, 0]
|
|
boxes[:, 1] = (predictions[:, 1] * anchors[:, 3]) + anchors[:, 1]
|
|
boxes[:, 2] = np.exp(predictions[:, 2]) * anchors[:, 2]
|
|
boxes[:, 3] = np.exp(predictions[:, 3]) * anchors[:, 3]
|
|
boxes = np.asarray([create_box(*box) for box in boxes])
|
|
|
|
# return np.insert(boxes, 4, predictions[:, class_index], axis=-1)
|
|
return np.concatenate([boxes, predictions[:, class_index:class_index + 1]], axis=1)
|
|
|
|
|
|
def fast_nms(boxes: np.ndarray, min_iou: float) -> np.ndarray:
|
|
# if there are no boxes, return an empty list
|
|
if len(boxes) == 0:
|
|
return []
|
|
|
|
# initialize the list of picked indexes
|
|
pick = []
|
|
|
|
# grab the coordinates of the bounding boxes
|
|
y_min = boxes[:, 0] - (boxes[:, 2] / 2)
|
|
y_max = boxes[:, 0] + (boxes[:, 2] / 2)
|
|
x_min = boxes[:, 1] - (boxes[:, 3] / 2)
|
|
x_max = boxes[:, 1] + (boxes[:, 3] / 2)
|
|
scores = boxes[:, 4]
|
|
|
|
# compute the area of the bounding boxes and sort the bounding boxes by the scores
|
|
areas = (x_max - x_min) * (y_max - y_min)
|
|
idxs = np.argsort(scores)
|
|
|
|
# keep looping while some indexes still remain in the indexes
|
|
# list
|
|
while len(idxs) > 0:
|
|
# grab the last index in the indexes list and add the
|
|
# index value to the list of picked indexes
|
|
last = len(idxs) - 1
|
|
i = idxs[last]
|
|
pick.append(i)
|
|
|
|
inter_tops = np.maximum(y_min[i], y_min[idxs[:last]])
|
|
inter_bottoms = np.minimum(y_max[i], y_max[idxs[:last]])
|
|
inter_lefts = np.maximum(x_min[i], x_min[idxs[:last]])
|
|
inter_rights = np.minimum(x_max[i], x_max[idxs[:last]])
|
|
inter_areas = (inter_rights - inter_lefts) * (inter_bottoms - inter_tops)
|
|
|
|
# compute the ratio of overlap
|
|
union_area = (areas[idxs[:last]] + areas[i]) - inter_areas
|
|
overlap = inter_areas / union_area
|
|
|
|
# delete all indexes from the index list that have less overlap than min_iou
|
|
idxs = np.delete(
|
|
idxs, np.concatenate(([last], np.where(overlap > min_iou)[0])))
|
|
|
|
# return only the bounding boxes that were picked using the
|
|
# integer data type
|
|
return boxes[pick]
|