Mercurial > hg > nsaunier > traffic-intelligence

comparison scripts/dltrack.py @ 1237:31a441efca6c

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
ped-bike grouping working, bike merging left todo
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Mon, 02 Oct 2023 16:51:43 -0400
parents 100fe098abe9
children b684135d817f
comparison
equal deleted inserted replaced
1236:100fe098abe9 1237:31a441efca6c
1 #! /usr/bin/env python3 1 #! /usr/bin/env python3
2 # from https://docs.ultralytics.com/modes/track/ 2 # from https://docs.ultralytics.com/modes/track/
3 import sys, argparse 3 import sys, argparse
4 from copy import copy 4 from copy import copy
5 from collections import Counter 5 from collections import Counter
6 import numpy as np
7 from scipy.optimize import linear_sum_assignment
6 from ultralytics import YOLO 8 from ultralytics import YOLO
7 from torch import cat 9 from torch import cat
8 from torchvision import ops 10 from torchvision.ops import box_iou
9 import cv2 11 import cv2
10 12
11 from trafficintelligence import cvutils, moving, storage, utils 13 from trafficintelligence import cvutils, moving, storage, utils
12 14
13 parser = argparse.ArgumentParser(description='The program tracks objects following the ultralytics yolo executable.')#, epilog = 'Either the configuration filename or the other parameters (at least video and database filenames) need to be provided.') 15 parser = argparse.ArgumentParser(description='The program tracks objects following the ultralytics yolo executable.')#, epilog = 'Either the configuration filename or the other parameters (at least video and database filenames) need to be provided.')
16 parser.add_argument('-m', dest = 'detectorFilename', help = 'name of the detection model file', required = True) 18 parser.add_argument('-m', dest = 'detectorFilename', help = 'name of the detection model file', required = True)
17 parser.add_argument('-t', dest = 'trackerFilename', help = 'name of the tracker file', required = True) 19 parser.add_argument('-t', dest = 'trackerFilename', help = 'name of the tracker file', required = True)
18 parser.add_argument('--display', dest = 'display', help = 'show the results (careful with long videos, risk of running out of memory)', action = 'store_true') 20 parser.add_argument('--display', dest = 'display', help = 'show the results (careful with long videos, risk of running out of memory)', action = 'store_true')
19 parser.add_argument('-f', dest = 'firstFrameNum', help = 'number of first frame number to process', type = int, default = 0) 21 parser.add_argument('-f', dest = 'firstFrameNum', help = 'number of first frame number to process', type = int, default = 0)
20 parser.add_argument('-l', dest = 'lastFrameNum', help = 'number of last frame number to process', type = int, default = float('Inf')) 22 parser.add_argument('-l', dest = 'lastFrameNum', help = 'number of last frame number to process', type = int, default = float('Inf'))
21 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) 23 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)
24 parser.add_argument('--cyclist-iou', dest = 'cyclistIou', help = 'IoU threshold to associate a bike and ped bounding box', type = float, default = 0.15)
25 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)
26 # mask!!
22 args = parser.parse_args() 27 args = parser.parse_args()
23 28
24 # required functionality? 29 # required functionality?
25 # # filename of the video to process (can be images, eg image%04d.png) 30 # # filename of the video to process (can be images, eg image%04d.png)
26 # video-filename = laurier.avi 31 # video-filename = laurier.avi
63 68
64 69
65 # check if one can go to specific frame https://docs.ultralytics.com/modes/track/#persisting-tracks-loop 70 # check if one can go to specific frame https://docs.ultralytics.com/modes/track/#persisting-tracks-loop
66 71
67 # Load a model 72 # Load a model
68 model = YOLO(args.detectorFilename, ) # seg yolov8x-seg.pt 73 model = YOLO(args.detectorFilename) # seg yolov8x-seg.pt
69 # seg could be used on cropped image... if can be loaded and kept in memory 74 # seg could be used on cropped image... if can be loaded and kept in memory
70 # model = YOLO('/home/nicolas/Research/Data/classification-models/yolo_nas_l.pt ') # AttributeError: 'YoloNAS_L' object has no attribute 'get' 75 # model = YOLO('/home/nicolas/Research/Data/classification-models/yolo_nas_l.pt ') # AttributeError: 'YoloNAS_L' object has no attribute 'get'
71 76
72 # Track with the model 77 # Track with the model
73 if args.display: 78 if args.display:
83 frameNum = args.firstFrameNum 88 frameNum = args.firstFrameNum
84 capture.set(cv2.CAP_PROP_POS_FRAMES, frameNum) 89 capture.set(cv2.CAP_PROP_POS_FRAMES, frameNum)
85 lastFrameNum = args.lastFrameNum 90 lastFrameNum = args.lastFrameNum
86 91
87 success, frame = capture.read() 92 success, frame = capture.read()
88 results = model.track(frame, tracker=args.trackerFilename, classes=list(moving.cocoTypeNames.keys()), persist=True) 93 results = model.track(frame, tracker=args.trackerFilename, classes=list(moving.cocoTypeNames.keys()), persist=True, verbose=False)
89 # create object with user type and list of 3 features (bottom ones and middle) + projection 94 # create object with user type and list of 3 features (bottom ones and middle) + projection
90 while capture.isOpened() and success and frameNum <= lastFrameNum: 95 while capture.isOpened() and success and frameNum <= lastFrameNum:
91 #for frameNum, result in enumerate(results): 96 #for frameNum, result in enumerate(results):
92 result = results[0] 97 result = results[0]
93 print(frameNum, len(result.boxes), 'objects') 98 print(frameNum, len(result.boxes), 'objects')
129 obj.setUserType(4) 134 obj.setUserType(4)
130 else: 135 else:
131 obj.setUserType(userTypeStats.most_common()[0][0]) 136 obj.setUserType(userTypeStats.most_common()[0][0])
132 137
133 # merge bikes and people 138 # merge bikes and people
134 #Construire graphe bipartite vélo/moto personne 139 twowheels = [num for num, obj in currentObjects.items() if obj.getUserType() in (3,4)]
135 #Lien = somme des iou / longueur track vélo 140 pedestrians = [num for num, obj in currentObjects.items() if obj.getUserType() == 2]
136 #Algo Hongrois
137 #Verif overlap piéton vélo : si long, changement mode (trouver exemples)
138 141
139 # for all cyclists and motorbikes 142 costs = []
143 for twInd in twowheels:
144 tw = currentObjects[twInd]
145 twCost = []
146 for pedInd in pedestrians:
147 ped = currentObjects[pedInd]
148 nmatches = 0
149 for t in tw.bboxes:
150 if t in ped.bboxes:
151 #print(tw.num, ped.num, t, box_iou(tw.bboxes[t], ped.bboxes[t]))
152 if box_iou(tw.bboxes[t], ped.bboxes[t]).item() > args.cyclistIou:
153 nmatches += 1
154 twCost.append(nmatches/len(tw.bboxes))
155 costs.append(twCost)
156
157 costs = -np.array(costs)
158 # before matching, scan for pedestrians with good non-overlapping temporal match with different bikes
159 for pedInd in costs.shape[1]:
160 if sum(costs[:,pedInd] < -args.cyclistMatchingProportion) >1:
161 twIndices = np.nonzero(costs[:,pedInd] < -args.cyclistMatchingProportion)
162 # 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
163 temporalOverlaps = np.zeros((len(twIndices),len(twIndices)))
164
165
166 twIndices, matchingPedIndices = linear_sum_assignment(costs)
167 for twInd, pedInd in zip(twIndices, matchingPedIndices): # caution indices in the cost matrix
168 if -costs[twInd, pedInd] >= args.cyclistMatchingProportion:
169 tw = currentObjects[twowheels[twInd]]
170 ped = currentObjects[pedestrians[pedInd]]
171 timeInstants = set(tw.bboxes).union(set(ped.bboxes))
172 for t in timeInstants:
173 if t in tw.bboxes and t in ped.bboxes:
174 tw.features[0].tmpPositions[t] = moving.Point(min(tw.features[0].tmpPositions[t].x, ped.features[0].tmpPositions[t].x),
175 min(tw.features[0].tmpPositions[t].y, ped.features[0].tmpPositions[t].y))
176 tw.features[1].tmpPositions[t] = moving.Point(max(tw.features[1].tmpPositions[t].x, ped.features[1].tmpPositions[t].x),
177 max(tw.features[1].tmpPositions[t].y, ped.features[1].tmpPositions[t].y))
178 elif t in ped.bboxes:
179 tw.features[0].tmpPositions[t] = ped.features[0].tmpPositions[t]
180 tw.features[1].tmpPositions[t] = ped.features[1].tmpPositions[t]
181 tw.timeInterval = moving.TimeInterval(min(tw.getFirstInstant(), ped.getFirstInstant()), max(tw.getLastInstant(), ped.getLastInstant()))
182 del currentObjects[pedestrians[pedInd]]
183 #Verif overlap piéton vélo : si long hors overlap, changement mode (trouver exemples)
140 184
141 # interpolate and generate velocity (?) before saving 185 # interpolate and generate velocity (?) before saving
142 for num, obj in currentObjects.items(): 186 for num, obj in currentObjects.items():
143 obj.features[0].timeInterval = copy(obj.getTimeInterval()) 187 obj.features[0].timeInterval = copy(obj.getTimeInterval())
144 obj.features[1].timeInterval = copy(obj.getTimeInterval()) 188 obj.features[1].timeInterval = copy(obj.getTimeInterval())