# -*- coding: utf-8 -*-
import os.path as osp
import shutil
import os
from tqdm import tqdm
import cv2

# 此函數將原始的紅色圖像轉換為黑白的二值圖像
def gt2png(folder_path="C:/Users/chenmingsong/Desktop/unetnnn/data/Training_GT", save_folder="C:/Users/chenmingsong/Desktop/unetnnn/data/Training_Labels"):
    # folder_path = osp.join(sys_path, "ISBI2016_ISIC_Part1_Training_GroundTruth")
    # save_folder = osp.join(sys_path, "labels")
    
    # 如果儲存資料夾已存在，則刪除後重新建立
    if osp.isdir(save_folder):
        shutil.rmtree(save_folder)
        os.makedirs(save_folder)
    else:
        os.makedirs(save_folder)
    
    images = os.listdir(folder_path)
    
    # 使用進度條顯示轉換進度
    with tqdm(total=len(images)) as pbar:
        for image in images:
            image_name = image.split(".")[0]
            src_path = osp.join(folder_path, image)
            # 讀取圖像，將其轉換為灰階模式
            img = cv2.imread(src_path, 0)
            # 將非零像素設為 255
            img[img > 0] = 255
            save_path = osp.join(save_folder, image_name + ".png")
            # 儲存轉換後的圖像
            cv2.imwrite(save_path, img)
            pbar.update(1)
    
    print("標籤轉換完成")


if __name__ == '__main__':
    gt2png()