How to run Keras model on RK3399Pro | DLology

convert_rknn_snippets.py

from rknn.api import RKNN


INPUT_NODE = ["input_1"]
OUTPUT_NODE = ["predictions/Softmax"]

img_height = 299

# Create RKNN object
rknn = RKNN()

# pre-process config
# channel_mean_value "0 0 0 255" while normalize the image data to range [0, 1]
# channel_mean_value "128 128 128 128" while normalize the image data to range [-1, 1]
# reorder_channel "0 1 2" will keep the color channel, "2 1 0" will swap the R and B channel,
# i.e. if the input is BGR loaded by cv2.imread, it will convert it to RGB for the model input.
# need_horizontal_merge is suggested for inception models (v1/v3/v4).
rknn.config(
    channel_mean_value="128 128 128 128",
    reorder_channel="0 1 2",
    need_horizontal_merge=True,
    quantized_dtype="asymmetric_quantized-u8",
)

# Load tensorflow model
ret = rknn.load_tensorflow(
    tf_pb="./model/frozen_model.pb",
    inputs=INPUT_NODE,
    outputs=OUTPUT_NODE,
    input_size_list=[[img_height, img_height, 3]],
)
if ret != 0:
    print("Load inception_v3 failed!")
    exit(ret)

# Build model
# dataset: A input data set for rectifying quantization parameters.
ret = rknn.build(do_quantization=True, dataset="./dataset.txt")
if ret != 0:
    print("Build inception_v3 failed!")
    exit(ret)

# Export rknn model
ret = rknn.export_rknn("./inception_v3.rknn")
if ret != 0:
    print("Export inception_v3.rknn failed!")
    exit(ret)