6-yolov3.py

#%% Initiate processing

# Init count
count = 0

# Create new window
cv2.namedWindow('stream')

while(vid.isOpened()):
    # Perform detection every 60 frames
    perform_detection = count % 60 == 0
    ok, frame = vid.read()

    if ok:
        if perform_detection: # perform detection
            blob = cv2.dnn.blobFromImage(frame, 1 / 255, (416, 416), [0,0,0], 1, crop=False)
            # Pass blob to model
            model.setInput(blob)
            # Execute forward pass
            outputs = model.forward(outputNames)
            bboxes, probs, class_ids = where_is_it(frame, outputs)
            
            if len(bboxes) > 0:
                # Init multitracker
                mtracker = cv2.MultiTracker_create()
                # Apply non-max suppression and pass boxes to the multitracker
                idxs = cv2.dnn.NMSBoxes(bboxes, probs, P_THRESH, NMS_THRESH)
                for i in idxs:
                    bbox = [int(v) for v in bboxes[i[0]]]
                    x, y, w, h = bbox
                    # Use median flow
                    mtracker.add(cv2.TrackerMedianFlow_create(), frame, (x, y, w, h))
                # Increase counter
                count += 1
            else: # declare failure
                cv2.putText(frame, 'Detection failed', (20, 80), 
                            cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0,0,255), 2)
        else: # perform tracking
            is_tracking, bboxes = mtracker.update(frame)
            if is_tracking:
                for i, bbox in enumerate(bboxes):
                    x, y, w, h = [int(val) for val in bbox]
                    class_id = classes[class_ids[idxs[i][0]]]
                    col = [int(c) for c in colors[class_ids[idxs[i][0]], :]]
                    # Mark tracking frame with corresponding color, write class name on top
                    cv2.rectangle(frame, (x, y), (x+w, y+h), col, 2)
                    cv2.putText(frame, class_id, (x, y - 15),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, col, 2)
                # Increase counter
                count += 1
            # If tracking fails, reset count to trigger detection
            else:
                count = 0

        # Display the resulting frame
        cv2.imshow('stream', frame)
        out.write(frame)
        # Press ESC to exit
        if cv2.waitKey(25) & 0xFF == 27:
            break
    # Break if capture read does not work
    else:
        print('Exhausted video capture.')
        break
out.release()
cv2.destroyAllWindows()