coco segmentation 转 二值mask

import cv2

from pycocotools import mask as maskUtils
import matplotlib.pyplot as plt
import json
import os
from PIL import Image
import numpy as np
import random


def annToRLE(segm, img_size):
    h, w = img_size
    rles = maskUtils.frPyObjects(segm, h, w)
    rle = maskUtils.merge(rles)
    return rle


def annToMask(segm, img_size):
    rle = annToRLE(segm, img_size)
    m = maskUtils.decode(rle)
    return m


def cocoSegm2BinaryMask(annFile, output_dir):
    no_ann = []
    no_ann_txt = './no_ann.txt'

    with open(os.path.join(annFile), "r") as f:
        records = json.load(f)

    segms = {}
    for ann in records["annotations"]:
        img_id = ann['image_id']
        if not segms.get(img_id, None):
            segms[img_id] = [ann['segmentation']]
        else:
            segms[img_id].append(ann['segmentation'])
    print('Num of imgs: ', len(segms))

    if len(segms) != len(records['images']):
        print(f'len of segms {len(segms)} != len of images {len(records["images"])}')
        
    cout = 1
    for img in records['images']:
        img_size = [img['height'], img['width']]
        img_name = img['file_name']
        img_id = img['id']
        print(f'[{cout}/{len(records["images"])}] Processing img >>> {img_name}', end='\r', flush=True)
        cout += 1
        annotation = segms.get(img_id, None)
        mask = np.zeros(img_size, dtype=np.uint8)
        output = np.zeros(img_size, dtype=np.uint8)
        if not annotation:
            print('Do not find annotation of img: ', img_id)
            no_ann.append(img_id)
        else:
            for ann in annotation:
                mask += annToMask(ann, img_size)
        output[mask == 0] = 1
        cv2.imwrite(os.path.join(output_dir, img_name), output)

    with open(no_ann_txt, 'w') as file:
        for item in no_ann:
            file.write(str(item) + '\n')
        f.close()


def sample_visualize(data_dir, show_dir, sample_num):
    file_list = os.listdir(data_dir)
    for img_name in random.sample(file_list, sample_num):
        img_path = os.path.join(data_dir, img_name)
        img = cv2.imread(img_path, -1)
        assert 1 in img, 'Image is not a binary mask'
        img[img == 1] = 255
        cv2.imwrite(os.path.join(show_dir, 'stufthingmaps_'+img_name), img)
    print('Done')


if __name__ == '__main__':
    root = './dataset/small_building/'
    splits = ['train', 'val']
    for split in splits:
        annotation_file = os.path.join(root, 'annotations', f'{split}_coco.json')
        out_dir = os.path.join(root, 'stuffthingmaps01', f'{split}')
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)
        cocoSegm2BinaryMask(annotation_file, out_dir)
        sample_visualize(out_dir, root, 3)