PyTorch 分割模型的转换与 OpenCV 的启动

目标

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

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

我们将通过 FCN ResNet-50 架构的示例来探讨以上列出的要点。

介绍

使用 OpenCV API 进行 PyTorch 分类 和分割模型转换流程中涉及的关键点是相同的。第一步是使用 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. 除以 std
input_blob[0] /= np.asarray(std, dtype=np.float32).reshape(3, 1, 1)

在此步骤中，我们读取图像并使用 cv2.dnn.blobFromImage 函数准备模型输入，该函数返回 4 维 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/")
    # 减少 val：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 的填充参数自主执行。

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

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