物体検出

obj_detect.py

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

from __future__ import print_function, division, unicode_literals

import cv2


def mask_ratio(mask):
    w, h = mask.shape
    p = cv2.countNonZero(mask) / (w * h)
    return int(p * 100)


def make_substractor(cap, history):
    substractor = cv2.createBackgroundSubtractorMOG2(history=history)

    # learn background
    init_counter = 0
    while(init_counter <= history):
        hasSucceeded, frame = cap.read()
        if not hasSucceeded:
            continue
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        substractor.apply(gray)
        init_counter = init_counter + 1

    return substractor


def main(ratio_threshold=5, counter_threshold=40, history=500):
    '''
    物体検出ループを実行する。
    このループは次のように動作する。

    1. MOG2法により背景とオブジェクトを区別する為のマスク画像を作る。
    2. マスク画像におけるオブジェクトの写っている部分の占めるピクセル数と全ピクセルの比率を計算する。
    3. counter_threshold回以上、ピクセル比率がratio_thresholdより大きい場合に、オブジェクトがあると判定する。

    Keyword arguments:
    ratio_threshold -- int, オブジェクトがあると判定する閾値
    counter_threshold -- 閾値以上のピクセル比率のマスク画像がこの値より大きい場合に、オブジェクトがいると判定する。
    '''

    cap = cv2.VideoCapture(0)
    substractor = make_substractor(cap, history)

    # repeat detection
    has_detected = False
    ratio_counter = 0

    while(True):
        hasSucceeded, frame = cap.read()
        if not hasSucceeded:
            continue
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        mask = substractor.apply(gray)
        ratio = mask_ratio(mask)

        # cv2.imshow('frame', frame)
        # cv2.imwrite('data/frame' + str(loop_counter) + '.jpg', frame)
        cv2.imshow('mask', mask)

        if ratio > ratio_threshold:
            ratio_counter = ratio_counter + 1
        else:
            ratio_counter = 0
            has_detected = False

        if ratio_counter > counter_threshold and not has_detected:
            has_detected = True

        results = [ratio, 50 if has_detected else 0]
        print(','.join([str(x) for x in results]))

        if cv2.waitKey(1) & 0xFF == ord(' '):
            break

if __name__ == '__main__':
    main()