動画ファイルに対して物体検出(yolo3_mobilenet1.0_coco)

初めに

以前Webカメラの動画に対してのリアルタイム物体検出をやった。
touch-sp.hatenablog.com

今回は動画ファイルに対しての物体検出をやってみた。
動画ファイルに対してすべてのフレームに対して物体検出モデルを適用しその結果を動画ファイルに保存するのは簡単である。
今回は動画ファイルに対してリアルタイムに物体検出することが目的である。すべてのフレームに対してモデルを適応すると処理が追い付かなので少し工夫が必要であった。

環境

Panasonic let's note CF-LX3
CPU Intel(R) Core(TM) i7-4600U CPU @ 2.10GHz 2.69GHz
RAM 8.00 GB
Windows10 Pro
GPUなし
Python 3.8.2

バージョンの確認(pip freeze)

インストールが必要なのは「mxnet」と「gluoncv」と「opencv-python」のみ。

pip install mxnet
pip install gluoncv
pip install opencv-python

その他は勝手についてくる。

certifi==2020.4.5.1
chardet==3.0.4
cycler==0.10.0
gluoncv==0.7.0
graphviz==0.8.4
idna==2.6
kiwisolver==1.2.0
matplotlib==3.2.1
mxnet==1.6.0
numpy==1.18.3
opencv-python==4.2.0.34
Pillow==7.1.2
portalocker==1.7.0
pyparsing==2.4.7
python-dateutil==2.8.1
pywin32==227
requests==2.18.4
scipy==1.4.1
six==1.14.0
tqdm==4.45.0
urllib3==1.22

実行ファイル

import mxnet as mx
import gluoncv

import time
import cv2

# Load the model
net = gluoncv.model_zoo.get_model('yolo3_mobilenet1.0_coco', pretrained=True)
# Compile the model for faster speed
net.hybridize()

# Load the webcam handler
cap = cv2.VideoCapture('traffic.mp4')
fps = cap.get(cv2.CAP_PROP_FPS)

start = time.time()

while(True):
    # Load frame from the camera
    end = time.time()
    sec = end-start
    cap.set(cv2.CAP_PROP_POS_FRAMES, round(fps * sec))
    ret, frame = cap.read()

    if(ret):
        # Image pre-processing
        frame = mx.nd.array(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).astype('uint8')
        rgb_nd, frame = gluoncv.data.transforms.presets.yolo.transform_test(frame,short=360)

        # Run frame through network
        class_IDs, scores, bounding_boxes = net(rgb_nd)

        # Display the result
        img = gluoncv.utils.viz.cv_plot_bbox(frame, bounding_boxes[0], scores[0], class_IDs[0], class_names=net.classes)
        gluoncv.utils.viz.cv_plot_image(img)
    else:
        break

    # escを押したら終了
    if cv2.waitKey(1) == 27:
        break

cap.release()
cv2.destroyAllWindows()