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

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=1src_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)