【ATU Book-i.MX8系列 - TFLite】 i.MX8 平台实现 !! Tensorflow Lite 手写识别

tf.saved_model.save(model, "Tensorflow_2_HandWrite")

model. save('tf2_keras_handwrite.h5')

#---------------------------------------------------------------------------------------------------------------
# 转成 Tensorflow Lite
#---------------------------------------------------------------------------------------------------------------
import tensorflow as tf
import numpy as np
def representative_dataset_gen():
    for _ in range(250):
        yield [np.random.uniform(0.0, 1.0, size=(1, 28, 28, 1)).astype(np.float32)] #需要输入端大小
        
converter = tf.lite.TFLiteConverter.from_saved_model("Tensorflow_2_HandWrite")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type  = tf.uint8
converter.inference_output_type = tf.uint8
converter.representative_dataset = representative_dataset_gen
tflite_model = converter.convert()
with open('tf2_handwrite_model_qunat_uint8.tflite','wb') as f:
    f.write(tflite_model)​​

#---------------------------------------------------------------------------------------------------------------
# 转成 Tensorflow Lite
#---------------------------------------------------------------------------------------------------------------
import tensorflow as tf
import numpy as np
def representative_dataset_gen():
    for _ in range(250):
        yield [np.random.uniform(0.0, 1.0, size=(1, 28, 28, 1)).astype(np.float32)] #需要输入端大小
        
converter = tf.compat.v1.lite.TFLiteConverter.from_keras_model_file('tf2_keras_handwrite.h5')
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type  = tf.uint8
converter.inference_output_type = tf.uint8
converter.representative_dataset = representative_dataset_gen
tflite_model = converter.convert()
with open('tf2_handwrite_model_qunat_uint8.tflite','wb') as f:
    f.write(tflite_model)

import os
import argparse
import cv2
import numpy as np
import sys
import time
import importlib.util
from tflite_runtime.interpreter import Interpreter
from tflite_runtime.interpreter import load_delegate

#取得终端机之输入资讯
def get_cmd():
    parser = argparse.ArgumentParser()
    parser.add_argument('--image', help='path to file image', required=True)
    parser.add_argument('--model', help='Path to tflite model.', required=True)
    return parser.parse_args()

#主要程序
def main():
    args = get_cmd() #读取终端机提供的资讯
    interpreter = Interpreter(args.model)
    interpreter.allocate_tensors() #读取模型架构资讯 Get model graph details
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    width = input_details[0]['shape'][2]
    height = input_details[0]['shape'][1]

    # 读取影像与处理成输入张量
    frame = cv2.imread(args.image)
    image_height,image_width=frame.shape[:2] 
    frame_counter = 0
    start = time.time()
    frame_counter += 1
    frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) 
    frame_resized = cv2.resize(frame_gray, (width, height))
    input_data = np.expand_dims(frame_resized, axis=0)
    input_data = np.expand_dims(input_data, axis=3)
    interpreter.set_tensor(input_details[0]['index'], input_data) # set frame as input tensors
    interpreter.invoke()# perform inference

    # 定义输出张量
    classes = interpreter.get_tensor(output_details[0]['index'])[0] 
    prob_max_num = np.where(classes==classes.max())[0] 
    print(prob_max_num)

if __name__ == "__main__":
    main()