Mercurial > hg > nsaunier > traffic-intelligence

comparison scripts/dltrack.py @ 1245:371c718e57d7

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
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))