# onnx_segmentation_demo.py

import cv2
import numpy as np
import onnxruntime as ort
import matplotlib.pyplot as plt


def preprocess(image_path, input_size=(640, 640)):
    """图像预处理：resize, normalize, HWC → CHW"""
    img = cv2.imread(image_path)
    img_resized = cv2.resize(img, input_size)  # resize
    img_rgb = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)

    # 归一化: [0,255] → [0,1] 或 ImageNet: (x - mean)/std
    img_norm = img_rgb.astype(np.float32) / 255.0
    # img_norm = (img_rgb.astype(np.float32) / 255.0 - 0.5) / 0.5  # 如果是 [-1,1]

    # HWC → CHW → 1xCxHxW
    img_chw = np.transpose(img_norm, (2, 0, 1))
    input_tensor = np.expand_dims(img_chw, axis=0)

    return input_tensor, img_rgb  # 返回输入张量和原始图像（用于显示）


def postprocess(output, original_shape, num_classes=1, threshold=0.5):
    """
    后处理 ONNX 输出
    :param output: ONNX 推理输出 (C, H, W) 或 (1, C, H, W)
    :param original_shape: 原图 (H, W, C)
    :param num_classes: 类别数（1: 二值分割；N: 多类分割）
    :param threshold: 二值化阈值
    :return: mask (H, W) 或 (H, W, C)
    """
    # 假设输出 shape: (1, C, H, W) 或 (C, H, W)
    if output.ndim == 4:
        output = output[0]  # 去掉 batch 维度

    h, w = original_shape[:2]

    if num_classes == 1:
        # 二值分割: Sigmoid 输出
        prob_map = output[0]  # (H, W)
        mask = (prob_map > threshold).astype(np.uint8) * 255
        mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
        return mask
    else:
        # 多类分割: Softmax 输出，shape (C, H, W)
        pred = np.argmax(output, axis=0)  # (H, W)
        mask = pred.astype(np.uint8)
        mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
        return mask


def visualize_segmentation(image, mask, alpha=0.5, colormap=cv2.COLORMAP_JET):
    """叠加 mask 到原图上"""
    if len(mask.shape) == 2 and mask.max() <= 1:
        mask = (mask * 255).astype(np.uint8)

    # 彩色化 mask
    colored_mask = cv2.applyColorMap(mask, colormap)
    # 融合
    overlay = cv2.addWeighted(image, 1 - alpha, colored_mask, alpha, 0)
    return overlay


# ==================== 主程序 ====================

if __name__ == "__main__":
    ONNX_MODEL = "best.onnx"
    TEST_IMAGE = "1.jpg"
    INPUT_SIZE = (640, 640)
    NUM_CLASSES = 1  # 1: 二值分割；>1: 多类分割
    THRESHOLD = 0.8

    # 1. 加载 ONNX 模型
    print("📦 加载 ONNX 模型...")
    ort_session = ort.InferenceSession(ONNX_MODEL, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

    # 获取输入名称
    input_name = ort_session.get_inputs()[0].name

    # 2. 预处理
    input_tensor, original_image = preprocess(TEST_IMAGE, INPUT_SIZE)
    print(f"📥 输入形状: {input_tensor.shape}")

    # 3. 推理
    print("🚀 开始推理...")
    outputs = ort_session.run(None, {input_name: input_tensor})
    logits = outputs[0]  # 假设只有一个输出
    print(f"📤 输出形状: {logits.shape}")

    # 4. 后处理
    mask = postprocess(logits, original_image.shape, num_classes=NUM_CLASSES, threshold=THRESHOLD)

    # 5. 可视化
    overlay = visualize_segmentation(original_image, mask, alpha=0.6)

    # 6. 显示
    plt.figure(figsize=(12, 6))
    plt.subplot(1, 3, 1)
    plt.title("原始图像")
    plt.imshow(original_image)
    plt.axis("off")

    plt.subplot(1, 3, 2)
    plt.title("预测 Mask")
    plt.imshow(mask, cmap="gray")
    plt.axis("off")

    plt.subplot(1, 3, 3)
    plt.title("叠加效果")
    plt.imshow(overlay)
    plt.axis("off")

    plt.tight_layout()
    plt.show()

    # （可选）保存结果
    # cv2.imwrite("overlay.jpg", cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
    print("✅ 验证完成！")