YOLOv5-Seg 深度解析:与 YOLOv5 检测模型的区别

YOLOv5-Seg 深度解析:与 YOLOv5 检测模型的区别

1. 概述

YOLOv5 是一个强大的目标检测框架,而 YOLOv5-Seg 是其扩展版本,增加了实例分割功能。

YOLOv5(目标检测) vs YOLOv5-Seg(实例分割)

版本 主要任务 输出结果
YOLOv5 目标检测(Object Detection) 目标类别、边界框 (Bounding Box)
YOLOv5-Seg 目标检测 + 实例分割(Instance Segmentation) 目标类别、边界框、掩码 (Mask)

2. YOLOv5-Seg 与 YOLOv5 的结构差异

2.1 网络结构对比

  • YOLOv5-Seg 在 YOLOv5 的基础上增加了 Mask 预测头,用于输出目标的掩码。
  • YOLOv5 仅有 类别预测头边界框预测头

2.2 主要结构变化

YOLOv5-Seg 中,主要修改了 model.yaml 文件,增加了一个 segmentation 相关的输出:

  • nc: 类别数,与目标检测一致。
  • mask: 32×32 的二值化掩码。
  • proto: 原始掩码特征,用于恢复高分辨率掩码。

在代码结构上,YOLOv5-Segmodels/yolo.pyforward() 中增加了额外的分割分支。

Mask 预测头的具体结构

YOLOv5-Seg 在 detect.py 代码中增加了 segmentation 分支,其核心是 mask 预测头,计算方式如下:

  • YOLOv5 目标检测头的基础上,增加了一个 ProtoNet 网络。
  • ProtoNet 负责提取 32×32 的掩码特征。
  • mask 预测头输出 32 维的 mask embedding,并与 ProtoNet 进行点积,生成最终掩码。

models/yolo.py 代码片段中,ProtoNet 计算方式如下:

python 复制代码
class ProtoNet(nn.Module):
    def __init__(self, c1, c2, k=3):
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, stride=1, padding=k//2)
    def forward(self, x):
        return self.conv(x)

ProtoNet 结构的作用是对特征图进行降维,使 mask 预测变得更加高效。

3. YOLOv5-Seg 输出格式解析

YOLOv5 目标检测模型中,输出格式如下:

python 复制代码
# YOLOv5 输出 (N, 6): [x1, y1, x2, y2, 置信度, 类别]

YOLOv5-Seg 额外输出了 maskproto,格式如下:

python 复制代码
# YOLOv5-Seg 输出 (N, 6+32): [x1, y1, x2, y2, 置信度, 类别, mask_1, mask_2, ..., mask_32]

其中:

  • mask_1 ~ mask_32 是目标的 32 维掩码向量,用于通过 proto 计算最终掩码。
  • proto 是一个 (32, H/4, W/4) 的张量,用于掩码解码。

4. 训练 YOLOv5-Seg

YOLOv5 训练类似,只需要指定 --task segment

bash 复制代码
python train.py --task segment --data coco128-seg.yaml --weights yolov5s-seg.pt --epochs 100

其中:

  • coco128-seg.yaml 是适用于分割任务的数据集配置。
  • yolov5s-seg.ptYOLOv5-Seg 预训练模型。

5. 推理与后处理

5.1 推理命令

bash 复制代码
python detect.py --task segment --weights yolov5s-seg.pt --source image.jpg

5.2 结果后处理

对于 YOLOv5-Seg,我们需要对 mask 进行解码:

python 复制代码
import torch
import cv2
import numpy as np

def process_mask(mask, proto, bbox, img_shape):
    """
    解析 YOLOv5-Seg 的 32×32 掩码,将其映射到原图
    """
    x1, y1, x2, y2 = map(int, bbox)  # 边界框坐标
    mask_decoded = (proto @ mask.T).sigmoid().cpu().numpy()  # 计算掩码
    mask_resized = cv2.resize(mask_decoded, (x2-x1, y2-y1))
    full_mask = np.zeros(img_shape[:2], dtype=np.uint8)
    full_mask[y1:y2, x1:x2] = (mask_resized > 0.5).astype(np.uint8)
    return full_mask

6. 总结

  • YOLOv5-SegYOLOv5 目标检测基础上,增加了掩码分支,实现了实例分割。
  • 输出增加了 mask,需要结合 proto 进行解码。
  • 训练方式与 YOLOv5 类似,推理时需要额外处理 mask
  • 适用于需要同时进行目标检测和实例分割的任务,如医学影像、自动驾驶等。

如果你对 YOLOv5-Seg 有任何问题,欢迎交流!🚀

相关推荐
whaosoft-14330 分钟前
51c自动驾驶~合集7
人工智能
刘晓倩4 小时前
Coze智能体开发实战-多Agent综合实战
人工智能·coze
石迹耿千秋4 小时前
迁移学习--基于torchvision中VGG16模型的实战
人工智能·pytorch·机器学习·迁移学习
路人蛃8 小时前
通过国内扣子(Coze)搭建智能体并接入discord机器人
人工智能·python·ubuntu·ai·aigc·个人开发
CV-杨帆8 小时前
论文阅读:arxiv 2025 A Survey of Large Language Model Agents for Question Answering
论文阅读·人工智能·语言模型
绝顶大聪明8 小时前
【深度学习】神经网络-part2
人工智能·深度学习·神经网络
加百力8 小时前
AI助手竞争白热化,微软Copilot面临ChatGPT的9亿下载挑战
人工智能·microsoft·copilot
Danceful_YJ9 小时前
16.使用ResNet网络进行Fashion-Mnist分类
人工智能·深度学习·神经网络·resnet
香蕉可乐荷包蛋9 小时前
AI算法之图像识别与分类
人工智能·学习·算法
张较瘦_10 小时前
[论文阅读] 人工智能 + 软件工程 | 当LLMs遇上顺序API调用:StateGen与StateEval如何破解测试难题?
论文阅读·人工智能