【物体検出2024】YOLOv9を試してみる　〜導入からデモまで〜

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

!git clone https://github.com/WongKinYiu/yolov9
%cd yolov9

!pip install -r requirements.txt -q
!pip install supervision -q

!mkdir -p weights

!wget -P weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-c.pt
!wget -P weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-e.pt

!wget "https://user0514.cdnw.net/shared/img/thumb/susi1123PAR50424312_TP_V4.jpg?&download=1" -O "sample.jpg"

import torch
import cv2
import numpy as np
from models.common import DetectMultiBackend
from utils.general import non_max_suppression, scale_boxes
from utils.torch_utils import select_device, smart_inference_mode
from utils.augmentations import letterbox
import PIL.Image
import supervision as sv
%matplotlib inline

def predict(image_path, weights, imgsz, conf_thres, iou_thres, device, data):
    device = select_device(device)
    model = DetectMultiBackend(weights, device=device, fp16=False, data=data)
    stride, names, pt = model.stride, model.names, model.pt
    image = PIL.Image.open(image_path)
    img0 = np.array(image)
    assert img0 is not None, f'Image Not Found {image_path}'
    img = letterbox(img0, imgsz, stride=stride, auto=True)[0]
    img = img[:, :, ::-1].transpose(2, 0, 1)
    img = np.ascontiguousarray(img)
    img = torch.from_numpy(img).to(device).float()
    img /= 255.0
    if img.ndimension() == 3:
        img = img.unsqueeze(0)

    pred = model(img, augment=False, visualize=False)

    pred = non_max_suppression(pred[0][0], conf_thres, iou_thres, classes=None, max_det=1000)

    bounding_box_annotator = sv.BoxAnnotator()
    label_annotator = sv.LabelAnnotator(text_position=sv.Position.TOP_LEFT)

    for i, det in enumerate(pred):
        if len(det):
            det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], img0.shape).round()
            for *xyxy, conf, cls in reversed(det):
                label_name = names[int(cls)]
                label = f'{label_name} {conf:.2f}'
                detections = sv.Detections(
                    xyxy=torch.stack(xyxy).cpu().numpy().reshape(1, -1),
                    class_id=np.array([int(cls)]),
                    confidence=np.array([float(conf)])
                )

                labels = [
                    f"{names[int(class_id)]} {confidence:0.2f}"
                    for class_id, confidence
                    in zip(detections.class_id, detections.confidence)
                ]

                img0 = bounding_box_annotator.annotate(img0, detections)
                img0 = label_annotator.annotate(img0, detections, labels)

    return img0[:, :, ::-1]

img = predict(image_path='sample.jpg', weights='weights/yolov9-e.pt', imgsz=640, conf_thres=0.5, iou_thres=0.7, device='0', data='data.yaml')
sv.plot_image(img)