上篇讲述如何把tensorflow模型转换成tflite模型,用于部署到移动端。
这篇分享如何在PC端对tflite模型进行预测,测试模型是否可用
首先,加载tflite模型,查看模型的输入输出
import numpy as np
import tensorflow as tf
import cv2
# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path="newModel.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print(input_details)
print(output_details)
打印结果:
[{'name': 'input_1', 'index': 115, 'shape': array([ 1, 224, 224, 3]), 'dtype': , 'quantization': (0.0, 0)}]
[{'name': 'activation_1/truediv', 'index': 6, 'shape': array([ 1, 12544, 2]), 'dtype': , 'quantization': (0.0, 0)}]
从input details
里可看到,需要输入的numpy数组 [1, 224, 224, 3],数据类型为float32。
index:115
的意思,我理解是传入数据放置的位置
于是我们需要这里处理下input的数据,为了方便,我把数据预存成csv文件了,取出来能直接使用
input_data = np.loadtxt('C:/Users\WIN10/input.csv',delimiter=',')
input_data = input_data.reshape(1,224,224,3)
input_data = input_data.astype(np.float32)
index = input_details[0]['index']
interpreter.set_tensor(index, input_data)
这样就把input数据设置好了,并且把数据传入网络模型
开始做预测:
interpreter.invoke()
读出预测结果
output_data = interpreter.get_tensor(output_details[0]['index'])
print('output_data shape:',output_data.shape)
output_data就是预测结果的源数据,print出来的shape就是上面output_detail 里的shape[ 1, 12544, 2]
最后还需要对预测的源数据进行后处理,解析出我们要的结果,这里每个工程的都不一样,仅供参考。
output_data = output_data.reshape(224,112)
pr = output_data.reshape(112,112,2).argmax( axis=2 )
seg_img = np.zeros( ( 112 , 112 , 3 ) )
seg_img[:,:,0] += ((pr[:,: ] == 1 )*200).astype('uint8')
seg_img[:,:,1] += ((pr[:,: ] == 1 )*200).astype('uint8')
seg_img[:,:,2] += ((pr[:,: ] == 1 )*200).astype('uint8')
cv2.imshow('img',seg_img)
cv2.waitKey(0)