Mercurial Hosting > traffic-intelligence
comparison scripts/dltrack.py @ 1245:371c718e57d7
interface updates
author | Nicolas Saunier <nicolas.saunier@polymtl.ca> |
---|---|
date | Thu, 08 Feb 2024 16:10:54 -0500 |
parents | 4cd8ace3552f |
children | 2397de73770d |
comparison
equal
deleted
inserted
replaced
1244:00b71da2baac | 1245:371c718e57d7 |
---|---|
11 import cv2 | 11 import cv2 |
12 | 12 |
13 from trafficintelligence import cvutils, moving, storage, utils | 13 from trafficintelligence import cvutils, moving, storage, utils |
14 | 14 |
15 parser = argparse.ArgumentParser(description='The program tracks objects using the ultralytics models and trakcers.') | 15 parser = argparse.ArgumentParser(description='The program tracks objects using the ultralytics models and trakcers.') |
16 parser.add_argument('-i', dest = 'videoFilename', help = 'name of the video file', required = True) | 16 parser.add_argument('--cfg', dest = 'configFilename', help = 'name of the configuration file') |
17 parser.add_argument('-d', dest = 'databaseFilename', help = 'name of the Sqlite database file', required = True) | 17 parser.add_argument('-d', dest = 'databaseFilename', help = 'name of the Sqlite database file (overrides the configuration file)') |
18 parser.add_argument('-i', dest = 'videoFilename', help = 'name of the video file (overrides the configuration file)') | |
18 parser.add_argument('-m', dest = 'detectorFilename', help = 'name of the detection model file', required = True) | 19 parser.add_argument('-m', dest = 'detectorFilename', help = 'name of the detection model file', required = True) |
19 parser.add_argument('-t', dest = 'trackerFilename', help = 'name of the tracker file', required = True) | 20 parser.add_argument('-t', dest = 'trackerFilename', help = 'name of the tracker file', required = True) |
20 parser.add_argument('-o', dest = 'homographyFilename', help = 'filename of the homography matrix', default = 'homography.txt') | 21 parser.add_argument('-o', dest = 'homographyFilename', help = 'filename of the homography matrix', default = 'homography.txt') |
21 parser.add_argument('-k', dest = 'maskFilename', help = 'name of the mask file') | 22 parser.add_argument('-k', dest = 'maskFilename', help = 'name of the mask file') |
22 parser.add_argument('--undistort', dest = 'undistort', help = 'undistort the video', action = 'store_true') | 23 parser.add_argument('--undistort', dest = 'undistort', help = 'undistort the video', action = 'store_true') |
23 parser.add_argument('--intrinsic', dest = 'intrinsicCameraMatrixFilename', help = 'name of the intrinsic camera file') | 24 parser.add_argument('--intrinsic', dest = 'intrinsicCameraMatrixFilename', help = 'name of the intrinsic camera file') |
24 parser.add_argument('--distortion-coefficients', dest = 'distortionCoefficients', help = 'distortion coefficients', nargs = '*', type = float) | 25 parser.add_argument('--distortion-coefficients', dest = 'distortionCoefficients', help = 'distortion coefficients', nargs = '*', type = float) |
25 parser.add_argument('--display', dest = 'display', help = 'show the raw detection and tracking results', action = 'store_true') | 26 parser.add_argument('--display', dest = 'display', help = 'show the raw detection and tracking results', action = 'store_true') |
27 parser.add_argument('--no-image-coordinates', dest = 'notSavingImageCoordinates', help = 'not saving the raw detection and tracking results', action = 'store_true') | |
26 parser.add_argument('-f', dest = 'firstFrameNum', help = 'number of first frame number to process', type = int, default = 0) | 28 parser.add_argument('-f', dest = 'firstFrameNum', help = 'number of first frame number to process', type = int, default = 0) |
27 parser.add_argument('-l', dest = 'lastFrameNum', help = 'number of last frame number to process', type = int, default = float('Inf')) | 29 parser.add_argument('-l', dest = 'lastFrameNum', help = 'number of last frame number to process', type = int, default = float('Inf')) |
28 parser.add_argument('--conf', dest = 'confindence', help = 'object confidence threshold for detection', type = float, default = 0.25) | 30 parser.add_argument('--conf', dest = 'confindence', help = 'object confidence threshold for detection', type = float, default = 0.25) |
29 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) | 31 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) |
30 parser.add_argument('--cyclist-iou', dest = 'cyclistIou', help = 'IoU threshold to associate a bike and ped bounding box', type = float, default = 0.15) | 32 parser.add_argument('--cyclist-iou', dest = 'cyclistIou', help = 'IoU threshold to associate a bike and ped bounding box', type = float, default = 0.15) |
31 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) | 33 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) |
32 parser.add_argument('--max-temp-overal', dest = 'maxTemporalOverlap', help = 'maximum proportion of time to merge 2 bikes associated with same pedestrian', type = float, default = 0.05) | 34 parser.add_argument('--max-temp-overal', dest = 'maxTemporalOverlap', help = 'maximum proportion of time to merge 2 bikes associated with same pedestrian', type = float, default = 0.05) |
35 | |
33 args = parser.parse_args() | 36 args = parser.parse_args() |
37 params, videoFilename, databaseFilename, invHomography, intrinsicCameraMatrix, distortionCoefficients, undistortedImageMultiplication, undistort, firstFrameNum = storage.processVideoArguments(args) | |
38 | |
39 if args.intrinsicCameraMatrixFilename is not None: | |
40 intrinsicCameraMatrix = loadtxt(args.intrinsicCameraMatrixFilename) | |
41 if args.distortionCoefficients is not None: | |
42 distortionCoefficients = args.distortionCoefficients | |
43 if args.firstFrameNum is not None: | |
44 firstFrameNum = args.firstFrameNum | |
45 if args.lastFrameNum is not None: | |
46 lastFrameNum = args.lastFrameNum | |
34 | 47 |
35 # TODO add option to refine position with mask for vehicles | 48 # TODO add option to refine position with mask for vehicles |
49 # TODO work with optical flow (farneback or RAFT) https://pytorch.org/vision/main/models/raft.html | |
36 | 50 |
37 # use 2 x bytetrack track buffer to remove objects from existing ones | 51 # use 2 x bytetrack track buffer to remove objects from existing ones |
38 | 52 |
39 # Load a model | 53 # Load a model |
40 model = YOLO(args.detectorFilename) # seg yolov8x-seg.pt | 54 model = YOLO(args.detectorFilename) # seg yolov8x-seg.pt |
55 | 69 |
56 success, frame = capture.read() | 70 success, frame = capture.read() |
57 if not success: | 71 if not success: |
58 print('Input {} could not be read. Exiting'.format(args.videoFilename)) | 72 print('Input {} could not be read. Exiting'.format(args.videoFilename)) |
59 import sys; sys.exit() | 73 import sys; sys.exit() |
74 | |
60 results = model.track(frame, tracker=args.trackerFilename, classes=list(moving.cocoTypeNames.keys()), persist=True, verbose=False) | 75 results = model.track(frame, tracker=args.trackerFilename, classes=list(moving.cocoTypeNames.keys()), persist=True, verbose=False) |
61 # create object with user type and list of 3 features (bottom ones and middle) + projection | |
62 while capture.isOpened() and success and frameNum <= lastFrameNum: | 76 while capture.isOpened() and success and frameNum <= lastFrameNum: |
63 #for frameNum, result in enumerate(results): | |
64 result = results[0] | 77 result = results[0] |
65 if frameNum %10 == 0: | 78 if frameNum %10 == 0: |
66 print(frameNum, len(result.boxes), 'objects') | 79 print(frameNum, len(result.boxes), 'objects') |
67 for box in result.boxes: | 80 for box in result.boxes: |
68 #print(box.cls, box.id, box.xyxy) | 81 #print(box.cls, box.id, box.xyxy) |
101 obj.setUserType(utils.mostCommon(obj.userTypes)) # improve? mix with speed? | 114 obj.setUserType(utils.mostCommon(obj.userTypes)) # improve? mix with speed? |
102 | 115 |
103 # add quality control: avoid U-turns | 116 # add quality control: avoid U-turns |
104 | 117 |
105 # merge bikes and people | 118 # merge bikes and people |
106 twowheels = [num for num, obj in objects.items() if obj.getUserType() in (3,4)] | 119 twowheels = [num for num, obj in objects.items() if obj.getUserType() in (moving.userType2Num['motorcyclist'],moving.userType2Num['cyclist'])] |
107 pedestrians = [num for num, obj in objects.items() if obj.getUserType() == 2] | 120 pedestrians = [num for num, obj in objects.items() if obj.getUserType() == moving.userType2Num['pedestrian']] |
108 | 121 |
109 def mergeObjects(obj1, obj2): | 122 def mergeObjects(obj1, obj2): |
110 obj1.features = obj1.features+obj2.features | 123 obj1.features = obj1.features+obj2.features |
111 obj1.featureNumbers = obj1.featureNumbers+obj2.featureNumbers | 124 obj1.featureNumbers = obj1.featureNumbers+obj2.featureNumbers |
112 obj1.timeInterval = moving.TimeInterval(min(obj1.getFirstInstant(), obj2.getFirstInstant()), max(obj1.getLastInstant(), obj2.getLastInstant())) | 125 obj1.timeInterval = moving.TimeInterval(min(obj1.getFirstInstant(), obj2.getFirstInstant()), max(obj1.getLastInstant(), obj2.getLastInstant())) |
132 # before matching, scan for pedestrians with good non-overlapping temporal match with different bikes | 145 # before matching, scan for pedestrians with good non-overlapping temporal match with different bikes |
133 for pedInd in range(costs.shape[1]): | 146 for pedInd in range(costs.shape[1]): |
134 nMatchedBikes = (costs[:,pedInd] < -args.cyclistMatchingProportion).sum() | 147 nMatchedBikes = (costs[:,pedInd] < -args.cyclistMatchingProportion).sum() |
135 if nMatchedBikes == 0: # peds that have no bike matching: see if they have been classified as bikes sometimes | 148 if nMatchedBikes == 0: # peds that have no bike matching: see if they have been classified as bikes sometimes |
136 userTypeStats = Counter(obj.userTypes) | 149 userTypeStats = Counter(obj.userTypes) |
137 if (4 in userTypeStats or (3 in userTypeStats and 4 in userTypeStats and userTypeStats[3]<=userTypeStats[4])) and userTypeStats[3]+userTypeStats[4] > args.bikeProportion*userTypeStats.total(): # 3 is motorcycle and 4 is cyclist (verif if not turning all motorbike into cyclists) | 150 if (moving.userType2Num['cyclist'] in userTypeStats or (moving.userType2Num['motorcyclist'] in userTypeStats and moving.userType2Num['cyclist'] in userTypeStats and userTypeStats[moving.userType2Num['motorcyclist']]<=userTypeStats[moving.userType2Num['cyclist']])) and userTypeStats[moving.userType2Num['motorcyclist']]+userTypeStats[moving.userType2Num['cyclist']] > args.bikeProportion*userTypeStats.total(): # verif if not turning all motorbike into cyclists |
138 obj.setUserType(4) | 151 obj.setUserType(moving.userType2Num['cyclist']) |
139 elif nMatchedBikes > 1: # try to merge bikes first | 152 elif nMatchedBikes > 1: # try to merge bikes first |
140 twIndices = np.nonzero(costs[:,pedInd] < -args.cyclistMatchingProportion)[0] | 153 twIndices = np.nonzero(costs[:,pedInd] < -args.cyclistMatchingProportion)[0] |
141 # we have to compute temporal overlaps of all 2 wheels among themselves, then remove the ones with the most overlap (sum over column) one by one until there is little left | 154 # we have to compute temporal overlaps of all 2 wheels among themselves, then remove the ones with the most overlap (sum over column) one by one until there is little left |
142 nTwoWheels = len(twIndices) | 155 nTwoWheels = len(twIndices) |
143 twTemporalOverlaps = np.zeros((nTwoWheels,nTwoWheels)) | 156 twTemporalOverlaps = np.zeros((nTwoWheels,nTwoWheels)) |