Add SSD

ssd/box.py

@@ -0,0 +1,86 @@
+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]