Update inference.py

inference.py

@@ -18,6 +18,7 @@ import torch
 import torchvision
 
 from lib.fish_eye.tracker import Tracker
+from lib.fish_eye.associative import Associate
 
 
 ### Configuration options
@@ -329,11 +330,24 @@ def do_tracking(all_preds, image_meter_width, image_meter_height, gp=None, max_a
 
     return json_data
 
-def do_associative_tracking(raw_detections, image_meter_width, image_meter_height, gp=None, max_age=MAX_AGE, iou_thres=IOU_THRES, min_hits=MIN_HITS, min_length=MIN_LENGTH, verbose=True):
+def do_associative_tracking(raw_detections, image_meter_width, image_meter_height, gp=None, conf_thresh=0.2, max_age=MAX_AGE, iou_thres=IOU_THRES, min_hits=MIN_HITS, min_length=MIN_LENGTH, verbose=True):
 
     if (gp): gp(0, "Tracking...")
 
-    print(len(raw_detections))
+    low_dets = []
+    high_dets = []
+    for batch in raw_detections:
+        for frame in batch:
+            low_frame = []
+            high_frame = []
+            for bbox in frame:
+                if bbox[4] > conf_thresh:
+                    high_frame.append(bbox)
+                else:
+                    low_frame.append(bbox)
+            low_dets.append(low_frame)
+            high_dets.append(high_frame)
+
 
     # Initialize tracker
     clip_info = {
@@ -342,14 +356,16 @@ def do_associative_tracking(raw_detections, image_meter_width, image_meter_heigh
         'image_meter_width': image_meter_width,
         'image_meter_height': image_meter_height
     }
-    tracker = Tracker(clip_info, args={ 'max_age': max_age, 'min_hits': 0, 'iou_threshold': iou_thres}, min_hits=min_hits)
+    tracker = Tracker(clip_info, algorithm=Associate, args={ 'max_age': max_age, 'min_hits': 0, 'iou_threshold': iou_thres}, min_hits=min_hits)
     
     # Run tracking
-    with tqdm(total=len(all_preds), desc="Running tracking", ncols=0, disable=not verbose) as pbar:
-        for i, key in enumerate(sorted(all_preds.keys())):
-            if gp: gp(i / len(all_preds), pbar.__str__())
-            boxes = all_preds[key]
-            if boxes is not None:
+    with tqdm(total=len(low_dets), desc="Running tracking", ncols=0, disable=not verbose) as pbar:
+        for i, key in enumerate(sorted(low_dets.keys())):
+            if gp: gp(i / len(low_dets), pbar.__str__())
+            low_boxes = low_dets[key]
+            high_boxes = high_dets[key]
+            boxes = (low_boxes, high_boxes)
+            if len(low_boxes) + len(high_boxes) > 0:
                 tracker.update(boxes)
             else:
                 tracker.update()