OpenVINO 日本語手書き文字認識 to TensorflowLite model test

example.py

import sys
import cv2
import numpy as np
import tensorflow as tf

def preprocess_input(image_name, height, width):
    src = cv2.imread(image_name, cv2.IMREAD_GRAYSCALE)
    ratio = float(src.shape[1]) / float(src.shape[0])
    tw = int(height * ratio)
    rsz = cv2.resize(src, (tw, height), interpolation=cv2.INTER_AREA).astype(np.float32)
    # [h,w] -> [c,h,w]
    img = rsz[None, :, :]
    _, h, w = img.shape
    pad_img = np.pad(img, ((0, 0), (0, height - h), (0, width -  w)), mode='edge')
    return pad_img

def main():
    
    imgf = 'test.jpg'
    modelf = 'handwritten_japanese_recognition_90x2000_float16_quant.tflite'

    interpreter = tf.lite.Interpreter(model_path=modelf)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    input_shape = input_details[0]['shape']
    input_batch_size, input_height, input_width, input_channel = input_shape

    input_image = preprocess_input(imgf, height=input_height, width=input_width)[None,:,:,:]
    input_image = input_image.transpose((0,2,3,1))

    # Input (1, 96, 2000, 1)
    input_data = np.array(input_image, dtype=np.float32)
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()

    # Output (186, 1, 4442)
    preds = interpreter.get_tensor(output_details[0]['index'])
    preds_index = np.argmax(preds, 2)
    preds_index = preds_index.transpose(1, 0).reshape(-1)

    print(preds_index)
    sys.exit()


if __name__ == '__main__':
    main()