add sorter and improve detections

app.py

@@ -1,6 +1,33 @@
 import gradio as gr
+import numpy as np
 import cv2
-from detector.utils import detect_plates, detect_chars
+from norfair import Detection, Tracker, Video
+from detector.utils import detect_plates, detect_chars, imcrop
+
+DISTANCE_THRESHOLD_BBOX: float = 0.7
+DISTANCE_THRESHOLD_CENTROID: int = 30
+MAX_DISTANCE: int = 10000
+
+
+def yolo_to_norfair(yolo_detections):
+    norfair_detections = []
+    detections_as_xyxy = yolo_detections.xyxy[0]
+    for detection_as_xyxy in detections_as_xyxy:
+        bbox = np.array(
+            [
+                [detection_as_xyxy[0].item(), detection_as_xyxy[1].item()],
+                [detection_as_xyxy[2].item(), detection_as_xyxy[3].item()],
+            ]
+        )
+        scores = np.array(
+            [detection_as_xyxy[4].item(), detection_as_xyxy[4].item()]
+        )
+        norfair_detections.append(
+            Detection(
+                points=bbox, scores=scores, label=int(detection_as_xyxy[-1].item())
+            )
+        )
+    return norfair_detections
 
 
 def fn_image(foto):
@@ -18,7 +45,10 @@ def fn_image(foto):
 
 
 def fn_video(video, initial_time, duration):
-    plates_text = []
+    tracker = Tracker(
+        distance_function="iou_opt",
+        distance_threshold=DISTANCE_THRESHOLD_BBOX,
+    )
     cap = cv2.VideoCapture(video)
     fps = cap.get(cv2.CAP_PROP_FPS)
     image_size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
@@ -26,6 +56,7 @@ def fn_video(video, initial_time, duration):
     num_frames = 0
     min_frame = int(initial_time * fps)
     max_frame = int((initial_time + duration) * fps)
+    plates = {}
     while cap.isOpened():
         try:
             ret, frame = cap.read()
@@ -34,24 +65,47 @@ def fn_video(video, initial_time, duration):
         except Exception as e:
             print(e)
             continue
+        # num_frames += 1
+        # if num_frames % 3 != 0:
+        #     continue
         if num_frames < min_frame:
             num_frames += 1
             continue
-        plates = detect_plates(frame)
-        for plate in plates:
-            p1, p2, crop = plate
-            if len(crop) > 0:
-                cv2.rectangle(frame, p1, p2, (0, 0, 255), 2)
-                text, crop = detect_chars(crop)
-                cv2.putText(frame, text, p1, cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 5)
-                plates_text.append(text)
+        yolo_detections = detect_plates(frame)
+        detections = yolo_to_norfair(yolo_detections)
+        tracked_objects = tracker.update(detections=detections)
+        for obj in tracked_objects:
+            if obj.last_detection is not None:
+                bbox = obj.last_detection.points
+                bbox = int(bbox[0][0]), int(bbox[0][1]), int(bbox[1][0]), int(bbox[1][1])
+                if obj.id not in plates.keys():
+                    crop = imcrop(frame, bbox)
+                    text = detect_chars(crop)
+                    plates[obj.id] = text
+
+                cv2.rectangle(
+                    frame,
+                    (bbox[0], bbox[1]),
+                    (bbox[2], bbox[3]),
+                    (0, 255, 0),
+                    2,
+                )
+                cv2.putText(
+                    frame,
+                    plates[obj.id],
+                    (bbox[0], bbox[1]),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    1,
+                    (0, 255, 0),
+                    2,
+                )
         final_video.write(frame)
         num_frames += 1
         if num_frames == max_frame:
             break
     cap.release()
     final_video.release()
-    return 'output.mp4', plates_text
+    return 'output.mp4', [plates[k] for k in plates.keys()]
 
 
 image_interface = gr.Interface(

detector/utils.py

@@ -30,29 +30,15 @@ def imcrop(img, bbox):
 
 
 def detect_plates(img):
-    detect = model_plates(img)
-    records = detect.pandas().xyxy[0].to_dict(orient='records')
-    plates = []
-    if records:
-        for plate in records:
-            xi, yi, xf, yf = int(plate['xmin']), int(plate['ymin']), int(plate['xmax']), int(plate['ymax'])
-            crop = imcrop(img, (xi, yi, xf, yf))
-            plates.append(((xi, yi), (xf, yf), crop))
-    return plates
+    return model_plates(img)
 
 
 def detect_chars(img):
     img = cv2.resize(img, (640, 320))
     detect = model_chars(img)
     records = detect.pandas().xyxy[0].to_dict(orient='records')
-    yolo = np.squeeze(detect.render())
     text = ''
     if records:
         records = sorted(records, key=lambda d: d['xmin'])
         text = ''.join([i.get('name') for i in records])
-    return text, yolo
-
-
-def save_plates(img):
-    detect = model_plates(img)
-    detect.crop(save=True)
+    return text