import cv2
import numpy as np
from rknnlite import RKNNLite

# 初始化 RKNN
rknn = RKNNLite()
rknn.load_rknn("yolov8n.rknn")
rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)

def preprocess(image):
    # 调整尺寸并转换为 BCHW 格式
    img = cv2.resize(image, (640, 640))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # BGR 转 RGB（根据模型配置）
    img = img.transpose(2, 0, 1)[np.newaxis, ...].astype(np.uint8)  # HWC -> BCHW
    return img

# 加载测试图像
img = cv2.imread("test.jpg")
input_data = preprocess(img)

# 推理
outputs = rknn.inference(inputs=[input_data])[0]  # 输出形状 [1, 84, 8400]

# 后处理（示例：提取检测框）
def postprocess(output, conf_thresh=0.5):
    # 输出格式: [batch, 84, 8400]
    # 84 = 4（坐标） + 80（类别）
    boxes = []
    for i in range(output.shape[2]):
        scores = output[0, 4:84, i]
        class_id = np.argmax(scores)
        confidence = scores[class_id]
        if confidence > conf_thresh:
            cx, cy, w, h = output[0, :4, i]
            # 转换为 x1,y1,x2,y2 格式
            x1 = int((cx - w/2) * 640)
            y1 = int((cy - h/2) * 640)
            x2 = int((cx + w/2) * 640)
            y2 = int((cy + h/2) * 640)
            boxes.append([x1, y1, x2, y2, confidence, class_id])
    return boxes

# 执行后处理
detections = postprocess(outputs)
print("检测结果:", detections)

rknn.release()