Preparation
In order to understand what is happening, and eventually solve our problem, we first need to improve the script a little.
I've added logging of the important steps of your algorithm, refactored the code a little, and added saving of the mask and processed images, added ability to run the script using the individual frame images, along with some other modifications.
This is what the script looks like at this point:
import logging
import logging.handlers
import os
import time
import sys
import cv2
import numpy as np
from vehicle_counter import VehicleCounter
# ============================================================================
IMAGE_DIR = "images"
IMAGE_FILENAME_FORMAT = IMAGE_DIR + "/frame_%04d.png"
# Support either video file or individual frames
CAPTURE_FROM_VIDEO = False
if CAPTURE_FROM_VIDEO:
IMAGE_SOURCE = "traffic.avi" # Video file
else:
IMAGE_SOURCE = IMAGE_FILENAME_FORMAT # Image sequence
# Time to wait between frames, 0=forever
WAIT_TIME = 1 # 250 # ms
LOG_TO_FILE = True
# Colours for drawing on processed frames
DIVIDER_COLOUR = (255, 255, 0)
BOUNDING_BOX_COLOUR = (255, 0, 0)
CENTROID_COLOUR = (0, 0, 255)
# ============================================================================
def init_logging():
main_logger = logging.getLogger()
formatter = logging.Formatter(
fmt='%(asctime)s.%(msecs)03d %(levelname)-8s [%(name)s] %(message)s'
, datefmt='%Y-%m-%d %H:%M:%S')
handler_stream = logging.StreamHandler(sys.stdout)
handler_stream.setFormatter(formatter)
main_logger.addHandler(handler_stream)
if LOG_TO_FILE:
handler_file = logging.handlers.RotatingFileHandler("debug.log"
, maxBytes = 2**24
, backupCount = 10)
handler_file.setFormatter(formatter)
main_logger.addHandler(handler_file)
main_logger.setLevel(logging.DEBUG)
return main_logger
# ============================================================================
def save_frame(file_name_format, frame_number, frame, label_format):
file_name = file_name_format % frame_number
label = label_format % frame_number
log.debug("Saving %s as '%s'", label, file_name)
cv2.imwrite(file_name, frame)
# ============================================================================
def get_centroid(x, y, w, h):
x1 = int(w / 2)
y1 = int(h / 2)
cx = x + x1
cy = y + y1
return (cx, cy)
# ============================================================================
def detect_vehicles(fg_mask):
log = logging.getLogger("detect_vehicles")
MIN_CONTOUR_WIDTH = 21
MIN_CONTOUR_HEIGHT = 21
# Find the contours of any vehicles in the image
contours, hierarchy = cv2.findContours(fg_mask
, cv2.RETR_EXTERNAL
, cv2.CHAIN_APPROX_SIMPLE)
log.debug("Found %d vehicle contours.", len(contours))
matches = []
for (i, contour) in enumerate(contours):
(x, y, w, h) = cv2.boundingRect(contour)
contour_valid = (w >= MIN_CONTOUR_WIDTH) and (h >= MIN_CONTOUR_HEIGHT)
log.debug("Contour #%d: pos=(x=%d, y=%d) size=(w=%d, h=%d) valid=%s"
, i, x, y, w, h, contour_valid)
if not contour_valid:
continue
centroid = get_centroid(x, y, w, h)
matches.append(((x, y, w, h), centroid))
return matches
# ============================================================================
def filter_mask(fg_mask):
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
# Fill any small holes
closing = cv2.morphologyEx(fg_mask, cv2.MORPH_CLOSE, kernel)
# Remove noise
opening = cv2.morphologyEx(closing, cv2.MORPH_OPEN, kernel)
# Dilate to merge adjacent blobs
dilation = cv2.dilate(opening, kernel, iterations = 2)
return dilation
# ============================================================================
def process_frame(frame_number, frame, bg_subtractor, car_counter):
log = logging.getLogger("process_frame")
# Create a copy of source frame to draw into
processed = frame.copy()
# Draw dividing line -- we count cars as they cross this line.
cv2.line(processed, (0, car_counter.divider), (frame.shape[1], car_counter.divider), DIVIDER_COLOUR, 1)
# Remove the background
fg_mask = bg_subtractor.apply(frame, None, 0.01)
fg_mask = filter_mask(fg_mask)
save_frame(IMAGE_DIR + "/mask_%04d.png"
, frame_number, fg_mask, "foreground mask for frame #%d")
matches = detect_vehicles(fg_mask)
log.debug("Found %d valid vehicle contours.", len(matches))
for (i, match) in enumerate(matches):
contour, centroid = match
log.debug("Valid vehicle contour #%d: centroid=%s, bounding_box=%s", i, centroid, contour)
x, y, w, h = contour
# Mark the bounding box and the centroid on the processed frame
# NB: Fixed the off-by one in the bottom right corner
cv2.rectangle(processed, (x, y), (x + w - 1, y + h - 1), BOUNDING_BOX_COLOUR, 1)
cv2.circle(processed, centroid, 2, CENTROID_COLOUR, -1)
log.debug("Updating vehicle count...")
car_counter.update_count(matches, processed)
return processed
# ============================================================================
def main():
log = logging.getLogger("main")
log.debug("Creating background subtractor...")
bg_subtractor = cv2.BackgroundSubtractorMOG()
log.debug("Pre-training the background subtractor...")
default_bg = cv2.imread(IMAGE_FILENAME_FORMAT % 119)
bg_subtractor.apply(default_bg, None, 1.0)
car_counter = None # Will be created after first frame is captured
# Set up image source
log.debug("Initializing video capture device #%s...", IMAGE_SOURCE)
cap = cv2.VideoCapture(IMAGE_SOURCE)
frame_width = cap.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)
frame_height = cap.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)
log.debug("Video capture frame size=(w=%d, h=%d)", frame_width, frame_height)
log.debug("Starting capture loop...")
frame_number = -1
while True:
frame_number += 1
log.debug("Capturing frame #%d...", frame_number)
ret, frame = cap.read()
if not ret:
log.error("Frame capture failed, stopping...")
break
log.debug("Got frame #%d: shape=%s", frame_number, frame.shape)
if car_counter is None:
# We do this here, so that we can initialize with actual frame size
log.debug("Creating vehicle counter...")
car_counter = VehicleCounter(frame.shape[:2], frame.shape[0] / 2)
# Archive raw frames from video to disk for later inspection/testing
if CAPTURE_FROM_VIDEO:
save_frame(IMAGE_FILENAME_FORMAT
, frame_number, frame, "source frame #%d")
log.debug("Processing frame #%d...", frame_number)
processed = process_frame(frame_number, frame, bg_subtractor, car_counter)
save_frame(IMAGE_DIR + "/processed_%04d.png"
, frame_number, processed, "processed frame #%d")
cv2.imshow('Source Image', frame)
cv2.imshow('Processed Image', processed)
log.debug("Frame #%d processed.", frame_number)
c = cv2.waitKey(WAIT_TIME)
if c == 27:
log.debug("ESC detected, stopping...")
break
log.debug("Closing video capture device...")
cap.release()
cv2.destroyAllWindows()
log.debug("Done.")
# ============================================================================
if __name__ == "__main__":
log = init_logging()
if not os.path.exists(IMAGE_DIR):
log.debug("Creating image directory `%s`...", IMAGE_DIR)
os.makedirs(IMAGE_DIR)
main()
This script is responsible for processing of the stream of images, and identifying all the vehicles in each frame -- I refer to them as matches
in the code.
The task of counting the detected vehicles is delegated to class VehicleCounter
. The reason why I chose to make this a class will become evident as we progress. I did not implement your vehicle counting algorithm, because it will not work for reasons that will again become evident as we dig into this deeper.
File vehicle_counter.py
contains the following code:
import logging
# ============================================================================
class VehicleCounter(object):
def __init__(self, shape, divider):
self.log = logging.getLogger("vehicle_counter")
self.height, self.width = shape
self.divider = divider
self.vehicle_count = 0
def update_count(self, matches, output_image = None):
self.log.debug("Updating count using %d matches...", len(matches))
# ============================================================================
Finally, I wrote a script that will stitch all the generated images together, so it's easier to inspect them:
import cv2
import numpy as np
# ============================================================================
INPUT_WIDTH = 160
INPUT_HEIGHT = 120
OUTPUT_TILE_WIDTH = 10
OUTPUT_TILE_HEIGHT = 12
TILE_COUNT = OUTPUT_TILE_WIDTH * OUTPUT_TILE_HEIGHT
# ============================================================================
def stitch_images(input_format, output_filename):
output_shape = (INPUT_HEIGHT * OUTPUT_TILE_HEIGHT
, INPUT_WIDTH * OUTPUT_TILE_WIDTH
, 3)
output = np.zeros(output_shape, np.uint8)
for i in range(TILE_COUNT):
img = cv2.imread(input_format % i)
cv2.rectangle(img, (0, 0), (INPUT_WIDTH - 1, INPUT_HEIGHT - 1), (0, 0, 255), 1)
# Draw the frame number
cv2.putText(img, str(i), (2, 10)
, cv2.FONT_HERSHEY_PLAIN, 0.7, (255, 255, 255), 1)
x = i % OUTPUT_TILE_WIDTH * INPUT_WIDTH
y = i / OUTPUT_TILE_WIDTH * INPUT_HEIGHT
output[y:y+INPUT_HEIGHT, x:x+INPUT_WIDTH,:] = img
cv2.imwrite(output_filename, output)
# ============================================================================
stitch_images("images/frame_%04d.png", "stitched_frames.png")
stitch_images("images/mask_%04d.png", "stitched_masks.png")
stitch_images("images/processed_%04d.png", "stitched_processed.png")
Analysis
In order to solve this problem, we should have some idea about what results we expect to get. We should also label all the distinct cars in the video, so it's easier to talk about them.
If we run our script, and stitch the images together, we get the a number of useful files to help us analyze the problem:
Upon inspecting those, a number