YOLOv5の認識プログラムを書いてみた

## 最新版をご覧ください。
https://qiita.com/ez_tera/items/e62987c66113c2e9d85e

## はじめに
* https://github.com/amirhosseinh77/JetsonYolo を参考にwebカメラを入力とした認識プログラムを書いてみました。
* 環境構築は https://qiita.com/ez_tera/items/c12662187164b8c06a80 を参照。

##前提条件
* ウェイトファイルには Google Colab上で学習を実行した kimetsu.pt を用意しました。
* クラス名には 0:inosuke、1:tanjiro、2:zenitsu の3つを指定してあります。

##プログラム
```:yolo.py
import torch
import cv2
import numpy as np
from models.experimental import attempt_load
from utils.general import non_max_suppression

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class OBJ_DETECTION():
    def __init__(self, model_path, classes):
        self.classes = classes
        self.yolo_model = attempt_load(weights=model_path, map_location=device)
        self.input_width = 320

    def detect(self, main_img):
        height, width = main_img.shape[:2]
        new_height = int((((self.input_width/width)*height)//32)*32)

        img = cv2.resize(main_img, (self.input_width, new_height))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = np.moveaxis(img, -1, 0)
        img = torch.from_numpy(img).to(device)
        img = img.float()/255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        pred = self.yolo_model(img, augment=False)[0]
        pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45, classes=None)
        items = []

        if pred[0] is not None and len(pred):
            for p in pred[0]:
                score = np.round(p[4].cpu().detach().numpy(), 2)
                label = self.classes[int(p[5])]
                xmin = int(p[0] * main_img.shape[1]/self.input_width)
                ymin = int(p[1] * main_img.shape[0]/new_height)
                xmax = int(p[2] * main_img.shape[1]/self.input_width)
                ymax = int(p[3] * main_img.shape[0]/new_height)

                item = {'label': label,
                        'bbox': [(xmin, ymin), (xmax, ymax)],
                        'score': score
                        }

                items.append(item)

        return items
```

```:JetsonYolo.py
import cv2
import numpy as np
from yolo import OBJ_DETECTION

Object_classes = ['inosuke', 'tanjiro', 'zenitsu']
Object_colors = list(np.random.rand(3, 3)*255)  # 3つ x 3要素(0.0-1.0)の乱数生成 *255
Object_detector = OBJ_DETECTION('kimetsu.pt', Object_classes)  # ウェイトファイル名を指定

cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

cv2.namedWindow("win", cv2.WINDOW_AUTOSIZE)

while True:
    ret, frame = cap.read()

    objs = Object_detector.detect(frame)

    for obj in objs:
        print(obj)
        label = obj['label']
        score = obj['score']

        [(xmin, ymin), (xmax, ymax)] = obj['bbox']
        color = Object_colors[Object_classes.index(label)]
        frame = cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), color, 2)
        frame = cv2.putText(frame, f'{label} ({str(score)})', (xmin, ymin),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, 4)

    cv2.imshow("win", frame)
    if cv2.waitKey(30) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()
```
##実行方法
python3 JetsonYolo.py