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

yolov5-seg分割后处理流程

article 2025/2/26 9:34:13

在 YOLOv5s-seg 模型中，分割后处理的目的是将模型的原始输出（包括检测框、类别概率和掩码系数）转换为最终的实例分割结果。以下是使用 Python 实现分割后处理的完整代码示例：

1. 分割后处理步骤

解析模型输出：
- 从模型输出中提取检测框、类别概率和掩码系数。
过滤低置信度的检测框：
- 使用置信度阈值过滤掉低置信度的检测框。
非极大值抑制（NMS）：
- 去除重叠的检测框，保留最优的检测结果。
生成实例分割掩码：
- 使用掩码系数和掩码原型生成每个检测框的分割掩码。
调整掩码分辨率并裁剪：
- 将掩码调整到检测框的实际大小，并裁剪到检测框范围。

2. 完整代码实现

import numpy as np
import torch
import cv2

def seg_postprocess(output, mask_prototypes, conf_thres=0.25, iou_thres=0.45, img_size=640):
    """
    YOLOv5s-seg 分割后处理函数。

    参数:
        output: 模型原始输出，形状为 (1, 25200, 117)。
        mask_prototypes: 掩码原型，形状为 (32, 160, 160)。
        conf_thres: 置信度阈值。
        iou_thres: NMS 的 IoU 阈值。
        img_size: 输入图像的尺寸。

    返回:
        boxes: 检测框，形状为 (N, 4)，格式为 (x1, y1, x2, y2)。
        scores: 置信度分数，形状为 (N,)。
        labels: 类别标签，形状为 (N,)。
        masks: 实例分割掩码，形状为 (N, H, W)。
    """
    # 1. 解析模型输出
    preds = output[0]  # (25200, 117)
    boxes = preds[:, :4]  # 检测框 (x_center, y_center, width, height)
    scores = preds[:, 4]  # 置信度分数
    class_probs = preds[:, 5:85]  # 类别概率 (80 个类别)
    mask_coefficients = preds[:, 85:]  # 掩码系数 (32 个)

    # 2. 过滤低置信度的检测框
    keep = scores > conf_thres
    boxes = boxes[keep]
    scores = scores[keep]
    class_probs = class_probs[keep]
    mask_coefficients = mask_coefficients[keep]

    # 3. 非极大值抑制 (NMS)
    keep = nms(boxes, scores, iou_thres)
    boxes = boxes[keep]
    scores = scores[keep]
    class_probs = class_probs[keep]
    mask_coefficients = mask_coefficients[keep]

    # 4. 生成实例分割掩码
    masks = generate_masks(mask_coefficients, mask_prototypes, boxes, img_size)

    # 5. 返回结果
    labels = np.argmax(class_probs, axis=1)  # 类别标签
    return boxes, scores, labels, masks


def nms(boxes, scores, iou_thres):
    """
    非极大值抑制 (NMS)。

    参数:
        boxes: 检测框，形状为 (N, 4)，格式为 (x1, y1, x2, y2)。
        scores: 置信度分数，形状为 (N,)。
        iou_thres: IoU 阈值。

    返回:
        keep: 保留的检测框索引。
    """
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    areas = (x2 - x1) * (y2 - y1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0, xx2 - xx1)
        h = np.maximum(0, yy2 - yy1)
        inter = w * h
        iou = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(iou <= iou_thres)[0]
        order = order[inds + 1]

    return keep


def generate_masks(mask_coefficients, mask_prototypes, boxes, img_size):
    """
    生成实例分割掩码。

    参数:
        mask_coefficients: 掩码系数，形状为 (N, 32)。
        mask_prototypes: 掩码原型，形状为 (32, 160, 160)。
        boxes: 检测框，形状为 (N, 4)，格式为 (x1, y1, x2, y2)。
        img_size: 输入图像的尺寸。

    返回:
        masks: 实例分割掩码，形状为 (N, H, W)。
    """
    masks = []
    for i in range(len(boxes)):
        # 1. 加权组合掩码原型
        mask = np.sum(mask_coefficients[i][:, None, None] * mask_prototypes, axis=0)  # (160, 160)

        # 2. 调整掩码分辨率到检测框大小
        x1, y1, x2, y2 = boxes[i]
        mask = cv2.resize(mask, (int(x2 - x1), int(y2 - y1)), interpolation=cv2.INTER_LINEAR)

        # 3. 将掩码裁剪到检测框范围
        full_mask = np.zeros((img_size, img_size))
        full_mask[int(y1):int(y2), int(x1):int(x2)] = mask
        masks.append(full_mask)

    return np.array(masks)


# 示例用法
if __name__ == "__main__":
    # 假设模型输出和掩码原型已经加载
    output = np.random.randn(1, 25200, 117)  # 模型原始输出
    mask_prototypes = np.random.randn(32, 160, 160)  # 掩码原型

    # 执行分割后处理
    boxes, scores, labels, masks = seg_postprocess(output, mask_prototypes)

    # 打印结果
    print("检测框:", boxes)
    print("置信度分数:", scores)
    print("类别标签:", labels)
    print("实例分割掩码形状:", masks.shape)