detectvehicles.py

"""
Detect vehicles in a video sequence writing the result to an hdf5 file.

Usage:
    detectvehicles (-h | --help)
    detectvehicles [options] <videofile> <hdf5output>

Options:
    -h, --help              Show brief usage summary.
    -q, --quiet             Reduce logging verbosity.

"""
import datetime
import logging
import os
import sys

import cv2
import docopt
import h5py
import numpy as np
from pkg_resources import Requirement, resource_filename

from srftraffic.util import video_frames

CLASSIFIER_FN = resource_filename(
    Requirement.parse('srftraffic'), 'srftraffic/data/classifier.xml'
)

def _main():
    opts = docopt.docopt(__doc__)
    logging.basicConfig(
        level=logging.WARN if opts['--quiet'] else logging.INFO,
    )

    cc = cv2.CascadeClassifier()
    if not cc.load(CLASSIFIER_FN):
        raise RuntimeError('Error loading classifier')

    video_fn = os.path.abspath(opts['<videofile>'])

    logging.info('Opening output %s', opts['<hdf5output>'])
    out = h5py.File(opts['<hdf5output>'], 'w')

    detections = out.create_dataset(
        'detections', (0, 5), chunks=(100, 5), maxshape=(None, 5), dtype='i32',
    )
    detections.attrs['video_filename'] = video_fn
    detections.attrs['created_at'] = datetime.datetime.utcnow().isoformat()

    logging.info('Processing...')
    for frame_idx, frame in enumerate(video_frames(video_fn)):
        if frame_idx % 20 == 0:
            logging.info('Frame %s...', frame_idx)
        n_row = detections.shape[0]

        d = cc.detectMultiScale(frame)
        if len(d) == 0:
            continue

        detections.resize((n_row + len(d), detections.shape[1]))
        detections[-len(d):, :] = np.asarray(list(
            [frame_idx] + list(o) for o in d
        ))

def main():
    sys.exit(_main())

if __name__ == '__main__':
    main()