YOLOv5 v7.0 Google Colabで学習し、Jetson NanoとPythonで物体検出する #ubuntu20.04

Jetson Nano 4GB
micro SDXC 64GB
logicool C270N
Ubuntu 20.04

Google Colab上でYOLOv5 v7.0を用いて学習する

Pythonのバージョンを確認します(2023/9/30:3.10.12)

!python3 -V

Google Driveをマウントします

from google.colab import drive
drive.mount('/content/drive')

YOLOv5用のディレクトリを作成します

%cd /content/drive/My Drive/
%mkdir Python3.10_YOLOv5v7.0
%cd Python3.10_YOLOv5v7.0

YOLOv5 v7.0をクローンし、ディレクトリ内に移動します

!git clone https://github.com/ultralytics/yolov5 -b v7.0
%cd yolov5

必要なモジュールをインストールします

!pip install -r requirements.txt

YOLOv5の動作確認をします

!python3 detect.py
from IPython.display import Image, clear_output
Image(filename='runs/detect/exp/zidane.jpg', width=600)

yolov5/data内にtrain、validフォルダとdata.yamlをアップロードし、学習を行います

!python train.py --img 640 --batch 8 --epochs 300 --data ./data.yaml --cfg ./models/yolov5s.yaml --weights ''

マイドライブ > ultralytics > yolov5 > runs > train > exp > weights に重みファイル best.ptが保存されます。

Jetson Nanoで YOLOv5 v7.0を動かす

Jetson Nano with Ubuntu 20.04 OS imageを入手します

※username：jetson、password：jetson

キーボード、タイムゾーンの設定

Settings
-> Region & Language -> Input Sources：Japanese
-> Data & Time -> Time Zone：JST(Tokyo, Japan)

モジュールの追加・更新

pip3 install numpy==1.20.3
pip3 install psutil==5.9.5
pip3 install matplotlib==3.6.0
pip3 install ultralytics==8.0.111
pip3 install python-dateutil==2.8.2
pip3 install IPython

YOLOv5 v7.0をクローンします

git clone https://github.com/ultralytics/yolov5.git -b v7.0
cd yolov5/
nano requirements.txt

requirements.tx 全ての行をコメントアウトします

# YOLOv5 requirements
# Usage: pip install -r requirements.txt

# Base ------------------------------------------------------------------------
# gitpython>=3.1.30
# matplotlib>=3.3

途中略

# albumentations>=1.0.3
# pycocotools>=2.0.6  # COCO mAP

YOLOv5の動作確認

detect.pyの実行

python3 detect.py

検出結果は「yolov5/runs/detect/exp」内に保存されます（bus.jpg、zidane.jpg）。

Google Colabで学習した重みファイルを使用して認識を実行

python3 detect.py --source 0 --imgsz=256 --nosave --weights best.pt

Jetson Nano：1 frame：43.0ms（23.3fps）

Pythonでプログラムを作成する

プログラム

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, device=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 = ['pcb', 'battery', 'nut', 'spacer']
Object_colors = list(np.random.rand(4, 3)*255)  # 4クラス x 3色(0.0-1.0)の乱数生成 *255
Object_detector = OBJ_DETECTION('best.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

YOLOv5 v7.0 Google Colabで学習し、Jetson NanoとPythonで物体検出する