PyTorch分割模型的转换和使用

目标

在本教程中，您将学习如何：

• 转换PyTorch分割模型
• 使用OpenCV运行转换后的PyTorch模型
• 评估PyTorch和OpenCV DNN模型

我们将以FCN ResNet-50架构为例，探讨以上几点。

简介

将PyTorch分类和分割模型转换为OpenCV API的关键步骤是相同的。第一步是使用PyTorch内置函数torch.onnx.export将模型转换为ONNX格式。然后，将生成的.onnx模型传递给cv.dnn.readNetFromONNX，该函数返回一个cv.dnn.Net对象，该对象已准备好进行DNN操作。

实践

在本部分中，我们将介绍以下几点：

1. 创建分割模型转换流程并进行推理
2. 评估和测试分割模型

如果您只想运行评估或测试模型流程，“模型转换流程”部分可以跳过。

模型转换流程

本小节中的代码位于dnn_model_runner模块中，可以使用以下命令执行：

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_fcnresnet50

以下代码描述了以下步骤：

1. 实例化PyTorch模型
2. 将PyTorch模型转换为.onnx
3. 使用OpenCV API读取转换后的网络
4. 准备输入数据
5. 进行推理
6. 从预测结果中获取彩色掩码
7. 可视化结果

# 初始化PyTorch FCN ResNet-50模型
original_model = models.segmentation.fcn_resnet50(pretrained=True)
 
# 获取转换为ONNX的PyTorch模型的路径
full_model_path = get_pytorch_onnx_model(original_model)
 
# 使用OpenCV API读取转换后的.onnx模型
opencv_net = cv2.dnn.readNetFromONNX(full_model_path)
print("OpenCV模型读取成功。层ID：\n", opencv_net.getLayerNames())
 
# 获取预处理图像
img, input_img = get_processed_imgs("test_data/sem_segm/2007_000033.jpg")
 
# 获取OpenCV DNN预测结果
opencv_prediction = get_opencv_dnn_prediction(opencv_net, input_img)
 
# 获取原始PyTorch ResNet50预测结果
pytorch_prediction = get_pytorch_dnn_prediction(original_model, input_img)
 
pascal_voc_classes, pascal_voc_colors = read_colors_info("test_data/sem_segm/pascal-classes.txt")
 
# 获取彩色分割掩码
opencv_colored_mask = get_colored_mask(img.shape, opencv_prediction, pascal_voc_colors)
pytorch_colored_mask = get_colored_mask(img.shape, pytorch_prediction, pascal_voc_colors)
 
# 获取PASCAL VOC颜色调色板
color_legend = get_legend(pascal_voc_classes, pascal_voc_colors)
 
cv2.imshow('PyTorch彩色掩码', pytorch_colored_mask)
cv2.imshow('OpenCV DNN彩色掩码', opencv_colored_mask)
cv2.imshow('颜色图例', color_legend)
 
cv2.waitKey(0)

为了进行模型推理，我们将使用来自PASCAL VOC验证数据集的以下图片：

目标分割结果是：

为了对PASCAL VOC颜色进行解码并将其与预测的掩码映射，我们还需要pascal-classes.txt文件，其中包含PASCAL VOC类的完整列表和相应的颜色。

让我们以预训练的PyTorch FCN ResNet-50为例，更深入地了解每个代码步骤。

• 实例化PyTorch FCN ResNet-50模型

# 初始化PyTorch FCN ResNet-50模型
original_model = models.segmentation.fcn_resnet50(pretrained=True)

• 将PyTorch模型转换为ONNX格式

# 定义转换后模型保存的目录
onnx_model_path = "models"
# 定义转换后模型的名称
onnx_model_name = "fcnresnet50.onnx"
 
# 创建转换后模型的目录
os.makedirs(onnx_model_path, exist_ok=True)
 
# 获取转换后模型的完整路径
full_model_path = os.path.join(onnx_model_path, onnx_model_name)
 
# 生成模型输入以构建图
generated_input = Variable(
torch.randn(1, 3, 500, 500)
)
 
# 将模型导出为ONNX格式
torch.onnx.export(
original_model,
generated_input,
full_model_path,
verbose=True,
input_names=["input"],
output_names=["output"],
opset_version=11
)

此步骤的代码与分类转换的情况没有区别。因此，在成功执行上述代码后，我们将得到models/fcnresnet50.onnx。

• 使用cv.dnn.readNetFromONNX读取转换后的网络，并将上一步中获得的ONNX模型传递给它。

# 使用OpenCV API读取转换后的.onnx模型
opencv_net = cv2.dnn.readNetFromONNX(full_model_path)

• 准备输入数据

# 读取图像
input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
input_img = input_img.astype(np.float32)
 
# 目标图像大小
img_height = input_img.shape[0]
img_width = input_img.shape[1]
 
# 定义预处理参数
mean = np.array([0.485, 0.456, 0.406]) * 255.0
scale = 1 / 255.0
std = [0.229, 0.224, 0.225]
 
# 准备输入blob以适应模型输入
# 1. 减去均值
# 2. 缩放以将像素值设置为0到1
input_blob = cv2.dnn.blobFromImage(
image=input_img,
scalefactor=scale,
size=(img_width, img_height), # 图像目标大小
mean=mean,
swapRB=True, # BGR -> RGB
crop=False # 中心裁剪
)
# 3. 除以标准差
input_blob[0] /= np.asarray(std, dtype=np.float32).reshape(3, 1, 1)

在此步骤中，我们读取图像并使用cv2.dnn.blobFromImage函数准备模型输入，该函数返回一个四维blob。需要注意的是，在cv2.dnn.blobFromImage中，首先减去均值，然后缩放像素值。因此，mean乘以255.0是为了重现原始图像预处理顺序。

img /= 255.0
img -= [0.485, 0.456, 0.406]
img /= [0.229, 0.224, 0.225]

• OpenCV cv.dnn_Net 推理

# 设置OpenCV DNN输入
opencv_net.setInput(preproc_img)
 
# OpenCV DNN推理
out = opencv_net.forward()
print("OpenCV DNN分割预测结果：\n")
print("* 形状：", out.shape)
 
# 获取预测类别的ID
out_predictions = np.argmax(out[0], axis=0)

执行上述代码后，我们将得到以下输出：

OpenCV DNN分割预测结果
* 形状：(1, 21, 500, 500)

21 个预测通道中的每一个（其中 21 代表 PASCAL VOC 类别的数量）都包含概率，这些概率表示像素属于 PASCAL VOC 类别的可能性。

• PyTorch FCN ResNet-50 模型推理

original_net.eval()
preproc_img = torch.FloatTensor(preproc_img)
 
with torch.no_grad()
    # 获取每个类别的未归一化概率
out = original_net(preproc_img)['out']
 
print("\nPyTorch 分割模型预测：\n")
print("* 形状：", out.shape)
 
# 获取预测类别的ID
out_predictions = out[0].argmax(dim=0)

运行上述代码后，我们将得到以下输出

PyTorch 分割模型预测
* 形状：torch.Size([1, 21, 366, 500])

PyTorch 预测也包含对应于每个类别预测的概率。

• 从预测结果中获取彩色掩码

# 将掩码值转换为 PASCAL VOC 颜色
processed_mask = np.stack([colors[color_id] for color_id in segm_mask.flatten()])
 
# 将掩码重塑为 3 通道图像
processed_mask = processed_mask.reshape(mask_height, mask_width, 3)
processed_mask = cv2.resize(processed_mask, (img_width, img_height), interpolation=cv2.INTER_NEAREST).astype(
np.uint8)
 
# 将彩色掩码从 BGR 转换为 RGB，以与 PASCAL VOC 颜色兼容
processed_mask = cv2.cvtColor(processed_mask, cv2.COLOR_BGR2RGB)

在此步骤中，我们将分割掩码中的概率与预测类别的相应颜色映射。让我们看看结果

为了对模型进行扩展评估，我们可以使用 `dnn_model_runner` 模块的 `py_to_py_segm` 脚本。此模块部分将在下一小节中介绍。

模型评估

在 `dnn/samples` 中提出的 `dnn_model_runner` 模块允许在 PASCAL VOC 数据集上运行完整的评估流程，并测试以下 PyTorch 分割模型的执行情况：

• FCN ResNet-50
• FCN ResNet-101

此列表还可以通过进一步的适当评估流程配置进行扩展。

评估模式

以下行表示在评估模式下运行模块

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_segm --model_name <pytorch_segm_model_name>

从列表中选择的分割模型将被读入 OpenCV `cv.dnn_Net` 对象。PyTorch 和 OpenCV 模型的评估结果（像素精度、平均 IoU、推理时间）将写入日志文件。推理时间值也将以图表的形式显示，以概括获得的模型信息。

必要的评估配置在 test_config.py 中定义

@dataclass
class TestSegmConfig
frame_size: int = 500
img_root_dir: str = "./VOC2012"
img_dir: str = os.path.join(img_root_dir, "JPEGImages/")
img_segm_gt_dir: str = os.path.join(img_root_dir, "SegmentationClass/")
    # 简化验证集：https://github.com/shelhamer/fcn.berkeleyvision.org/blob/master/data/pascal/seg11valid.txt
segm_val_file: str = os.path.join(img_root_dir, "ImageSets/Segmentation/seg11valid.txt")
colour_file_cls: str = os.path.join(img_root_dir, "ImageSets/Segmentation/pascal-classes.txt")

这些值可以根据选择的模型流程进行修改。

要启动 PyTorch FCN ResNet-50 的评估，请运行以下命令：

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_segm --model_name fcnresnet50

测试模式

以下行表示在测试模式下运行模块，该模式提供了模型推理的步骤

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_segm --model_name <pytorch_segm_model_name> --test True --default_img_preprocess <True/False> --evaluate False

这里 `default_img_preprocess` 键定义了您是否希望使用某些特定值或默认值来参数化模型测试过程，例如 `scale`、`mean` 或 `std`。

测试配置在 test_config.py 的 `TestSegmModuleConfig` 类中表示

@dataclass
class TestSegmModuleConfig
segm_test_data_dir: str = "test_data/sem_segm"
test_module_name: str = "segmentation"
test_module_path: str = "segmentation.py"
input_img: str = os.path.join(segm_test_data_dir, "2007_000033.jpg")
model: str = ""
 
frame_height: str = str(TestSegmConfig.frame_size)
frame_width: str = str(TestSegmConfig.frame_size)
scale: float = 1.0
mean: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
std: List[float] = field(default_factory=list)
crop: bool = False
rgb: bool = True
classes: str = os.path.join(segm_test_data_dir, "pascal-classes.txt")

默认图像预处理选项在 `default_preprocess_config.py` 中定义

pytorch_segm_input_blob = {
    "mean": ["123.675", "116.28", "103.53"],
    "scale": str(1 / 255.0),
    "std": ["0.229", "0.224", "0.225"],
    "crop": "False",
    "rgb": "True"
}

模型测试的基础在 `samples/dnn/segmentation.py` 中。`segmentation.py` 可以使用在 `--input` 中提供的已转换模型和为 `cv2.dnn.blobFromImage` 填充的参数自主执行。

要使用 `dnn_model_runner` 从头开始重现“模型转换流程”中描述的 OpenCV 步骤，请执行以下命令：

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_segm --model_name fcnresnet50 --test True --default_img_preprocess True --evaluate False