The code behind the picture

3 teaching cells · Python standard library · Read-only

1. 1. Prepare the inputs

   The lesson's nine decoded candidates as [letter, x, y, w, h, score, class] (one decimal), with its two separate settings: the score threshold and the overlap cutoff. The raw outputs behind them were written by hand for the lesson.

   In [ ]:

   import json
   
   candidates = [
       ["A", 52.9, 58.1, 54.1, 51.0, 0.88, "watermelon"],
       ["B", 61.4, 151.1, 45.1, 38.8, 0.825, "mug"],
       ["C", 205.9, 160.1, 60.3, 35.8, 0.82, "watermelon"],
       ["D", 154.9, 54.0, 30.0, 32.2, 0.714, "apple"],
       ["E", 228.1, 140.9, 35.9, 34.2, 0.698, "watermelon"],
       ["F", 52.6, 55.4, 57.9, 54.0, 0.677, "watermelon"],
       ["G", 188.2, 158.6, 64.0, 38.0, 0.569, "watermelon"],
       ["H", 74.0, 57.9, 50.5, 50.5, 0.368, "watermelon"],
       ["I", 248.0, 24.0, 48.0, 48.0, 0.167, "mug"],
   ]
   score_min = 0.50
   iou_max = 0.45

2. 2. Threshold, then class-aware greedy NMS

   Keep scores at or above 0.50, then repeatedly keep the highest-scoring waiting box and drop later boxes of the same class whose IoU with it is greater than 0.45. Boxes of different classes are never compared.

   In [ ]:

   def iou(a, b):
       w = max(0, min(a[1] + a[3], b[1] + b[3]) - max(a[1], b[1]))
       h = max(0, min(a[2] + a[4], b[2] + b[4]) - max(a[2], b[2]))
       shared = w * h
       return shared / (a[3] * a[4] + b[3] * b[4] - shared)
   
   waiting = sorted([c for c in candidates if c[5] >= score_min], key=lambda c: -c[5])
   result = []
   while waiting:
       best = waiting.pop(0)
       result.append(best[0])
       waiting = [c for c in waiting if c[6] != best[6] or iou(best, c) <= iou_max]

3. 3. Inspect the result

   Print the detections as JSON so both languages use the same output format. Try iou_max = 0.15 to see the small watermelon E suppressed.

   In [ ]:

   print(json.dumps(result))

[
  "A",
  "B",
  "C",
  "D",
  "E"
]