пока нет моделей S и M, тестируем модель С на cuda.
import torch,time
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
start_time = time.time()
image_path='street.jpg'
@smart_inference_mode()
def predict(image_path, weights='yolov9-c.pt', imgsz=640, conf_thres=0.1, iou_thres=0.45, device='0', data='data/coco.yaml'):
Initialize
device = select_device(device)
model = DetectMultiBackend(weights, device=device, fp16=False, data=data)
stride, names, pt = model.stride, model.names, model.pt
Load image
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)
Init bounding box annotator and label annotator
bounding_box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
Inference
pred = model(img, augment=False, visualize=False)
Apply NMS
pred = non_max_suppression(pred[0][0], conf_thres, iou_thres, classes=None, max_det=1000)
Process detections
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 = f'{names[int(cls)]} {conf:.2f}'
Transform detections to supervisions detections
detections = sv.Detections(
xyxy=torch.stack(xyxy).cpu().numpy().reshape(1, -1),
class_id=np.array([int(cls)]),
confidence=np.array([float(conf)])
)
Labels
labels = [
f"{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]
image=predict(image_path, weights='yolov9-c.pt', imgsz=640, conf_thres=0.1, iou_thres=0.45, device='0', data='data/coco.yaml')
cv2.imwrite('out.jpg', image)
print("--- %s seconds ---" % (time.time() - start_time))