vision-multiprocess-example.py

import cv2
import time
import numpy as np
import multiprocessing as mp


log_const = 30


class FPS:
    def __init__(self, length: int, title: str):
        self.length = length
        self.counter = 0
        self.start = time.time_ns()
        self.title = title

    def tick(self):
        self.counter += 1
        if (self.counter % self.length == 0):
            print(self.title, round(self.counter/round((time.time_ns() -
                                                        self.start)/1e9)), "FPS", round((time.time_ns() -
                                                                                         self.start)/(self.counter*1e3)), 'us')


def scale_frame(frame, wanted_frame_width: int):
    '''Scales frame to wanted width if smaller than frame or returns original'''
    raw_frame_height, raw_frame_width, _ = frame.shape
    if wanted_frame_width >= raw_frame_width:
        return frame
    scale_factor = wanted_frame_width / raw_frame_width
    width = int(raw_frame_width * scale_factor)
    height = int(raw_frame_height * scale_factor)
    dim = (width, height)
    frame_scaled = cv2.resize(
        frame, dim, interpolation=cv2.INTER_AREA)
    return frame_scaled


def cpu_bound():
    face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
    cpu_bound_ticker = FPS(5, 'CPU Bound')

    def inner(frame):
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        face_cascade.detectMultiScale(gray, 1.01, 5)
        cpu_bound_ticker.tick()

    return inner


def camera_producer():
    camera_ticker = FPS(log_const, 'Camera')

    cam = cv2.VideoCapture(0)
    while True:
        ok, frame = cam.read()
        if ok:
            camera_ticker.tick()
            yield frame


display_ticker = FPS(log_const, 'Display')


def display(image_frame):
    cv2.imshow('display', scale_frame(image_frame, 200))
    cv2.waitKey(1)
    display_ticker.tick()


def display_thread(input_queue: mp.Queue):
    while True:
        image = input_queue.get()
        display(image)


def camera_thread(display_queue: mp.Queue, cpu_bound_queue: mp.Queue):
    camera = camera_producer()
    for frame in camera:
        display_queue.put(frame)
        try:
            cpu_bound_queue.put(frame, False)
        except mp.queues.Full:
            pass


def cpu_bound_thread(input_queue: mp.Queue):
    cpu_bound_op = cpu_bound()
    while True:
        frame = input_queue.get()
        cpu_bound_op(frame)


def main():
    ctx = mp.get_context('spawn')
    display_input_queue = ctx.Queue(1)
    cpu_bound_input_queue = ctx.Queue(1)

    PROCESSES = [
        ctx.Process(name="camera_thread", target=camera_thread, args=(
            display_input_queue, cpu_bound_input_queue)),
        ctx.Process(name="cpu_bound_thread", target=cpu_bound_thread, args=(
            cpu_bound_input_queue,)),
    ]

    for process in PROCESSES:
        process.start()

    display_thread(display_input_queue)

    for process in PROCESSES:
        process.join()


if __name__ == '__main__':
    main()