OpenCV:Example:ObjectAngleTracking

# -*- coding: utf-8 -*-

from math import pi, atan, atan2, floor, sqrt, degrees
import math
import time

import cv2
import numpy as np
from typing import Tuple

SHOW_CONTOUR = False
SHOW_HULL = False
SHOW_APPROX_POLYGON = False
SHOW_MIN_AREA_BOX = True

FILTER_MIN_AREA_SIZE = 4000
FILTER_MAX_AREA_SIZE = 10000
MAX_MISSING_COUNT = 20

LEFT_KEY = 63234
RIGHT_KEY = 63235

Point = Tuple[float, float]


def norm(p1: Point, p2: Point) -> float:
    x1, y1 = p1
    x2, y2 = p2
    dx = x2 - x1
    dy = y2 - y1
    return sqrt((dx ** 2) + (dy ** 2))


def angle(a: Point, b: Point, c: Point) -> float:
    ang = degrees(
        math.atan2(c[1] - b[1], c[0] - b[0]) - math.atan2(a[1] - b[1], a[0] - b[0])
    )
    return ang + 360 if ang < 0 else ang


def find_nearest_point(pivot: Point, ps: Tuple[Point, Point, Point, Point]) -> Point:
    p1, p2, p3, p4 = ps
    n1 = norm(pivot, p1)
    n2 = norm(pivot, p2)
    n3 = norm(pivot, p3)
    n4 = norm(pivot, p4)
    ns = [n1, n2, n3, n4]
    return ps[ns.index(min(ns))]


def find_nearest_point_from_rect(pivot: Point, rect: Tuple[Point, Point]) -> Point:
    p1, p2 = rect
    x1, y1 = p1
    x2, y2 = p2
    ps = (x1, y1), (x2, y1), (x2, y2), (x1, y2)
    return find_nearest_point(pivot, ps)


def get_center(rect: Tuple[Point, Point]) -> Point:
    p1, p2 = rect
    x1, y1 = p1
    x2, y2 = p2
    return (x2 + x1) / 2, (y2 + y1) / 2


def normalize_point(center: Point, point0: Point) -> Point:
    p1, p2 = center, point0
    x1, y1 = p1
    x2, y2 = p2
    return x2 - x1, y2 - y1


def calc_angle(
    prev: Tuple[Point, Point],
    current: Tuple[Point, Point],
) -> float:
    n0 = normalize_point(prev[0], prev[1])
    n1 = normalize_point(current[0], current[1])
    return angle(n0, (0.0, 0.0), n1)


def main():
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # mp4v
    file_name = 'dongkuk_angle'

    # 파일 불러오기
    # data_list = glob.glob('data/*/*/*.avi')
    # data_list = glob.glob('data/*/*/*.mp4')

    Vid = cv2.VideoCapture('1-short.avi')
    # Vid = cv2.VideoCapture('curve_movie.mp4')
    w = round(Vid.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = round(Vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = Vid.get(cv2.CAP_PROP_FPS)
    speed = 2  # play speed

    # 동영상 저장
    out = cv2.VideoWriter(f'video_result/{file_name}.mp4', fourcc, fps, (w, h))

    # color
    green = (0, 255, 0)
    red = (0, 0, 255)
    blue = (255, 0, 0)
    black = (0, 0, 0)
    white = (255, 255, 255)

    missing_count = 0
    result_angle = 0.0
    pivot_angle = 0.0

    pivot_rect = None
    pivot_point0 = None
    pivot_center = None
    current_point0 = None
    current_center = None

    # 동영상 재생 시작
    while Vid.isOpened():
        ret: bool
        frame: np.ndarray
        ret, frame = Vid.read()
        if not ret:
            break

        # 전처리(전 기능 공통)
        img_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        ret, threshold_image = cv2.threshold(img_gray, 252, 255, cv2.THRESH_BINARY)
        dst2 = np.zeros(frame.shape, np.uint8)

        contours: Tuple[np.ndarray, ...]
        hierarchy: np.ndarray
        contours, hierarchy = cv2.findContours(threshold_image, cv2.RETR_EXTERNAL,
                                               cv2.CHAIN_APPROX_SIMPLE)

        c_list = []
        for index in range(len(contours)):
            c_list.append(cv2.contourArea(contours[index]))

        max_index = c_list.index(max(c_list))
        largest_contour_size = c_list[max_index]

        largest_contour = contours[max_index]
        # 사각형 들어가기전 1개로 추려낸 테두리
        # cv2.drawContours(show_frame, [contours[max_index]], 0, red, 2)

        if SHOW_CONTOUR:
            contour_preview = frame.copy()
            cv2.drawContours(contour_preview, contours, max_index, (0, 0, 255), 2)
            cv2.imshow('Contour', contour_preview)

        # 화면 오버레이 구현부

        font_size = 1.2
        line_thick = 2

        txt_x = 88
        txt_y = 100
        txt_margin = 40

        line1_x = 290
        line2_x = 430
        line_curve = 8

        cv2.putText(frame, 'AG:', (txt_x, txt_y), cv2.FONT_HERSHEY_SIMPLEX, font_size,
                    (0, 0, 0), line_thick + 7, cv2.LINE_AA)
        cv2.putText(frame, 'AG:', (txt_x, txt_y), cv2.FONT_HERSHEY_SIMPLEX, font_size,
                    (255, 255, 255), line_thick, cv2.LINE_AA)

        # cv2.putText(frame,'BL:',(txt_x,txt_y),cv2.FONT_HERSHEY_SIMPLEX,font_size,(0,0,0),line_thick+7,cv2.LINE_AA)
        # cv2.putText(frame,'BL:',(txt_x,txt_y),cv2.FONT_HERSHEY_SIMPLEX,font_size,(255,255,255),line_thick,cv2.LINE_AA)

        # cv2.putText(frame,'GR:',(txt_x,txt_y+txt_margin),cv2.FONT_HERSHEY_SIMPLEX,font_size,(0,0,0),line_thick+7,cv2.LINE_AA)
        # cv2.putText(frame,'GR:',(txt_x,txt_y+txt_margin),cv2.FONT_HERSHEY_SIMPLEX,font_size,(255,255,255),line_thick,cv2.LINE_AA)

        # 각도 기능 구현부
        hull = cv2.convexHull(largest_contour)
        if SHOW_HULL:
            cv2.drawContours(frame, [hull], 0, (0, 255, 0), 3)

        hull_length = cv2.arcLength(hull, True)
        epsilon = hull_length * 0.02  # 7
        approx_polygon = cv2.approxPolyDP(hull, epsilon, closed=True, approxCurve=None)
        if SHOW_APPROX_POLYGON:
            cv2.polylines(frame [approx_polygon], True, blue, 2)

        min_area_rect = cv2.minAreaRect(approx_polygon)
        min_area_box = cv2.boxPoints(min_area_rect)
        if SHOW_MIN_AREA_BOX:
            cv2.drawContours(frame, [np.int0(min_area_box)], 0, (0, 255, 0), 3)

        # lines = get_4point(min_area_rect[:2])
        # angle = get_angle()
        #
        # min_area_rect_angle = min_area_rect[2]
        # min_area_rect_center = min_area_rect[0]

        approach = FILTER_MIN_AREA_SIZE <= largest_contour_size <= FILTER_MAX_AREA_SIZE

        if approach:
            missing_count = 0
            if pivot_rect is None:
                pivot_rect = min_area_rect
                current_point0 = float(min_area_box[0][0]), float(min_area_box[0][1])
                current_center = float(min_area_rect[0][0]), float(min_area_rect[0][1])
                pivot_point0 = current_point0
                pivot_center = current_center
                result_angle = 0
            else:
                current_point0 = find_nearest_point(
                    current_point0,
                    (
                        (min_area_box[0][0], min_area_box[0][1]),
                        (min_area_box[1][0], min_area_box[1][1]),
                        (min_area_box[2][0], min_area_box[2][1]),
                        (min_area_box[3][0], min_area_box[3][1]),
                    ),
                )
                current_center = float(min_area_rect[0][0]), float(min_area_rect[0][1])

                result_angle = calc_angle(
                    (pivot_center, pivot_point0),
                    (current_center, current_point0),
                )
                if pivot_angle == 0:
                    pivot_angle = result_angle
        else:
            missing_count += 1
            if missing_count >= MAX_MISSING_COUNT:
                missing_count = 0
                result_angle = 0
                pivot_rect = None
                pivot_angle = 0

        if current_point0:
            cv2.circle(frame, np.int0(current_point0), 4, red, 3)
            cv2.circle(frame, np.int0(current_center), 4, blue, 3)

        printable_angle = str(floor(result_angle))
        cv2.putText(frame,  printable_angle, (txt_x + 90, txt_y),
                    cv2.FONT_HERSHEY_SIMPLEX, font_size, (0, 0, 0),
                    line_thick + 7, cv2.LINE_AA)
        cv2.putText(frame,  printable_angle, (txt_x + 90, txt_y),
                    cv2.FONT_HERSHEY_SIMPLEX, font_size, (255, 255, 255),
                    line_thick, cv2.LINE_AA)

        # cv2.putText(frame,  f"Pivot: {floor(pivot_angle)}", (txt_x, txt_y + 30),
        #             cv2.FONT_HERSHEY_SIMPLEX, font_size, (0, 0, 0),
        #             line_thick + 7, cv2.LINE_AA)
        # cv2.putText(frame,  f"Pivot: {floor(pivot_angle)}", (txt_x, txt_y + 30),
        #             cv2.FONT_HERSHEY_SIMPLEX, font_size, (255, 255, 255),
        #             line_thick, cv2.LINE_AA)

        cv2.imshow("Result", frame)
        out.write(frame)
        key = cv2.waitKeyEx(5)
        if key == ord('q'):
            break
        elif key == RIGHT_KEY:
            speed *= 2
            if speed == 8:
                speed = 8
        elif key == LEFT_KEY:  # left
            speed /= 2
            if speed == 0.5: speed = 0.5

        time.sleep(1. / fps / speed)

    Vid.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()