import onnxruntime as ort
import numpy as np
from PIL import Image
import cv2

# === 1. 載入 ONNX 模型 ===
onnx_path = "work_dirs/meconfig8/latest.onnx"
session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])

# === 2. 前處理輸入圖像（724x362） ===
def preprocess(img_path):
    image = Image.open(img_path).convert("RGB")
    image = image.resize((724, 362), Image.BILINEAR)
    img = np.array(image) / 255.0
    img = np.transpose(img, (2, 0, 1))  # HWC → CHW
    img = np.expand_dims(img, 0).astype(np.float32)  # (1, 3, 362, 724)
    return img

img_path = "test.png"
input_tensor = preprocess(img_path)

# === 3. 執行推論 ===
input_name = session.get_inputs()[0].name
output = session.run(None, {input_name: input_tensor})  # list of np.array

# === 4. 後處理 + 預測 Mask ===
output_tensor = output[0][0]           # shape: (num_classes, H, W)
pred_mask = np.argmax(output_tensor, axis=0).astype(np.uint8)  # (H, W)

# === 5. 可視化結果 ===
colors = [
    [128, 0, 0],    # 0: bunker
    [0, 0, 128],    # 1: car
    [0, 128, 0],    # 2: grass
    [0, 255, 0],    # 3: greenery
    [255, 0, 0],    # 4: person
    [255, 165, 0],  # 5: road
    [0, 255, 255],  # 6: tree
]

color_mask = np.zeros((pred_mask.shape[0], pred_mask.shape[1], 3), dtype=np.uint8)
for cls_id, color in enumerate(colors):
    color_mask[pred_mask == cls_id] = color

# 儲存可視化圖片
cv2.imwrite("onnx_pred_mask.png", color_mask)
print("✅ 預測結果已儲存為：onnx_pred_mask.png")