TensorFlow 分类模型的转换和使用 OpenCV Python 进行启动

原作者	Anastasia Murzova
兼容性	OpenCV >= 4.5

目标

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

• 获取 TensorFlow (TF) 分类模型的冻结图
• 使用 OpenCV Python API 运行转换后的 TensorFlow 模型
• 评估 TensorFlow 和 OpenCV DNN 模型

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

简介

让我们简要了解一下 TensorFlow 模型与 OpenCV API 交互流程中涉及的关键概念。将 TensorFlow 模型转换为cv.dnn.Net 的第一步是获取冻结的 TF 模型图。冻结图定义了模型图结构与所需变量（例如权重）的保留值的组合。通常，冻结图保存在protobuf (.pb) 文件中。生成模型 .pb 文件后，可以使用cv.dnn.readNetFromTensorflow 函数读取。

需求

要能够使用下面的代码进行实验，您需要安装一系列库。我们将为此使用 python3.7+ 的虚拟环境

virtualenv -p /usr/bin/python3.7 <env_dir_path>
source <env_dir_path>/bin/activate

对于从源码构建 OpenCV-Python，请遵循OpenCV 简介中的相应说明。

在开始安装库之前，您可以自定义requirements.txt，排除或包含（例如，opencv-python）某些依赖项。以下命令将启动将需求安装到先前激活的虚拟环境中

pip install -r requirements.txt

实践

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

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

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

模型转换管道

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

python -m dnn_model_runner.dnn_conversion.tf.classification.py_to_py_mobilenet

以下代码包含对以下步骤的描述：

1. 实例化 TF 模型
2. 创建 TF 冻结图
3. 使用 OpenCV API 读取 TF 冻结图
4. 准备输入数据
5. 提供推理

# 初始化 TF MobileNet 模型
original_tf_model = MobileNet(
include_top=True,
weights="imagenet"
)
 
# 获取 TF 冻结图路径
full_pb_path = get_tf_model_proto(original_tf_model)
 
# 使用 OpenCV API 读取冻结图
opencv_net = cv2.dnn.readNetFromTensorflow(full_pb_path)
print("OpenCV 模型已成功读取。模型层：\n", opencv_net.getLayerNames())
 
# 获取预处理图像
input_img = get_preprocessed_img("../data/squirrel_cls.jpg")
 
# 获取 ImageNet 标签
imagenet_labels = get_imagenet_labels("../data/dnn/classification_classes_ILSVRC2012.txt")
 
# 获取 OpenCV DNN 预测结果
get_opencv_dnn_prediction(opencv_net, input_img, imagenet_labels)
 
# 获取 TF 模型预测结果
get_tf_dnn_prediction(original_tf_model, input_img, imagenet_labels)

为了进行模型推理，我们将使用下面的松鼠照片（在CC0 许可下），对应于 ImageNet 类 ID 335

狐松鼠，东部狐松鼠，Sciurus niger

分类模型输入图像

对于获得的预测结果的标签解码，我们还需要 imagenet_classes.txt 文件，其中包含 ImageNet 类的完整列表。

让我们通过预训练的 TF MobileNet 例子深入了解每个步骤

• 实例化 TF 模型

# 初始化 TF MobileNet 模型
original_tf_model = MobileNet(
include_top=True,
weights="imagenet"
)

• 创建 TF 冻结图

# 定义 .pb 模型的目录
pb_model_path = "models"
 
# 定义 .pb 模型的名称
pb_model_name = "mobilenet.pb"
 
# 为转换后的模型创建目录
os.makedirs(pb_model_path, exist_ok=True)
 
# 获取模型 TF 图
tf_model_graph = tf.function(lambda x: tf_model(x))
 
# 获取具体函数
tf_model_graph = tf_model_graph.get_concrete_function(
tf.TensorSpec(tf_model.inputs[0].shape, tf_model.inputs[0].dtype))
 
# 获取冻结的具体函数
frozen_tf_func = convert_variables_to_constants_v2(tf_model_graph)
# 获取冻结图
frozen_tf_func.graph.as_graph_def()
 
# 保存完整的 tf 模型
tf.io.write_graph(graph_or_graph_def=frozen_tf_func.graph,
logdir=pb_model_path,
name=pb_model_name,
as_text=False)

成功执行上述代码后，我们将在 models/mobilenet.pb 中得到一个冻结图。

• 使用cv.dnn.readNetFromTensorflow 读取 TF 冻结图，并将上一步中获得的 mobilenet.pb 传递给它

# 获取 TF 冻结图路径
full_pb_path = get_tf_model_proto(original_tf_model)

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

# 读取图像
input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
input_img = input_img.astype(np.float32)
 
# 定义预处理参数
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=(224, 224), # 图片目标大小
mean=mean,
swapRB=True, # BGR -> RGB
crop=True # 中心裁剪
)
print("输入 blob 形状：{}\n".format(input_blob.shape))

请注意 cv2.dnn.blobFromImage 函数中的预处理顺序。首先，减去均值，然后将像素值乘以定义的比例因子。因此，为了重现来自 TF mobilenet.preprocess_input 函数的图像预处理管道，我们将 mean 乘以 127.5。

结果，得到了一个 4 维的 input_blob

输入 blob 形状： (1, 3, 224, 224)

• 提供 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
imagenet_class_id = np.argmax(out)
 
# 获取置信度
confidence = out[0][imagenet_class_id]
print("* 类 ID：{}, 标签：{}".format(imagenet_class_id, imagenet_labels[imagenet_class_id]))
print("* 置信度：{:.4f}\n".format(confidence))

执行上述代码后，我们将得到以下输出：

OpenCV DNN 预测结果
* 形状： (1, 1000)
* 类 ID：335，标签：狐松鼠，东部狐松鼠，Sciurus niger
* 置信度：0.9525

• 提供 TF MobileNet 推理

# 推理
preproc_img = preproc_img.transpose(0, 2, 3, 1)
print("TF 输入 blob 形状: {}\n".format(preproc_img.shape))
 
out = original_net(preproc_img)
 
print("\nTensorFlow 模型预测: \n")
print("* 形状：", out.shape)
 
# 获取预测的类 ID
imagenet_class_id = np.argmax(out)
print("* 类 ID：{}, 标签：{}".format(imagenet_class_id, imagenet_labels[imagenet_class_id]))
 
# 获取置信度
confidence = out[0][imagenet_class_id]
print("* 置信度: {:.4f}".format(confidence))

为了适应 TF 模型输入，input_blob 进行了转置

TF 输入 blob 形状: (1, 224, 224, 3)

TF 推理结果如下

TensorFlow 模型预测
* 形状： (1, 1000)
* 类 ID：335，标签：狐松鼠，东部狐松鼠，Sciurus niger
* 置信度：0.9525

从实验结果可以看出，OpenCV 和 TF 推理结果相同。

模型评估

在 dnn/samples 中提出的 dnn_model_runner 模块允许在 ImageNet 数据集上运行完整的评估流程，并测试以下 TensorFlow 分类模型的执行情况：

• vgg16
• vgg19
• resnet50
• resnet101
• resnet152
• densenet121
• densenet169
• densenet201
• inceptionresnetv2
• inceptionv3
• mobilenet
• mobilenetv2
• nasnetlarge
• nasnetmobile
• xception

此列表还可以通过进一步的适当评估流程配置进行扩展。

评估模式

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

python -m dnn_model_runner.dnn_conversion.tf.classification.py_to_py_cls --model_name <tf_cls_model_name>

从列表中选择的分类模型将被读入 OpenCV cv.dnn_Net 对象。TF 和 OpenCV 模型的评估结果（准确率、推理时间、L1）将写入日志文件。推理时间值也将以图表的形式显示，以概括获得的模型信息。

必要的评估配置定义在 test_config.py 中，可以根据实际数据位置路径进行修改：

@dataclass
class TestClsConfig
batch_size: int = 50
frame_size: int = 224
img_root_dir: str = "./ILSVRC2012_img_val"
    # 图像类别匹配位置
img_cls_file: str = "./val.txt"
bgr_to_rgb: bool = True

可以根据所选模型自定义 TestClsConfig 中的值。

要启动 TensorFlow MobileNet 的评估，请运行以下命令：

python -m dnn_model_runner.dnn_conversion.tf.classification.py_to_py_cls --model_name mobilenet

脚本启动后，将在 dnn_model_runner/dnn_conversion/logs 中生成包含评估数据的日志文件。

===== 使用以下参数运行模型评估
* 验证数据位置：./ILSVRC2012_img_val
* 日志文件位置：dnn_model_runner/dnn_conversion/logs/TF_mobilenet_log.txt

测试模式

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

python -m dnn_model_runner.dnn_conversion.tf.classification.py_to_py_cls --model_name <tf_cls_model_name> --test True --default_img_preprocess <True/False> --evaluate False

这里 default_img_preprocess 键定义是否要使用某些特定值或默认值来参数化模型测试过程，例如 scale、mean 或 std。

测试配置在 test_config.py 的 TestClsModuleConfig 类中表示。

@dataclass
class TestClsModuleConfig
cls_test_data_dir: str = "../data"
test_module_name: str = "classification"
test_module_path: str = "classification.py"
input_img: str = os.path.join(cls_test_data_dir, "squirrel_cls.jpg")
model: str = ""
 
frame_height: str = str(TestClsConfig.frame_size)
frame_width: str = str(TestClsConfig.frame_size)
scale: str = "1.0"
mean: List[str] = field(default_factory=lambda: ["0.0", "0.0", "0.0"])
std: List[str] = field(default_factory=list)
crop: str = "False"
rgb: str = "True"
rsz_height: str = ""
rsz_width: str = ""
classes: str = os.path.join(cls_test_data_dir, "dnn", "classification_classes_ILSVRC2012.txt")

默认图像预处理选项定义在 default_preprocess_config.py 中。例如，对于 MobileNet：

tf_input_blob = {
    "mean": ["127.5", "127.5", "127.5"],
    "scale": str(1 / 127.5),
    "std": [],
    "crop": "True",
    "rgb": "True"
}

模型测试的基础在 samples/dnn/classification.py 中。classification.py 可以使用在 --input 中提供的已转换模型和为 cv.dnn.blobFromImage 填充的参数独立执行。

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

python -m dnn_model_runner.dnn_conversion.tf.classification.py_to_py_cls --model_name mobilenet --test True --default_img_preprocess True --evaluate False

网络预测显示在输出窗口的左上角。

TF MobileNet OpenCV 推理输出