obj_detect.py

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

from __future__ import print_function, division, unicode_literals

import cv2
import urllib
import urllib2
from apiclient import discovery
from apiclient import http
from oauth2client.client import GoogleCredentials
import json
import io
from datetime import datetime


class Notifier:
    SLACK_TEST_TOKEN = '***********************************************'
    SLACK_IOTSAMPLE_CHANNEL_ID = '*************************************'
    SLACK_POST_MESSAGE_URL = 'https://slack.com/api/chat.postMessage'
    GCS_BUCKET_NAME = '******************'

    def __init__(self):
        credentials = GoogleCredentials.get_application_default()
        self.gcs = discovery.build('storage', 'v1', credentials=credentials)

    def post_to_slack(self, text):
        params = {'token': self.SLACK_TEST_TOKEN,
                  'channel': self.SLACK_IOTSAMPLE_CHANNEL_ID,
                  'text': text}
        params = urllib.urlencode(params)
        req = urllib2.Request(self.SLACK_POST_MESSAGE_URL)
        req.add_header('Content-Type', 'application/x-www-form-urlencoded')
        req.add_data(params)
        res = urllib2.urlopen(req)
        body = res.read()
        print(body)

    def notify(self, bytesio):
        now = datetime.now()
        file_name = now.strftime('%Y-%m-%d-%H%M%S') + '.jpg'
        media = http.MediaIoBaseUpload(bytesio, mimetype='image/jpeg')
        req = self.gcs.objects().insert(bucket=self.GCS_BUCKET_NAME,
                                        name=file_name,
                                        media_body=media,
                                        predefinedAcl='publicRead')
        resp = req.execute()
        print(json.dumps(resp, indent=2))
        message = 'detected!\nhttp://storage.googleapis.com/{bucket}/{file}' \
            .format(bucket=self.GCS_BUCKET_NAME, file=file_name)
        self.post_to_slack(message)


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 -- 閾値以上のピクセル比率のマスク画像がこの値より大きい場合に、オブジェクトがいると判定する。
    '''

    observer = Notifier()
    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
            _, stream = cv2.imencode('.jpg', frame)
            observer.notify(io.BytesIO(stream.tostring()))

        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()