【Realtime Object Detection】動画に対して物体検出（ssd_512_mobilenet1.0）

MXNet GluonCV 物体検出 opencv-pythonで動画を扱う

2020年12月24日記事を更新しました。

初めに

GluonCVの学習済みモデルを使ってWebカメラの動画に対して物体検出をやってみました。

環境

OSはWindowsです。
GPUあり/なしの二つの環境で動作確認しました。

インストールが必要なのは「mxnet」と「gluoncv」のみです。
「opencv-python」を使用しますが「gluoncv」をインストールする際に一緒にインストールされます。

「mxnet」のインストールはこちらを参照して下さい。
「gluoncv」のインストールはこちらを参照して下さい。

Windows 10
GTX 1080
CUDA Toolkit 10.1
Python 3.7.9

gluoncv==0.9.0
mxnet-cu101==1.7.0
opencv-python==4.4.0.46

Windows10
GPUなし
Python 3.8.6

gluoncv==0.9.0
mxnet==1.7.0.post1
opencv-python==4.4.0.46

Pythonスクリプト（GPUあり/なしどちらでも実行可能）

import mxnet as mx
import gluoncv

import time
import cv2

ctx = mx.gpu() if mx.context.num_gpus() >0 else mx.cpu()

# Load the model
net = gluoncv.model_zoo.get_model('ssd_512_mobilenet1.0_voc', pretrained=True, ctx=ctx, root='./models')
# Compile the model for faster speed
net.hybridize()

# Load the webcam handler
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
# letting the camera autofocus
time.sleep(1)

while(True):
    # Load frame from the camera
    ret, frame = cap.read()

    # Image pre-processing
    frame = mx.nd.array(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).astype('uint8')
    rgb_nd, frame = gluoncv.data.transforms.presets.ssd.transform_test(frame, short=512, max_size=700)

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

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

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

cap.release()
cv2.destroyAllWindows()

補足

カメラがうまく検出されない場合には実行ファイルの中の次の部分を少し変えてみる。

# Load the webcam handler
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)

数字を「0」から「1」や「2」に変えてみる。