yolov5-python调用精简源码

import time

import numpy as np
import torch

from models.common import DetectMultiBackend
from utils.augmentations import letterbox
from utils.general import (check_img_size, cv2,
                           non_max_suppression, scale_coords)
from utils.plots import Annotator, colors
from utils.torch_utils import select_device


class YOLOv5Detector:
    """ YOLOv5 object detection """

    def __init__(self, weights='yolov5s.pt', conf_thres=0.25, iou_thres=0.45, imgsz=640, data='data/coco128.yaml'):
        """ Initialization """
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        self.device = select_device('0')
        self.model = DetectMultiBackend(weights, device=self.device, dnn=False, data=data, fp16=False)
        self.stride, self.names, self.pt = self.model.stride, self.model.names, self.model.pt
        self.imgsz = check_img_size(imgsz, s=self.stride)  # check image size

    def image_preprocess(self, image):
        im0 = image.copy()
        im = letterbox(im0, self.imgsz, stride=32, auto=True)[0]  # padded resize
        im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        im = np.ascontiguousarray(im)  # contiguous
        im = torch.from_numpy(im).to(self.device)
        im = im.half() if self.model.fp16 else im.float()  # uint8 to fp16/32
        im /= 255  # 0 - 255 to 0.0 - 1.0
        if len(im.shape) == 3:
            im = im[None]  # expand for batch dim
            # Dataloader
        return im

    def __call__(self, image, *args, **kwargs):
        im = self.image_preprocess(image)
        pred = self.model(im, augment=False, visualize=False)
        pred = non_max_suppression(pred,
                                   conf_thres=0.25,
                                   iou_thres=0.45,
                                   classes=None,
                                   agnostic=False,
                                   multi_label=False,
                                   labels=(),
                                   max_det=1000)

        for i, det in enumerate(pred):  # per image
            annotator = Annotator(im0, example=str(self.names))
            if len(det):
                det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
                for *xyxy, conf, cls in reversed(det):
                    label = f'{self.names[int(cls)]} {conf:.2f}'
                    annotator.box_label(xyxy, label, color=colors(2, True))

        return im0


yolov5_detector = YOLOv5Detector(weights='best.pt')
img = r'C:\Users\Administrator\Desktop\000000011244.jpg'
while True:
    im0 = cv2.imread(img)

    t0 = time.time()
    im0 = yolov5_detector(im0)
    print(f'Done. ({time.time() - t0:.3f}s)')
    # print(time.time() - t0)
    cv2.imshow("123456", im0)
    cv2.waitKey(1)  # 1 millisecond