【DAMO-YOLO】PyTorch＋Webカメラでリアルタイム推論してみた
宣伝
こんなコミュニティもやっているので良ければご参加ください！
connpassはこちらから！
はじめに
こんにちはToshikiHaraguchです。今回は勉強会【DAMO-YOLO + ONNX】物体検出AIをノートPCで動かそう！に合わせて、いろいろしたことをつらつらと書き連ねていこうと思います。
まずはDAMO-YOLOの元のリポジトリからクローンしたものを実行＆リアルタイム推論できるように改造していきます！
実験環境
今回の実験環境は次の通りです。
- Windows11（12th Gen Intel(R) Core(TM) i5-12400F 2.50 GHz）
- conda
いつもはクラウド（Google Colaboratory）で行うのですが、リアルタイム推論をするとなるとやはりローカルのほうが・・・。ということで、自前のPCで行っています。
CPUで推論するのでCPUの型番も併せて載せています。CPU推論ではCPUの馬力がものをいうので、速度が出ない場合は自身のCPUと筆者のCPUを比較してもらえたらと思います。
Anaconda環境を作る
Windowsにおいてconda環境を作ることはそこまで難しくないかな・・・？と思っています。
すこし手順を追って説明しますね。
Step1：インストーラーをダウンロードする
AnacondaをWindowsにインストールには、インストーラーを利用するのが一番手っ取り早いです。
ここから以下の画像を参考にダウンロードします。
めちゃ簡単ですね・・・。
続いてインストールに移りましょう。
Step2：インストールする
続いてインストールをしていきましょう。ダウンロードしたインストーラーをダブルクリックするとしたのような画面が出るので「Next」を選択します。
次の画面ではライセンス系のお話をされます。とりあえず「I Agree」しておきましょう。
アグリーできるとインストールするユーザーを聞かれます。「Just Me」を選びましょう。
ユーザを決定すると、インストール先を聞かれます。特に問題なさそうでしたらそのまま「Next」決めちゃいましょう！
最後はパスについてです。Anacondaは「add Anaconda3 to my PATh environment variable」することを非推奨としていますが、今回はガンガン入れていきます。これを入れるとAnaconda専用のコマンドプロンプト以外でもAnacondaを使えるようになるので、入れておいて損はないです。準備できれば「Install」でPythonの世界へひとっとびしましょう。
Step3：PowerShellの設定をしよう
最後にPowerShellでAnacondaが使えるように設定しましょう。
PowerShellを起動して次のコマンドを実行します。
conda init
これでOKです。今開いているPowerShellを落として、新しいPowerShellを起動しましょう。
(base) PS C:\Users\T
こんな感じで表示されていれば準備完了です。続いてGithubからクローンするためにGitをインストールしましょう。
Gitの準備
続いてGitのインストールをしましょう。ここから「Donwload」を選択してGitインストーラーをダウンロードしましょう。
インストーラーを実行すると「install」ボタンが出てくるので、それをクリック。準備は完了です。
DAMO-YOLOのインストール
ではお待ちかねのDAMO-YOLOをインストールしていきましょう。以下のコマンドからDAMO-YOLOのデータをクローンします。
git clone https://github.com/tinyvision/DAMO-YOLO
DAMO-YOLOフォルダができていれば準備完了です。
では公式のインストールに従って進めていきましょう。
conda create -n DAMO-YOLO python=3.7 -y
conda activate DAMO-YOLO
conda install pytorch==1.7.0 torchvision==0.8.0 torchaudio==0.7.0 -c pytorch
pip install -r requirements.txt
DEMOの実行
リアルタイム推論を行う前に、公式が準備してくれているデモを実行してみましょう。今回はTモデルを利用しました。学習済みモデルはここからダウンロードできます。その他のモデルは公式ページの「Model Zoo」を確認してください。
ダウンロードした学習済みモデルはDAMO-YOLOフォルダ直下に配置しています。
では実行です。
python tools/torch_inference.py -f configs/damoyolo_tinynasL20_T.py --ckpt ./damoyolo_tinynasL20_T_418.pth --path assets/dog.jpg
ModuleNotFoundError: No module named 'damo.base_models.core'
・・・？ということで、そのままするとエラーで止まります。悲しいね・・・。
解決方法として、tools/torch_inference.pyをDAMO-YOLO直下にコピーして、
python torch_inference.py -f configs/damoyolo_tinynasL20_T.py --ckpt ./damoyolo_tinynasL20_T_418.pth --path assets/dog.jpg
にすると実行できます。
結果は・・・？
| INFO | __main__:main:120 - saved torch inference result into ./demo\dog.jpg
いつものYOLOベンチマーク画像が保存されています。いい感じに囲めてますね。
左奥のバイクっぽいのが取れてないのは内緒だよ
とりあえず公式の動作は確認できました。ここからWebカメラを使った推論に書き換えていきましょう。
Webカメラ画像に対する推論
コードを作るために少しtorch_inference.pyを眺めてみましょう・・・。
origin_img = np.asarray(Image.open(args.path).convert('RGB'))
PILのImageを利用して読み込んでいます。色空間はRGBでOpenCVと若干違うため注意が必要そうです。
モデルの読み込みは
config = parse_config(args.config_file)
model = build_local_model(config, device)
ckpt = torch.load(ckpt_file, map_location=device)
new_state_dict = {}
for k, v in ckpt['model'].items():
new_state_dict[k] = v
model.load_state_dict(new_state_dict, strict=False)
を組み合わせるとよさそうです。
入力する画像は
img = transform_img(origin_img, args.img_size, **config.test.augment.transform)
で変換されています。この関数はdemo_utils.py内で定義されており、
def transform_img(origin_img, size_divisibility, image_max_range, flip_prob,
image_mean, image_std):
transform = [
T.Resize(image_max_range),
T.RandomHorizontalFlip(flip_prob),
T.ToTensor(),
T.Normalize(mean=image_mean, std=image_std),
]
transform = T.Compose(transform)
img, _ = transform(origin_img)
img = to_image_list(img, size_divisibility)
return img
で構成されています。T.Resizeで画像中の短辺が指定サイズになるようにしています。Normalizeでは画像の平均・分散を所定の数値になるように変換していますね。
平均・分散の指定はaugmentation.pyにて指定されています。
'image_mean': [0.0, 0.0, 0.0],
'image_std': [1.0, 1.0, 1.0],
これを用いればよさそうです。（今回は特に意味がない数字が入ってますね・・・。）
ではこのあたりを参考に実装を進めてみましょう。
実装した結果が次の通りです。
import cv2
from damo.base_models.core.ops import RepConv
from damo.config.base import parse_config
from damo.detectors.detector import build_local_model
from damo.utils import get_model_info, vis
from damo.utils.demo_utils import transform_img
import torch
import numpy as np
from time import time
cap = cv2.VideoCapture(0)
COCO_CLASSES = []
for i in range(80):
COCO_CLASSES.append(str(i))
COCO_CLASSES = tuple(COCO_CLASSES)
config = parse_config("./configs/damoyolo_tinynasL20_T.py")
model = build_local_model(config, "cpu")
ckpt = torch.load("damoyolo_tinynasL20_T_418.pth", map_location="cpu")
new_state_dict = {}
for k, v in ckpt['model'].items():
new_state_dict[k] = v
model.load_state_dict(new_state_dict, strict=False)
model.eval()
def resize(img, size):
if img.shape[0] > img.shape[1]:
return cv2.resize(img, (size, int(img.shape[0]/img.shape[1]*size/32)*32))
else:
return cv2.resize(img, (int(img.shape[1]/img.shape[0]*size/32)*32, size))
with torch.no_grad():
while True:
ret, frame = cap.read()
frame = frame[:, ::-1, :]
frame = resize(frame, 640)
frame_z = frame.transpose(2, 0, 1)[np.newaxis]
frame_z = torch.from_numpy(frame_z.copy()).to(torch.float32)
start = time()
output = model(frame_z)
end = time()
bboxes = output[0].bbox
scores = output[0].get_field('scores')
cls_inds = output[0].get_field('labels')
out_img = vis(frame,
bboxes,
scores,
cls_inds,
conf=0.7,
class_names=COCO_CLASSES)
cv2.putText(out_img, f"time:{end-start:4.3f}s",(0,20), cv2.FONT_HERSHEY_SIMPLEX,1.0, (0,255,0), 2 , cv2.LINE_4)
cv2.imshow("test",out_img)
cv2.waitKey(1)
def resize(img, size):
if img.shape[0] > img.shape[1]:
return cv2.resize(img, (int(size/32)*32, int(img.shape[0]/img.shape[1]*size/32)*32))
else:
return cv2.resize(img, (int(img.shape[1]/img.shape[0]*size/32)*32, int(size/32)*32))
の処理では画像を32の倍数になるように制限しています。YOLO系統は画像サイズが32の倍数でないとうまく処理ができない構造になっています。（モデル深層部に行くと画像サイズが小さくなっていくため、32*Nのサイズがないと最後までうまく畳み込みが実行できないためです。）
上記コードを実行するとリアルタイム推論が実行可能です。・・・・が、やはり重い・・・。PyTorch+CPUだとなんでこんなに重いんでしょうね。
次回は中身のスクリプトを調べながらONNXへと変換し、高速化を図りたいと思います！
Discussion