def main(): ''' input ''' # choose the input stream #captureSource = "video/video_116.mp4" #captureSource = "video/video_205.mp4" captureSource = "video/video_white.mp4" #captureSource = 0 # webcam cap = cv2.VideoCapture(captureSource) ''' trackers typology ''' # choose the tracker trackerName = "CSRT" # "MOSSE" | "KCF" | "CSRT" tm = TrackerManager(trackerName, maxFailures=20) ''' parameters ''' # try to change these parameters period = 1 # length of the period: only on the first frame of the period we detect objects (instead, we track them in every frame) maintainDetected = True # True if in transition frames, in case of overlapping bboxes, we want to keep those of the detector (False if we want to keep those of the tracker) frameWidth = 512 ''' some background subtractor with default params ''' # bgSubtractor = cv2.bgsegm.createBackgroundSubtractorMOG(history=200, nmixtures=5, backgroundRatio=0.7, noiseSigma=0) # bgSubtractor = cv2.createBackgroundSubtractorMOG2(history=500, varThreshold=16, detectShadows=True) # bgSubtractor = cv2.bgsegm.createBackgroundSubtractorGMG(initializationFrames=120, decisionThreshold=0.8) # bgSubtractor = cv2.createBackgroundSubtractorKNN(history=500, dist2Threshold=400.0, detectShadows=True) # bgSubtractor = CompositeBackgroundSubtractor(bgSubtractor1, bgSubtractor2, ...) ''' background subtractor ''' # define a background subtractor to be used for object detection #bgSubtractor = cv2.createBackgroundSubtractorMOG2(history=200, varThreshold=25, detectShadows=True) # good for video/video_116.mp4 bgSubtractor = cv2.createBackgroundSubtractorKNN( history=500, dist2Threshold=500.0, detectShadows=True) # good for video/video_205.mp4 # bgSubtractor = CompositeBackgroundSubtractor( # cv2.bgsegm.createBackgroundSubtractorMOG(history=600, nmixtures=3, backgroundRatio=0.2, noiseSigma=2.3), # cv2.createBackgroundSubtractorMOG2(history=200, varThreshold=14, detectShadows=True)) # good for video/video_white.mp4 ''' pipeline ''' # define the pipeline of functions to be executed on the b/w image, after the background subtraction and before getting bounding rects of contours kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7)) pipeline = ProcessPipeline() pipeline \ .add(cv2.medianBlur, ksize=5) \ .add(cv2.dilate, kernel=kernel) \ .add(cv2.dilate, kernel=kernel) \ .add(cv2.dilate, kernel=kernel) \ .add(cv2.dilate, kernel=kernel) \ .add(cv2.dilate, kernel=kernel) \ .add(fillHoles) \ .add(cv2.erode, kernel=kernel) \ .add(cv2.erode, kernel=kernel) \ ''' create object detector and face detector ''' od = ObjectDetector(bgSubtractor, pipeline) fd = FaceDetector() ''' auto-definition of output folder ''' outputDir = "output" if captureSource == 0: outputDir = os.path.join(outputDir, "webcam") else: outputDir = os.path.join(outputDir, captureSource[:captureSource.find(".")]) outputDir = os.path.join(outputDir, trackerName) print("Tracking video '%s' with tracker %s" % (captureSource, trackerName)) ''' cycle begins ''' frameNumber = 0 frames = 0 seconds = 0 eta = 0.05 totalTime = 0 show = True oneSkipOnly = False while True: ''' handle input: esc to quit; space to pause/start; "n" to go one frame at a time ''' k = cv2.waitKey(30) & 0xff if k == 27: break elif k == ord(' ') or oneSkipOnly: show = not show oneSkipOnly = False elif k == ord('n'): show = True oneSkipOnly = True if not show: if oneSkipOnly: show = False continue start = timer() ''' reading next frame ''' ret, frameOrig = cap.read() if not ret: break frameOrig = cv2.flip(frameOrig, 1) frame = imutils.resize(frameOrig, width=frameWidth) scale = frameOrig.shape[1] / frameWidth detectedObjects = [] if frameNumber % period == 0: ''' detection by background subtraction ''' detectedObjects = od.detect(frame) ''' objects tracking, faces detection''' ''' tracking ''' success, objects = tm.update(frame, detectedObjects, maintainDetected=maintainDetected) objIDs = tm.getIDs() tm.removeDeadTrackers() failed_objects = [obj for suc, obj in zip(success, objects) if not suc] failed_objIDs = [ objID for suc, objID in zip(success, objIDs) if not suc ] succ_objIDs = [objID for suc, objID in zip(success, objIDs) if suc] objects = [obj for suc, obj in zip(success, objects) if suc] ''' detection of faces ''' faces_bboxes = fd.detectFaces(frameOrig, objects, objIDs, scale=scale) ''' images merging and show ''' frameOrig = draw_bboxes(frameOrig, objects, (255, 0, 0), succ_objIDs, scale=scale) frameOrig = draw_bboxes(frameOrig, failed_objects, (0, 0, 255), failed_objIDs, scale=scale) frameOrig = draw_bboxes(frameOrig, faces_bboxes, (0, 255, 0)) frameOrig = cv2.resize(frameOrig, (640, 640)) cv2.imshow('frame', frameOrig) ''' some stats ''' frameNumber += 1 end = timer() frames = eta + (1 - eta) * frames seconds = eta * (end - start) + (1 - eta) * seconds print( "\rFrame: %04d FPS: %03d Active trackers: %02d Failed trackers: %02d " % (frameNumber, int( frames // seconds), len(objects), len(failed_objects)), end="") totalTime += end - start cap.release() cv2.destroyAllWindows() ''' save on disk ''' fd.dump(outputDir) avgFPS = str(round(frameNumber / totalTime, 2)) print("\rAverage FPS: " + avgFPS) with open(os.path.join(outputDir, "info.txt"), "w") as file: bgSubClsName = str(bgSubtractor.__class__) bgSubClsName = bgSubClsName[bgSubClsName.index("'") + 1:bgSubClsName.rindex("'")] file.write("tracker: " + trackerName + "\n") file.write("background subtractor: " + bgSubClsName + "\n") file.write("average FPS: " + avgFPS + "\n")