Save video only object detected using YOLOv5

yolov5 YOLOv5

YOLOv5を使って特定の物体が検出された場合だけビデオを保存する方法を解説します。またビデオをAWSクラウドに保存する方法も解説します。

目的

監視カメラの映像を後から見る場合、監視対象が写っていない映像を自動で削除できれば、監視対象を検索する手間を軽減できたり、保存容量を小さくすることが可能です。この記事ではYOLOv5を使用し監視対象が写っていない映像を自動で削除し、監視対象が写っている映像だけをクラウドに保存するシステムについて解説します。今回、クラウドのストレージとしてAWS S3を使用します。

方針

以下の方針で実装をしようと思います。

  • YOLOv5で物体検出を行いながら映像を1分ごとにローカルに保存。
  • ローカルに保存された映像に人(=Person)が含まれていたらクラウドに保存
  • ローカルに保存された映像に人(=Person)が含まれていなければ削除

実装

今回はこちらの記事で使用したpython版のOpenCVを使ったYOLOv5を使用して実装することとします。

なおソースコードはこちらのレポジトリにあります。

以下ではmainの関数を解説します。

def main(is_cuda, video, class_list, object_ids, s3_bucket):
    executor = ThreadPoolExecutor()

    colors = [(255, 255, 0), (0, 255, 0), (0, 255, 255), (255, 0, 0)]

    net = build_model(is_cuda)
    capture = cv2.VideoCapture(video)

    start = time.time_ns()
    frame_count = 0
    total_frames = 0
    fps = -1

    is_save = False
    writer = None

    timestamp = datetime.now().strftime("%Y%m%d%H%M")
    video_name = f"{timestamp}.mp4"

    while True:

        _, frame = capture.read()
        if frame is None:
            print("End of stream")
            break

        inputImage = format_yolov5(frame)
        outs = detect(inputImage, net)

        class_ids, confidences, boxes = wrap_detection(inputImage, outs[0])

        if len([c for c in class_ids if c in object_ids]) != 0:
            is_save = True

        frame_count += 1
        total_frames += 1

        for (classid, confidence, box) in zip(class_ids, confidences, boxes):
            color = colors[int(classid) % len(colors)]
            cv2.rectangle(frame, box, color, 2)
            cv2.rectangle(
                frame, (box[0], box[1] - 20), (box[0] + box[2], box[1]), color, -1
            )
            cv2.putText(
                frame,
                class_list[classid],
                (box[0], box[1] - 10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (0, 0, 0),
            )

        if frame_count >= 30:
            end = time.time_ns()
            fps = 1000000000 * frame_count / (end - start)
            frame_count = 0
            start = time.time_ns()

        if fps > 0:
            fps_label = "FPS: %.2f" % fps
            cv2.putText(
                frame, fps_label, (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2
            )

        cv2.imshow("output", frame)

        if writer is None and fps > 0:
            w, h = frame.shape[1], frame.shape[0]
            writer = cv2.VideoWriter(
                video_name, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)
            )

        if writer:
            writer.write(frame)

            # update timestamp
            new_timestamp = datetime.now().strftime("%Y%m%d%H%M")
            if timestamp != new_timestamp:
                writer.release()
                writer = None
                if is_save:
                    print(f"save {video_name}")
                    executor.submit(upload, s3_bucket, video_name)
                else:
                    print(f"remove {video_name}")
                    os.remove(video_name)
                is_save = False
                timestamp = new_timestamp

        if cv2.waitKey(1) & 0xFF == ord("q"):
            print("finished by user")
            running = False
            if writer:
                writer.release()
            break
Code language: Python (python)

build_modelはYOLOv5モデルのオブジェクトを返します。is_cudaがTrueかつCUDAが使用可能な場合、CUDAを使用するオブジェクトを、それ以外の場合はCPUを使用するオブジェクトを返します。

def build_model(is_cuda):
    net = cv2.dnn.readNet("config_files/yolov5s.onnx")
    if is_cuda:
        print("Attempty to use CUDA")
        net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
        net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA_FP16)
    else:
        print("Running on CPU")
        net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
        net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
    return netCode language: Python (python)

cv2.VideoCaptureでweb cameraのデバイス番号をvideoで指定しデバイスをオープンします。web cameraのデバイスが/dev/video0の場合0を指定します。

    capture = cv2.VideoCapture(video)Code language: Python (python)

capture.read()でビデオフレームを取得します。

        _, frame = capture.read()Code language: Python (python)

format_yolov5でframeを変換し、detectで物体を検出し、wrap_detectionで検出された物体のID(class_id)、検出物体の信頼度(confidences)、バウンディングボックス(boxes)に変換します。

        inputImage = format_yolov5(frame)
        outs = detect(inputImage, net)

        class_ids, confidences, boxes = wrap_detection(inputImage, outs[0])Code language: Python (python)

mainの引数で渡されたobject_idsに検出された物体のID(class_id)があれば作成中のビデオは保存するためis_saveにTrueを設定しています。なおobject_idsに人(Person)のID(=0)を指定すれば、そのIDが検出された場合is_saveがTrueに設定され、ビデオが保存されます。IDの番号はconfig_files/classes.txtを参照してください。

        if len([c for c in class_ids if c in object_ids]) != 0:
            is_save = TrueCode language: Python (python)

cv2.VideoWriterでビデオ生成オブジェクトを作成します。ビデオのフォーマットはcv2.VideoWriter_fourccで指定します。ここでは*"mp4v"を指定しています。指定できるビデオフォーマットはOSやサポートするコーデック、OpenCVのffmpegサポートに依存するようです。こちらの記事が参考になるかもしれません。fps(Frame Per Second)は直近のfpsを指定しています。

            writer = cv2.VideoWriter(
                video_name, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)
            )Code language: Python (python)

timestampはYYYYmmddHHMMのフォーマットで1分ごとに更新されます。

            # update timestamp
            new_timestamp = datetime.now().strftime("%Y%m%d%H%M")
            if timestamp != new_timestamp:
                writer.release()
                writer = NoneCode language: Python (python)

is_saveがTrueの場合、uploadスレッドをexecutor.submitで起動し、ローカルに保存されたビデオ(video_name)をクラウドにアップロードします。Falseの場合は単にローカルに保存されたビデオを削除します。

                if is_save:
                    print(f"save {video_name}")
                    executor.submit(upload, s3_bucket, video_name)
                else:
                    print(f"remove {video_name}")
                    os.remove(video_name)Code language: Python (python)

uploadではAWS APIのpythonライブラリboto3を使ってBucket.upload_fileを呼び出しファイルをS3にアップロードします。アップロードが終わったファイルは不要ですのでos.removeで削除します。

def upload(s3_bucket, file):
    if s3_bucket:
        s3 = boto3.resource("s3")
        bucket = s3.Bucket(s3_bucket)
        # upload
        bucket.upload_file(file, file)
        print(f"upload {file} done.")
    # remove
    os.remove(file)Code language: Python (python)

AWS

以下の操作はAWSアカウントが必要です。

また管理者権限でAWS CLIを実行できる必要があります。

AWS S3にファイルをアップロードする為、S3のバケットとそのS3バケットにファイルをアップロードすることができる専用のIAM Userを用意します。以下はAWS CloudFormationのテンプレートです。

AWSTemplateFormatVersion: 2010-09-09
Description: Create IAM Users

Parameters:
  UserName:
    Type: String
    Description: username to upload files to s3
  BucketName:
    Type: String
    Description: bucket name to create

Resources:
  S3Bucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Ref BucketName

  S3UploadUser:
    Type: AWS::IAM::User
    Properties:
      UserName: !Ref UserName
      Policies:
      - PolicyName: S3UploadOnly
        PolicyDocument:
          Version: '2012-10-17'
          Statement:
          - Effect: Allow
            Action:
            - s3:PutObject
            Resource:
            - !Sub arn:aws:s3:::${BucketName}/*
    DependsOn: S3BucketCode language: YAML (yaml)

boto3のBucket.upload_fileに必要な権限は、s3:PutObjectですのでその権限だけを与えます。S3のバケット名はグローバルですのですでに存在するバケット名を指定した場合バケットの作成に失敗します。バケットの作成に成功した場合のみS3UploadUserを作成するようにDependsOn: S3Bucketを指定しています。

以下を実行し、S3バケットyolov5-opencv-save-detectedと、IAMユーザーmy-s3-uploaderを作成します。

 aws cloudformation deploy --template-file template.yaml --stack-name my-stack-name --parameter-overrides UserName=my-s3-uploader BucketName=yolov5-opencv-save-detected --capabilities CAPABILITY_NAMED_IAMCode language: Bash (bash)

上記でIAMユーザーmy-s3-uploaderを作成したら以下のコマンドでアクセスキーを取得します。

$ aws iam create-access-key --user-name my-s3-uploader
{
    "AccessKey": {
        "UserName": "my-s3-uploader",
        "AccessKeyId": "xxxxxxxxxxxxxxxxxxxx",
        "Status": "Active",
        "SecretAccessKey": "yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy",
        "CreateDate": "2022-10-13T12:41:00+00:00"
    }
}Code language: Bash (bash)

得られたAccessKeyIdとSecretAccessKeyをぞれぞれAWS_ACCESS_KEY_ID、AWS_SECRET_ACCESS_KEYに指定し作成したプログラムを以下のように実行します。--s3-bucketには作成したバケット名を指定します。

実行

AWS_ACCESS_KEY_ID=xxxxxxxxxxxxxxxxxxxx \
AWS_SECRET_ACCESS_KEY=yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy \
python3 python/yolo.py --s3-bucket yolov5-opencv-save-detectedCode language: Bash (bash)

人(=Person)を検出した映像はS3バケットyolov5-opencv-save-detectedにアップロードされていることを確認します。

$ aws s3 ls s3://yolov5-opencv-save-detected/
2022-10-13 22:10:01     908341 202210132209.mp4
2022-10-13 22:14:01     666566 202210132213.mp4Code language: Bash (bash)

すべてのコードは https://github.com/otamajakusi/yolov5-opencv-save-detected から参照できます。

以上です。

参照