Yolov5 最简推理代码

Yolov5模型加载与推理的简化代码，只需要一个文件即可，不依赖其他文件。

如果报下列错误，请下载精简版的models文件夹。下载地址：Yolov5最简推理代码-深度学习文档类资源-CSDN下载。

File "D:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py", line 853, in _load
result = unpickler.load()
ModuleNotFoundError: No module named 'models'

这个错误是因为torch保存了模型存储路径和结构，但是并没有保存计算过程。因此，需要models文件中对模型的定义。

File "D:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py", line 853, in _load
result = unpickler.load()
ModuleNotFoundError: No module named 'models'

Yolov5最简推理代码如下所示：

# -*- coding: utf-8 -*-
"""
Created on Fri Apr  1 16:12:34 2022@author: suiyingy
"""import cv2
import torch
import torchvision
import numpy as np
import time
import osdef xywh2xyxy(x):# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-righty = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left xy[:, 1] = x[:, 1] - x[:, 3] / 2  # top left yy[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right xy[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right yreturn ydef xyxy2xywh(x):# Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-righty = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x centery[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y centery[:, 2] = x[:, 2] - x[:, 0]  # widthy[:, 3] = x[:, 3] - x[:, 1]  # heightreturn ydef box_iou(box1, box2):# https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py"""Return intersection-over-union (Jaccard index) of boxes.Both sets of boxes are expected to be in (x1, y1, x2, y2) format.Arguments:box1 (Tensor[N, 4])box2 (Tensor[M, 4])Returns:iou (Tensor[N, M]): the NxM matrix containing the pairwiseIoU values for every element in boxes1 and boxes2"""def box_area(box):# box = 4xnreturn (box[2] - box[0]) * (box[3] - box[1])area1 = box_area(box1.T)area2 = box_area(box2.T)# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)return inter / (area1[:, None] + area2 - inter)  # iou = inter / (area1 + area2 - inter)def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,labels=(), max_det=300):"""Runs Non-Maximum Suppression (NMS) on inference resultsReturns:list of detections, on (n,6) tensor per image [xyxy, conf, cls]"""nc = prediction.shape[2] - 5  # number of classesxc = prediction[..., 4] > conf_thres  # candidates# Checksassert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'# Settingsmin_wh, max_wh = 2, 7680  # (pixels) minimum and maximum box width and heightmax_nms = 30000  # maximum number of boxes into torchvision.ops.nms()time_limit = 10.0  # seconds to quit afterredundant = True  # require redundant detectionsmulti_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)merge = False  # use merge-NMSt = time.time()output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]for xi, x in enumerate(prediction):  # image index, image inference# Apply constraintsx[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-heightx = x[xc[xi]]  # confidence# Cat apriori labels if autolabellingif labels and len(labels[xi]):lb = labels[xi]v = torch.zeros((len(lb), nc + 5), device=x.device)v[:, :4] = lb[:, 1:5]  # boxv[:, 4] = 1.0  # confv[range(len(lb)), lb[:, 0].long() + 5] = 1.0  # clsx = torch.cat((x, v), 0)# If none remain process next imageif not x.shape[0]:continue# Compute confx[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf# Box (center x, center y, width, height) to (x1, y1, x2, y2)box = xywh2xyxy(x[:, :4])# Detections matrix nx6 (xyxy, conf, cls)if multi_label:i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).Tx = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)else:  # best class onlyconf, j = x[:, 5:].max(1, keepdim=True)x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]# Filter by classif classes is not None:x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]# Apply finite constraint# if not torch.isfinite(x).all():#     x = x[torch.isfinite(x).all(1)]# Check shapen = x.shape[0]  # number of boxesif not n:  # no boxescontinueelif n > max_nms:  # excess boxesx = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence# Batched NMSc = x[:, 5:6] * (0 if agnostic else max_wh)  # classesboxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scoresi = torchvision.ops.nms(boxes, scores, iou_thres)  # NMSif i.shape[0] > max_det:  # limit detectionsi = i[:max_det]if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)# update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrixweights = iou * scores[None]  # box weightsx[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxesif redundant:i = i[iou.sum(1) > 1]  # require redundancyoutput[xi] = x[i]if (time.time() - t) > time_limit:break  # time limit exceededreturn outputdef scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):# Rescale coords (xyxy) from img1_shape to img0_shapeif ratio_pad is None:  # calculate from img0_shapegain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / newpad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh paddingelse:gain = ratio_pad[0][0]pad = ratio_pad[1]coords[:, [0, 2]] -= pad[0]  # x paddingcoords[:, [1, 3]] -= pad[1]  # y paddingcoords[:, :4] /= gainclip_coords(coords, img0_shape)return coordsdef clip_coords(boxes, shape):# Clip bounding xyxy bounding boxes to image shape (height, width)if isinstance(boxes, torch.Tensor):  # faster individuallyboxes[:, 0].clamp_(0, shape[1])  # x1boxes[:, 1].clamp_(0, shape[0])  # y1boxes[:, 2].clamp_(0, shape[1])  # x2boxes[:, 3].clamp_(0, shape[0])  # y2else:  # np.array (faster grouped)boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, shape[1])  # x1, x2boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0])  # y1, y2def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):# Resize and pad image while meeting stride-multiple constraintsshape = im.shape[:2]  # current shape [height, width]if isinstance(new_shape, int):new_shape = (new_shape, new_shape)# Scale ratio (new / old)r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])if not scaleup:  # only scale down, do not scale up (for better val mAP)r = min(r, 1.0)# Compute paddingratio = r, r  # width, height ratiosnew_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh paddingif auto:  # minimum rectangledw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh paddingelif scaleFill:  # stretchdw, dh = 0.0, 0.0new_unpad = (new_shape[1], new_shape[0])ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratiosdw /= 2  # divide padding into 2 sidesdh /= 2if shape[::-1] != new_unpad:  # resizeim = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))left, right = int(round(dw - 0.1)), int(round(dw + 0.1))im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add borderreturn im, ratio, (dw, dh)def preprocess(path, img_size, stride, auto):img_bgr = cv2.imread(path)  # BGRassert img_bgr is not None, f'Image Not Found {path}'# Padded resizeimg_rgb = letterbox(img_bgr, img_size, stride=stride, auto=auto)[0]# Convertimg_rgb = img_rgb.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGBimg_rgb = np.ascontiguousarray(img_rgb)return img_rgb, img_bgrclass Detect():def __init__(self, weights='yolov5s.pt'):self.device = 'cpu'self.weights =  weightsself.model = Noneself.imgsz = (640, 640)self.conf_thres=0.25self.iou_thres=0.45self.save_txt=Trueself.save_conf=Trueself.save_dir = './detect/'if not os.path.exists(self.save_dir):os.makedirs(self.save_dir) if torch.cuda.is_available() and torch.cuda.device_count() > 1:self.device = torch.device('cuda:0')self.init_model()self.stride = max(int(self.model.stride.max()), 32)def init_model(self):ckpt = torch.load(self.weights, map_location=self.device)  # loadckpt = (ckpt.get('ema') or ckpt['model']).float()  # FP32 modelfuse = Trueself.model = ckpt.fuse().eval() if fuse else ckpt.eval() # fused or un-fused model in eval modeself.model.float()def infer_image(self, image_path):im, im0 = preprocess(image_path, img_size=self.imgsz, stride=self.stride, auto=True)im = torch.from_numpy(im).to(self.device).float() / 255if len(im.shape) == 3:im = im[None]  # expand for batch dim# Inferencepred = self.model(im, augment=False, visualize=False)[0]# NMSpred = non_max_suppression(pred, self.conf_thres, self.iou_thres, None, False, max_det=1000)det  = pred[0]fn =  image_path.rsplit('/', 1)[-1]txt_path  = self.save_dir + fn.split('.', 1)[0] + '.txt'gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwhif len(det):# Rescale boxes from img_size to im0 sizedet[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()# Write resultsfor *xyxy, conf, cls in reversed(det):if self.save_txt:  # Write to filexywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywhline = (cls, *xywh, conf) if self.save_conf else (cls, *xywh)  # label formatwith open(txt_path, 'a') as f:f.write(('%g ' * len(line)).rstrip() % line + '\n')if __name__ == "__main__":detect = Detect('best.pt')detect.infer_image('test.png')

Yolov5模型加载与推理的简化代码，只需要一个文件即可，不依赖其他文件。

如果报下列错误，请下载精简版的models文件夹。下载地址：Yolov5最简推理代码-深度学习文档类资源-CSDN下载。

File "D:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py", line 853, in _load
result = unpickler.load()
ModuleNotFoundError: No module named 'models'

这个错误是因为torch保存了模型存储路径和结构，但是并没有保存计算过程。因此，需要models文件中对模型的定义。

File "D:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py", line 853, in _load
result = unpickler.load()
ModuleNotFoundError: No module named 'models'

更多三维、二维感知算法和金融量化分析算法请关注“乐乐感知学堂”微信公众号，并将持续进行更新。

Yolov5 最简推理代码相关推荐

【yolov5检测代码简化】Yolov5 detect.py推理代码简化，输入图片，输出图片和结果
前言最近的项目里有yolov5的嵌入,需求是只需要推理,模型文件是已有的,输入需要是图片(原yolov5是输入路径),输出结果的图片和标签.这样的话需要对原来的代码进行一些简化和变更. 路径模型这 ...
OpenVINO-yolov5推理代码
文章目录前言一.OpenVINO2021.3代码二.OpenVINO2022.3代码三.结果展示总结前言使用OpenVINO实现对yolov5s口罩模型推理,其他版本yolov5可以 ...
yolov5的detect.py代码详解
目标检测系列之yolov5的detect.py代码详解前言哈喽呀!今天又是小白挑战读代码啊!所写的是目标检测系列之yolov5的detect.py代码详解.yolov5代码对应的是官网v6.1版本 ...
PointPillars点云检测在OpenPCDet推理代码详解
在之前的文章中已经详细解析了PointPillars的论文和训练代码的实现和详解,可以参考之前的博客:PointPillars论文解析和OpenPCDet代码解析_NNNNNathan的博客-CSDN ...
一个极简操作系统的代码实现
一个极简操作系统的代码实现在网上看的demo OS实现时,发现一个名为Hurlex的demo OS project,实现精简,麻雀虽小五脏俱全,挺适合对OS实现进行代码级别的快速粗略了解一下的. 当 ...
win10,vs2015深度学习目标检测YOLOV5+deepsort C++多目标跟踪代码实现，源码注释，拿来即用。
打死不用CSDN,整改的太恶心了,发什么都审核不过,各种图片和链接不让发.人如果没有立场那还是人吗?不用CSDN并且博客园就很好! DeepSort纯C++ Yolov5[s,l,m系列],详细讲解- ...
Yolov5部署训练及代码解读
5.Yolov5实操训练(重点) 一.前言 1.集成的资源,包括我自己做成的成品,可以直接train与detect.需要加qq群:938162384 2.本文目的主要是能够让读者复现,直接使用,而且少 ...
这25条极简Python代码，你还不知道
作者 | 小F 来源 | 法纳斯特头图 | 下载于视觉中国自从我用Python编写第一行代码以来,就被它的简单性.出色的可读性和特别流行的一行代码所吸引. 下面,我将给大家介绍一些Python一行 ...
30 段极简 Python 代码：这些小技巧你都 Get 了么？
选自 | towardsdatascienc 编译 | 机器之心学 Python 怎样才最快,当然是实战各种小项目,只有自己去想与写,才记得住规则.本文是 30 个极简任务,初学者可以尝试着自己实现 ...

Yolov5 最简推理代码

Yolov5 最简推理代码相关推荐

最新文章

热门文章