Mercurial Hosting > traffic-intelligence
comparison scripts/dltrack.py @ 1237:31a441efca6c
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()) |