coco语义分割标注json转png

参考COPY FROM：

coco语义分割标注json转png_coco格式的语义分割标签_Coding的叶子的博客-CSDN博客

repair the up file's bug ,after get my file:

follow code:

import shutil
import cv2
from PIL import Image
import imgviz
import argparse
import os
import tqdm
import numpy as np
from pycocotools.coco import COCO


def save_colored_mask(mask, save_path):
    lbl_pil = Image.fromarray(mask.astype(np.uint8), mode="P")
    colormap = imgviz.label_colormap()
    lbl_pil.putpalette(colormap.flatten())
    lbl_pil.save(save_path)


def main(args):
    json_path = 'annotations.json'
    annotation_file = os.path.join(args.input_dir, json_path)
    os.makedirs(os.path.join(args.input_dir, 'maskValueAnnotations'), exist_ok=True)
    os.makedirs(os.path.join(args.input_dir, 'segmentationClass'), exist_ok=True)
    #os.makedirs(os.path.join(args.input_dir, 'JPEGImages'), exist_ok=True)
    coco = COCO(annotation_file)
    catIds = coco.getCatIds()
    #print("catIds")#[0, 1]
    imgIds = coco.getImgIds()
    print("catIds len:{}, imgIds len:{}".format(len(catIds), len(imgIds)))
    for imgId in tqdm.tqdm(imgIds, ncols=100):
        img = coco.loadImgs(imgId)[0]
        annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
        anns = coco.loadAnns(annIds)
        if len(annIds) > 0:
            pre_cat_mask = coco.annToMask(anns[0])
            mask = pre_cat_mask * (catIds.index(anns[0]['category_id']) + 1)
            print(mask)
            for i in range(len(anns) - 1):
                binary_mask = coco.annToMask(anns[i + 1])
                mask += binary_mask * (catIds.index(anns[i + 1]['category_id']) + 1)
                mask_area = pre_cat_mask + binary_mask
                mask_intersection_area = np.where(mask_area == 2)
                if len(mask_intersection_area[0]) > 0:
                    mask[mask_intersection_area] = catIds.index(anns[i + 1]['category_id']) + 1
                    mask_area[mask_intersection_area] = 1
                # for j in range(len(mask_intersection_area[0])):
                #     mask[mask_intersection_area[0][j]][mask_intersection_area[1][j]] = \
                #         catIds.index(anns[i + 1]['category_id']) + 1
                #     mask_area[mask_intersection_area[0][j]][mask_intersection_area[1][j]] = 1
                pre_cat_mask = mask_area
            #img_origin_path = os.path.join(args.input_dir, img['file_name'])
            #img_output_path = os.path.join(args.input_dir, 'segclass', img['file_name'].split("/")[-1])
            seg_output_path = os.path.join(args.input_dir, 'maskValueAnnotations',
                                           img['file_name'].split("/")[-1].replace('.jpg', '.png'))
            seg_output_path_show = os.path.join(args.input_dir, 'segmentationClass',
                                                img['file_name'].split("/")[-1].replace('.jpg', '.png'))
            if len(np.where(mask > 80)[0]) > 0:
                print("error")
            #shutil.copy(img_origin_path, img_output_path)
            cv2.imwrite(seg_output_path, mask)
            save_colored_mask(mask, seg_output_path_show)
    print("process end")


def get_args():
    parser = argparse.ArgumentParser()
    coco_dir = "/home/ubuntu/Projects/LISI_Pointer/SAM_HALFAUTO_LABELEIMG/sam_halfauto_lableimgs_data_type_6/"
    parser.add_argument("--input_dir", default=coco_dir, type=str,
                        help="input dataset directory")
    parser.add_argument("--split", default="train2014", type=str,
                        help="train2014 or val2014")
    return parser.parse_args()


if __name__ == '__main__':
    args = get_args()
    main(args)