yolov7简化pt调用示例

import time

import cv2

import numpy as np
import torch
from numpy import random

from models.experimental import attempt_load
from utils.datasets import letterbox
from utils.general import non_max_suppression, scale_coords, xyxy2xywh
from utils.plots import plot_one_box
from utils.torch_utils import select_device


class YOLOv7Detector:

    def __init__(self, weights='best.pt', conf_thres=0.3, iou_thres=0.45):
        """ Initialization """
        self.half = False
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        self.device = select_device('0')
        self.model = attempt_load(weights, map_location=self.device)  # load FP32 model

        if self.half:
            self.model.half()  # to FP16

        # stride = int(model.stride.max())  # model stride
        # imgsz = check_img_size(imgsz, s=stride)  # check img_size

        # Get names and colors
        self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names
        self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in self.names]
     
    def image_preprocess(self, image):
        im0 = image.copy()

        img = letterbox(im0, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True,
                       stride=32)[0]
        # Convert
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(self.device)
        img = img.half() if self.half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        return img, im0

    def __call__(self, image, *args, **kwargs):
        img, img0 = self.image_preprocess(image)
        pred = self.model(img, augment=False)[0]
        pred = non_max_suppression(pred, conf_thres=self.conf_thres, iou_thres=self.iou_thres, classes=None, agnostic=False)

        for i, det in enumerate(pred):  # detections per image
            # gn = torch.tensor(img0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            if len(det):
                # # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()

                # Write results
                for *xyxy, conf, cls in reversed(det):
                    # xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                    label = f'{self.names[int(cls)]} {conf:.2f}'
                    plot_one_box(xyxy, img0, label=label, color=self.colors[int(cls)], line_thickness=1)

        return img0

if __name__ == '__main__':
    yolov7_detector = YOLOv7Detector(weights='best.pt')
    img = 'E:/yolov7_ui_qt/VOCdevkit/VOC2007/JPEGImages/000000016761.jpg'
    while True:
        im0 = cv2.imread(img)
        t0 = time.time()
        im = yolov7_detector(im0)
        print(f'Done. ({time.time() - t0:.3f}s)')
        cv2.imshow("123456", im)
        cv2.waitKey(1)  # 1 millisecond