当前位置：首页 > article >正文

YOLOv5-Seg 完全指南：从训练到后处理

article 2025/2/24 3:50:59

YOLOv5-Seg 是 Ultralytics 官方基于 YOLOv5 目标检测模型的分割分支，能够在 目标检测 的基础上进行 实例分割。相比传统的分割模型（如 Mask R-CNN），YOLOv5-Seg 具备 速度快、结构轻量、容易部署 等优势。

本教程将详细介绍 YOLOv5-Seg 的 安装、训练、推理、输出格式解析以及后处理方法，帮助你快速掌握 YOLOv5-Seg。

1. YOLOv5-Seg 简介

YOLOv5-Seg 采用 Anchor-Free 机制，输出 目标的边界框 (Bounding Box)、类别 (Class) 和分割掩码 (Segmentation Mask)。它的核心思想是在目标检测的基础上增加一个额外的分割头，从而实现 实例分割。

主要特点：

端到端实例分割：不需要额外的后处理步骤，直接输出目标的掩码。
轻量级：相比 Mask R-CNN，推理速度更快，适合实时应用。
与 YOLOv5 兼容：使用相同的数据格式和训练方式，迁移成本低。

2. YOLOv5-Seg 安装

首先，克隆 YOLOv5 仓库并安装依赖项：

# 克隆 YOLOv5 仓库
git clone https://github.com/ultralytics/yolov5.git
cd yolov5

# 安装依赖
pip install -r requirements.txt

3. 训练 YOLOv5-Seg

YOLOv5-Seg 的数据格式与 YOLOv5 目标检测类似，但需要额外提供 分割掩码 (masks)。

3.1 数据格式

YOLOv5-Seg 采用 COCO 格式或 YOLO 格式的数据集，数据组织方式如下：

dataset/
  ├── images/
  │   ├── train/
  │   ├── val/
  │   ├── test/
  ├── labels/
  │   ├── train/
  │   ├── val/
  │   ├── test/

每个 labels/*.txt 文件包含目标的信息，格式如下：

<class_id> <x_center> <y_center> <width> <height> <polygon_points>

polygon_points 是目标的 归一化分割轮廓点，用于生成分割掩码。

3.2 训练命令

python segment/train.py --weights yolov5s-seg.pt --data coco.yaml --epochs 300

其中：

yolov5s-seg.pt：预训练模型
coco.yaml：数据集配置文件
epochs：训练轮数

4. 推理（Inference）

使用训练好的模型进行推理，处理单张图片：

python segment/predict.py --weights yolov5s-seg.pt --source image.jpg

处理视频：

python segment/predict.py --weights yolov5s-seg.pt --source video.mp4

5. YOLOv5-Seg 输出格式解析

YOLOv5-Seg 的推理结果主要包含 目标类别、边界框、置信度和掩码，通常返回一个 torch.Tensor 数组，格式如下：

[array([x1, y1, x2, y2, conf, class_id, mask1, mask2, ...])]

其中：

(x1, y1, x2, y2)：目标的 边界框坐标
conf：目标的 置信度
class_id：目标的 类别 ID
mask：目标的掩码，是一个固定大小的 32x32 归一化分割掩码，需要进行反向映射恢复到原图大小。

6. YOLOv5-Seg 后处理

为了将 YOLOv5-Seg 的输出转换成可用的 二值掩码 (Binary Mask)，需要进行 插值和阈值处理。

6.1 还原掩码到原图尺寸

YOLOv5-Seg 采用 32x32 的小尺寸掩码，需要插值恢复到目标的 真实边界框尺寸。

意思是输出的mask是检测框bbox的mask，表明了检测框范围内哪些像素是分割项目。所以需要先将32*32缩放到bbox实际shape，然后再转换到全图中。

import torch
import cv2
import numpy as np

def process_mask(mask, bbox, img_shape):
    """
    处理 YOLOv5-Seg 的 32x32 掩码，恢复到原图大小
    """
    x1, y1, x2, y2 = bbox  # 边界框
    mask = mask.reshape(32, 32)  # 转换为 32x32
    mask = cv2.resize(mask, (x2 - x1, y2 - y1))  # 插值放大
    binary_mask = (mask > 0.5).astype(np.uint8)  # 二值化处理
    
    # 创建全图掩码
    full_mask = np.zeros(img_shape[:2], dtype=np.uint8)
    full_mask[y1:y2, x1:x2] = binary_mask
    
    return full_mask

6.2 叠加掩码到原图

def overlay_mask(image, mask):
    """在原图上叠加分割掩码"""
    colored_mask = np.zeros_like(image)
    colored_mask[:, :, 1] = mask * 255  # 绿色掩码
    overlayed_image = cv2.addWeighted(image, 0.7, colored_mask, 0.3, 0)
    return overlayed_image

6.3 运行完整的后处理流程

# 读取图片
image = cv2.imread("image.jpg")

# 解析 YOLOv5-Seg 结果
for det in results.pred[0]:
    x1, y1, x2, y2, conf, class_id, *mask = det.cpu().numpy()
    mask = np.array(mask)
    full_mask = process_mask(mask, (int(x1), int(y1), int(x2), int(y2)), image.shape)
    image = overlay_mask(image, full_mask)

# 显示结果
cv2.imshow("Segmented Image", image)
cv2.waitKey(0)