分割图像的json文件，多个分割目标选择自己需要的分割对象，生成伪颜色mask训练数据

from PIL import Image, ImageDraw
import json
import numpy as np
import shutil
import os
import glob
import tqdm


def pro():
    # 图片有，json 没有有，遍历删除多余的图片
    # 定义JSON文件所在的目录,所有json 数据提取有内中膜和血管的json 分割到另一个文件夹
    
    save_json_dir = "/home/syy/data/carotid_all/vessel_intima_jsons/"
    json_dir = "/home/syy/data/carotid_all/jsons"
    
    # 存储所有标签名称的列表
    all_labels = []
    n=0
    # 遍历JSON文件
    for json_file in os.listdir(json_dir):
        list_set = set()
        src_path = ""
        dst_path=""
        src_img_path=""
        dst_img_path=""
        
        yes=0
        if json_file.endswith(".json"):
            with open(os.path.join(json_dir, json_file), "r") as f:
                data = json.load(f)
                imageName = data["imagePath"]
                # 获取该JSON文件中的标签信息
                # "label": "intima",
                # "label": "vessel",
                
                if "shapes" in data:
                    for shape in data["shapes"]:
                        label = shape["label"]
                        list_set.add(label)
                
                if "intima" in list_set and "vessel" in list_set   :
                    print(list_set)
                    print(json_file)
                    yes=1
                    
                    src_path = os.path.join(json_dir, json_file)
                    dst_path = os.path.join(save_json_dir, json_file)
                    
                    src_img_path = "/home/syy/data/carotid_all/images/" + imageName
                    dst_img_path = "/home/syy/data/carotid_all/vessel_intima_imgs/" + imageName
        if yes:
            shutil.copy(src_path, dst_path)
            shutil.copy(src_img_path,dst_img_path)
            n=n+1
            print(n)
    
   
                
                
    
    # # 打印所有标签名称
    # for label in all_labels:
    #     print(label)


def get_all_json_label():
    # 分割的json 文件得到所有的label
    json_dir = "/home/syy/data/carotid_all/vessel_intima_jsons"

    # 存储所有标签名称的列表，不重复
    list_set = set()

    # 遍历JSON文件
    for json_file in os.listdir(json_dir):
        if json_file.endswith(".json"):
            with open(os.path.join(json_dir, json_file), "r") as f:
                data = json.load(f)
            
                # 获取该JSON文件中的标签信息
                # "label": "intima",
                # "label": "vessel",
                
                if "shapes" in data:
                    for shape in data["shapes"]:
                        label = shape["label"]
                        list_set.add(label)
    print(list_set)
    
    
def get_color_map_list(num_classes):
    """
    Returns the color map for visualizing the segmentation mask,
    which can support arbitrary number of classes.
    Args:
        num_classes (int): Number of classes.
    Returns:
        (list). The color map.
    """

    num_classes += 1
    color_map = num_classes * [0, 0, 0]
    for i in range(0, num_classes):
        j = 0
        lab = i
        while lab:
            color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
            color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
            color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
            j += 1
            lab >>= 3
    color_map = color_map[3:]
    return color_map



def get_indexes_np(arr, target):
    return np.where(arr == target)[0].tolist()


def json_to_mask(json_file,save_mask_root,label_to_id ,concat_label):
    color_map = get_color_map_list(256)
    #提供 json_file ，和img_file   绝对路径
    # json文件中的部分label , 或者全部label
    # 读取JSON文件
    with open(json_file, 'r') as json_file:
        data = json.load(json_file)
    
    # 获取图像的尺寸
    width = data['imageWidth']
    height = data['imageHeight']
    imageName = data["imagePath"]
    
    # 创建一个全零的数组作为伪颜色 mask，使用uint8数据类型
    mask = np.zeros((height, width), dtype=np.uint8)
    # 两个同步，用于大目标小目标，优先级谁，小目标黏贴大目标中，
    mask_list = []
    mask_label_list = []
    # 遍历分割数据
    
    for segment in data['shapes']:  # 一个json 文件多个分割标注，遍历
        label = segment['label']
        points = segment['points']
        polygon = np.array(points, dtype=np.int32)
    
    
        # 有些分割对象是包含关系需要，先做大的再做小的分割对象，不然就覆盖了,这里是每个分割对象有了对应唯一的值，之后合并出来的,没有这个问题
        
        if label in label_to_id:  # 字典中没有这个label 跳出
            label_value = label_to_id.get(label, -1)  # 获取键'intima'对应的值,不存在返回-1  # 已经是 0， 1 2这种label映射了
   
            # 使用 ImageDraw.Draw 来填充多边形区域
            mask_image = Image.fromarray(np.zeros((height, width), dtype=np.uint8))
            image_draw = ImageDraw.Draw(mask_image)
            
            points_list = [tuple(point) for point in points]
            assert len(points_list) > 2, 'Polygon must have points more than 2'
            image_draw.polygon(xy=points_list, outline=label_value, fill=label_value)
    
            # # 合并 'vessel' 和 'intima' 到一个伪颜色 mask
            # mask = mask + mask_image  # 0 , 1, 2, 如果有重叠的，说明有3，把3变成1
            mask_label_list.append(label)
            mask_list.append(mask_image)
            
            # 这种拼接没有交集可以这么拼接，有交集，以最大的像素为主，不适应多种情况
            # mask = np.maximum(mask, np.array(mask_image))  # 它用于比较两个数组（或数组和标量）的元素，并返回一个新的数组，该数组包含了两个输入数组中对应位置元素的较大值
        
    # 拼接目标的操作
    # mask_list.sort(key=lambda x: np.max(x), reverse=True)  # 大目标在前面，或者指定哪个目标在前面  2是血管
    # 拼接操作，大的在下面，小目标重叠黏贴到前一个图片上，这种操作
    # concat_label = ["vessel","intima"]  # 拼接图片优先级别
    for b in range(len(concat_label)):
        lab = concat_label[b]
        arr = np.array(mask_label_list)   # ["vessel","intima","intima","intima","vessel]
        target = lab
        index_list = get_indexes_np(arr, target) # 【0,4】
        # print("=====>", target,mask_label_list,index_list)
        for ind in index_list:
            mask_data = np.array(mask_list[ind])
            vv = np.max(np.unique(mask_data))
            # print("*******",lab,vv,np.unique(mask_data))
            
            # 大目标在前，按照小目标，特定顺序， 后者黏贴到前者上，保证小目标不会被覆盖
            # 后一张图片黏贴到前一张图片当中，两张图片大小是一致的
            mask = np.where(mask_data == vv, vv, mask)  # mask_data 中的像素值等于1的，区域黏贴到 mask 图片中并且像素值也是1
     
    all_label_id = np.unique(np.array(mask))
    if len(all_label_id)<len(label_to_id):# 背景  内中膜，血管  ，一定有3个
        print("分割图像和label,和给定的label 不一致")
        print(all_label_id,label_to_id)
        return
    
    # 保存伪颜色 mask  ，多个对象拼接好了之后，像素值不变还是0,1,2
    mask_all = Image.fromarray(mask.astype(np.uint8), mode='P')
    
    mask_all.putpalette(color_map)
    mask_all.save(save_mask_root+imageName)
    


if __name__ == "__main__":
    # 打印所有的lab_set信息
    # get_all_json_label()
    # {'intima', 'vessel', 'plaque'}
    
    # 1、对于 images  labels  文件夹，找到同时有内中膜和血管的图片和json 复制到新的文件
    # pro()  #符合要求的，同时有血管内中膜的图片+json，复制出来
    
    
    # 2、设置lab 顺序，内中膜1， 血管2，  传递参数到jsontomask 函数中，保存  可以训练的mask 伪颜色图片
    label_sort = ["__background__","intima","vessel"]  # 顺序不要错了，第一个一定是背景0  
    # 拼接顺序
    concat_label = ["vessel","intima"] #大目标覆盖小目标， 拼接多个分割对象，大目标在前，小的目标在后，后者直接黏贴到这个位置

    # 创建一个标签名称到标签ID的映射字典
    label_to_id = {label: i for i, label in enumerate(label_sort)}
    print(label_to_id)

    save_mask_root = "/home/syy/data/carotid_all/vessel_intima_mask/"
    for json_file in tqdm(glob.glob("/home/syy/data/carotid_all/vessel_intima_jsons/*.json")):
        json_to_mask(json_file, save_mask_root,label_to_id,concat_label)