Mercurial > hg > nsaunier > traffic-intelligence

comparison scripts/classify-objects.py @ 901:753a081989e2

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
factorized some argument handling code
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Thu, 22 Jun 2017 12:02:34 -0400
parents 1466a63dd1cf
children c69a8defe5c3
comparison
equal deleted inserted replaced
900:85b81c46c526 901:753a081989e2
17 parser.add_argument('-n', dest = 'nObjects', help = 'number of objects to classify', type = int, default = None) 17 parser.add_argument('-n', dest = 'nObjects', help = 'number of objects to classify', type = int, default = None)
18 parser.add_argument('--plot-speed-distributions', dest = 'plotSpeedDistribution', help = 'simply plots the distributions used for each user type', action = 'store_true') 18 parser.add_argument('--plot-speed-distributions', dest = 'plotSpeedDistribution', help = 'simply plots the distributions used for each user type', action = 'store_true')
19 parser.add_argument('--max-speed-distribution-plot', dest = 'maxSpeedDistributionPlot', help = 'if plotting the user distributions, the maximum speed to display (km/h)', type = float, default = 50.) 19 parser.add_argument('--max-speed-distribution-plot', dest = 'maxSpeedDistributionPlot', help = 'if plotting the user distributions, the maximum speed to display (km/h)', type = float, default = 50.)
20 20
21 args = parser.parse_args() 21 args = parser.parse_args()
22 params = storage.ProcessParameters(args.configFilename) 22 params, videoFilename, databaseFilename, invHomography, intrinsicCameraMatrix, distortionCoefficients, undistortedImageMultiplication, undistort, firstFrameNum = storage.processVideoArguments(args)
23
23 classifierParams = storage.ClassifierParameters(params.classifierFilename) 24 classifierParams = storage.ClassifierParameters(params.classifierFilename)
24 classifierParams.convertToFrames(params.videoFrameRate, 3.6) # conversion from km/h to m/frame 25 classifierParams.convertToFrames(params.videoFrameRate, 3.6) # conversion from km/h to m/frame
25
26 if args.videoFilename is not None:
27 videoFilename = args.videoFilename
28 else:
29 videoFilename = params.videoFilename
30 if args.databaseFilename is not None:
31 databaseFilename = args.databaseFilename
32 else:
33 databaseFilename = params.databaseFilename
34
35 if params.homography is not None:
36 invHomography = np.linalg.inv(params.homography)
37 else:
38 invHomography = None
39 26
40 if classifierParams.speedAggregationMethod == 'median': 27 if classifierParams.speedAggregationMethod == 'median':
41 speedAggregationFunc = np.median 28 speedAggregationFunc = np.median
42 elif classifierParams.speedAggregationMethod == 'mean': 29 elif classifierParams.speedAggregationMethod == 'mean':
43 speedAggregationFunc = np.mean 30 speedAggregationFunc = np.mean
89 capture = cv2.VideoCapture(videoFilename) 76 capture = cv2.VideoCapture(videoFilename)
90 width = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)) 77 width = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH))
91 height = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)) 78 height = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT))
92 79
93 pastObjects = [] 80 pastObjects = []
94 if params.undistort: # setup undistortion 81 if undistort: # setup undistortion
95 [map1, map2] = cvutils.computeUndistortMaps(width, height, params.undistortedImageMultiplication, params.intrinsicCameraMatrix, params.distortionCoefficients) 82 [map1, map2] = cvutils.computeUndistortMaps(width, height, undistortedImageMultiplication, intrinsicCameraMatrix, distortionCoefficients)
96 if capture.isOpened(): 83 if capture.isOpened():
97 ret = True 84 ret = True
98 frameNum = timeInterval.first 85 frameNum = timeInterval.first
99 capture.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, frameNum) 86 capture.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, frameNum)
100 lastFrameNum = timeInterval.last 87 lastFrameNum = timeInterval.last
102 while ret and frameNum <= lastFrameNum: 89 while ret and frameNum <= lastFrameNum:
103 ret, img = capture.read() 90 ret, img = capture.read()
104 if ret: 91 if ret:
105 if frameNum%50 == 0: 92 if frameNum%50 == 0:
106 print('frame number: {}'.format(frameNum)) 93 print('frame number: {}'.format(frameNum))
107 if params.undistort: 94 if undistort:
108 img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR) 95 img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR)
109 currentObjects = [] 96 currentObjects = []
110 for obj in objects: 97 for obj in objects:
111 inter = obj.getTimeInterval() 98 inter = obj.getTimeInterval()
112 if inter.contains(frameNum): 99 if inter.contains(frameNum):