Mercurial Hosting > traffic-intelligence
comparison python/cvutils.py @ 868:1fdafa9f6bf4
added colors more friendly for color blind people (thanks Ryan Louie!)
author | Nicolas Saunier <nicolas.saunier@polymtl.ca> |
---|---|
date | Wed, 01 Feb 2017 11:58:04 -0500 |
parents | 8f28b4fcf129 |
children | ff92801e5c54 |
comparison
equal
deleted
inserted
replaced
867:003445db1e30 | 868:1fdafa9f6bf4 |
---|---|
28 | 28 |
29 | 29 |
30 #import aggdraw # agg on top of PIL (antialiased drawing) | 30 #import aggdraw # agg on top of PIL (antialiased drawing) |
31 | 31 |
32 | 32 |
33 cvRed = (0,0,255) | 33 cvRed = {'default': (0,0,255), |
34 cvGreen = (0,255,0) | 34 'colorblind': (0,114,178)} |
35 cvBlue = (255,0,0) | 35 cvGreen = {'default': (0,255,0), |
36 cvCyan = (255, 255, 0) | 36 'colorblind': (0,158,115)} |
37 cvYellow = (0, 255, 255) | 37 cvBlue = {'default': (255,0,0), |
38 cvMagenta = (255, 0, 255) | 38 'colorblind': (213,94,0)} |
39 cvWhite = (255, 255, 255) | 39 cvCyan = {'default': (255, 255, 0), |
40 cvBlack = (0,0,0) | 40 'colorblind': (240,228,66)} |
41 cvColors3 = utils.PlottingPropertyValues([cvRed, | 41 cvYellow = {'default': (0, 255, 255), |
42 cvGreen, | 42 'colorblind': (86,180,233)} |
43 cvBlue]) | 43 cvMagenta = {'default': (255, 0, 255), |
44 cvColors = utils.PlottingPropertyValues([cvRed, | 44 'colorblind': (204,121,167)} |
45 cvGreen, | 45 cvWhite = {k: (255, 255, 255) for k in ['default', 'colorblind']} |
46 cvBlue, | 46 cvBlack = {k: (0,0,0) for k in ['default', 'colorblind']} |
47 cvCyan, | 47 |
48 cvYellow, | 48 cvColors3 = {k: utils.PlottingPropertyValues([cvRed[k], cvGreen[k], cvBlue[k]]) for k in ['default', 'colorblind']} |
49 cvMagenta, | 49 cvColors = {k: utils.PlottingPropertyValues([cvRed[k], cvGreen[k], cvBlue[k], cvCyan[k], cvYellow[k], cvMagenta[k], cvWhite[k], cvBlack[k]]) for k in ['default', 'colorblind']} |
50 cvWhite, | |
51 cvBlack]) | |
52 | 50 |
53 def quitKey(key): | 51 def quitKey(key): |
54 return chr(key&255)== 'q' or chr(key&255) == 'Q' | 52 return chr(key&255)== 'q' or chr(key&255) == 'Q' |
55 | 53 |
56 def saveKey(key): | 54 def saveKey(key): |
144 newCameraMatrix = deepcopy(intrinsicCameraMatrix) | 142 newCameraMatrix = deepcopy(intrinsicCameraMatrix) |
145 newCameraMatrix[0,2] = newImgSize[0]/2. | 143 newCameraMatrix[0,2] = newImgSize[0]/2. |
146 newCameraMatrix[1,2] = newImgSize[1]/2. | 144 newCameraMatrix[1,2] = newImgSize[1]/2. |
147 return cv2.initUndistortRectifyMap(intrinsicCameraMatrix, array(distortionCoefficients), identity(3), newCameraMatrix, newImgSize, cv2.CV_32FC1) | 145 return cv2.initUndistortRectifyMap(intrinsicCameraMatrix, array(distortionCoefficients), identity(3), newCameraMatrix, newImgSize, cv2.CV_32FC1) |
148 | 146 |
149 def playVideo(filenames, windowNames = None, firstFrameNums = None, frameRate = -1, interactive = False, printFrames = True, text = None, rescale = 1., step = 1): | 147 def playVideo(filenames, windowNames = None, firstFrameNums = None, frameRate = -1, interactive = False, printFrames = True, text = None, rescale = 1., step = 1, colorBlind = False): |
150 '''Plays the video(s)''' | 148 '''Plays the video(s)''' |
149 if colorBlind: | |
150 colorType = 'colorblind' | |
151 else: | |
152 colorType = 'default' | |
151 if len(filenames) == 0: | 153 if len(filenames) == 0: |
152 print('Empty filename list') | 154 print('Empty filename list') |
153 return | 155 return |
154 if windowNames is None: | 156 if windowNames is None: |
155 windowNames = ['frame{}'.format(i) for i in xrange(len(filenames))] | 157 windowNames = ['frame{}'.format(i) for i in xrange(len(filenames))] |
179 if array(rets).all(): | 181 if array(rets).all(): |
180 if printFrames: | 182 if printFrames: |
181 print('frame shown {0}'.format(nFramesShown)) | 183 print('frame shown {0}'.format(nFramesShown)) |
182 for i in xrange(len(filenames)): | 184 for i in xrange(len(filenames)): |
183 if text is not None: | 185 if text is not None: |
184 cv2.putText(images[i], text, (10,50), cv2.FONT_HERSHEY_PLAIN, 1, cvRed) | 186 cv2.putText(images[i], text, (10,50), cv2.FONT_HERSHEY_PLAIN, 1, cvRed[colorType]) |
185 cvImshow(windowNames[i], images[i], rescale) # cv2.imshow('frame', img) | 187 cvImshow(windowNames[i], images[i], rescale) # cv2.imshow('frame', img) |
186 key = cv2.waitKey(wait) | 188 key = cv2.waitKey(wait) |
187 if saveKey(key): | 189 if saveKey(key): |
188 cv2.imwrite('image-{}.png'.format(frameNum), img) | 190 cv2.imwrite('image-{}.png'.format(frameNum), img) |
189 nFramesShown += step | 191 nFramesShown += step |
300 else: | 302 else: |
301 croppedImg = None | 303 croppedImg = None |
302 return croppedImg, yCropMin, yCropMax, xCropMin, xCropMax | 304 return croppedImg, yCropMin, yCropMax, xCropMin, xCropMax |
303 | 305 |
304 | 306 |
305 def displayTrajectories(videoFilename, objects, boundingBoxes = {}, homography = None, firstFrameNum = 0, lastFrameNumArg = None, printFrames = True, rescale = 1., nFramesStep = 1, saveAllImages = False, nZerosFilenameArg = None, undistort = False, intrinsicCameraMatrix = None, distortionCoefficients = None, undistortedImageMultiplication = 1., annotations = [], gtMatches = {}, toMatches = {}): | 307 def displayTrajectories(videoFilename, objects, boundingBoxes = {}, homography = None, firstFrameNum = 0, lastFrameNumArg = None, printFrames = True, rescale = 1., nFramesStep = 1, saveAllImages = False, nZerosFilenameArg = None, undistort = False, intrinsicCameraMatrix = None, distortionCoefficients = None, undistortedImageMultiplication = 1., annotations = [], gtMatches = {}, toMatches = {}, colorBlind = False): |
306 '''Displays the objects overlaid frame by frame over the video ''' | 308 '''Displays the objects overlaid frame by frame over the video ''' |
309 if colorBlind: | |
310 colorType = 'colorblind' | |
311 else: | |
312 colorType = 'default' | |
313 | |
307 capture = cv2.VideoCapture(videoFilename) | 314 capture = cv2.VideoCapture(videoFilename) |
308 width = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)) | 315 width = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)) |
309 height = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)) | 316 height = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)) |
310 | 317 |
311 windowName = 'frame' | 318 windowName = 'frame' |
346 if not hasattr(obj, 'projectedPositions'): | 353 if not hasattr(obj, 'projectedPositions'): |
347 if homography is not None: | 354 if homography is not None: |
348 obj.projectedPositions = obj.positions.project(homography) | 355 obj.projectedPositions = obj.positions.project(homography) |
349 else: | 356 else: |
350 obj.projectedPositions = obj.positions | 357 obj.projectedPositions = obj.positions |
351 cvPlot(img, obj.projectedPositions, cvColors[obj.getNum()], frameNum-obj.getFirstInstant()) | 358 cvPlot(img, obj.projectedPositions, cvColors[colorType][obj.getNum()], frameNum-obj.getFirstInstant()) |
352 if frameNum not in boundingBoxes.keys() and obj.hasFeatures(): | 359 if frameNum not in boundingBoxes.keys() and obj.hasFeatures(): |
353 imgcrop, yCropMin, yCropMax, xCropMin, xCropMax = imageBox(img, obj, frameNum, homography, width, height) | 360 imgcrop, yCropMin, yCropMax, xCropMin, xCropMax = imageBox(img, obj, frameNum, homography, width, height) |
354 cv2.rectangle(img, (xCropMin, yCropMin), (xCropMax, yCropMax), cvBlue, 1) | 361 cv2.rectangle(img, (xCropMin, yCropMin), (xCropMax, yCropMax), cvBlue[colorType], 1) |
355 objDescription = '{} '.format(obj.num) | 362 objDescription = '{} '.format(obj.num) |
356 if moving.userTypeNames[obj.userType] != 'unknown': | 363 if moving.userTypeNames[obj.userType] != 'unknown': |
357 objDescription += moving.userTypeNames[obj.userType][0].upper() | 364 objDescription += moving.userTypeNames[obj.userType][0].upper() |
358 if len(annotations) > 0: # if we loaded annotations, but there is no match | 365 if len(annotations) > 0: # if we loaded annotations, but there is no match |
359 if frameNum not in toMatches[obj.getNum()]: | 366 if frameNum not in toMatches[obj.getNum()]: |
360 objDescription += " FA" | 367 objDescription += " FA" |
361 cv2.putText(img, objDescription, obj.projectedPositions[frameNum-obj.getFirstInstant()].asint().astuple(), cv2.FONT_HERSHEY_PLAIN, 1, cvColors[obj.getNum()]) | 368 cv2.putText(img, objDescription, obj.projectedPositions[frameNum-obj.getFirstInstant()].asint().astuple(), cv2.FONT_HERSHEY_PLAIN, 1, cvColors[colorType][obj.getNum()]) |
362 # plot object bounding boxes | 369 # plot object bounding boxes |
363 if frameNum in boundingBoxes.keys(): | 370 if frameNum in boundingBoxes.keys(): |
364 for rect in boundingBoxes[frameNum]: | 371 for rect in boundingBoxes[frameNum]: |
365 cv2.rectangle(img, rect[0].asint().astuple(), rect[1].asint().astuple(), cvColors[obj.getNum()]) | 372 cv2.rectangle(img, rect[0].asint().astuple(), rect[1].asint().astuple(), cvColors[colorType][obj.getNum()]) |
366 # plot ground truth | 373 # plot ground truth |
367 if len(annotations) > 0: | 374 if len(annotations) > 0: |
368 for gt in annotations: | 375 for gt in annotations: |
369 if gt.existsAtInstant(frameNum): | 376 if gt.existsAtInstant(frameNum): |
370 if frameNum in gtMatches[gt.getNum()]: | 377 if frameNum in gtMatches[gt.getNum()]: |
371 color = cvColors[gtMatches[gt.getNum()][frameNum]] # same color as object | 378 color = cvColors[colorType][gtMatches[gt.getNum()][frameNum]] # same color as object |
372 else: | 379 else: |
373 color = cvRed | 380 color = cvRed[colorType] |
374 cv2.putText(img, 'Miss', gt.topLeftPositions[frameNum-gt.getFirstInstant()].asint().astuple(), cv2.FONT_HERSHEY_PLAIN, 1, cvRed) | 381 cv2.putText(img, 'Miss', gt.topLeftPositions[frameNum-gt.getFirstInstant()].asint().astuple(), cv2.FONT_HERSHEY_PLAIN, 1, color) |
375 cv2.rectangle(img, gt.topLeftPositions[frameNum-gt.getFirstInstant()].asint().astuple(), gt.bottomRightPositions[frameNum-gt.getFirstInstant()].asint().astuple(), color) | 382 cv2.rectangle(img, gt.topLeftPositions[frameNum-gt.getFirstInstant()].asint().astuple(), gt.bottomRightPositions[frameNum-gt.getFirstInstant()].asint().astuple(), color) |
376 # saving images and going to next | 383 # saving images and going to next |
377 if not saveAllImages: | 384 if not saveAllImages: |
378 cvImshow(windowName, img, rescale) | 385 cvImshow(windowName, img, rescale) |
379 key = cv2.waitKey() | 386 key = cv2.waitKey() |