Mercurial > hg > nsaunier > traffic-intelligence

comparison python/cvutils.py @ 868:1fdafa9f6bf4

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