rect.py

#! /usr/bin/env python
from turtle import circle
import cv2
import sys
import numpy as np
import typing
from typing import Tuple
import collections

# See also
# https://github.com/microsoft/python-type-stubs/blob/main/cv2/__init__.pyi
Mat = np.ndarray

def detect_draw_circle(detect_image: Mat, paint_image: Mat, pts: collections.deque) -> Mat:
    circles = cv2.HoughCircles(
        detect_image,
        method=cv2.HOUGH_GRADIENT,
        dp=2,
        minDist=detect_image.shape[0] * 2,
        param1=300,
        param2=0.9,
        minRadius=5,
        maxRadius=10)
    output_frame = paint_image.copy()
    # draw circles
    if circles is not None:
        circles = np.round(circles[0, :]).astype("int")
        for circle in circles:
            x, y, r, *_rest = circle
            pt = (x, y)
            cv2.circle(output_frame, center=pt, radius=r,
                       color=(0, 255, 0), thickness=2)
            pts.appendleft(pt)
    # draw tail
    for i in np.arange(1, len(pts)):
        LINE_MAX_THICKNESS = 5
        thickness = int(np.sqrt(LINE_MAX_THICKNESS / float(i + 1)) * 2.5)
        cv2.line(output_frame, pts[i - 1], pts[i], (0, 0, 255), thickness)
    return output_frame


# preprocess image with mog2
def handle_image_mog2(image: Mat, back_sub: cv2.BackgroundSubtractorMOG2) -> Mat:
    blur = cv2.GaussianBlur(image, (5, 5), 2, 2)
    fg_mask = back_sub.apply(blur)
    blur_again = cv2.GaussianBlur(fg_mask, (5, 5), 2, 2)
    return blur_again

# preprocess image with color range
def handle_image_color(image: Mat,
                       color_min: Tuple[int, int, int],
                       color_max: Tuple[int, int, int],
                       ) -> Mat:
    blur = cv2.GaussianBlur(image, (5, 5), 2, 2)
    lab = cv2.cvtColor(blur, cv2.COLOR_BGR2LAB)
    l_min, a_min, b_min = color_min
    l_max, a_max, b_max = color_max
    lab_red = cv2.inRange(lab, np.array(
        [l_min, a_min, b_min]), np.array([l_max, a_max, b_max]))
    blur_again = cv2.GaussianBlur(lab_red, (5, 5), 2, 2)
    return blur_again


def main():
    color_min = (0, 137, 126)
    color_max = (245, 245, 135)
    q = collections.deque(maxlen=25)
    back_sub = cv2.createBackgroundSubtractorMOG2(varThreshold=16)

    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    if (cap.isOpened() == False):
        print("Error opening video stream or file")
    while (cap.isOpened()):
        ret: bool
        frame: Mat
        ret, frame = cap.read()
        if ret == True:
            # flip frame horizontally
            frame = cv2.flip(frame, 1)
            prepreocessed_image = handle_image_mog2(frame, back_sub)
            # prepreocessed_image = handle_image_color(frame, color_min, color_max)
            res = detect_draw_circle(prepreocessed_image, frame, q)
            cv2.imshow('res', res)
            cv2.imshow('blur', prepreocessed_image)
            if cv2.waitKey(25) & 0xFF == ord('q'):
                break
            if cv2.waitKey(25) & 0xFF == ord('s'):
                # save frame
                cv2.imwrite('test.png', frame)
                break
        else:
            break
    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()