import os
import sys
import time
import cv2
import numpy as np
import pandas as pd
from rknnlite.api import RKNNLite

pd.set_option("display.max_rows", 1000)

dim = 448
thresh = 0.5

# 初始化RKNNLite
rknn_lite = RKNNLite(verbose=False)

# 加载RKNN模型
ret = rknn_lite.load_rknn("model.rknn")
if ret != 0:
    print('加载RKNN模型失败')
    exit(ret)

# 初始化运行时环境
ret = rknn_lite.init_runtime()
if ret != 0:
    print('初始化运行时环境失败')
    exit(ret)

label_names = pd.read_csv("selected_tags.csv")

target_img = "input.jpg" if len(sys.argv) < 2 else sys.argv[1]

try:
    # 图像预处理
    # 1. 读取图像并转RGB
    img = cv2.imread(target_img)
    if img is None:
        print(f"无法读取图像: {target_img}")
        exit(1)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    # 2. 处理透明通道
    if img.shape[2] == 4:  # RGBA图像
        img = img[:, :, :3]  # 只保留RGB通道
    elif len(img.shape) == 2:  # 灰度图
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
    
    # 3. 填充为正方形
    h, w = img.shape[:2]
    if h > w:
        diff = h - w
        pad_left = diff // 2
        pad_right = diff - pad_left
        img = cv2.copyMakeBorder(img, 0, 0, pad_left, pad_right, 
                                cv2.BORDER_CONSTANT, value=(255, 255, 255))
    elif w > h:
        diff = w - h
        pad_top = diff // 2
        pad_bottom = diff - pad_top
        img = cv2.copyMakeBorder(img, pad_top, pad_bottom, 0, 0, 
                                cv2.BORDER_CONSTANT, value=(255, 255, 255))
    
    # 4. resize到目标尺寸
    img = cv2.resize(img, (dim, dim), interpolation=cv2.INTER_AREA)
    
    # 5. 转换为float32并添加batch维度
    img = img.astype(np.float32)
    img = np.expand_dims(img, 0)
    print(img.shape) # 1, 448, 448, 3
    # NHWC -> NWCH
    # img = np.transpose(img, (0, 2, 3, 1))
    # 执行推理
    # np.save("img.npy",img)
    start = time.time()
    outputs = rknn_lite.inference(inputs=[img], data_format="nhwc")
    end = time.time()
    print(f"Inference time: {end - start} seconds")
    probs = outputs[0]  # 获取第一个输出
    print(probs.shape)

    # 后处理
    label_names["probs"] = probs[0]
    found_tags = label_names[label_names["probs"] > thresh][["tag_id", "name", "probs"]]
    print(found_tags)

finally:
    # 释放资源
    rknn_lite.release()