如何将Yolov5的detect.py修改为可以直接调用的函数详解
作者:guluC 发布时间:2021-12-12 22:21:28
标签:yolov5,detect.py,函数
前几天学习了Yolov5,当我想实际将Yolov5实际运用的时候却不知道怎么办了
然后我决定对Yolov5的detect.py修改为可以直接调用的函数
因为我只需要识别图片,所以我将detect.py修改为只要传入一张图片他就可以返回坐标
ps:我这里用的是Yolov5(6.0版本)
# Copyright (c) 2022 guluC
#导入需要的库
import os
import sys
from pathlib import Path
import numpy as np
import cv2
import torch
import torch.backends.cudnn as cudnn
#初始化目录
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # 定义YOLOv5的根目录
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # 将YOLOv5的根目录添加到环境变量中(程序结束后删除)
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
from models.common import DetectMultiBackend
from utils.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr,
increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
from utils.plots import Annotator, colors, save_one_box
from utils.torch_utils import select_device, time_sync
#导入letterbox
from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterbox, mixup, random_perspective
weights=ROOT / 'yolov5s.pt' # 权重文件地址 .pt文件
source=ROOT / 'data/images' # 测试数据文件(图片或视频)的保存路径
data=ROOT / 'data/coco128.yaml' # 标签文件地址 .yaml文件
imgsz=(640, 640) # 输入图片的大小 默认640(pixels)
conf_thres=0.25 # object置信度阈值 默认0.25 用在nms中
iou_thres=0.45 # 做nms的iou阈值 默认0.45 用在nms中
max_det=1000 # 每张图片最多的目标数量 用在nms中
device='0' # 设置代码执行的设备 cuda device, i.e. 0 or 0,1,2,3 or cpu
classes=None # 在nms中是否是只保留某些特定的类 默认是None 就是所有类只要满足条件都可以保留 --class 0, or --class 0 2 3
agnostic_nms=False # 进行nms是否也除去不同类别之间的框 默认False
augment=False # 预测是否也要采用数据增强 TTA 默认False
visualize=False # 特征图可视化 默认FALSE
half=False # 是否使用半精度 Float16 推理 可以缩短推理时间 但是默认是False
dnn=False # 使用OpenCV DNN进行ONNX推理
# 获取设备
device = select_device(device)
# 载入模型
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data)
stride, names, pt, jit, onnx, engine = model.stride, model.names, model.pt, model.jit, model.onnx, model.engine
imgsz = check_img_size(imgsz, s=stride) # 检查图片尺寸
# Half
# 使用半精度 Float16 推理
half &= (pt or jit or onnx or engine) and device.type != 'cpu' # FP16 supported on limited backends with CUDA
if pt or jit:
model.model.half() if half else model.model.float()
def detect(img):
# Dataloader
# 载入数据
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
# Run inference
# 开始预测
model.warmup(imgsz=(1, 3, *imgsz), half=half) # warmup
dt, seen = [0.0, 0.0, 0.0], 0
#对图片进行处理
im0 = img
# Padded resize
im = letterbox(im0, imgsz, stride, auto=pt)[0]
# Convert
im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
im = np.ascontiguousarray(im)
t1 = time_sync()
im = torch.from_numpy(im).to(device)
im = im.half() if half 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
t2 = time_sync()
dt[0] += t2 - t1
# Inference
# 预测
pred = model(im, augment=augment, visualize=visualize)
t3 = time_sync()
dt[1] += t3 - t2
# NMS
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
dt[2] += time_sync() - t3
#用于存放结果
detections=[]
# Process predictions
for i, det in enumerate(pred): # per image 每张图片
seen += 1
# im0 = im0s.copy()
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
# Write results
# 写入结果
for *xyxy, conf, cls in reversed(det):
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4))).view(-1).tolist()
xywh = [round(x) for x in xywh]
xywh = [xywh[0] - xywh[2] // 2, xywh[1] - xywh[3] // 2, xywh[2],
xywh[3]] # 检测到目标位置,格式:(left,top,w,h)
cls = names[int(cls)]
conf = float(conf)
detections.append({'class': cls, 'conf': conf, 'position': xywh})
#输出结果
for i in detections:
print(i)
#推测的时间
LOGGER.info(f'({t3 - t2:.3f}s)')
return detections
path = 'C://Users//25096//Desktop//yoloV5//yolov5//yolov5-master//data//images//zidane.jpg'
img = cv2.imread(path)
#传入一张图片
detect(img)
我这里用的是Yolov5自带的zidane.jpg
这是输出结果
class:标签的名称
conf:置信度
position:xywh ( 左上角x,左上角y,宽,高 )
来源:https://blog.csdn.net/guluC/article/details/122683269
0
投稿
猜你喜欢
- 关于SQL Server数据库的一切信息都保存在它的系统表格里。我怀疑你是否花过比较多的时间来检查系统表格,因为你总是忙于用户表格。但是,你
- 简介今天试着用ptyhon做了一个抓取网页内容,并生成word文档的功能,功能很简单,做一下记录以备以后用到。生成word用到了第三方组件p
- MySQL使用于认证目的的用户名,与Unix用户名(登录名字)或Windows用户名无关。缺省地,大多数MySQL客户尝试使用当前Unix用
- 导言如我们在之前的教程里讨论的那样,分页可以通过两种方法来实现:1.默认分页– 你仅仅只用选中data Web control的 智能标签的
- 这篇文章主要介绍了python线程信号量semaphore使用解析,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价
- 网上的关于django-scrapy的介绍比较少,该博客只在本人查资料的过程中学习的,如果不对之处,希望指出改正;以后的博客可能不会再出关于
- 下面是代码,如果看不懂,建议先把表格的一些<tr><td>的表格原理弄清楚了,就可以了代码如下:<table&
- 本文进通过实例,讲述通过 prototype 自定义方法的过程,旨在抛砖引玉,如果不对的地方,欢迎指正!prototype 是在 IE 4
- 用语言实现 好处: 1、可以减少对数据库的访问。 2、可移植性好。 坏处: 1、操作起来考虑的东西较多,修改一处就要修改别一处。也就是说是相
- matplotlib及相关cmap参数的取值在matplotlib中对于图片的显示有如下方法(这不是重点), 其中有cmap=&ls
- EcmaScript正則表達式( 深入淺出系列之淺出 ^_^ )在线正则表达式测试:http://www.aspxhome.com/RegE
- Python元字典字典(dictionary)是除列表以外python之中最灵活的内置数据结构类型。列表是有序的对象结合,字典是无序的对象集
- 本文实例讲述了python flask框架实现重定向功能。分享给大家供大家参考,具体如下:flask 重定向:from flask impo
- 好几年前我在做一些自动化的脚本时,脑子里也闪过这样的想法:能不能直接把浏览器的cookies取出来用呢?直到昨天看到代码《python模拟发
- 1、requests 的常见用法requests 除了 url 之外,还有 params, data 和 files 三个参数,用于和服务器
- MFCC梅尔倒谱系数(Mel-scaleFrequency Cepstral Coefficients,简称MFCC)。MFCC通常有以下之
- 如下所示:import osDIR = "/home/serho/workspace/lisp"def compare(
- SSI是英文Server Side Includes的缩写,翻译成中文就是服务器端包含的意思。从技术角度上说,SSI就是在HTML文件中,可
- 1、subprocess这个模块来产生子进程,并且可以连接到子进程的标准输入、输出、错误中,还可以获得子进程的返回值。2、subproces
- Image and text elements that appear in another element are called floa