Template module for video frame processing.

process_video.py

"""Template code for video frame processing.

Features:
    - video source from files or webcams.
    - automatically flip the frame if reading from webcam.
    - built in video writer to output the processed result.
    - built in FPS metter.

For more: https://github.com/yinguobing
"""
from argparse import ArgumentParser

import cv2

# Take arguments from user input.
parser = ArgumentParser()
parser.add_argument("--video", type=str, default=None,
                    help="Video file to be processed.")
parser.add_argument("--cam", type=int, default=None,
                    help="The webcam index.")
parser.add_argument("--write_video", type=bool, default=False,
                    help="Write output video.")
args = parser.parse_args()


if __name__ == "__main__":
    # Initialization. This is a good point to initialize your objects like
    # detectors or something else.

    # Setup the video source. If no video file provided, the default webcam will
    # be used.
    video_src = args.cam if args.cam is not None else args.video
    if video_src is None:
        print("Warning: video source not assigned, default webcam will be used.")
        video_src = 0

    cap = cv2.VideoCapture(video_src)

    # If reading frames from a webcam, try setting the camera resolution.
    if video_src == 0:
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)

    # Get the real frame resolution.
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    frame_rate = cap.get(cv2.CAP_PROP_FPS)

    # Video output by video writer.
    if args.write_video:
        video_writer = cv2.VideoWriter(
            'output.avi', cv2.VideoWriter_fourcc(*'avc1'), frame_rate,
            (frame_width, frame_height))

    # Introduce a metter to measure the FPS.
    tm = cv2.TickMeter()

    # Loop through the video frames.
    while True:
        # Start the metter.
        tm.start()

        # Read frame, crop it, flip it, suits your needs.
        frame_got, frame = cap.read()
        if not frame_got:
            break

        # Crop it if frame is larger than expected.
        # frame = frame[0:480, 300:940]

        # If frame comes from webcam, flip it so it looks like a mirror.
        if video_src == 0:
            frame = cv2.flip(frame, 2)

        # Do your processing here.

        # Stop the metter. You can use `tm.getFPS` to get the FPS.
        tm.stop()

        # Show the result in windows.
        cv2.imshow('Preview', frame)

        # Write video file.
        if args.write_video:
            video_writer.write(frame)

        if cv2.waitKey(1) == 27:
            break