yolov5 python推理代码

isyoungboy

于 2024-07-16 11:05:48 发布

阅读量256

点赞数 5

文章标签： YOLO python 开发语言

本文链接：https://blog.csdn.net/isyoungboy/article/details/140460973

版权

code

import os
import cv2
import argparse
import time
import torch
import numpy as np
from models.common import DetectMultiBackend
from utils.torch_utils import select_device, smart_inference_mode
from utils.augmentations import (Albumentations, augment_hsv, classify_albumentations, classify_transforms, copy_paste,
                                 letterbox, mixup, random_perspective)
from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
                           increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh)
from ultralytics.utils.plotting import Annotator, colors, save_one_box

class Yolov5:
    def __init__(self, model_path, img_size, conf_thres, iou_thres):
        self.img_size=img_size
        self.conf_thres=conf_thres
        self.iou_thres=iou_thres
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = DetectMultiBackend(model_path, device=self.device , dnn=False, data='data/VOC.yaml', fp16=False)
        self.stride, self.names, self.pt = self.model.stride, self.model.names, self.model.pt

    def preprocess(self, img):
        im = letterbox(img, self.img_size, stride=self.stride, auto=True)[0]  # padded resize
        im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        im = np.ascontiguousarray(im)  # contiguous
        return img

    def inference(self, img):
        im0=img
        imShape=img.shape
        gn = torch.tensor(img.shape)[[1, 0, 1, 0]]
        img = self.preprocess(img)
        with torch.no_grad():
            im = torch.from_numpy(img).to(self.device)
            im=im.float()
            im = im.unsqueeze(0) 
            im=im.permute(0, 3, 1, 2)
            im /= 255
            print(im.shape)
            pred = self.model(im, augment=False, visualize=False)
            conf_thres=0.25  # confidence threshold
            iou_thres=0.45  # NMS IOU threshold
            max_det=1000  # maximum detections per image)
            classes=[0]
            print(conf_thres)
            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, False, max_det=max_det)
            for i, det in enumerate(pred): 
                annotator = Annotator(im0, line_width=3, example=str(self.names))
                if len(det):
                    # Rescale boxes from img_size to im0 size
                    det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], imShape).round()


                    # Write results
                    for *xyxy, conf, cls in reversed(det):
                        c = int(cls)  # integer class
                        confidence = float(conf)
                        confidence_str = f'{confidence:.2f}'
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()
                        annotator.box_label(xyxy,'C', color=colors(c, True))
            im0 = annotator.result()
            cv2.imwrite('res.jpg',im0)
        return ''