TensorFlow 分割模型的转换及使用 OpenCV 启动

目标

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

• 转换 TensorFlow (TF) 分割模型
• 使用 OpenCV 运行转换后的 TensorFlow 模型
• 获取 TensorFlow 和 OpenCV DNN 模型的评估结果

我们将通过 DeepLab 架构的示例来探讨上述要点。

简介

通过 OpenCV API 转换 TensorFlow 分类和分割模型所涉及的关键概念几乎相同，除了图优化阶段。将 TensorFlow 模型转换为 cv.dnn.Net 的初始步骤是获取冻结的 TF 模型图。冻结图定义了模型图结构与所需变量的保留值的组合，例如权重。通常，冻结图保存在 protobuf (.pb) 文件中。要使用 cv.dnn.readNetFromTensorflow 读取生成的分割模型 .pb 文件，需要使用 TF 图转换工具修改该图。

实践

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

1. 创建 TF 分类模型转换管道并提供推理
2. 评估和测试 TF 分类模型

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

模型转换管道

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

python -m dnn_model_runner.dnn_conversion.tf.segmentation.py_to_py_deeplab

TensorFlow 分割模型可以在 TensorFlow 研究模型部分中找到，其中包含基于已发布的研究论文的模型实现。我们将从以下链接检索包含预训练 TF DeepLabV3 的存档

http://download.tensorflow.org/models/deeplabv3_mnv2_pascal_trainval_2018_01_29.tar.gz

完整的冻结图获取管道在 deeplab_retrievement.py 中描述

def get_deeplab_frozen_graph()
    # 定义要下载的模型路径
models_url = 'http://download.tensorflow.org/models/'
mobilenetv2_voctrainval = 'deeplabv3_mnv2_pascal_trainval_2018_01_29.tar.gz'
 
    # 构建要下载的模型链接
model_link = models_url + mobilenetv2_voctrainval
 
    try:
urllib.request.urlretrieve(model_link, mobilenetv2_voctrainval)
    except Exception
print("TF DeepLabV3 未检索：{}".format(model_link))
        return
 
tf_model_tar = tarfile.open(mobilenetv2_voctrainval)
 
    # 迭代获得的模型存档
    for model_tar_elem in tf_model_tar.getmembers()
        # 检查模型存档是否包含冻结图
        if TF_FROZEN_GRAPH_NAME in os.path.basename(model_tar_elem.name)
            # 提取冻结图
tf_model_tar.extract(model_tar_elem, FROZEN_GRAPH_PATH)
 
tf_model_tar.close()

运行此脚本后

python -m dnn_model_runner.dnn_conversion.tf.segmentation.deeplab_retrievement

我们将在 deeplab/deeplabv3_mnv2_pascal_trainval 中获得 frozen_inference_graph.pb。

在使用 OpenCV 加载网络之前，需要优化提取的 frozen_inference_graph.pb。为了优化图，我们使用带有默认参数的 TF TransformGraph

DEFAULT_OPT_GRAPH_NAME = "optimized_frozen_inference_graph.pb"
DEFAULT_INPUTS = "sub_7"
DEFAULT_OUTPUTS = "ResizeBilinear_3"
DEFAULT_TRANSFORMS = "remove_nodes(op=Identity)" \
                     " merge_duplicate_nodes" \
                     " strip_unused_nodes" \
                     " fold_constants(ignore_errors=true)" \
                     " fold_batch_norms" \
                     " fold_old_batch_norms"
 
 
def optimize_tf_graph(
in_graph,
out_graph=DEFAULT_OPT_GRAPH_NAME,
inputs=DEFAULT_INPUTS,
outputs=DEFAULT_OUTPUTS,
transforms=DEFAULT_TRANSFORMS,
is_manual=True,
was_optimized=True
):
    # ...
 
tf_opt_graph = TransformGraph(
tf_graph,
inputs,
outputs,
transforms
    )

要运行图优化过程，执行以下行

python -m dnn_model_runner.dnn_conversion.tf.segmentation.tf_graph_optimizer --in_graph deeplab/deeplabv3_mnv2_pascal_trainval/frozen_inference_graph.pb

结果，deeplab/deeplabv3_mnv2_pascal_trainval 目录将包含 optimized_frozen_inference_graph.pb。

获得模型图后，让我们检查以下列出的步骤

1. 读取 TF frozen_inference_graph.pb 图
2. 使用 OpenCV API 读取优化的 TF 冻结图
3. 准备输入数据
4. 提供推理
5. 从预测中获取彩色掩码
6. 可视化结果

# 从获得的冻结图获取 TF 模型图
deeplab_graph = read_deeplab_frozen_graph(deeplab_frozen_graph_path)
 
# 使用 OpenCV API 读取 DeepLab 冻结图
opencv_net = cv2.dnn.readNetFromTensorflow(opt_deeplab_frozen_graph_path)
print("OpenCV 模型已成功读取。模型层：\n", opencv_net.getLayerNames())
 
# 获取处理后的图像
original_img_shape, tf_input_blob, opencv_input_img = get_processed_imgs("test_data/sem_segm/2007_000033.jpg")
 
# 获取 OpenCV DNN 预测
opencv_prediction = get_opencv_dnn_prediction(opencv_net, opencv_input_img)
 
# 获取 TF 模型预测
tf_prediction = get_tf_dnn_prediction(deeplab_graph, tf_input_blob)
 
# 获取 PASCAL VOC 类和颜色
pascal_voc_classes, pascal_voc_colors = read_colors_info("test_data/sem_segm/pascal-classes.txt")
 
# 获取彩色分割掩码
opencv_colored_mask = get_colored_mask(original_img_shape, opencv_prediction, pascal_voc_colors)
tf_colored_mask = get_tf_colored_mask(original_img_shape, tf_prediction, pascal_voc_colors)
 
# 获取 PASCAL VOC 颜色的调色板
color_legend = get_legend(pascal_voc_classes, pascal_voc_colors)
 
cv2.imshow('TensorFlow 彩色掩码', tf_colored_mask)
cv2.imshow('OpenCV DNN 彩色掩码', opencv_colored_mask)
 
cv2.imshow('颜色图例', color_legend)

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

目标分割结果是

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

让我们以预训练的 TF DeepLabV3 MobileNetV2 为例，更深入地了解每个步骤

• 读取 TF frozen_inference_graph.pb 图

# 初始化 deeplab 模型图
model_graph = tf.Graph()
 
# 获取
with tf.io.gfile.GFile(frozen_graph_path, 'rb') as graph_file
tf_model_graph = GraphDef()
tf_model_graph.ParseFromString(graph_file.read())
 
with model_graph.as_default()
tf.import_graph_def(tf_model_graph, name='')

• 使用 OpenCV API 读取优化的 TF 冻结图

# 使用 OpenCV API 读取 DeepLab 冻结图
opencv_net = cv2.dnn.readNetFromTensorflow(opt_deeplab_frozen_graph_path)

• 使用 cv2.dnn.blobFromImage 函数准备输入数据

# 读取图像
input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
input_img = input_img.astype(np.float32)
 
# 预处理 TF 模型输入的图像
tf_preproc_img = cv2.resize(input_img, (513, 513))
tf_preproc_img = cv2.cvtColor(tf_preproc_img, cv2.COLOR_BGR2RGB)
 
# 定义 OpenCV DNN 的预处理参数
mean = np.array([1.0, 1.0, 1.0]) * 127.5
scale = 1 / 127.5
 
# 准备输入 blob 以适应模型输入
# 1. 减去均值
# 2. 缩放以将像素值设置为 0 到 1
input_blob = cv2.dnn.blobFromImage(
image=input_img,
scalefactor=scale,
size=(513, 513), # 图像目标大小
mean=mean,
swapRB=True, # BGR -> RGB
crop=False # 中心裁剪
)

请注意 cv2.dnn.blobFromImage 函数中的预处理顺序。首先，减去平均值，然后将像素值乘以定义的比例。因此，为了重现 TF 图像预处理管道，我们将 mean 乘以 127.5。另一个重点是 TF DeepLab 的图像预处理。要将图像传递到 TF 模型中，我们只需要构建一个合适的形状，其余的图像预处理在 feature_extractor.py 中描述，并将自动调用。

• 提供 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 分割预测
* shape: (1, 21, 513, 513)

21 个预测通道中的每一个，其中 21 代表 PASCAL VOC 类的数量，都包含概率，指示像素对应于 PASCAL VOC 类的可能性。

• 提供 TF 模型推理

preproc_img = np.expand_dims(preproc_img, 0)
 
# 初始化 TF 会话
tf_session = Session(graph=model_graph)
 
input_tensor_name = "ImageTensor:0",
output_tensor_name = "SemanticPredictions:0"
 
# 运行推理
out = tf_session.run(
output_tensor_name,
feed_dict={input_tensor_name: [preproc_img]}
)
 
print("TF 分割模型预测：\n")
print("* 形状：", out.shape)

TF 推理结果如下

TF 分割模型预测
* shape: (1, 513, 513)

TensorFlow 预测包含相应 PASCAL VOC 类的索引。

• 将 OpenCV 预测转换为彩色掩码

mask_height = segm_mask.shape[0]
mask_width = segm_mask.shape[1]
 
img_height = original_img_shape[0]
img_width = original_img_shape[1]
 
# 将掩码值转换为 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
processed_mask = cv2.cvtColor(processed_mask, cv2.COLOR_BGR2RGB)

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

• 将 TF 预测转换为彩色掩码

colors = np.array(colors)
processed_mask = colors[segm_mask[0]]
 
img_height = original_img_shape[0]
img_width = original_img_shape[1]
 
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/samples 中提出的 dnn_model_runner 模块允许在 PASCAL VOC 数据集上运行完整的评估管道，并测试 DeepLab MobileNet 模型的执行情况。

评估模式

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

python -m dnn_model_runner.dnn_conversion.tf.segmentation.py_to_py_segm

该模型将被读入 OpenCV cv.dnn_Net 对象。TF 和 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")

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

测试模式

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

python -m dnn_model_runner.dnn_conversion.tf.segmentation.py_to_py_segm --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 中定义

tf_segm_input_blob = {
    "scale": str(1 / 127.5),
    "mean": ["127.5", "127.5", "127.5"],
    "std": [],
    "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.tf.segmentation.py_to_py_segm --test True --default_img_preprocess True --evaluate False