Update tracker.py

lib/fish_eye/tracker.py

@@ -6,6 +6,7 @@ import numpy as np
 
 from fish_length import Fish_Length
 from sort import Sort
+from associative import Associate
 
 class Tracker:
     def __init__(self, clip_info, algorithm=Sort, args={'max_age':1, 'min_hits':0, 'iou_threshold':0.05}, min_hits=3):
@@ -24,16 +25,23 @@ class Tracker:
             # Match confidence with correct track
             conf = 0
             min_score = 1000000
-            for det in dets[0]:
-                score = sum(abs(det[0:4] - track[0:4]))
-                if (score < min_score):
-                    min_score = score
-                    conf = det[4]
-            for det in dets[1]:
-                score = sum(abs(det[0:4] - track[0:4]))
-                if (score < min_score):
-                    min_score = score
-                    conf = det[4]
+            if self.algorithm == Sort:
+                for det in dets:
+                    score = sum(abs(det[0:4] - track[0:4]))
+                    if (score < min_score):
+                        min_score = score
+                        conf = det[4]
+            elif self.algorith == Associate:
+                for det in dets[0]:
+                    score = sum(abs(det[0:4] - track[0:4]))
+                    if (score < min_score):
+                        min_score = score
+                        conf = det[4]
+                for det in dets[1]:
+                    score = sum(abs(det[0:4] - track[0:4]))
+                    if (score < min_score):
+                        min_score = score
+                        conf = det[4]
             
             # Assign Track
             self.fish_ids[int(track[4])] += 1