Source code for rcpcv.demo

# useful machine vision routines shared by multiple samples

import argparse
import numpy as np
import cv2 as cv

#================================================================
def open_video_output(frame, args):
    # Open a video output file sized for the given sample frame.  The frame is not written.

    codec_code = cv.VideoWriter.fourcc(*'MJPG')  # motion JPEG
    frame_width = frame.shape[1]   # cols
    frame_height = frame.shape[0]  # rows
    frame_rate = args.fps
    return cv.VideoWriter(args.out, codec_code, frame_rate, (frame_width, frame_height))

#================================================================


def parse_args(parser):
    """Add additional common arguments and process the command line.  Returns an argparse results object."""
    parser.add_argument('-q','--quiet', action='store_true', help="run without viewer window or console output")
    parser.add_argument('--fps', default=30, type=int, help="output file frame rate  (default: %(default)s)")
    parser.add_argument('--verbose', action='store_true', help="enable even more detailed console output" )
    parser.add_argument('--live', action='store_true', help="use camera for input (ignores -i) (default: %(default)s).")
    parser.add_argument('--camera', default=0, type=int, help="camera number (default: %(default)s).")
    parser.add_argument('-i', '--input', default="frames.avi", type=str, help="input AVI file name (default: %(default)s)")
    parser.add_argument('out', default="processed.avi", type=str, nargs='?', help="output AVI file name (default: %(default)s)")
    return parser.parse_args()


#================================================================
def prepare_inputs(args):
    if args.live:
        capture = cv.VideoCapture(args.camera)
    else:
        capture = cv.VideoCapture(args.input)

    if not capture.isOpened():
        print("Unable to open video input source, exiting.")
        exit(1)

    if not args.quiet:
        cv.namedWindow("capture")

    return capture

#================================================================


def run_filter(args, callback):
    """Process a video input stream through a filter callback function."""

    capture = prepare_inputs(args)

    # the output file will be created after the frame size is known
    output = None

    try:
        while True:
            success, frame = capture.read()
            if not success or frame is None:
                if args.live:
                    continue # when using a camera, ignore empty frames
                else:
                    break # when reading from a file, empty frame is the end

            if args.verbose:
                print("read image of shape %s" % (str(frame.shape)))

            # apply script-specific processing
            frame = callback(frame, args)

            if not args.quiet:
                cv.imshow("capture", frame) # display the frame

            if output is None:
                output = open_video_output(frame, args)

            if output is not None:
                output.write(frame)
                if args.verbose:
                    print("wrote frame to file.")

            # If using the GUI, calling waitKey allows the window system event loop to run and update
            # the display.  The waitKey argument is in milliseconds.
            if not args.quiet:
                key = cv.waitKey(33) & 0xFF
                if key == ord('q'):
                    break

    except KeyboardInterrupt:
        print ("User quit.")

    output.release()
    capture.release()

#================================================================