1import numpy as np
2
3S, B, C = 7, 2, 20 # grid size, boxes per cell, classes
4
5
6def iou(box_a, box_b):
7 """Intersection over Union of two [x1, y1, x2, y2] boxes."""
8 x1 = max(box_a[0], box_b[0]); y1 = max(box_a[1], box_b[1])
9 x2 = min(box_a[2], box_b[2]); y2 = min(box_a[3], box_b[3])
10 inter = max(0, x2 - x1) * max(0, y2 - y1)
11 area_a = (box_a[2] - box_a[0]) * (box_a[3] - box_a[1])
12 area_b = (box_b[2] - box_b[0]) * (box_b[3] - box_b[1])
13 return inter / (area_a + area_b - inter)
14
15
16def decode(pred, conf_thresh=0.25):
17 """pred: (S, S, B*5 + C) tensor -> list of scored boxes."""
18 boxes = []
19 for row in range(S):
20 for col in range(S):
21 cell = pred[row, col]
22 class_probs = cell[B * 5:] # 20 class probabilities
23 for b in range(B):
24 x, y, w, h, conf = cell[b * 5 : b * 5 + 5]
25 # class-specific score = Pr(Class|Obj) * Pr(Obj) * IoU
26 scores = class_probs * conf
27 cls = int(np.argmax(scores))
28 if scores[cls] < conf_thresh:
29 continue # too unsure -> discard
30 # (x, y) are relative to the cell; (w, h) to the image
31 cx, cy = (col + x) / S, (row + y) / S
32 boxes.append({
33 "box": [cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2],
34 "score": float(scores[cls]),
35 "class": cls,
36 })
37 return boxes
38
39
40def nms(boxes, iou_thresh=0.45):
41 """Keep the best box per object; drop overlapping duplicates."""
42 boxes = sorted(boxes, key=lambda d: d["score"], reverse=True)
43 kept = []
44 for cand in boxes:
45 duplicate = any(
46 k["class"] == cand["class"]
47 and iou(k["box"], cand["box"]) > iou_thresh
48 for k in kept
49 )
50 if not duplicate:
51 kept.append(cand)
52 return kept
53
54
55# --- Try it on random 'network output' ---
56pred = np.random.rand(S, S, B * 5 + C)
57detections = nms(decode(pred))
58print(f"{len(detections)} objects after decode + NMS")