import os
import cv2
import numpy as np
from rknnlite.api import RKNNLite

objectThresh = 0.5  # 置信度阈值，可调整

def letterbox_resize(image, size, bg_color=114):
    target_width, target_height = size
    h, w, _ = image.shape
    scale = min(target_width / w, target_height / h)
    new_w, new_h = int(w * scale), int(h * scale)
    resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
    canvas = np.ones((target_height, target_width, 3), dtype=np.uint8) * bg_color
    offset_x, offset_y = (target_width - new_w) // 2, (target_height - new_h) // 2
    canvas[offset_y:offset_y+new_h, offset_x:offset_x+new_w] = resized
    return canvas, scale, offset_x, offset_y

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def detect_top2_single_masks(model_path, image_path):
    img = cv2.imread(image_path)
    if img is None:
        print(f"❌ 无法读取图像 {image_path}")
        return []

    img_resized, scale, offset_x, offset_y = letterbox_resize(img, (640, 640))
    infer_img = img_resized[..., ::-1]  # BGR->RGB
    infer_img = np.expand_dims(infer_img, 0)

    rknn = RKNNLite(verbose=True)
    rknn.load_rknn(model_path)
    rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
    outputs = rknn.inference([infer_img])
    rknn.release()

    mask_coeffs_list = [outputs[3], outputs[7], outputs[11]]  # mask coefficients
    conf_list = [outputs[1], outputs[5], outputs[9]]           # object confidence
    proto = outputs[12][0]                                     # proto [32,160,160]

    # 遍历所有候选，生成(置信度, 尺度索引, 目标索引)
    candidates = []
    for scale_idx, conf_map in enumerate(conf_list):
        conf_map_flat = conf_map.flatten()
        for idx, conf in enumerate(conf_map_flat):
            if conf > objectThresh:
                candidates.append((conf, scale_idx, idx))

    if not candidates:
        print(f"⚠️ 未检测到目标，置信度低于 {objectThresh}")
        return []

    # 按置信度排序，取前两个
    candidates.sort(key=lambda x: x[0], reverse=True)
    top2 = candidates[:2]

    masks = []
    for i, (conf, scale_idx, idx) in enumerate(top2):
        coeff = mask_coeffs_list[scale_idx].reshape(mask_coeffs_list[scale_idx].shape[1], -1)
        mask_flat = np.matmul(coeff[:, idx], proto.reshape(proto.shape[0], -1))
        mask = sigmoid(mask_flat).reshape(proto.shape[1], proto.shape[2])
        mask_resized = cv2.resize(mask, (img.shape[1], img.shape[0]))
        mask_bin = (mask_resized > 0.5).astype(np.uint8) * 255
        masks.append(mask_bin)
        cv2.imwrite(f"mask_single_{i+1}.png", mask_bin)
        print(f"✅ mask {i+1} 已保存: mask_single_{i+1}.png")

    return masks

if __name__ == "__main__":
    model_path = "seg.rknn"
    image_path = "2.jpg"
    detect_top2_single_masks(model_path, image_path)