YOLO系列后处理 自定义推理颜色

时间：2025-04-28 19:26:03

YOLOv5的predict文件中的

#Mask plotting
(masks,
                colors=[colors(x, True) for x in det[:, 5]],
                im_gpu=None if retina_masks else im[i])

# Write results
for j, (*xyxy, conf, cls) in enumerate(reversed(det[:, :6])):
    if save_txt:  # Write to file
        segj = segments[j].reshape(-1)  # (n,2) to (n*2)
        line = (cls, *segj, conf) if save_conf else (cls, *segj)  # label format
        with open(f'{txt_path}.txt', 'a') as f:
            (('%g ' * len(line)).rstrip() % line + '\n')

    if save_img or save_crop or view_img:  # Add bbox to image
        c = int(cls)  # integer class
        label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
        annotator.box_label(xyxy, label, color=colors(c, True))
        # (segments[j], outline=colors(c, True), width=3)
    if save_crop:
        save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{}.jpg', BGR=True)

替换为以下代码，我已经将其注释为中文：

# 假设有一个颜色字典，其中键是类别的索引，值是相应的颜色
#custom_colors = {0: 'red', 1: 'blue', 2: 'green'}
# 自定义颜色映射
custom_colors = {0: (255, 0, 0),  # 红色
                 1: (0, 255, 0),  # 绿色
                 2: (0, 0, 255)}  # 蓝色
default_color = (0, 0, 0)
# Mask plotting
# 使用自定义颜色而不是colors函数
(masks,
                colors=[custom_colors.get(x, default_color) for x in det[:, 5]],
                im_gpu=None if retina_masks else im[i])

# Write results
for j, (*xyxy, conf, cls) in enumerate(reversed(det[:, :6])):
    if save_txt:  # Write to file
        segj = segments[j].reshape(-1)  # (n,2) to (n*2)
        line = (cls, *segj, conf) if save_conf else (cls, *segj)
        with open(f'{txt_path}.txt', 'a') as f:
            (('%g ' * len(line)).rstrip() % line + '\n')

    if save_img or save_crop or view_img:  # Add bbox to image
        c = int(cls)  # integer class
        label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
        # 使用自定义颜色而不是colors函数
        annotator.box_label(xyxy, label, color=custom_colors.get(c, default_color))
        # annotator.box_label(xyxy, label, color=custom_colors.get(c, default_color))
        # (segments[j], outline=custom_colors.get(c, default_color), width=3)

    if save_crop:
        save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{}.jpg', BGR=True)

在基于 PyTorch 的 YOLOv8 架构中自定义实例分割推理后的掩码颜色，你需要在后处理步骤中对推理结果进行操作。这里有一些具体的步骤：

1. **获取模型输出**：模型的输出通常包含掩码、类别以及置信度。确保你可以访问这些输出。

2. **定义颜色映射**：创建一个映射，将类别ID映射到你想要的颜色。

3. **处理掩码**：遍历每个掩码，并根据类别ID应用颜色映射。

以下是一个简化的代码示例，演示了如何在 PyTorch 中实现这个过程：

```python
import torch
import numpy as np

# 假设 model 是你的 YOLOv8 模型
# data 是模型的输入数据
# output 是模型的输出，包含了掩码和其他信息

# 模型推理
# output = model(data)

# 定义你的颜色映射，这里我们为三个类别定义了颜色
color_map = {
    0: (255, 0, 0),    # 类别 0 的颜色为红色
    1: (0, 255, 0),    # 类别 1 的颜色为绿色
    2: (0, 0, 255)     # 类别 2 的颜色为蓝色
}

# 假设 output['masks'] 是模型输出的掩码，output['classes'] 是对应的类别ID
# masks = output['masks']
# classes = output['classes']

# 这里我们使用随机数据来模拟掩码和类别
masks = ((3, 480, 640))  # 3个掩码，假设尺寸为480x640
classes = ([0, 1, 2])   # 3个类别

# 应用颜色映射
colored_masks = ((3, [1], [2], 3), dtype=torch.uint8)
for i, mask in enumerate(masks):
    class_id = classes[i].item()
    color = color_map[class_id]
    binary_mask = mask > 0.5  # 假设掩码的阈值为0.5来二值化掩码
    for j, c in enumerate(color):
        colored_masks[i, :, :, j] = binary_mask * c

# colored_masks 现在包含应用了颜色的三个掩码

# 如果需要将颜色掩码添加到原始图像上，可以使用以下代码：
# image 是原始图像的 PyTorch tensor，形状为 (3, H, W)
# alpha 是掩码的透明度系数

# for i in range(3):
#     image[:, colored_masks[i] > 0] = (1 - alpha) * image[:, colored_masks[i] > 0] + alpha * colored_masks[i]
```

请注意，以上代码是一个示例，可能需要根据你的具体场景进行调整。确保根据你的数据格式和模型输出正确地调整代码。

相关文章