@对coco格式的分割标注生成二值mask
对coco格式的分割标注生成二值mask
给定_annotations.coco.json文件,将里面的特定类别解码批量生成二值mask,代码如下:
import json from pathlib import Path import numpy as np from PIL import Image from pycocotools import mask as maskUtils def decode_coco_rle(rle: dict) -> np.ndarray: """ 使用 pycocotools 官方实现解码 COCO RLE,返回 0/255 的二值掩码。 rle 形如 {"counts": "...", "size": [height, width]}。 """ # 有些版本要求 counts 为 bytes,这里兼容处理一下 if isinstance(rle.get("counts"), str): rle = rle.copy() rle["counts"] = rle["counts"].encode("utf-8") mask = maskUtils.decode(rle) # H x W,0/1 uint8 # 统一转为 0/255 mask = (mask.astype(np.uint8)) * 255 return mask def main(): root = Path(__file__).resolve().parent coco_path = root / "_annotations.coco.json" images_dir = root / "images" masks_dir = root / "mask" masks_dir.mkdir(exist_ok=True) # 只保留的类别 id TARGET_CATEGORY_ID = 1 with coco_path.open("r", encoding="utf-8") as f: coco = json.load(f) # image_id -> image_info image_map = {img["id"]: img for img in coco["images"]} # image_id -> 该图像下、且 category_id == 1 的所有标注 annos_by_image = {} for anno in coco["annotations"]: if anno.get("category_id") != TARGET_CATEGORY_ID: # 忽略 category_id 为 0 或其它的标注 continue img_id = anno["image_id"] annos_by_image.setdefault(img_id, []).append(anno) # 对每张图,合成只包含 category_id == 1 的 mask for image_id, info in image_map.items(): height = info["height"] width = info["width"] mask = np.zeros((height, width), dtype=np.uint8) for anno in annos_by_image.get(image_id, []): seg = anno["segmentation"] # 只处理 RLE 字典形式 if not isinstance(seg, dict) or "counts" not in seg: continue decoded = decode_coco_rle(seg) # 多个实例取最大,合并为一张前景=255 的二值图 mask = np.maximum(mask, decoded) # 保存为灰度图(0/255) out_path = masks_dir / info["file_name"] Image.fromarray(mask, mode="L").save(out_path) print(f"saved: {out_path}") print("全部处理完成。") if __name__ == "__main__": main()coco格式的存储形式示例如下: