はじめに
前回までの記事では、Yolov5 の github からダウンロードしたプログラムのみを使用していました。
今回は、Yolov5 の推論を自作コードに組み込む方法について説明します。
前提条件
前提条件は、以下の通りです。
- numpy>1.22.0, opencv-python>4.5.0, pytorch>1.8.0 がインストールされている(CPUとGPUはどちらでも問題ありません)
- Visual Studio Code がインストールされている
- コマンドはUbuntuのものですが、Windowsでもほとんど動作します(python3, pip3 → python, pip に置き換える必要があります)
プログラムの準備
まずは、 git clone で yolov5 をダウンロードしていきます。
mkdir ~/yolov5-dev
cd yolov5-dev
git clone https://github.com/ultralytics/yolov5.git
mv yolov5 yolov5_object
touch yolov5_detection.py作成した yolov5_detection.py のコードを以下のようにしてください。
import numpy as np
import cv2
import torch
from yolov5_object.models.common import DetectMultiBackend
from yolov5_object.utils.augmentations import letterbox
from yolov5_object.utils.general import non_max_suppression
# ======================= yolov5 ========================= #
frame = cv2.imread("./output.png")
h, w, _ = frame.shape
device = torch.device("cuda")
model = DetectMultiBackend("./yolov5/yolov5s.pt", device=device, dnn=False, data='./yolov5/data/coco128.yaml', fp16=True)
stride, names, pt = model.stride, model.names, model.pt
conf_thres = 0.45
iou_thres = 0.25
classes = None
# Padded resize
img = letterbox(frame, 640, stride=stride, auto=True)[0]
# Convert
img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
img = np.ascontiguousarray(img)
im = torch.from_numpy(img).to(device)
im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
im /= 255 # 0 - 255 to 0.0 - 1.0
if len(im.shape) == 3:
im = im[None] # expand for batch dim
pred = model(im, augment=False, visualize=False)
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, False, max_det=50)
pred_cpu = pred[0][0].cpu().detach().numpy()
index = int(pred_cpu[5])
print(pred_cpu, index, names[index])コードの説明は別途行います。
まずは、動かすための準備をしていきます。
Visual Studio Code で yolov5-dev を開いて、以下のようになっていることを確認してください。
Yolov5 の依存関係を解消
import の依存関係を解消していきます。
from yolov5_object.models.common import DetectMultiBackendこちらは、yolov5_object フォルダの models フォルダ内の common.py ファイルのDetectMultiBackend() 関数を import することを表しています。
実際に yolov5_object/models/common.py ファイルを確認してみます。
from utils import TryExcept上記の一文がありますが、これは先ほどの例と同じで utils フォルダ内の __init__.py から TryExcept を import しています。
ひとまず、このまま動かしてみましょう。
python3 yolov5_detection.pyすると、以下のようなエラーが出力されます。
Traceback (most recent call last):
File "yolov5_detection.py", line 5, in <module>
from yolov5_object.models.common import DetectMultiBackend
File "~/yolov5-dev/yolov5_object/models/common.py", line 28, in <module>
from utils import TryExcept
ModuleNotFoundError: No module named 'utils'~/yolov5-dev/yolov5_object/models/common.py において、
utils が見つからないというエラーが発生しています。
これは、yolov5-dev フォルダにある yolov5_detection.py から utils を参照しようとしているので、見つからないエラーが発生しています。
なので、utils をすべて yolov5_object.utils に変更します。
変更したら、再度動かしてみましょう。
python3 yolov5_detection.py以下のようなエラーが出ます。
Traceback (most recent call last):
File "yolov5_detection.py", line 5, in <module>
from yolov5_object.models.common import DetectMultiBackend
File "~/yolov5-dev/yolov5_object/models/common.py", line 29, in <module>
from yolov5_object.utils.dataloaders import exif_transpose, letterbox
File "~/yolov5-dev/yolov5_object/utils/dataloaders.py", line 31, in <module>
from utils.augmentations import (Albumentations, augment_hsv, classify_albumentations, classify_transforms, copy_paste,
ModuleNotFoundError: No module named 'utils'今度は、 utils/dataloaders.py にて同様のエラーが発生しています。
utils を yolov5_object.utils に変更します。
以降、これをエラーがなくなるまで続けてください。
- models/common.py の先頭import と 332, 603, 604行目
- models/experimental.py
- utils/dataloaders.py
- utils/augmentations.py
- utils/general.py
- utils/metrics.py
- utils/torch_utils.py
- utils/plots.py
の models, utils, export の import文を yolov5_object.models, yolov5_object.utils といったように変更してください。
yolov5 オブジェクトの実行
推論に使用する画像はこちらです。
プログラムを実行します。
python3 yolov5_detection.py出力は以下のようになります。
Fusing layers...
YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients
[ 249.5 373.25 373.5 470.75 0.93408 64] 64 mousemouse が検出されていれば成功です。
おわりに
今回は Yolov5 をオブジェクト化して簡単に再利用する方法について説明しました。
次回は、推論コードの詳細について説明していきます。
コメント