【yolov5改编系列之一】yolov5-6.0中detect.py检测单张图片,官方文档内容太多,看着眼花缭乱,我直接给它删掉一大堆。

引言

各位同仁大家好,网上的资源众多但是感觉用起来有点麻烦,干脆自己整理一下。

yolov5官方的detect.py文件集成度比较高。一般情况下我们不需要同时实现这么多的功能,比如:检测图像、视频、摄像头等,包括支持pytorch、tensorflow等模型框架。

本节内容是基于yolov5-6.0的detect.py文件改写的。我相信也适用于相近的yolov5版本。

本节代码功能支撑:windows系统、pytorch模型框架、检测单张图像、基于python语言

使用方法:直接用本节代码替换detect.py文件或者在官方detect.py同级目录下新建一个py文件,配置好环境后直接运行。

代码文件.py
import cv2
import numpy as np
import torch

from models.experimental import attempt_load
from utils.general import check_img_size, non_max_suppression, scale_coords
from utils.plots import Annotator, colors
from utils.augmentations import letterbox
from utils.torch_utils import select_device

img_size = 640
stride = 32
weights = 'yolov5s.pt'  # 模型权重文件路径
device = '0'  # 设置设备类型
image_path = 'data/images/zidane.jpg'  # 输入图像路径(也可以是绝对路径)
save_path = 'run1/11.jpg'  # 输出图像保存路径(也可以是绝对路径)
view_img = True  # 是否显示检测结果的图像窗口
half = False

device = select_device(device)
half &= device.type != 'cpu'  # half precision only supported on CUDA

# 导入模型
model = attempt_load(weights, map_location=device)  # 加载模型
img_size = check_img_size(img_size, s=stride)  # 检查图像尺寸是否符合规范
names = model.names  # 获取模型输出的类别名称(people、bus等)

# Padded resize
img0 = cv2.imread(image_path)  # 读取输入图像
img = letterbox(img0, img_size, stride=stride, auto=True)[0]  # 对输入图像进行填充和调整大小

# Convert
img = img.transpose((2, 0, 1))[::-1]  # 图像通道转换和颜色通道转换
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(device)
img = img.float() / 255.0   # 归一化图像数据
img = img[None]     # [h w c] -> [1 h w c]

# inference
pred = model(img)[0]  # 进行目标检测
pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45, max_det=1000)  # 非最大抑制,保留置信度最高的目标框

# plot label
det = pred[0]
annotator = Annotator(img0, line_width=3, example=str(names))  # 创建绘制标签的对象
if len(det):
    det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()  # 将检测结果的坐标转换到原始图像坐标系中
    for *xyxy, conf, cls in reversed(det):  # 遍历每个目标框
        c = int(cls)  # 目标类别
        label = f'{names[c]} {conf:.2f}'  # 标签文字内容
        annotator.box_label(xyxy, label, color=colors(c, True))  # 绘制标签框和文字

# write image
im0 = annotator.result()  # 获取带有标签的图像
cv2.imwrite(save_path, im0)  # 保存图像
if view_img:
    im0 = cv2.resize(im0, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_CUBIC) # 按比例修改需要展示的图像大小
    cv2.imshow(str('image'), im0)  # 显示图像
    cv2.waitKey(1000)  # 等待按键(1000毫秒)
print(f'Inference {image_path} finish, save to {save_path}') # 打印图片来源和输出保存路径

请添加图片描述

结语

大家一起努力,一步一个脚印,加油!!!

参考资料

链接: YOLOv5的Tricks | 【Trick13】YOLOv5的detect.py脚本的解析与简化

  • 2
    点赞
  • 4
    收藏
    觉得还不错? 一键收藏
  • 2
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值