【yolov5改编系列之一】yolov5-6.0中detect.py检测单张图片，官方文档内容太多，看着眼花缭乱，我直接给它删掉一大堆。_yolo5-6.0 petect.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}') # 打印图片来源和输出保存路径
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

结语

大家一起努力，一步一个脚印，加油！！！

参考资料

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