ラズパイカメラによるMNIST手書き文字分類 with TensorFlow lite

import tensorflow as tf

localpb = 'frozen_inference_graph_frcnn.pb' #変換したいpbファイルの名前
tflite_file = 'frcnn_od.lite' #変換後のtfliteファイルの名前。デフォルトでは.liteになってるが正しくは.tfliteなので要修正

print("{} -> {}".format(localpb, tflite_file))

converter = tf.lite.TFLiteConverter.from_frozen_graph(
    localpb, 
    ["image_tensor"], #変換対象モデルの入力ノードの名前
    ['detection_boxes'] #変換対象モデルの出力ノードの名前
)

tflite_model = converter.convert()

open(tflite_file,'wb').write(tflite_model)

interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

import numpy as np
from tflite_runtime.interpreter import Interpreter
import time
import cv2

camera = cv2.VideoCapture(0) # カメラCh.(ここでは0)を指定

# TFLiteモデルの読み込み
interpreter = Interpreter(model_path="lenet_mnist.tflite")
# メモリ確保
interpreter.allocate_tensors()

# 学習モデルの入力層・出力層のプロパティをGet.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

def detect(image):
    frame = cv2.resize(image, (28,28)) #28*28ピクセルに画像縮小
    frame = frame / 255 #0.0~1.0に正規化

    #モデルに入力するために次元を合わせる
    #この操作により(1,1,28,28)の配列になる
    frame = np.expand_dims(frame, 0)
    frame = np.expand_dims(frame, 0)

    #float32形式にキャストする
    frame = frame.astype(np.float32)

    # indexにテンソルデータのポインタをセット
    interpreter.set_tensor(input_details[0]['index'], frame)
    start_time = time.time()
    interpreter.invoke() #推論実行
    stop_time = time.time()
    print("time: ", stop_time - start_time) #推論時間の表示

    # get results
    label = interpreter.get_tensor(output_details[0]['index'])
    return label


if __name__ == '__main__':

    while True:
        ret, frame = camera.read()  #フレームを取得
        frame = frame[ : , 80:560 ]  #正方形にトリミング。640*480解像度なので、横方向を80~560(480)トリミング
        frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) #グレースケール化
        ret, frame_bin = cv2.threshold(frame_gray, 50, 255, cv2.THRESH_BINARY_INV)
        #↑白黒反転二値化。MNISTデータセットが白黒反転のため、取得画像も反転させる必要あり。50は二値化の閾値。255は最大値。

        cv2.imshow('camera', frame)  # フレームを画面に表示。ここでは反転前の生画像を表示している。

        num = detect(frame_bin)  #推論結果の保存
        print("label: " + str(np.argmax(num)) + " accuracy: " + str(np.amax(num)))  #推論結果の表示

        # キー操作があればwhileループを抜ける
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # 撮影用オブジェクトとウィンドウの解放
    camera.release()
    cv2.destroyAllWindows()