import os
import torch
from mmseg.apis import inference_segmentor, init_segmentor

def main():
    # 設定路徑
    config_file = 'configs/stdc/kn_stdc1_in1k-pre_512x1024_80k_cityscapes.py'
    checkpoint_file = 'work_dirs/kn_stdc1_in1k-pre_512x1024_80k_cityscapes/latest.pth'
    img_dir = 'data/cityscapes/leftImg8bit/val'
    out_dir = 'work_dirs/vis_results'

    # 初始化模型
    model = init_segmentor(config_file, checkpoint_file, device='cuda:0')
    print('CLASSES：', model.CLASSES)
    print('PALETTE：', model.PALETTE)
    # 建立輸出資料夾
    os.makedirs(out_dir, exist_ok=True)

    # 找出所有圖片檔
    img_list = []
    for root, _, files in os.walk(img_dir):
        for f in files:
            if f.endswith('.png') or f.endswith('.jpg'):
                img_list.append(os.path.join(root, f))

    # 推論每一張圖片
    for img_path in img_list:
        result = inference_segmentor(model, img_path)
        filename = os.path.basename(img_path)
        out_path = os.path.join(out_dir, filename)
        model.show_result(img_path, result, out_file=out_path, opacity=0.5)

    print(f'✅ 推論完成，共處理 {len(img_list)} 張圖片，結果已輸出至：{out_dir}')

if __name__ == '__main__':
    main()