import os
import cv2
import numpy as np
import pandas as pd
from pathlib import Path
from tqdm import tqdm

# ✅ 設定資料夾
datasets = [
    {
        "name": "train",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\0507stdcgrass\train",
        "output_img_dir": r"C:\Users\rd_de\kneronstdc\data\cityscapes\leftImg8bit\train",
        "output_mask_dir": r"C:\Users\rd_de\kneronstdc\data\cityscapes\gtFine\train",
    },
    {
        "name": "val",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\0507stdcgrass\valid",
        "output_img_dir": r"C:\Users\rd_de\kneronstdc\data\cityscapes\leftImg8bit\val",
        "output_mask_dir": r"C:\Users\rd_de\kneronstdc\data\cityscapes\gtFine\val",
    },
    {
        "name": "test",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\0507stdcgrass\test",
        "output_img_dir": r"C:\Users\rd_de\kneronstdc\data\cityscapes\leftImg8bit\test",
        "output_mask_dir": r"C:\Users\rd_de\kneronstdc\data\cityscapes\gtFine\test",
    }
]

# ✅ 只保留 grass: label 1 → 0，其餘變成 ignore (255)
for dataset in datasets:
    name = dataset["name"]
    input_dir = dataset["input_dir"]
    output_img_dir = dataset["output_img_dir"]
    output_mask_dir = dataset["output_mask_dir"]

    os.makedirs(output_img_dir, exist_ok=True)
    os.makedirs(output_mask_dir, exist_ok=True)

    print(f"\n📂 資料集: {name}")

    for file in tqdm(os.listdir(input_dir), desc=f"轉換 {name}"):
        if not file.endswith("_mask.png"):
            continue

        mask_path = os.path.join(input_dir, file)
        image_name = file.replace("_mask.png", ".jpg")
        image_path = os.path.join(input_dir, image_name)

        if not os.path.exists(image_path):
            print(f"⚠️ 找不到對應圖片: {image_name}")
            continue

        img = cv2.imread(image_path)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

        if img is None or mask is None:
            print(f"❌ 無法讀取圖像或 mask: {file}")
            continue

        # ✅ 將草地 (1) → 0，其餘設為 ignore (255)
        remapped_mask = np.full_like(mask, 255, dtype=np.uint8)
        remapped_mask[mask == 1] = 0

        # ✅ 輸出
        stem = Path(file).stem.replace("_mask", "")
        out_img_path = os.path.join(output_img_dir, f"{stem}_leftImg8bit.png")
        out_mask_path = os.path.join(output_mask_dir, f"{stem}_gtFine_labelIds.png")

        cv2.imwrite(out_img_path, img)
        cv2.imwrite(out_mask_path, remapped_mask)

print("\n🎉 轉換完成，只保留 grass 類別（label=0）！")