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

狐松鼠，东狐松鼠，小黑

分类模型输入图像

为了解码获得的预测结果标签，我们还需要 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 中获得一个冻结图。

• 将获取的mobilenet.pb传递给cv.dnn.readNetFromTensorflow，使用该函数读取以 TensorFlow 框架格式存储的网络模型

# 获取 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), # img 目标大小
mean=mean,
swapRB=True, # BGR -> RGB
crop=True # 中心裁剪
)
print("输入Blob形状: {}\n".format(input_blob.shape))

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

这样就获得了四维的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，标签：fox squirrel，eastern fox squirrel，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，标签：fox squirrel，eastern fox squirrel，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
类 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 推理输出