main.py

import argparse

from traffic_counter import TrafficCounter


def CLI():
    # Define default values here to make documentation self-updating
    minArea_default = 200
    direction_default = ['H', '0.5']
    numCount_default = 10
    videoWidth_default = 640
    videoParams_default = ['mjpg', 'avi']
    startingFrame_default = 10

    parser = argparse.ArgumentParser(
        description='Finds the contours on a video file')  # creates a parser object
    parser.add_argument('-p', '--path', type=str, help="""A video filename or path. Works better with .avi files. If no path or name is provided, the camera will be used instead.""")  # instead of using metavar='--path', just type '--path'. For some reason the metavar argument was causing problems
    parser.add_argument('-a', '--minArea', type=int, help=f'The minimum area (in pixels) to draw a bounding box (default is {minArea_default})',
                        default=minArea_default)
    parser.add_argument('-d', '--direction', type=str, default=direction_default, nargs=2, help=f"""A character: H or V
    representing the orientation of the count line. H is horizontal, V is vertical.
    If not provided, the default is {direction_default[0]},{direction_default[1]}. The second parameter
    is a float number from 0 to 1 indicating the place at which the
    line should be drawn.""")
    parser.add_argument('-n', '--numCount', type=int, default=numCount_default,
                        help=f"""The number of contours to be detected by the program (default is {numCount_default}).""")
    parser.add_argument('-w', '--webcam', type=int, nargs='+',
                        help="""Allows the user to specify which to use as the video source""")
    parser.add_argument('--rgb', action='store_true',
                        help="Boolean flag to use rbg colors.")
    parser.add_argument('-vo', '--video_out', type=str,
                        default="", help="Provide a video filename to output")
    parser.add_argument('-vw', '--video_width', type=int, default=videoWidth_default,
                        help=f"Videos will be resized to this width (default is {videoWidth_default}). Height will be computed automatically to preserve aspect ratio")
    parser.add_argument('-vp', '--video_params', type=str, default=videoParams_default, nargs=2,
                        help=f"Provide video codec and extension (in that order) for the output video. Example: `--video_params mjpg avi`. Default values are {videoParams_default[0]} {videoParams_default[1]}")
    parser.add_argument('-sf', '--starting_frame', type=int, default=startingFrame_default,
                        help=f"Select the starting frame for video analysis (default is {startingFrame_default}). All frames before that will still be used for the background average")
    args = parser.parse_args()
    return args


def make_video_params_dict(video_params):
    codec = video_params[0]
    extension = video_params[1]

    params_dict = {
        'codec': codec,
        'extension': extension,
    }
    return params_dict


def main(args):
    video_source = args.path
    line_direction = args.direction[0]
    line_position = float(args.direction[1])
    video_width = args.video_width
    min_area = int(args.minArea)
    video_out = args.video_out
    numCnts = int(args.numCount)
    video_params = make_video_params_dict(args.video_params)
    starting_frame = args.starting_frame
    tc = TrafficCounter(video_source,
                        line_direction,
                        line_position,
                        video_width,
                        min_area,
                        video_out,
                        numCnts,
                        video_params,
                        starting_frame,)

    tc.main_loop()


if __name__ == '__main__':
    args = CLI()
    main(args)