深度神经网络——转换 PyTorch 分割模型并使用 OpenCV 启动 OpenCV v4.8.0

目标

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

• 转换 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', color_legend)
cv2.waitKey(0)

为了提供模型推理，我们将使用 PASCAL VOC 验证数据集中的以下图片：

PASCAL VOC img（此图片已损坏）

目标分割结果为

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 维 blob。需要注意的是，cv2.dnn.blobFromImage 首先要减去平均值，然后才对像素值进行缩放。因此，平均值乘以 255.0，以重现原始图像的预处理顺序：

img /= 255.0
img -= [0.485, 0.456, 0.406] (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 segmentation prediction: \n")
print("* shape: ", 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("* shape: ", out.shape)
# 获取预测类别的 ID
out_predictions = out[0].argmax(dim=0)

启动上述代码后，我们将得到以下输出结果：

PyTorch 分割模型预测：
* shape: 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)

在这一步中，我们将分割掩码中的概率与预测类别的适当颜色进行映射。让我们来看看结果：

OpenCV 彩色掩码（此图片已损坏）

为了对模型进行扩展评估，我们可以使用 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:
    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")
    color_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