COCO 2 mask

import os
import numpy as np
import pycocotools.mask as mask_utils
from pycocotools.coco import COCO
import cv2
from tqdm import tqdm
 
# 根据自己情况设置图片、json以及mask保存路径
img_path = ''
json_path = ''
save_mask = ''
 
os.makedirs(save_mask, exist_ok=True)
 
# 通过pycocotools加载json文件
coco = COCO(json_path)
 
# 获取json中的图像id
images_ids = coco.getImgIds()
 
# 逐图处理
for img_id in tqdm(images_ids):
    # 根据图像唯一id获取对应信息
    img_info = coco.loadImgs(img_id)[0]
    img_files_path = os.path.join(img_path, img_info['file_name'])
    height, width = img_info['height'], img_info['width'] 
   
    # 根据图像id获取对应标签信息
    ann_ids = coco.getAnnIds(imgIds=img_id)    
    anns = coco.loadAnns(ann_ids)
    
    # 逐实例处理（一个图像存在多个实例）
    for ann in anns:
        rle = coco.annToRLE(ann)
        mask = mask_utils.decode(rle)
        mask[mask == 1] = 255   # 调整mask中的像素值
 
        # 保存mask
        mask_name = img_info['file_name'].replace('.jpg', f'_{ann['category_id']}.jpg')
        cv2.imwrite(os.path.join(save_mask, mask_name), mask)