训练一个图像分类器并在OpenCV中调用

首页 分类 标签 链接 留言 关于 订阅 2020-12-04

|

分类 技术随笔 

单纯的记录。

目标：训练一个判断图片内是否包含文字的分类器，并让ImageTrans能通过OpenCV进行调用

1. 使用TensorFlow训练图片分类模型

   使用ImageTrans导出包含文字的图像和不包含文字的图像用于训练。使用TensorFlow提供的make_image_classifier脚本进行训练，会生成tflite和savedmodel格式的模型文件。

2. 转换savedmodel格式的模型为frozen graph

    import tensorflow as tf
   
    from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
   
    loaded = tf.saved_model.load('my_model')
    infer = loaded.signatures['serving_default']
   
    f = tf.function(infer).get_concrete_function(input_1=tf.TensorSpec(shape=[None, 224, 224, 3], dtype=tf.float32))
    f2 = convert_variables_to_constants_v2(f)
    graph_def = f2.graph.as_graph_def()
   
    # Export frozen graph
    with tf.io.gfile.GFile('frozen_graph.pb', 'wb') as f:
       f.write(graph_def.SerializeToString())
   

   代码来自：Import SavedModel from TensorFlow

3. 使用OpenCV的DNN模块进行调用

    import numpy as np
    import cv2
   
    net = cv2.dnn.readNet('frozen_graph.pb')
    image = cv2.imread(filename)
    img_tensor =  cv2.dnn.blobFromImage(image, 1 / 255.0, (224, 224), swapRB=True, crop=False) 
    net.setInput(img_tensor)
    out = net.forward()
    print(out[0][1])
    print(out.shape)