MMSegmentation 图像分割推理教程：使用预训练模型进行预测

MMSegmentation 图像分割推理教程：使用预训练模型进行预测

mmsegmentation OpenMMLab Semantic Segmentation Toolbox and Benchmark. 项目地址: https://gitcode.com/gh_mirrors/mm/mmsegmentation

前言

在计算机视觉领域，语义分割是一项基础且重要的任务，它旨在为图像中的每个像素分配一个类别标签。MMSegmentation 作为一个强大的语义分割开源工具箱，提供了丰富的预训练模型和便捷的推理接口。本文将详细介绍如何使用 MMSegmentation 中的预训练模型进行图像分割推理。

准备工作

在开始之前，请确保已经正确安装了 MMSegmentation 及其依赖项。本文假设读者已经完成了环境配置工作。

使用 MMSegInferencer 进行推理

MMSegInferencer 是 MMSegmentation 提供的高级接口，它封装了模型加载、预处理、推理和后处理的完整流程，使用起来非常便捷。

基础用法

最简单的推理只需要三行代码：

from mmseg.apis import MMSegInferencer
inferencer = MMSegInferencer(model='deeplabv3plus_r18-d8_4xb2-80k_cityscapes-512x1024')
result = inferencer('demo/demo.png', show=True)

这段代码会：

1. 加载指定的预训练模型（这里使用的是 DeepLabV3+ 模型）
2. 对输入图像进行推理
3. 显示分割结果的可视化效果

批量处理图像

MMSegInferencer 同样支持批量处理图像：

# 处理图像列表
images = ['img1.jpg', 'img2.jpg', 'img3.jpg']
results = inferencer(images)

# 处理整个目录下的图像
image_dir = 'path/to/image_folder'
results = inferencer(image_dir)

结果保存

如果需要保存推理结果，可以指定输出目录：

inferencer(images, 
           out_dir='outputs',
           img_out_dir='vis',    # 保存可视化结果
           pred_out_dir='pred')  # 保存预测掩码

结果格式控制

默认情况下，推理结果以字典形式返回，包含可视化结果和预测掩码：

result = inferencer('demo.png')
print(result.keys())  # 输出：dict_keys(['visualization', 'predictions'])

如果需要更详细的结果数据，可以设置 return_datasamples=True，返回 SegDataSample 对象：

result = inferencer('demo.png', return_datasamples=True)
print(type(result))  # 输出：<class 'mmseg.structures.seg_data_sample.SegDataSample'>

模型初始化参数详解

创建 MMSegInferencer 时可以指定多个参数：

inferencer = MMSegInferencer(
    model='模型名称或配置文件路径',
    weights='自定义权重路径',  # 可选，不指定则使用预训练权重
    classes=['道路', '建筑', '天空'],  # 自定义类别标签
    palette=[[255,0,0], [0,255,0], [0,0,255]],  # 自定义颜色映射
    dataset_name='cityscapes',  # 使用内置数据集配置
    device='cuda:0',  # 指定推理设备
    scope='mmseg'  # 模型作用域
)

可视化选项

推理时可以控制结果的可视化效果：

inferencer('demo.png',
           show=True,      # 是否显示结果
           wait_time=1,    # 显示时间(秒)，0表示持续显示
           opacity=0.6)    # 掩码透明度(0-1]

模型列表查询

可以通过以下方式查询所有可用的预训练模型：

from mmseg.apis import MMSegInferencer
models = MMSegInferencer.list_models('mmseg')

底层 API 使用指南

除了便捷的 MMSegInferencer，MMSegmentation 也提供了更底层的 API，适合需要更精细控制的场景。

模型初始化

from mmseg.apis import init_model

config = 'configs/pspnet/pspnet_r50-d8_4xb2-40k_cityscapes-512x1024.py'
checkpoint = 'checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes.pth'

model = init_model(config, checkpoint, device='cuda:0')

执行推理

from mmseg.apis import inference_model

result = inference_model(model, 'demo.png')

结果可视化

from mmseg.apis import show_result_pyplot

vis_image = show_result_pyplot(
    model,
    'demo.png',
    result,
    opacity=0.7,
    title='分割结果',
    out_file='result.png'  # 保存结果
)

结果数据结构

MMSegmentation 使用 SegDataSample 作为标准的数据结构，它包含以下主要属性：

1. gt_sem_seg: 真实分割标签（PixelData 格式）
2. pred_sem_seg: 预测分割结果（PixelData 格式）
3. seg_logits: 预测的 logits（PixelData 格式）

PixelData 是 MMEngine 中定义的像素级数据结构，支持多种格式的存储和转换。

最佳实践建议

1. 批量处理：对于大量图像，建议使用批量处理而非单张处理，可以提高效率
2. 内存管理：处理大尺寸图像时注意内存消耗，必要时可以先调整图像尺寸
3. 结果保存：在生产环境中，建议关闭显示窗口（show=False）直接保存结果
4. 自定义配置：根据实际场景调整类别标签和颜色映射，使可视化效果更符合需求
5. 设备选择：GPU 推理速度远快于 CPU，对于实时性要求高的场景建议使用 GPU

结语

通过本文的介绍，相信读者已经掌握了使用 MMSegmentation 进行图像分割推理的多种方法。无论是简单的单张图像处理，还是复杂的批量推理任务，MMSegmentation 都提供了灵活且高效的解决方案。在实际应用中，可以根据具体需求选择合适的方法和参数配置。

mmsegmentation OpenMMLab Semantic Segmentation Toolbox and Benchmark. 项目地址: https://gitcode.com/gh_mirrors/mm/mmsegmentation