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

目标

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

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

我们将以 DeepLab 架构为例，探讨上述要点。

简介

将TensorFlow 分类和分割模型与 OpenCV API 集成的关键概念几乎相同，只是图优化阶段除外。将 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分割预测
* 形状： (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分割模型预测
* 形状： (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/")
    # 缩减验证集：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。 使用提供的已转换模型作为--input参数，并填充cv2.dnn.blobFromImage所需的参数，即可独立运行segmentation.py。

要根据“模型转换流程”中描述的步骤，使用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