#!/usr/bin/env python3
import cv2
import numpy as np
from rknnlite.api import RKNNLite
from dataclasses import dataclass
from typing import List, Tuple, Optional

@dataclass
class Detection:
    class_id: int
    class_name: str
    confidence: float
    bbox: Tuple[int, int, int, int]

LABEL_MAP = {0: '安全帽', 4: '安全衣', 3: '人'}
INPUT_SIZE = (640, 640)

def nms(dets, iou_threshold=0.45):
    if len(dets) == 0:
        return []
    
    boxes = np.array([[d.bbox[0], d.bbox[1], d.bbox[2], d.bbox[3], d.confidence] for d in dets])
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    scores = boxes[:, 4]
    
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]
    
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        
        inds = np.where(ovr <= iou_threshold)[0]
        order = order[inds + 1]
    
    return [dets[i] for i in keep]

def letterbox(image, input_size=(640, 640)):
    h0, w0 = image.shape[:2]
    ih, iw = input_size
    scale = min(iw / w0, ih / h0)
    new_w, new_h = int(w0 * scale), int(h0 * scale)
    pad_w = (iw - new_w) // 2
    pad_h = (ih - new_h) // 2
    resized = cv2.resize(image, (new_w, new_h))
    canvas = np.full((ih, iw, 3), 114, dtype=np.uint8)
    canvas[pad_h:pad_h+new_h, pad_w:pad_w+new_w] = resized
    return canvas, scale, pad_w, pad_h, h0, w0

def test_model():
    model_path = "yolo11m_safety.rknn"
    conf_threshold_map = {3: 0.8, 0: 0.5, 4: 0.5}
    
    rknn = RKNNLite()
    
    ret = rknn.load_rknn(model_path)
    if ret != 0:
        print("[ERROR] load_rknn failed")
        return
    
    ret = rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2)
    if ret != 0:
        print("[ERROR] init_runtime failed")
        return
    
    image = cv2.imread("b.jpg")
    if image is None:
        print("无法读取测试图片")
        return
    
    canvas, scale, pad_w, pad_h, h0, w0 = letterbox(image)
    
    # RKNN expects NHWC input: (1, H, W, C), RGB, float32 normalized 0-1
    img = canvas[..., ::-1].astype(np.float32) / 255.0
    blob = img[None, ...]  # (1, 640, 640, 3)
    
    outputs = rknn.inference(inputs=[blob])
    
    if outputs:
        output = outputs[0]
        if len(output.shape) == 3:
            output = output[0]
        
        num_classes = 5
        dets = []
        for i in range(output.shape[1]):
            x_center = float(output[0, i])
            y_center = float(output[1, i])
            width = float(output[2, i])
            height = float(output[3, i])
            
            class_probs = output[4:4+num_classes, i]
            best_class = int(np.argmax(class_probs))
            confidence = float(class_probs[best_class])
            
            if best_class not in LABEL_MAP:
                continue
            
            conf_threshold = conf_threshold_map.get(best_class, 0.5)
            if confidence < conf_threshold:
                continue
            
            # Remove padding and scale to original image
            x1 = int(((x_center - width / 2) - pad_w) / scale)
            y1 = int(((y_center - height / 2) - pad_h) / scale)
            x2 = int(((x_center + width / 2) - pad_w) / scale)
            y2 = int(((y_center + height / 2) - pad_h) / scale)
            
            x1 = max(0, min(w0, x1))
            y1 = max(0, min(h0, y1))
            x2 = max(0, min(w0, x2))
            y2 = max(0, min(h0, y2))
            
            det = Detection(
                class_id=best_class,
                class_name=LABEL_MAP[best_class],
                confidence=confidence,
                bbox=(x1, y1, x2, y2)
            )
            dets.append(det)
        
        dets = nms(dets, iou_threshold=0.45)
        
        print(f"检测结果: {len(dets)} 个目标")
        for d in dets:
            print(f"  {d.class_name}: conf={d.confidence:.3f}, box={d.bbox}")
    
    rknn.release()

if __name__ == "__main__":
    test_model()