--- a/scripts/dltrack.py	Tue Sep 19 17:04:30 2023 -0400
+++ b/scripts/dltrack.py	Mon Oct 02 16:51:43 2023 -0400
@@ -3,9 +3,11 @@
 import sys, argparse
 from copy import copy
 from collections import Counter
+import numpy as np
+from scipy.optimize import linear_sum_assignment
 from ultralytics import YOLO
 from torch import cat
-from torchvision import ops
+from torchvision.ops import box_iou
 import cv2
 
 from trafficintelligence import cvutils, moving, storage, utils
@@ -18,7 +20,10 @@
 parser.add_argument('--display', dest = 'display', help = 'show the results (careful with long videos, risk of running out of memory)', action = 'store_true')
 parser.add_argument('-f', dest = 'firstFrameNum', help = 'number of first frame number to process', type = int, default = 0)
 parser.add_argument('-l', dest = 'lastFrameNum', help = 'number of last frame number to process', type = int, default = float('Inf'))
-parser.add_argument('--bike-pct', dest = 'bikeProportion', help = 'percent of time a person classified as bike or motorbike to be classified as cyclist', type = float, default = 0.2)
+parser.add_argument('--bike-prop', dest = 'bikeProportion', help = 'minimum proportion of time a person classified as bike or motorbike to be classified as cyclist', type = float, default = 0.2)
+parser.add_argument('--cyclist-iou', dest = 'cyclistIou', help = 'IoU threshold to associate a bike and ped bounding box', type = float, default = 0.15)
+parser.add_argument('--cyclist-match-prop', dest = 'cyclistMatchingProportion', help = 'minimum proportion of time a bike exists and is associated with a pedestrian to be merged as cyclist', type = float, default = 0.3)
+# mask!!
 args = parser.parse_args()
 
 # required functionality?
@@ -65,7 +70,7 @@
 # check if one can go to specific frame https://docs.ultralytics.com/modes/track/#persisting-tracks-loop
 
 # Load a model
-model = YOLO(args.detectorFilename, ) # seg yolov8x-seg.pt
+model = YOLO(args.detectorFilename) # seg yolov8x-seg.pt
 # seg could be used on cropped image... if can be loaded and kept in memory
 # model = YOLO('/home/nicolas/Research/Data/classification-models/yolo_nas_l.pt ') # AttributeError: 'YoloNAS_L' object has no attribute 'get'
 
@@ -85,7 +90,7 @@
 lastFrameNum = args.lastFrameNum
 
 success, frame = capture.read()
-results = model.track(frame, tracker=args.trackerFilename, classes=list(moving.cocoTypeNames.keys()), persist=True)
+results = model.track(frame, tracker=args.trackerFilename, classes=list(moving.cocoTypeNames.keys()), persist=True, verbose=False)
 # create object with user type and list of 3 features (bottom ones and middle) + projection
 while capture.isOpened() and success and frameNum <= lastFrameNum:
 #for frameNum, result in enumerate(results):
@@ -131,12 +136,51 @@
         obj.setUserType(userTypeStats.most_common()[0][0])
 
 # merge bikes and people
-#Construire graphe bipartite vélo/moto personne
-#Lien = somme des iou / longueur track vélo 
-#Algo Hongrois
-#Verif overlap piéton vélo : si long, changement mode (trouver exemples)
+twowheels = [num for num, obj in currentObjects.items() if obj.getUserType() in (3,4)]
+pedestrians = [num for num, obj in currentObjects.items() if obj.getUserType() == 2]
+
+costs = []
+for twInd in twowheels:
+    tw = currentObjects[twInd]
+    twCost = []
+    for pedInd in pedestrians:
+        ped = currentObjects[pedInd]
+        nmatches = 0
+        for t in tw.bboxes:
+            if t in ped.bboxes:
+                #print(tw.num, ped.num, t, box_iou(tw.bboxes[t], ped.bboxes[t]))
+                if box_iou(tw.bboxes[t], ped.bboxes[t]).item() > args.cyclistIou:
+                    nmatches += 1
+        twCost.append(nmatches/len(tw.bboxes))
+    costs.append(twCost)
 
-# for all cyclists and motorbikes
+costs = -np.array(costs)
+# before matching, scan for pedestrians with good non-overlapping temporal match with different bikes
+for pedInd in costs.shape[1]:
+    if sum(costs[:,pedInd] < -args.cyclistMatchingProportion) >1:
+        twIndices = np.nonzero(costs[:,pedInd] < -args.cyclistMatchingProportion)
+        # we have to compute temporal overlaps with everyone else, then remove the ones with the most overlap (sum over column) one by one until there is little left
+        temporalOverlaps = np.zeros((len(twIndices),len(twIndices)))
+        
+
+twIndices, matchingPedIndices = linear_sum_assignment(costs)
+for twInd, pedInd in zip(twIndices, matchingPedIndices): # caution indices in the cost matrix
+    if -costs[twInd, pedInd] >= args.cyclistMatchingProportion:
+        tw = currentObjects[twowheels[twInd]]
+        ped = currentObjects[pedestrians[pedInd]]
+        timeInstants = set(tw.bboxes).union(set(ped.bboxes))
+        for t in timeInstants:
+            if t in tw.bboxes and t in ped.bboxes:
+                tw.features[0].tmpPositions[t] = moving.Point(min(tw.features[0].tmpPositions[t].x, ped.features[0].tmpPositions[t].x),
+                                                              min(tw.features[0].tmpPositions[t].y, ped.features[0].tmpPositions[t].y))
+                tw.features[1].tmpPositions[t] = moving.Point(max(tw.features[1].tmpPositions[t].x, ped.features[1].tmpPositions[t].x),
+                                                              max(tw.features[1].tmpPositions[t].y, ped.features[1].tmpPositions[t].y))
+            elif t in ped.bboxes:
+                tw.features[0].tmpPositions[t] = ped.features[0].tmpPositions[t]
+                tw.features[1].tmpPositions[t] = ped.features[1].tmpPositions[t]
+        tw.timeInterval = moving.TimeInterval(min(tw.getFirstInstant(), ped.getFirstInstant()), max(tw.getLastInstant(), ped.getLastInstant()))
+        del currentObjects[pedestrians[pedInd]]
+#Verif overlap piéton vélo : si long hors overlap, changement mode (trouver exemples)
 
 # interpolate and generate velocity (?) before saving
 for num, obj in currentObjects.items():