import numpy as np
import ktc
import cv2
from PIL import Image

# === 1. 前處理 + 推論 ===
def run_e2e_simulation(img_path, onnx_path):
    # 圖片前處理（724x362）
    image = Image.open(img_path).convert("RGB")
    image = image.resize((724, 362), Image.BILINEAR)
    img_data = np.array(image) / 255.0
    img_data = np.transpose(img_data, (2, 0, 1))  # HWC → CHW
    img_data = np.expand_dims(img_data, 0)        # → NCHW (1,3,362,724)

    input_data = [img_data]
    inf_results = ktc.kneron_inference(
        input_data,
        onnx_file=onnx_path,
        input_names=["input"]
    )

    return inf_results

# === 2. 呼叫推論 ===
image_path = "test.png"
onnx_path = "work_dirs/meconfig8/latest_optimized.onnx"
result = run_e2e_simulation(image_path, onnx_path)

print("推論結果 shape:", np.array(result).shape)  # (1, 1, 7, 46, 91)

# === 3. 提取與處理輸出 ===
output_tensor = np.array(result)[0][0]        # shape: (7, 46, 91)
pred_mask = np.argmax(output_tensor, axis=0)  # shape: (46, 91)

print("預測的 segmentation mask：")
print(pred_mask)

# === 4. 上採樣回 724x362 ===
upsampled_mask = cv2.resize(pred_mask.astype(np.uint8), (724, 362), interpolation=cv2.INTER_NEAREST)

# === 5. 上色（簡單使用固定 palette）===
# 根據你的 7 類別自行定義顏色 (BGR)
colors = np.array([
    [0, 0, 0],          # 0: 背景
    [0, 255, 0],        # 1: 草地
    [255, 0, 0],        # 2: 車子
    [0, 0, 255],        # 3: 人
    [255, 255, 0],      # 4: 道路
    [255, 0, 255],      # 5: 樹
    [0, 255, 255],      # 6: 其他
], dtype=np.uint8)

colored_mask = colors[upsampled_mask]  # shape: (362, 724, 3)
colored_mask = np.asarray(colored_mask, dtype=np.uint8)

# === 6. 檢查並儲存 ===
if colored_mask.shape != (362, 724, 3):
    raise ValueError(f"❌ mask shape 不對: {colored_mask.shape}")

cv2.imwrite("pred_mask_resized.png", colored_mask)
print("✅ 已儲存語意遮罩圖：pred_mask_resized.png")