location_detection.py

import numpy as np
import cv2

FILE_ORG = "calibrated_output_background2.m4v"
WINDOW_ORG = "Orginal_frame"
FRAME_RATE = 30  # fps


class LocationDetector:
    def __init__(self,
                 file_org="calibrated_output_background2.m4v",
                 window_org="Orginal_frame"):
        self.file_org = file_org
        self.window_org = window_org

    def guess_location(self, src):
        index = np.array([i for i in enumerate(src)])
        index_sum_sort10 = index[index[:, -1].argsort()][::-1][0:40].T
        location = int(np.average(index_sum_sort10[0]))
        return location

    def calculate_location(self, src):
        row_max_location = self.guess_location(np.sum(src, axis=0))
        col_max_location = self.guess_location(np.sum(src, axis=1))
        return(row_max_location, col_max_location)

    def detect_contour(self, src):
        # グレースケール画像へ変換
        gray = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY)

        # 2値化
        retval, bw = cv2.threshold(gray, 50, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)

        image, contours, hierarchy = cv2.findContours(bw, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
        contours = cv2.findContours(bw, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)

        # 矩形検出された数（デフォルトで0を指定）
        detect_count = 0

        # 各輪郭に対する処理
        for i in range(0, len(contours)):

            print("contours[i]", contours[i])
            # 輪郭の領域を計算
            area = cv2.contourArea(contours[i])

            print("area", area)

            # ノイズ（小さすぎる領域）と全体の輪郭（大きすぎる領域）を除外
            if area < 1e1 or 1e5 < area:
                continue

            # 外接矩形
            if len(contours[i]) > 0:
                rect = contours[i]
                x, y, w, h = cv2.boundingRect(rect)
                cv2.rectangle(src, (x, y), (x + w, y + h), (0, 255, 0), 2)
 
                # 外接矩形毎に画像を保存
                cv2.imwrite(str(detect_count) + '.jpg', src[y:y + h, x:x + w])
 
                detect_count = detect_count + 1

        # 外接矩形された画像を表示
        cv2.imshow('output', src)
        # cv2.waitKey(0)
        # cv2.imwrite("output.jpg", img=src)

    def movie_process(self):
        mov_org = cv2.VideoCapture(self.file_org)
        has_next, i_frame = mov_org.read()
        # cv2.imwrite("one_scene.jpg", i_frame)

        self.detect_contour(i_frame)

        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        out_BINARY = cv2.VideoWriter("output_location.m4v",int(fourcc), FRAME_RATE, \
                             (int(mov_org.get(cv2.CAP_PROP_FRAME_WIDTH)),int(mov_org.get(cv2.CAP_PROP_FRAME_HEIGHT))))

        while has_next == True:
            gray_frame = cv2.cvtColor(i_frame, cv2.COLOR_RGB2GRAY)
            location = self.calculate_location(gray_frame)
            self.detect_contour(i_frame)
            print(location)

            cv2.circle(i_frame, location, 10, (0,0,255), -1)
            cv2.imshow(self.window_org, i_frame)
            cv2.waitKey(0)
            # cv2.waitKey(50) & 0xFF
            out_BINARY.write(i_frame.astype(np.uint8))

        mov_org.release()

if __name__ == "__main__":
    location_detector = LocationDetector()
    location_detector.movie_process()