YOLOで特定のオブジェクトの写真のみ保存

2024-03-01 linux raspi yolo

前回まで色々遊んでみた結果を応用して、写真に「人(0:person)」が映っていたら、サーバに画像をアップロードし、Slackのチャネルに投稿する、というコードにしてみました。プライベートな箇所は伏せているので、書き換えが必要ですが、SlackのAPIのURLと、画像をアップロード(&公開)できるサーバがあれば実装可能かと思います。

連投を避けるため、最大で10秒に1回としています（それでも多いかw）。サーバはVPSでrsyncコマンドを使っていますが、AWS S3とかにアップロードしても良いかもしれません。ただ、どちらも（容易には見つけにくいとはいえ）インターネットに公開前提なので、秘匿性の高い場面では使えないかもしれません。

import cv2
import json, os, requests, time
SLACK_WEBHOOK = "https://hooks.slack.com/services/(YOUR_API_ADDRESS)"

if __name__ == '__main__':
    last_taken = 0.0
    cv2.namedWindow('find-on-video')
    cam = cv2.VideoCapture(0)

    from ultralytics import YOLO
    model = YOLO('yolov8n.pt')
    while True:
        ret, image = cam.read()
        results = model(image)
        cv2.imshow('find-on-video',results[0].plot())
        boxes = results[0].boxes
        if len(boxes) > 0:
            cls = list(map(lambda b: int(b.cls),boxes))
            if time.time() - last_taken > 10.0 and 0 in cls: # 0:person
                last_taken = time.time()
                fname = "people/"+time.strftime("%y%m%d%H%M%S")+".jpg"
                cv2.imwrite(fname, results[0].plot())
                os.system("rsync -avz people/ (YOUR_DOMAIN):(YOUR_PATH)/people/")
                url = 'https://(YOUR_DOMAIN)'+fname
                print(url)
                requests.post(SLACK_WEBHOOK, data = json.dumps({
                    'blocks': [{
                    'type': 'image',
                    'alt_text': fname,
                    'image_url': url
                }]}))
        k = cv2.waitKey(1)
        if k != -1:
            break

    cam.release()
    cv2.destroyAllWindows()