yolov13推理示例
yolov13 推理示例
from ultralytics import YOLO
import torch
import os
import cv2
from numpy import randomnames = {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear', 22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard', 32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle', 40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli', 51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table', 61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard', 67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink', 72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'}def draw_image(x, img, color=None, label=None, line_thickness=None):tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thicknesscolor = color or [random.randint(0, 255) for _ in range(3)]c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)if label:tf = max(tl - 1, 1) # font thicknesst_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled backgroundcv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)return imgif __name__ == '__main__':pt_model_path = 'weights/yolov13s.pt'device = torch.device("cpu")model = YOLO(pt_model_path) image_path = "images/bus.jpg"image = cv2.imread(image_path)conf_threshold = 0.5iou_threshold = 0.45results = model(image, conf=conf_threshold, iou=iou_threshold)detections = []if len(results) > 0 and results[0].boxes is not None:boxes = results[0].boxesfor i in range(len(boxes)):box = boxes[i]detection = {"bbox": box.xyxy[0].cpu().numpy().tolist(), # [x1, y1, x2, y2]"confidence": float(box.conf[0]),"class_id": int(box.cls[0])}detections.append(detection)draw_image_data = imagefor detection in detections:label = '%s %.2f' % (names[detection["class_id"]], detection["confidence"])draw_image_data = draw_image(detection["bbox"], draw_image_data, label=label, color=(255, 0, 0))image_name = os.path.basename(image_path)cv2.imwrite('result_{}.jpg'.format(image_name[:-4]), draw_image_data)print("Done!")
说明
- yolo内部已经集成了图像预处理和推理后处理nms合并,然后输出外包了一层。