前言

最近的项目里有yolov5的嵌入，需求是只需要推理，模型文件是已有的，输入需要是图片（原yolov5是输入路径），输出结果的图片和标签。这样的话需要对原来的代码进行一些简化和变更。

路径

模型这块路径的结构我是这样的，因为项目里还会加别的模型，所以有点套娃

加入__init__.py是为了让文件夹可以被当成包来import
models和utils是yolo原来的文件，至于需要哪个不需要哪个大家可以自己看着弄

import torch
import numpy as np
from model.yolov5.utils.augmentations import letterbox
from model.yolov5.models.experimental import attempt_load
from model.yolov5.utils.general import check_img_size, non_max_suppression, scale_coords, set_logging, \xyxy2xywh
from model.yolov5.utils.torch_utils import select_device, time_sync
from model.yolov5.utils.plots import Annotator, colorsclass Yolov5:device = ''weights = 'model/yolov5/screw.pt' # model.pt path(s)imgsz = 640  # inference size (pixels)save_img = Truedef __init__(self):set_logging()self.device = select_device(self.device)# Load modelw = str(self.weights[0] if isinstance(self.weights, list) else self.weights)self.stride, self.names = 64, [f'class{i}' for i in range(1000)]  # assign defaultsself.model = torch.jit.load(w) if 'torchscript' in w else attempt_load(self.weights, map_location=self.device)self.stride = int(self.model.stride.max())  # model strideself.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names  # get class namesself.imgsz = [self.imgsz]self.imgsz *= 2self.imgsz = check_img_size(self.imgsz, s=self.stride)  # check image sizeif self.device.type != 'cpu':self.model(torch.zeros(1, 3, *self.imgsz, ).to(self.device).type_as(next(self.model.parameters())))  # run once@torch.no_grad()def run(self, im0s,  # HWC图片imgsz=640,  # inference size (pixels)conf_thres=0.25,  # confidence thresholdiou_thres=0.45,  # NMS IOU thresholdmax_det=1000,  # maximum detections per imageview_img=False,  # show resultssave_txt=False,  # save results to *.txtsave_conf=False,  # save confidences in --save-txt labelssave_crop=False,  # save cropped prediction boxesclasses=None,  # filter by class: --class 0, or --class 0 2 3agnostic_nms=False,  # class-agnostic NMSaugment=False,  # augmented inferenceline_thickness=3,  # bounding box thickness (pixels)hide_labels=False,  # hide labelshide_conf=False,  # hide confidences):# Load imageassert im0s is not None, 'Image Not Available 'img = letterbox(im0s, imgsz, stride=self.stride)[0]# Convertimg = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGBimg = np.ascontiguousarray(img)dt = [0.0, 0.0, 0.0]t1 = time_sync()img = torch.from_numpy(img).to(self.device)img = img.float()  # uint8 to fp16/32img = img / 255.0  # 0 - 255 to 0.0 - 1.0if len(img.shape) == 3:img = img[None]  # expand for batch dimt2 = time_sync()dt[0] += t2 - t1# Inferencevisualize = Falsepred = self.model(img, augment=augment, visualize=visualize)[0]t3 = time_sync()dt[1] += t3 - t2# NMSpred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)dt[2] += time_sync() - t3# Process predictionsresult = []result_lb = 0det = pred[0] # per images = ''im0 = im0s.copy()s += '%gx%g ' % img.shape[2:]  # print stringgn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwhimc = im0.copy() if save_crop else im0 # for save_cropannotator = Annotator(im0, line_width=line_thickness, example=str(self.names))if len(det):print('det', det)result_lb = 1# Rescale boxes from img_size to im0 sizedet[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()# Print resultsfor c in det[:, -1].unique():n = (det[:, -1] == c).sum()  # detections per classs += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "  # add to string# Write resultsfor *xyxy, conf, cls in reversed(det):if save_txt:  # Write to filexywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywhline = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label formatif self.save_img or save_crop or view_img:  # Add bbox to imagec = int(cls)  # integer classlabel = None if hide_labels else (self.names[c] if hide_conf else f'{self.names[c]} {conf:.2f}')annotator.box_label(xyxy, label, color=colors(c, True))im_result = annotator.result()result.append(im_result)  # 图片result.append(result_lb)# Print time (inference-only)print(f'{s}Done. ({t3 - t2:.3f}s)')return result  # [图片， 标签]

有几个点说一下，
现在这里device不手动指定，默认选择；
读取的模型文件一定是.pt，这个就是根据我自己这边定制的，原来支持各种格式，我把那些都删掉了简化；
输入图片是BGR也就是opencv读入的默认格式，因为后面有转化的代码；
imgsz输入一个int，后面会转化成yolo需要的格式（之前在这里debug好久）；
我这里是做成一个类，实例化的时候会加载模型并且运行一次warming up以减少后面的初次运行时间（我猜的）（原作者就是这样做的，正式推理之前先run once）
需要检测的时候执行run就可以了，返回值是一个图片和标签组成的列表，当然用元组也可以

代码还是有一些冗余的我懒得再仔细看了。

【yolov5检测代码简化】Yolov5 detect.py推理代码简化，输入图片，输出图片和结果相关推荐

Yolov5代码详解——detect.py
首先执行扩展包的导入: import argparse import os import platform import sys from pathlib import Path import t ...
YoLoV5学习（4）--detect.py程序（预测图片、视频、网络流）逐段讲解~
本章博客主要分析YoloV5代码中的detect程序代码,按照程序运行步骤顺序主要分为3大部分. 1.包与库的导入 1.1 导入安装好的python库.torch库等等其中:argparse模块.o ...
YOLOv5的Tricks | 【Trick13】YOLOv5的detect.py脚本的解析与简化
如有错误,恳请指出. 在之前介绍了一堆yolov5的训练技巧,train.py脚本也介绍得差不多了.之后还有detect和val两个脚本文件,还想把它们总结完. 在之前测试yolov5训练好的模型时, ...
OpenVINO-yolov5推理代码
文章目录前言一.OpenVINO2021.3代码二.OpenVINO2022.3代码三.结果展示总结前言使用OpenVINO实现对yolov5s口罩模型推理,其他版本yolov5可以 ...
yolov5-5.0版本代码详解----augmentations.py的augment_hsv函数
yolov5-5.0版本代码详解----augmentations.py的augment_hsv函数 1.用途图片的hsv色域增强模块 2.调用位置在datasets.py的LoadImagesA ...
YOLOV5检测代码detect.py注释与解析
YOLOv5代码注释版更新啦,注释的是最近的2021.07.14的版本,且注释更全 github: https://github.com/Laughing-q/yolov5_annotations Y ...
砍掉九成代码，重构并简化YOLOv5图像目标检测推理实现
YOLOv5官方开源代码给出了完成的推理实现,但过于封装,只能通过修改配置参数对指定文件夹下图像和视频进行推理,而且三百多行的推理代码也显得过于冗长.如果想要在项目上进行部署应用,显然需要更高的灵活性 ...
yolov5——detect.py代码【注释、详解、使用教程】
yolov5--detect.py代码[注释.详解.使用教程] yolov5--detect.py代码[注释.详解.使用教程] 1. 函数parse_opt() 2. 函数main() 3. 函数ru ...
将yolov5的detect.py改写成可以供其他程序调用的方式，并实现低时延（＜0.5s）直播推理
将yolov5的推理代码改成可供其它程序调用的方式,并实现低时延(<0.5s)直播推理 yolov5的代码具有高度的模块化,对于初学者十分友好,但是如果咱们要做二次开发,想直接调用其中一些函数, ...

【yolov5检测代码简化】Yolov5 detect.py推理代码简化，输入图片，输出图片和结果

前言

路径

【yolov5检测代码简化】Yolov5 detect.py推理代码简化，输入图片，输出图片和结果相关推荐

最新文章

热门文章