YOLOv7のAPIを使用してみた

概要

※2022/7/17 記述

最新バージョンであるYOLOv7についての投稿です。
論文によると、推論速度も精度も以前のバージョンを上回っています。
この記事はYOLOv7のAPIを使用し、私好みに実装した例を記述しております。
現時点では、学習済みモデルを使用した推論のみです。

※URL
Github : https://github.com/WongKinYiu/yolov7
Paper : https://arxiv.org/pdf/2207.02696.pdf

環境構築

$ mkdir [任意のディレクトリ]
$ cd [任意のディレクトリ]
$ conda create -n yolov7 python=3.7
$ conda activate yolov7
$ conda install git
$ git clone https://github.com/WongKinYiu/yolov7
$ cd yolov7
$ pip install -r requirements.txt

準備

以下のディレクトリ構造になるよう①～③の準備をする。
① src/ を作成し、入力画像を用意して格納する。
② model_weight/ を作成し、学習済みモデルをGithubよりダウンロードして格納する。
③ my_yolov7.ipynb を作成。

[任意のディレクトリ]
     └─ yolov7/
          ├─ cfg/
          ├─ data/
          ├─ figure/
          ├─ inference/
          ├─ models/
          ├─ scripts/
          ├─ tools/
          ├─ utils/
          ├─ src/ : 入力画像格納用に作成したディレクトリ
               └─ dog.jpg (今回は[https://github.com/pjreddie/darknet/blob/master/data/dog.jpg]を使用)
          ├─ dst/ : 出力画像格納用に作成したディレクトリ
               └─ yolov7_dog.jpg
          ├─ model_weight/ : 自分で作成したディレクトリ(モデルの重みをGithubよりダウンロードし、格納)
               └─ yolov7.pt ([https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt]よりダウンロード)
          ├─ LICENSE.md
          ├─ README.md
          ├─ detect.py
          ├─ hubconf.py
          ├─ requirement.txt
          ├─ test.py
          ├─ train.py
          ├─ train_aux.py
          └─ my_yolov7.ipynb : 今回使用しているファイル

推論

# my_yolov7.ipynb
import torch
import cv2
import random
import os
from utils.general import non_max_suppression, xyxy2xywh
from models.experimental import attempt_load
from utils.plots import plot_one_box

# 乱数固定
random.seed(0)
# 出力ディレクトリを作成
os.makedirs("./dst", exist_ok=True)

# モデルの読み込み
model = attempt_load("./model_weight/yolov7.pt", map_location="cpu")  
model.eval()

# 画像の読み込み
original_image = cv2.imread("./src/dog.jpg")
image = original_image.copy()
image = cv2.resize(image, [640,640])

# 前処理
inputs = image.copy()
inputs = inputs/255
inputs = inputs.transpose(2,0,1)
inputs = torch.tensor(inputs).unsqueeze(0).float()

# 推論
outputs = model(inputs)[0]

# 後処理
conf_thres = 0.5 # deafult:0.25
iou_thres = 0.45
classes = range(len(model.names))
agnostic_nms = True
prediction = non_max_suppression(outputs, conf_thres, iou_thres, classes=classes, agnostic=agnostic_nms)
detection = prediction[0]


# 画像に出力
gn = torch.tensor(image.shape)[[1, 0, 1, 0]]
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
save_conf = False
for *xyxy, conf, cls in reversed(detection):
    xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
    line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
    label = f'{names[int(cls)]} {conf:.2f}'
    plot_one_box(xyxy, image, label=label, color=colors[int(cls)], line_thickness=1)
image = cv2.resize(image, [original_image.shape[1], original_image.shape[0]])
cv2.imwrite("./dst/yolov7_dog.jpg", image)

# matplotで画像を表示
# import matplotlib
# import matplotlib.pyplot as plt
# %matplotlib inline
# plt.figure(figsize=(10,10))
# plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
# plt.show()

推論結果

学習

今後、記述予定。

最後に

以上、YOLOv7のAPIを使用した実装のご紹介でした。

今後、行っていきたいこと。
・モデルの学習
・OpenCVを使用した映像推論
・ONNXRuntimeでの推論

最後までご覧いただきありがとうございました。