import json
from labelme import utils
from PIL import Image
import numpy as np
from pathlib import Path
import glob

import matplotlib.pyplot as plt
# 定义文件夹路径
bad_folder_path = Path('bad')
mask_folder_path = Path('mask')

# 创建保存 Mask 的文件夹（如果不存在）
mask_folder_path.mkdir(parents=True, exist_ok=True)

def process_json_file(json_file_path, mask_folder_path):
    """处理单个 JSON 文件，生成并保存 Mask 图片"""
    try:
        # 加载 JSON 文件
        with open(json_file_path, 'r') as f:
            data = json.load(f)

        # 获取图像尺寸
        img_height = data['imageHeight']
        img_width = data['imageWidth']

        # 解析标注信息
        shapes = data['shapes']
        label_name_to_value = {'_background_': 0}
        for shape in shapes:
            label_name = shape['label']
            if label_name not in label_name_to_value:
                label_name_to_value[label_name] = len(label_name_to_value)

        # 生成 Mask
        mask = utils.shapes_to_label(
            img_shape=(img_height, img_width),
            shapes=shapes,
            label_name_to_value=label_name_to_value
        )

        # 提取 Mask 部分（假设 mask 是一个元组，第一个元素是 Mask）
        if isinstance(mask, tuple):
            mask = mask[0]  # 提取 Mask 数组

        # 将 Mask 转换为图像
        mask_image = Image.fromarray((mask * 255).astype(np.uint8)) # 缩放

        # 保存 Mask 图片
        mask_save_path = mask_folder_path / (Path(json_file_path).stem + '.png')
        mask_image.save(mask_save_path, 'PNG')
        # plt.imshow(mask_image, cmap='gray')
        # plt.show()
        print(f"Saved: {mask_save_path}")

    except Exception as e:
        print(f"Error processing {json_file_path}: {e}")

# 获取文件夹内所有 JSON 文件路径
json_files_path = glob.glob(str(bad_folder_path / '*.json'))

# 处理所有 JSON 文件
for json_file_path in json_files_path:
    process_json_file(json_file_path, mask_folder_path)