转换 PyTorch 分割模型并使用 OpenCV 启动

目标

在此教程中，你将学习如何

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

我们将通过 FCN ResNet-50 架构的示例来探讨上述各点。

简介

OpenCV API 中 PyTorch 分类和分割模型转换管道涉及的关键点是相同的。第一步是将模型转换为 ONNX 格式，方法是使用 PyTorch 的内置函数 torch.onnx.export。然后将获得的 .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 model was successfully read. Layer IDs: \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 函数准备模型输入，此函数会返回 4 维块。应当注意，首先在 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] = 字段(default_factory=lambda: [0.0, 0.0, 0.0])
std: List[float] = 字段(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