train.py文件中网络结构部分
YOLOv5网络结构包括Input、Backbone、Neck、Head四部分，其网络结构图如图所示。Input用
于对输入的图片进行预处理。Backbone用于对处理后的图片提取特征。Neck模块用于对特征图
进行混合或者组合，形成一些更为复杂的特征。Head模块用来预测图像的特征，生成边界框并预
测类别。
train.py中加载模型的代码部分
# ============================================== # 步骤1：判断是否使用【预训练权重】 # ============================================== # 判断权重文件是否以 .pt 结尾（PyTorch 模型文件格式） # 如果是，代表使用预训练权重；否则代表从零开始训练 pretrained = weights.endswith('.pt') # ============================================== # 步骤2：如果有预训练权重 → 加载权重 # ============================================== if pretrained: # 多GPU训练时的同步操作：保证只有主进程先下载权重，其他等待 with torch_distributed_zero_first(rank): attempt_download(weights) # 如果权重文件不存在，自动从云端下载 # 加载权重文件（.pt 文件）到当前设备 ckpt = torch.load(weights, map_location=device) # 创建 YOLO 模型结构 # 优先使用命令行 --cfg 指定的模型配置，否则使用权重里自带的配置 # ch=3 表示输入是 RGB 3 通道图像 # nc=nc 表示数据集类别数量 # anchors= 传入锚框（如果自动计算过就用新的） model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # 定义加载权重时要【排除】的键 # 如果是自定义模型/自定义锚框，不加载原始权重里的 anchor 参数 exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else [] # 把权重转成 FP32 高精度格式 state_dict = ckpt['model'].float().state_dict() # 筛选权重：只保留模型中存在的键（解决类别数不同导致的不匹配） state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # 把筛选后的权重加载进模型 # strict=False：允许缺少/多余的键，不报错 model.load_state_dict(state_dict, strict=False) # 打印日志：成功加载了多少个参数 logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights)) # ============================================== # 步骤3：如果没有预训练权重 → 从零创建模型 # ============================================== else: # 直接创建一个全新的、随机初始化的模型 model = Model(opt.cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # ============================================== # 步骤4：检查数据集是否正确（路径、标签格式） # ============================================== with torch_distributed_zero_first(rank): check_dataset(data_dict) # 检查数据集路径、图片、标签是否正常
核心代码
一般情况下，我们都是使用预训练模型的，这样网络才会更好的收敛。
model = Model( opt.cfg or ckpt['model'].yaml, # 第 1 个参数:模型结构配置 ch=3, # 第 2 个参数:输入通道数 nc=nc, # 第 3 个参数:类别数 anchors=hyp.get('anchors') # 第 4 个参数:anchor 配置 ).to(device) # 把模型转移到指定设备
参数部分
opt.cfg:用户通过 --cfg models/yolov5s.yaml 显式指定的配置路径,通常用于从零训练或改变结
构后重新训练
ckpt['model'].yaml:加载预训练权重时,checkpoint 文件里保存的原始模型结构配置,通常用于断
点续训或微调
Model 是 YOLOv5 官方写在 models/yolo.py 文件里的一个 “类”，专门用来创建 YOLO 检测模型
。from models.yolo import Model
yolo文件中的代码
Detect函数部分
- 输入图像：640×640
- 模型输出 3 个检测层：
- 80×80（小目标）
- 40×40（中目标）
- 20×20（大目标）
- 每个层 3 个锚框
- 类别数 nc = 80（COCO 数据集）
输入 = 3 个特征图（list 格式）
来自网络的上一层输出，是3 个不同大小的特征图。
x = [ # 第 1 个检测层：80x80 torch.Size([1, 256, 80, 80]), # 第 2 个检测层：40x40 torch.Size([1, 512, 40, 40]), # 第 3 个检测层：20x20 torch.Size([1, 1024, 20, 20]) ]
【输出】分两种：训练输出 / 推理输出
【情况 A：训练模式】输出
[ torch.Size([1, 3, 80, 80, 85]), torch.Size([1, 3, 40, 40, 85]), torch.Size([1, 3, 20, 20, 85]) ]
[1, 3, 80, 80, 85] 1：batch 3：每个格子 3 个锚框 80：特征图高 80：特征图宽 85：5 + 80（x,y,w,h,置信度 + 80类）
【情况 B：推理模式】输出（你看到的框）
torch.Size([1, 25200, 85])
1：一张图 25200：所有层的框总数（80×80×3 + 40×40×3 + 20×20×3） 85：x, y, w, h, confidence, class0, class1, ..., class79
编辑
Model 类初始化
parse_model
def parse_model(d, ch): # model_dict, input_channels(3)
d:从 YAML 解析得到的字典(就是上一题中的data_dict类似结构,但这里是模型配置)ch:输入通道数列表,初始为[3](RGB 三通道)
anchors = [[10,13, 16,30, 33,23], [30,61, 62,45, 59,119], [116,90, 156,198, 373,326]] nc = 80 gd = 0.33 # depth_multiple gw = 0.50 # width_multiple # 计算 anchor 数量 # anchors[0] = [10,13, 16,30, 33,23] 共 6 个数 → 3 对 (w,h) → na = 3 na = len(anchors[0]) // 2 = 6 // 2 = 3 # 计算每个 anchor 的输出维度 # no = 3 × (80 + 5) = 255 no = na * (nc + 5) = 3 * 85 = 255
初始化容器
layers, save, c2 = [], [], ch[-1]
layers:存放所有构建好的 PyTorch 模块save:存放需要"保存输出"的层索引(供后续跨层连接使用)c2:当前层的输出通道数,初始为ch[-1]= 3
我们取 yaml 里的真实一行（以第一层为例）
这一行的意思：
-1：输入来自上一层1：重复 1 次Conv：卷积层[64,6,2,2]：输出通道 64，卷积核 6x6，步长 2，padding2
# yolov5s.yaml 片段 backbone: # [from, repeats, module, args] - [-1, 1, Conv, [64, 6, 2, 2]] # 第0层
i, (f, n, m, args) = (0, -1, 1, 'Conv', [64,6,2,2])
m = eval('Conv') # 现在 m = 你定义的 Conv 类
n = round(1 * 0.33) = 1
c1 = ch[f] # f=-1 → ch[-1] = 3（输入通道） c2 = args[0] = 64 # 输出通道 # 宽度因子缩放：64 * 0.5 = 32 c2 = make_divisible(64 * 0.5, 8) = 32 # 最终参数变成： args = [3, 32, 6, 2, 2]
m_ = Conv(3, 32, 6, 2, 2)
m_.i = 0 # 第0层
m_.f = -1 # 输入来自上一层
ch.append(32) # 现在 ch = [3, 32]，下一层输入通道就是32
layers.append(Conv(3,32,6,2,2))
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, MblNetV3_Blk, CBAM, ConvBNHswish, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3, C3TR, SELayer, ASPP, simam_module, CoordAtt]:
from models.common import * 导入了所有基础网络层（Conv、C3、SPP、Bottleneck 等）
from models.experimental import * 导入了实验性层
所以 parse_model 才能直接识别并创建这些层
yolo.yaml文件
# anchors anchors: - [10,13, 16,30, 33,23] # P3/8 - [30,61, 62,45, 59,119] # P4/16 - [116,90, 156,198, 373,326] # P5/32
[10,13, 16,30, 33,23] ↓ ↓ ↓ ↓ ↓ ↓ w h w h w h 第1个 第2个 第3个 anchor anchor anchor
| 尺度 | 下采样倍数 | 特征图大小(640输入) | 适合检测 | Anchor 大小 |
| P3/8 | 8 倍 | 80 × 80 | 小目标 | 10~33 像素 |
| P4/16 | 16 倍 | 40 × 40 | 中目标 | 30~119 像素 |
| P5/32 | 32 倍 | 20 × 20 | 大目标 | 116~373 像素 |
# YOLOv5 backbone backbone: # [from, number, module, args] [[-1, 1, Focus, [64, 3]], # 0-P1/2 ← 320×320 [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 ← 160×160 [-1, 3, C3, [128]], [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 ← 80×80 ⭐ 小目标 [-1, 9, C3, [256]], [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 ← 40×40 ⭐ 中目标 [-1, 9, C3, [512]], [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 ← 20×20 ⭐ 大目标 [-1, 1, SPP, [1024, [5, 9, 13]]], [-1, 3, C3, [1024, False]]] # 9
第 0 层采用Focus模块，以 3×3 卷积核将特征图下采样至320×320×64，实现无信息损失的空间下采样与通道扩充；
第 1 层通过Conv卷积模块（128 通道、3×3 卷积核、步长 2）将特征图尺寸压缩至
160×160×128；
第 2 层堆叠3 个 C3模块，在160×160×128尺度完成残差特征融合与强化；
第 3 层通过Conv模块（256 通道、步长 2）下采样至80×80×256，形成小目标检测基础特征层；
第 4 层堆叠9 个 C3模块，在80×80×256尺度完成深层特征提取；
第 5 层经Conv模块（512 通道、步长 2）下采样至40×40×512，构成中目标特征层；
第 6 层再次堆叠9 个 C3模块，在40×40×512尺度进一步增强特征表达能力；
第 7 层通过Conv模块（1024 通道、步长 2）下采样至20×20×1024，形成大目标高语义特征层；
第 8 层接入SPP 空间金字塔池化模块，采用 5×5、9×9、13×13 池化核在20×20×1024尺度融合多
感受野全局特征；
第 9 层通过3 个 C3模块（无 shortcut）完成最终特征整合，输出20×20×1024的深层语义特征，为
后续多尺度检测提供强表征能力的特征支撑。
head: [[-1, 1, Conv, [512, 1, 1]], [-1, 1, nn.Upsample, [None, 2, 'nearest']], [[-1, 6], 1, Concat, [1]], # cat backbone P4 [-1, 3, C3, [512, False]], # 13 [-1, 1, Conv, [256, 1, 1]], [-1, 1, nn.Upsample, [None, 2, 'nearest']], [[-1, 4], 1, Concat, [1]], # cat backbone P3 [-1, 3, C3, [256, False]], # 17 (P3/8-small) [-1, 1, Conv, [256, 3, 2]], [[-1, 14], 1, Concat, [1]], # cat head P4 [-1, 3, C3, [512, False]], # 20 (P4/16-medium) [-1, 1, Conv, [512, 3, 2]], [[-1, 10], 1, Concat, [1]], # cat head P5 [-1, 3, C3, [1024, False]], # 23 (P5/32-large) [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) ]
该 YOLOv5 网络Head（检测头）基于 PANet 结构构建，负责多尺度特征融合、上采样、下采样
与目标检测，整体分为自上而下上采样融合和自下而下下采样融合两个阶段，最终输出三个尺度的
检测特征图并接入 Detect 层完成预测。
第10层采用Conv模块（512通道、1×1卷积核），将输入的20×20×1024特征图压缩为20×20×512，实现通道降维；
第11层通过nn.Upsample上采样模块，将特征图放大2倍，尺寸变为40×40×512；
第12层将上采样特征与主干网络第6层的40×40×512特征进行Concat拼接，输出
40×40×1024融合特征；
第13层堆叠3个C3模块（无shortcut），在40×40×512尺度完成特征精炼；
第14层使用Conv模块（256通道、1×1卷积），将特征压缩为40×40×256；
第15层通过上采样放大2倍，得到80×80×256特征图；
第16层将其与主干网络第4层80×80×256特征Concat拼接，输出80×80×512；
第17层堆叠3个C3模块，输出80×80×256特征，对应P3小目标检测分支；
第18层采用**Conv**模块（256通道、3×3卷积、步长2）下采样，得到40×40×256；
第19层与第14层40×40×512特征Concat拼接，输出40×40×768；
第20层堆叠3个C3模块，输出40×40×512特征，对应P4中目标检测分支；
第21层通过Conv模块（512通道、3×3卷积、步长2）下采样至20×20×512；
第22层与第10层20×20×512特征Concat拼接，输出20×20×1024；
第23层堆叠3个C3模块，输出20×20×1024特征，对应P5大目标检测分支；
第24层为Detect检测层，接收P3(80×80×256)、P4(40×40×512)、P5(20×20×1024)三个尺度特
征，完成目标框坐标、置信度与类别的预测输出。
common.py
关于Concat层的定义
class Concat(nn.Module): # Concatenate a list of tensors along dimension def __init__(self, dimension=1): super(Concat, self).__init__() self.d = dimension def forward(self, x): return torch.cat(x, self.d)
该 Concat 模块接收多个同高宽、不同通道的特征图作为输入，在通道维度（dim=1）进行拼接融
合，输出通道数叠加、空间尺寸保持不变的融合特征图，用于增强多尺度上下文信息与特征表达能
力。
编辑
编辑
打标签部分
这张图展示了YOLO 目标检测任务里，如何给图片里的物体打标签（写进 txt 文件），是训练模
型前最基础的一步。 特征是这张图片，标签是：
16 0.412 0.530 0.234 0.661 # 狗 2 0.812 0.620 0.180 0.290 # 车
| 数字 | 含义 | 大白话翻译 |
16 |
class_id | 这个物体是第 0 类，也就是 “狗” |
| 0.412 | x_center | 狗的中心，在图片宽度方向的相对位置 |
| 0.530 | y_center | 狗的中心，在图片高度方向的相对位置 |
| 0.234 | width | 狗的框，占图片总宽度的比例 |
| 0.661 | height | 狗的框，占图片总高度的比例 |
编辑
train.py文件中数据集读取部分
数据配置文件读取
这段代码的功能是读取一个 YAML 格式的配置文件,并把它解析成 Python 字典,以便后续训练
脚本能够访问其中的配置项(比如训练集路径、类别数、类别名称等)。
# 打开由命令行参数 --data 指定的 YAML 配置文件 # opt.data 通常是一个字符串路径,例如 'data/coco128.yaml' # 使用 with 上下文管理器,可以在读取完毕后自动关闭文件,避免资源泄漏 with open(opt.data) as f: # 使用 PyYAML 库的 load 函数,把文件内容解析成 Python 字典 # 参数 f:刚刚打开的文件对象,作为输入流 # 参数 Loader=yaml.SafeLoader:指定安全加载器 # - SafeLoader 只解析 YAML 标准的基本数据类型(字典、列表、字符串、数字等) # - 它会拒绝执行 YAML 中可能包含的任意 Python 代码,防止恶意代码注入 # - 相比之下,默认的 yaml.Loader 存在安全风险,官方已不推荐使用 # 解析结果赋值给 data_dict,这是一个 Python 字典对象 data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict
以 YOLOv5 自带的 coco128.yaml 为例:
# 训练集和验证集的路径 train: ../datasets/coco128/images/train2017/ val: ../datasets/coco128/images/train2017/ # 类别数量 nc: 80 # 类别名称列表(索引与 class_id 对应) names: ['person', 'bicycle', 'car', 'motorcycle', ...]
经过解析,上面的文件会变成这样一个 Python 字典:
data_dict = { 'train': '../datasets/coco128/images/train2017/', 'val': '../datasets/coco128/images/train2017/', 'nc': 80, 'names': ['person', 'bicycle', 'car', 'motorcycle', ...] }
获取数据集路径
train_path = data_dict['train'] # 训练集路径 test_path = data_dict['val'] # 验证集路径
创建训练数据加载器
# ===================== 核心功能 ===================== # 创建【训练集】的数据加载器(DataLoader)和数据集(Dataset)对象 # 作用：批量、多线程加载训练图片+标签，是模型训练的"数据搬运工" # ===================== dataloader, dataset = create_dataloader( train_path, # 1. 训练集图片/标签的根路径 imgsz, # 2. 模型输入的统一图像尺寸 (如 640) batch_size, # 3. 批次大小：每次喂给模型多少张图片 gs, # 4. 网格步长(Grid Stride)：YOLO下采样倍数(8/16/32)，保证尺寸整除 opt, # 5. 命令行参数对象(所有运行参数都在这里) hyp=hyp, # 6. 训练超参数字典(学习率、损失函数系数等) augment=True, # 7. 开启训练数据增强(翻转、缩放、色域变换等，提升泛化能力) cache=opt.cache_images, # 8. 是否缓存图片到内存/磁盘：加速训练(避免重复读取硬盘) rect=opt.rect, # 9. 矩形训练模式：按图片原始比例缩放，节省显存 rank=rank, # 10. 多GPU训练时：当前进程的GPU编号 world_size=opt.world_size, # 11. 多GPU训练时：总GPU数量 workers=opt.workers, # 12. 数据加载线程数：多线程并行读图片，速度更快 image_weights=opt.image_weights, # 13. 样本权重：给困难样本/小目标加权训练 quad=opt.quad, # 14. 四边形数据加载：适配超大批次训练的加速选项 prefix=colorstr('train: ') # 15. 日志打印前缀：终端输出带颜色的"train: "标识 )
| 返回值 | 含义 | 作用 |
dataset |
数据集对象 | 管理所有训练图片、标签路径，负责读取单张图 + 解析标签 |
dataloader |
数据加载器 | 核心搬运工：批量、多线程、乱序把数据递给模型训练 |
然后这个函数是从utils/datasets.py导入进去的
数据预处理与标签处理
# 仅在【不恢复训练】时执行（从头训练才走这里，断点续训跳过） if not opt.resume: # 1. 拼接训练集中所有图片的标签数据 (形状: [总目标数, 5])，5列=类别+中心x+中心y+宽+高 labels = np.concatenate(dataset.labels, 0) # 2. 提取所有目标的【类别ID】，转为PyTorch张量 c = torch.tensor(labels[:, 0]) # classes # 以下两行是注释掉的代码：根据类别频率初始化模型偏置，官方默认关闭 # cf = torch.bincount(c.long(), minlength=nc) + 1. # 统计每个类别的目标数量 # model._initialize_biases(cf.to(device)) # 3. 如果开启可视化(plots=True)，绘制标签分布图（类别、框大小、位置） if plots: plot_labels(labels, names, save_dir, loggers) # 如果启用TensorBoard，将类别分布写入日志，方便可视化查看 if tb_writer: tb_writer.add_histogram('classes', c, 0) # 4. 锚框(Anchor)优化：核心功能！自动计算适配数据集的最优锚框 # 如果没有关闭自动锚框（opt.noautoanchor=False），就执行检查/计算锚框 if not opt.noautoanchor: check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) # 5. 模型精度预处理：先转半精度再转回单精度，优化锚框计算精度（不影响训练） model.half().float()
我们设定：
图片 1：有 2 个物体（狗 = 0，车 = 1）→ 标签 2 行
图片 2：有 1 个物体（人 = 2）→ 标签 1 行
总目标数：3 个
labels = np.concatenate(dataset.labels, 0) # 合并所有标签
原始标签文件
img1.txt（一张图两个物体）
0 0.412 0.530 0.234 0.661 # 狗 1 0.812 0.620 0.180 0.290 # 车
img2.txt
2 0.500 0.500 0.100 0.200 # 人
# 🔥 第一步：数据集存储的标签（关键！） # dataset.labels 是列表： # 第0个元素 = 图片1的所有目标（2行，2个物体） # 第1个元素 = 图片2的所有目标（1行，1个物体） dataset.labels = [ # 图片1：2个物体 → 形状 (2, 5) np.array([ [0, 0.412, 0.530, 0.234, 0.661], [1, 0.812, 0.620, 0.180, 0.290] ]), # 图片2：1个物体 → 形状 (1, 5) np.array([ [2, 0.500, 0.500, 0.100, 0.200] ]) ]
labels = [ [0, 0.412, 0.530, 0.234, 0.661], # 图1-狗 [1, 0.812, 0.620, 0.180, 0.290], # 图1-车 [2, 0.500, 0.500, 0.100, 0.200] # 图2-人 ] # 最终形状：(3, 5) → 3个物体，5列参数
c = tensor([0, 1, 2]) # 3个物体的类别：狗、车、人
c = torch.tensor(labels[:, 0]) # 类别索引
网络最后的输出
import torch from models.yolo import Model model = Model('models/yolov5m.yaml') model.eval() x = torch.randn(2, 3, 640, 640) # 推理模式 with torch.no_grad(): out = model(x) print(type(out)) # tuple print(out[0].shape) # torch.Size([2, 25200, 85]) print(len(out[1])) # 3 print([o.shape for o in out[1]]) # [torch.Size([2, 3, 80, 80, 85]), # torch.Size([2, 3, 40, 40, 85]), # torch.Size([2, 3, 20, 20, 85])] # 训练模式 model.train() out = model(x) print(type(out)) # list print([o.shape for o in out]) # [torch.Size([2, 3, 80, 80, 85]), # torch.Size([2, 3, 40, 40, 85]), # torch.Size([2, 3, 20, 20, 85])]
return _VF.meshgrid(tensors, **kwargs) # type: ignore[attr-defined] <class 'tuple'> torch.Size([2, 25200, 85]) 3 [torch.Size([2, 3, 80, 80, 85]), torch.Size([2, 3, 40, 40, 85]), torch.Size([2, 3, 20, 20, 85])] <class 'list'> [torch.Size([2, 3, 80, 80, 85]), torch.Size([2, 3, 40, 40, 85]), torch.Size([2, 3, 20, 20, 85])]
[2, 3, 80, 80, 85]
小目标分支，负责检测小物体
[2, 3, 40, 40, 85]
中目标分支，负责检测中等物体
[2, 3, 20, 20, 85]
大目标分支，负责检测大物体
编辑
推理阶段
获取类别名称和颜色
names = model.names colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
# 假设 names 有 3 个类别 names = ['cat', 'dog', 'bird'] # 生成的 colors 可能是： colors = [ [123, 45, 67], # cat 的颜色：RGB(123, 45, 67) [255, 100, 200], # dog 的颜色：RGB(255, 100, 200) [50, 200, 150] # bird 的颜色：RGB(50, 200, 150) ]
图像预处理（直接调用项目中的 letterbox 函数）
输入图像 (720×1280×3, uint8, BGR) │ ▼ ┌──────────────────────────────────┐ │ letterbox(img0, new_shape=640) │ └──────────────────┬───────────────┘ │ ▼ (640×640×3, uint8, BGR) │ ▼ ┌──────────────────────────────────┐ │ [:, :, ::-1] → BGR→RGB │ │ .transpose(2, 0, 1) → HWC→CHW │ └──────────────────┬───────────────┘ │ ▼ (3×640×640, uint8, RGB) │ ▼ ┌──────────────────────────────────┐ │ torch.from_numpy() → tensor │ │ .float() / 255 → 归一化 │ │ .to(device) → 移动到设备 │ └──────────────────┬───────────────┘ │ ▼ (3×640×640, float32, RGB, 0~1) │ ▼ ┌──────────────────────────────────┐ │ .unsqueeze(0) → 添加批次维度 │ └──────────────────┬───────────────┘ │ ▼ 最终输出: (1×3×640×640, float32, RGB, 0~1, device)
推理
print("正在推理...") with torch.no_grad(): pred = model(img)[0] # 直接获取检测结果
output = model(img) # output = (detections, raw_features) # detections: 解码后的检测框，形状 [batch, 25200, 85] # raw_features: 原始特征图列表，包含3个张量
pred = model(img)[0] # 获取元组的第一个元素：detections
print(pred.shape) # torch.Size([1, 25200, 85])
| 维度 | 含义 | 值 | 说明 |
| 第0维 | 批次大小 | 1 | 我们只输入了一张图片 |
| 第1维 | 检测框总数 | 25200 | 3个检测层 × 8400个网格点 |
| 第2维 | 每个检测框的输出 | 85 | 5个坐标 + 80个类别概率 |
3个检测层： - 第一层：80×80 = 6400 个网格 - 第二层：40×40 = 1600 个网格 - 第三层：20×20 = 400 个网格 每个网格有3个锚点，所以： (6400 + 1600 + 400) × 3 = 8400 × 3 = 25200
假设输入一张包含"人"和"车"的图片
图像尺寸：640×640
内容：一个人站在车旁边
pred[0] = [ # 检测框 0：检测到一个人 [0.5, # x_center: 图像水平中心位置 0.3, # y_center: 图像上部（30%高度处） 0.2, # width: 宽度占图像的20% 0.4, # height: 高度占图像的40% 0.95, # conf: 95%置信度 1.0, # cls_0: person类别概率100% 0.0, # cls_1: bicycle类别概率0% 0.0, # cls_2: car类别概率0% ...], # 其他77个类别概率都是0% # 检测框 1：检测到一辆车 [0.7, # x_center: 图像右侧（70%宽度处） 0.6, # y_center: 图像中部偏下（60%高度处） 0.3, # width: 宽度占图像的30% 0.5, # height: 高度占图像的50% 0.88, # conf: 88%置信度 0.0, # cls_0: person类别概率0% 0.0, # cls_1: bicycle类别概率0% 1.0, # cls_2: car类别概率100% ...], # 其他类别概率都是0% # 检测框 2~25199：背景或低置信度检测 [0.1, 0.1, 0.05, 0.05, 0.01, 0.0, 0.0, ...], # 置信度只有1% ... ]
我们可以把它想象成一个25200*85 的矩阵，每一行代表一个检测框对应的85维度的向量
NMS过滤
print("正在进行NMS过滤...") pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)
| 参数 | 含义 | 值 |
pred |
输入检测结果 | [1, 25200, 85] |
conf_thres |
置信度阈值 | 0.25 |
iou_thres |
IOU阈值 | 0.45 |
pred_output = [ tensor([[256, 64, 384, 320, 0.95, 0]]), # person tensor([[352, 256, 544, 576, 0.88, 2]]) # car ]
| 检测框 | x1 | y1 | x2 | y2 | conf | cls | 含义 |
| 0 | 256 | 64 | 384 | 320 | 0.95 | 0 | person（类别0） |
| 1 | 352 | 256 | 544 | 576 | 0.88 | 2 | car（类别2） |
处理检测结果
det = pred[0]
det = tensor([ [256, 64, 384, 320, 0.95, 0], # person [352, 256, 544, 576, 0.88, 2] # car ]) det.shape = (2, 6) # 2个检测框，每个6个维度
转换坐标
# 转换后 det = tensor([ [512, 0, 768, 360, 0.95, 0], # 原始图像坐标 [704, 232, 1088, 872, 0.88, 2] ])
每个检测框的6个值 索引 名称 含义 示例值 0 x1 检测框左上角的 x 坐标（像素） 512 1 y1 检测框左上角的 y 坐标（像素） 0 2 x2 检测框右下角的 x 坐标（像素） 768 3 y2 检测框右下角的 y 坐标（像素） 360 4 conf 目标存在的 置信度（0~1） 0.95 5 cls 目标的 类别索引（0~79） 0
图像左上角为原点 (0, 0) ┌─────────────────────────────────────────────────────────────────────┐ │ (0,0) ← x轴方向 (向右) │ │ ↓ │ │ y轴方向 (向下) │ │ │ │ 检测框0 │ │ ┌─────────────────────────────┐ │ │ │ (512,0) │ │ │ │ person │ │ │ │ 0.95 │ │ │ └─────────────────────────────┘ (768,360) │ │ │ │ 检测框1 │ │ ┌─────────────────────────────┐ │ │ │ (704,232) │ │ │ │ car │ │ │ │ 0.88 │ │ │ └─────────────────────────────┘ (1088,872) │ │ │ └─────────────────────────────────────────────────────────────────────┘ 原始图像尺寸：1280 × 720 像素
names = model.names # ['person', 'bicycle', 'car', ...] # 检测框0 cls_idx = int(det[0, 5]) # 0 class_name = names[cls_idx] # 'person' # 检测框1 cls_idx = int(det[1, 5]) # 2 class_name = names[cls_idx] # 'car'
最后的检测效果
原始图像 (1280×720) ┌─────────────────────────────────────────────────────────────────────┐ │ │ │ ┌─────────────────────────────────────────────────────────────┐ │ │ │ person 0.95 │ │ │ └─────────────────────────────────────────────────────────────┘ │ │ │ │ ┌─────────────────────────────────────────┐ │ │ │ car 0.88 │ │ │ └─────────────────────────────────────────┘ │ │ │ └─────────────────────────────────────────────────────────────────────┘
