LoginSignup
1
4

More than 1 year has passed since last update.

@camultaneous

【Python】pytorchによる複数の物体検知とその追跡

Last updated at Posted at 2022-05-15

物体検知と追跡

物体検知とは、画像や動画に移ったものが何なのかをオブジェクト事に分類することである。物体検知の一つとして、機械学習の教師あり学習の分類(classification)があり、ラベリングされた多数の画像を学習(アノテーション)して、与えられた画像に写っているオブジェクトが、どの学習済み画像と近いのかをそれぞれ計算し分類する。今回は機械学習ライブラリPytorchと物体追跡ライブラリmotpyを使って、以下のような物体検知＆追跡のやり方を紹介する。

以下のコードで生成した物体検知＆追跡の動画
Animation.gif

YOLO

代表的な物体検知の分類アルゴリズムの一つとして、YOLO(You Look Only Once)がある。YOLOは検出窓をスライドさせるようなことはせず、一度のCNN(畳み込みニューラルネットワーク)で検出するリアルタイム物体検知アルゴリズムである。

YOLOのアルゴリズムの流れは大まかに以下のようである。

  1. 画像をグリッドに分ける
  2. それぞれのセルに決められた数、大きさのバウンディングボックスを置く
  3. 各バウンディングボックスで各クラス(ラベリングされた分類)にオブジェクトが存在する確率を計算する
  4. 各グリッドでどのクラスが最も高い確率を示したかを決め、各グリッドに当てはめる
  5. NMS(Non-Maxmum-Suppression)でIntersection over Union（IoU, どれくらい重なっているかの指標)の最も確率の高いグリッドを選び、それを物体の位置とする

YOLOと並ぶ代表的な物体検知アルゴリズムとしてSSD(Single Shot MultiBox Detector)がある。YOLOとの大きな違いは、YOLOがBounding Boxの出力を出力層だけで行っていたのに対し、SSDではCNNの複数の層から物体のBounding Boxを出力する。

SSDは、YOLOと同様にバウンディングボックスを推定していくのですが、YOLOとの違いは、YOLOが元画像を等間隔に区切ったgrid cellにより、バウンディングボックスを推定するのに対し、SSDでは、特徴マップに対してデフォルトボックスと呼ばれる矩形パターンを配置することでバウンディングボックスを推定する点になります。

古い物体検知として2001年に発表されたカスケード型識別器（Cascade detector）があり、複数の強識別器を連結した識別器である。カスケード型識別器では、各強識別器により順番に判別を行う。最初の識別器1で「正解」と判別されると、その次の識別器2に進む。これを繰り返し、強識別器Nまで一貫して「正解」になった場合のみ、結果を正解として出力する。

準備

今回、pytorchで学習済のYOLOの画像分類モデルを使用する。そのため、pytorchのコードレポジトリ(GitHub)から学習済モデルデータをダウンロードし、適当なディレクトリに置く。

また、必要なpython ライブラリを以下のようにインストールする。

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
pip install -r requirements.txt

GPUを利用する場合、ほとんどがNVIDIAのGPUなので、CUDA Tool KitとcuDNNをインストールする必要がある。
下記リンクからそれぞれインストール・ダウンロードする。

cuDNNは下記のフォルダにコピペする。

C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.0

Pythonコード

よく使う標準ライブラリ群と画像処理ライブラリであるcv2(opencv), 機械学習ライブラリpytorchを使う。

import os
import sys
import random
import time
import datetime
import glob
import re
import shutil
import numpy as np
import matplotlib.pyplot as plt
#動画編集ライブラリ
import cv2 as cv
#機械学習ライブラリ
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data as data
from torch.utils.data import DataLoader
import torchvision
import torchvision.datasets
import torchvision.transforms

機械学習の画像分類をCPUとGPUのどちらで行うのか選ぶ。torch.cuda.is_available()がYesなら実行環境ではGPUが使用可能なので、os.environ['CUDA_VISIBLE_DEVICES'] = '0'でGPUを使う(0)を環境変数入れ込む。CPUの場合は1を入れ込む。

# GPU情報確認
print(torch.cuda.is_available())
print(torch.cuda.current_device())
print(torch.cuda.device_count())
print(torch.cuda.get_device_name())
print(torch.cuda.get_device_capability())
print(torch.cuda.get_device_name('cuda:0'))

# 環境変数でどのGPUを使うか、またGPUを使わないかを選べる
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

torch.hub.list()でダウンロードしたpytorchの学習済みデータソース('ultralytics/yolov5')を読み込み、データソース内で使用可能な学習済モデルのバージョンを表示する。表示された学習済モデルのバージョンの中から適切なものを、torch.hub.load()で読み込む。

# 指定した学習データ内のversionをリスト表示する
entrypoints = torch.hub.list('ultralytics/yolov5', force_reload=True)
print(entrypoints)
# torch hubから学習済みモデルをダウンロード＆読み込む
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
print(torch.hub.get_dir())

Opencv(cv2)で接続しているWebカメラ(もしくは動画パス)の設定情報(fps,立幅,横幅,明るさ,コントラスト...)を取得する。また、保存する動画の形式やfps、画面サイズを設定する。最後の行のfps*0.4の0.4でfpsを調整して、保存した際に滑らかな動画になるようにする。

# OpenCVでカメラ情報取得
cap = cv.VideoCapture(0) # Webカメラ or 動画パス
fps = int(cap.get(cv.CAP_PROP_FPS))
w = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
brightness = int(cap.get(cv.CAP_PROP_BRIGHTNESS))
contrast = int(cap.get(cv.CAP_PROP_CONTRAST))
auto_exposure = int(cap.get(cv.CAP_PROP_AUTO_EXPOSURE))
auto_focus = int(cap.get(cv.CAP_PROP_AUTOFOCUS))
auto_WB = int(cap.get(cv.CAP_PROP_AUTO_WB))
fourcc = cv.VideoWriter_fourcc(*'mp4v')
out = cv.VideoWriter('output.mp4', fourcc, fps*0.4, (w,  h))

Webカメラ(もしくは動画)を1フレームずつwhile文で読み込み、1フレームずつ処理する。result = model(frame)で読み込んだフレームを学習済モデルに渡して、result.render()で分類結果をフレームに反映している。以上が、物体検知のコードである。

# Webカメラからフレームごとに読み込む
while True:
    ret, frame = cap.read()
    # yoloクラス分類モデル適応
    result = model(frame)
    # 分類結果をDataFrame形式で取得できる
    df = result.pandas().xywh[0] # .xyxy[0]
    result.render() # 画面にボンディングボックスを表示する
    cv.imshow("image show",result.imgs[0])
    out.write(result.imgs[0])
    if cv.waitKey(1) == ord('q'):
        break
cap.release()
out.release()
cv.destroyAllWindows()

サンプルコード全文

上記の物体検知のサンプルコード全文を以下に記載する。

Main.py
import os
import sys
import random
import time
import datetime
import glob
import re
import shutil
import numpy as np
import matplotlib.pyplot as plt
#動画編集ライブラリ
import cv2 as cv
#機械学習ライブラリ
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data as data
from torch.utils.data import DataLoader
import torchvision
import torchvision.datasets
import torchvision.transforms

# GPU情報確認
print(torch.cuda.is_available())
print(torch.cuda.current_device())
print(torch.cuda.device_count())
print(torch.cuda.get_device_name())
print(torch.cuda.get_device_capability())
print(torch.cuda.get_device_name('cuda:0'))

# 環境変数でどのGPUを使うか、またGPUを使わないかを選べる
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

# 指定した学習データ内のversionをリスト表示する
entrypoints = torch.hub.list('ultralytics/yolov5', force_reload=True)
print(entrypoints)
# torch hubから学習済みモデルをダウンロード＆読み込む
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
print(torch.hub.get_dir())

# OpenCVでカメラ情報取得
cap = cv.VideoCapture(0) # Webカメラ or 動画パス
fps = int(cap.get(cv.CAP_PROP_FPS))
w = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
brightness = int(cap.get(cv.CAP_PROP_BRIGHTNESS))
contrast = int(cap.get(cv.CAP_PROP_CONTRAST))
auto_exposure = int(cap.get(cv.CAP_PROP_AUTO_EXPOSURE))
auto_focus = int(cap.get(cv.CAP_PROP_AUTOFOCUS))
auto_WB = int(cap.get(cv.CAP_PROP_AUTO_WB))
fourcc = cv.VideoWriter_fourcc(*'mp4v')
out = cv.VideoWriter('output.mp4', fourcc, fps*0.4, (640,  480))

# Webカメラからフレームごとに読み込む
while True:
    ret, frame = cap.read()
    # yoloクラス分類モデル適応
    result = model(frame)
    # 分類結果をDataFrame形式で取得できる
    df = result.pandas().xywh[0] # .xyxy[0]
    result.render() # 画面にボンディングボックスを表示する
    cv.imshow("image show",result.imgs[0])
    out.write(result.imgs[0])
    if cv.waitKey(1) == ord('q'):
        break
cap.release()
out.release()
cv.destroyAllWindows()

MOT(Multi-Objects-Tracking)

MOTは、動画のあるフレームで複数の物体を検知したら、次のフレームではどこに移動したかを各物体を識別して追跡する手法の総称である。物体追跡の手法には、はカルマン・フィルタやパーティクル・フィルタ、SORTやオプティカル・フローなどがある。以下は、pythonライブラリのmotpyを使って、pytorchで検知した物体を前後のフレームで個別に識別IDを割り当て、追跡している。

motpyライブラリをインストールする。

pip install motpy

import os
import sys
import random
from random import randint
import time
import datetime
import glob
import re
import shutil
import numpy as np
import cv2 as cv
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data as data
from torch.utils.data import DataLoader
import torchvision
import torchvision.datasets
import torchvision.transforms
from motpy import Detection, MultiObjectTracker
from motpy.testing_viz import draw_track

class MOT:
    def __init__(self):
        self.tracker = MultiObjectTracker(dt=0.1)

    def track(self, outputs, ratio):
        try:
            outputs = np.array(outputs)
            outputs = [Detection(box=box[:4] / ratio, score=float(box[4] * box[5]), class_id=box[6]) for box in outputs]
        except:
            outputs = []

        self.tracker.step(detections=outputs)
        tracks = self.tracker.active_tracks()
        return tracks

# 環境変数でどのGPUを使うか、またGPUを使わないかを選べる
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

# 指定した学習データ内のversionをリスト表示する
entrypoints = torch.hub.list('ultralytics/yolov5', force_reload=True)
# torch hubから学習済みモデルをダウンロード＆読み込む
model = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True)

cap = cv.VideoCapture(0)
fps = int(cap.get(cv.CAP_PROP_FPS))
w = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
brightness = int(cap.get(cv.CAP_PROP_BRIGHTNESS))
contrast = int(cap.get(cv.CAP_PROP_CONTRAST))
auto_exposure = int(cap.get(cv.CAP_PROP_AUTO_EXPOSURE))
auto_focus = int(cap.get(cv.CAP_PROP_AUTOFOCUS))
auto_WB = int(cap.get(cv.CAP_PROP_AUTO_WB))
fourcc = cv.VideoWriter_fourcc(*'mp4v')
writer = cv.VideoWriter('output.mp4', fourcc, fps*0.4, (w,  h))

mot = MOT()

while True:
    ret, frame = cap.read()
    size = frame.shape[:2]
    if ret == False:
        break
    result = model(frame)
    df = result.pandas().xyxy[0] # .xyxy[0]
    df_person = df[df['name'].isin(['person'])]
    outputs = []
    for index, row in df_person.iterrows():
        name = row[6]
        xmin = int(float(row[0]))
        ymin = int(float(row[1]))
        xmax = int(float(row[2]))
        ymax = int(float(row[3]))    
        confidence = row[4]
        _box = [xmin, ymin, xmax, ymax, confidence, 1, randint(0, 255)]
        outputs.append(_box)
    tracks = mot.track(outputs, 1)
    for trc in tracks:
        draw_track(frame, trc, thickness=1)
    cv.imshow('Video', frame)
    writer.write(frame)
    if cv.waitKey(1) & 0xFF == ord('q'):
        break
cap.release()
writer.release()
cv.destroyAllWindows()

まとめ

今回はpytorchを用いた物体検知とmotpyを用いた追跡を紹介した。今回は前後のフレームでの追跡のみで、対象物体が一度隠れたりすると、追跡するために割り当てた識別IDが変わってしまうため、厳密には追跡できていない。そのため前後のフレームのみだけでなく、複数のフレーム間で追跡できるように改良できると、より実用的になると思う。

付録1: pytorch チートシート

pytorch-cheat.jpg

1
4
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
1
4