rknn_classify_image.py

import numpy as np
import time
from PIL import Image
from rknnlite.api import RKNNLite

def load_labels(label_file):
    labels = {}
    with open(label_file, 'r') as f:
        for index, line in enumerate(f.readlines()):
            labels[index] = line.strip()
    return labels

if __name__ == '__main__':
    labels = load_labels('labels.txt')

    rknn_lite = RKNNLite(verbose=False)

    print('--> Loading RKNN model')
    ret = rknn_lite.load_rknn('./mobilenet_v2_1.0_224_inat_bird_quant.rknn')
    if ret != 0:
        print('Load RKNN model failed!')
        exit(ret)
    print('done')

    print('--> Init runtime environment')
    ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_ALL)
    if ret != 0:
        print('Init runtime environment failed!')
        exit(ret)
    print('done')

    image = Image.open('parrot.jpg')
    size = (224, 224)
    image = image.convert('RGB').resize(size, Image.LANCZOS)

    image_data = np.asarray(image, dtype=np.uint8)
    image_data = np.expand_dims(image_data, axis=0)

    print('-----INFERENCE TIME-----')
    for _ in range(100):
        start = time.perf_counter()
        outputs = rknn_lite.inference(inputs=[image_data])
        inference_time = time.perf_counter() - start
        print('%.1fms' % (inference_time * 1000))

    print('-------RESULTS--------')
    predictions = outputs[0].squeeze()
    max_index = np.argmax(predictions)
    max_prob = predictions[max_index]

    predicted_label = labels.get(max_index, "Unknown")
    print(f'{predicted_label}: {max_prob:.5f}')

    rknn_lite.release()