detectron2 中的关于 coco 数据集的一些函数 - 主要用于解析 COCO-format 标注数据,转换为 Detectron2 format;语义分割数据集加载;detectron2 格式数据转换为 COCO json 格式.

其中,COCO 格式数据解析为 detectron2 格式数据,可用于检查非 COCO 数据转换为 COCO 数据集后的情况.

1. COCO 实例分割数据可视化

测试 COCO json 数据集.

import contextlib
import datetime
import io
import json
import logging
import numpy as np
import os
import pycocotools.mask as mask_util
from fvcore.common.file_io import PathManager, file_lock
from fvcore.common.timer import Timer
from PIL import Image

from detectron2.structures import Boxes, BoxMode, PolygonMasks

from .. import DatasetCatalog, MetadataCatalog # coco 数据集注册

logger = logging.getLogger(__name__)

__all__ = ["load_coco_json", "load_sem_seg", "convert_to_coco_json"]

def load_coco_json(json_file, image_root, dataset_name=None, extra_annotation_keys=None):
    加载 COCO 实例分割格式的 json 文件. 
    当前支持的标注数据: instance detection, instance segmentation, person keypoints

        json_file (str): full path to the json file in COCO instances annotation format.
        image_root (str or path-like): the directory where the images in this json file exists.
        dataset_name (str): the name of the dataset (e.g., coco_2017_train).
            If provided, this function will also put "thing_classes" into
            the metadata associated with this dataset.
        extra_annotation_keys (list[str]): list of per-annotation keys that should also be
            loaded into the dataset dict (besides "iscrowd", "bbox", "keypoints",
            "category_id", "segmentation"). The values for these keys will be returned as-is.
            For example, the densepose annotations are loaded in this way.

        list[dict]: a list of dicts in Detectron2 standard dataset dicts format. (See
        `Using Custom Datasets </tutorials/datasets.html>`_ )

        该函数不读取图像文件,results 里不包含 'image' 字段.
    from pycocotools.coco import COCO

    timer = Timer()
    json_file = PathManager.get_local_path(json_file)
    with contextlib.redirect_stdout(io.StringIO()):
        coco_api = COCO(json_file) # COCO 类
    if timer.seconds() > 1:"Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds()))

    id_map = None
    if dataset_name is not None: # COCO 数据集元数据
        meta = MetadataCatalog.get(dataset_name)
        cat_ids = sorted(coco_api.getCatIds())
        cats = coco_api.loadCats(cat_ids)
        # The categories in a custom json file may not be sorted.
        thing_classes = [c["name"] for c in sorted(cats, key=lambda x: x["id"])]
        meta.thing_classes = thing_classes

        # In COCO, certain category ids are artificially removed,
        # and by convention they are always ignored.
        # We deal with COCO's id issue and translate
        # the category ids to contiguous ids in [0, 80).

        # It works by looking at the "categories" field in the json, therefore
        # if users' own json also have incontiguous ids, we'll
        # apply this mapping as well but print a warning.
        if not (min(cat_ids) == 1 and max(cat_ids) == len(cat_ids)):
            if "coco" not in dataset_name:
                    Category ids in annotations are not in [1, #categories]! 
                    We'll apply a mapping for you.
        id_map = {v: i for i, v in enumerate(cat_ids)}
        meta.thing_dataset_id_to_contiguous_id = id_map

    # sort indices for reproducible results
    img_ids = sorted(coco_api.imgs.keys())
    imgs = coco_api.loadImgs(img_ids)
    # imgs 是 dicts 组成的 list(a list of dicts),每个 dict 的格式如下:
    # {'license': 4,
    #  'url': '',
    #  'file_name': 'COCO_val2014_000000001268.jpg',
    #  'height': 427,
    #  'width': 640,
    #  'date_captured': '2013-11-17 05:57:24',
    #  'id': 1268}
    anns = [coco_api.imgToAnns[img_id] for img_id in img_ids]
    # anns 是 list[list[dict]], 每个 dict 为一个物体的标注数据. 
    # list[dict] 表示单张图片中所有物体的标注数据.
    # anns[0] 的示例如:
    # [{'segmentation': [[192.81,
    #     247.09,
    #     ...
    #     219.03,
    #     249.06]],
    #   'area': 1035.749,
    #   'iscrowd': 0,
    #   'image_id': 1268,
    #   'bbox': [192.81, 224.8, 74.73, 33.43],
    #   'category_id': 16,
    #   'id': 42986},
    #  ...]

    if "minival" not in json_file:
        # The popular valminusminival & minival annotations for COCO2014 contain this bug.
        # However the ratio of buggy annotations there is tiny and does not affect accuracy.
        # Therefore we explicitly white-list them.
        ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
        assert len(set(ann_ids)) == len(ann_ids), "Annotation ids in '{}' are not unique!".format( json_file)
    imgs_anns = list(zip(imgs, anns))"Loaded {} images in COCO format from {}".format(len(imgs_anns), json_file))

    dataset_dicts = []
    ann_keys = ["iscrowd", "bbox", "keypoints", "category_id"] + (extra_annotation_keys or [])

    num_instances_without_valid_segmentation = 0
    for (img_dict, anno_dict_list) in imgs_anns:
        record = {}
        record["file_name"] = os.path.join(image_root, img_dict["file_name"])
        record["height"] = img_dict["height"]
        record["width"] = img_dict["width"]
        image_id = record["image_id"] = img_dict["id"]

        objs = []
        for anno in anno_dict_list:
            # 比对标注数据中的 image_id 与上面 image_id 是否一致.
            # 如果不一致,则可能是数据解析或者标注文件存在问题.
            # The original COCO valminusminival2014 & minival2014 annotation files
            # actually contains bugs that, together with certain ways of using COCO API,
            # can trigger this assertion.
            assert anno["image_id"] == image_id
            assert anno.get("ignore", 0) == 0, '"ignore" in COCO json file is not supported.'

            obj = {key: anno[key] for key in ann_keys if key in anno}

            segm = anno.get("segmentation", None)
            if segm:  # 分割数据格式为 list[list[float]] 或 dict(RLE)
                if not isinstance(segm, dict):
                    # 过滤无效多边形 (< 3 points)
                    segm = [poly for poly in segm if len(poly) % 2 == 0 and len(poly) >= 6]
                    if len(segm) == 0:
                        num_instances_without_valid_segmentation += 1
                        continue  # ignore this instance
                obj["segmentation"] = segm
            # keypoints
            keypts = anno.get("keypoints", None)
            if keypts:  # list[int]
                for idx, v in enumerate(keypts):
                    if idx % 3 != 2:
                        # COCO's segmentation coordinates are floating points in [0, H or W],
                        # but keypoint coordinates are integers in [0, H-1 or W-1]
                        # Therefore we assume the coordinates are "pixel indices" and
                        # add 0.5 to convert to floating point coordinates.
                        keypts[idx] = v + 0.5
                obj["keypoints"] = keypts

            obj["bbox_mode"] = BoxMode.XYWH_ABS
            if id_map:
                obj["category_id"] = id_map[obj["category_id"]]
        record["annotations"] = objs

    if num_instances_without_valid_segmentation > 0:
            "Filtered out {} instances without valid segmentation. ".format(
            + "There might be issues in your dataset generation process. "
            "A valid polygon should be a list[float] with even length >= 6."
    return dataset_dicts

if __name__ == "__main__":
        python -m \
            path/to/json path/to/image_root dataset_name output_folder
    from detectron2.utils.logger import setup_logger
    from detectron2.utils.visualizer import Visualizer
    import  # noqa # add pre-defined metadata
    import sys

    logger = setup_logger(name=__name__)
    assert sys.argv[3] in DatasetCatalog.list()
    meta = MetadataCatalog.get(sys.argv[3])

    dicts = load_coco_json(sys.argv[1], sys.argv[2], sys.argv[3])"Done loading {} samples.".format(len(dicts)))

    dirname = sys.argv[4]
    os.makedirs(dirname, exist_ok=True)
    fp = open(os.path.join(dirname, 'error.txt'), 'w')
    for d in dicts:
            img = np.array(["file_name"]))
            visualizer = Visualizer(img, metadata=meta)
            vis = visualizer.draw_dataset_dict(d)
            fpath = os.path.join(dirname, os.path.basename(d["file_name"]))


python -m \
                    /path/to/instances_train2014.json \ 
                      /path/to/images/coco_train2014 \ 
                    coco_2014_val \ 

2. detectron2 format 转 coco json format

导入 Python 包代码段与上共用.

def convert_to_coco_dict(dataset_name):
    将 detectron2 标准格式的 instance detection/segmentation 和 keypoint detection 数据集
    转换为 COCO json 格式.

    Generic dataset description can be found here:

    COCO data format description can be found here:

        dataset_name (str):
            name of the source dataset
            Must be registered in DatastCatalog and in detectron2's standard format.
            Must have corresponding metadata "thing_classes"
        coco_dict: serializable dict in COCO json format

    dataset_dicts = DatasetCatalog.get(dataset_name)
    metadata = MetadataCatalog.get(dataset_name)

    # unmap the category mapping ids for COCO
    if hasattr(metadata, "thing_dataset_id_to_contiguous_id"):
        reverse_id_mapping = {v: k for k, v in metadata.thing_dataset_id_to_contiguous_id.items()}
        reverse_id_mapper = lambda contiguous_id: reverse_id_mapping[contiguous_id]  # noqa
        reverse_id_mapper = lambda contiguous_id: contiguous_id  # noqa

    categories = [
        {"id": reverse_id_mapper(id), "name": name}
        for id, name in enumerate(metadata.thing_classes)
    ]"Converting dataset dicts into COCO format")
    coco_images = []
    coco_annotations = []

    for image_id, image_dict in enumerate(dataset_dicts):
        coco_image = {
            "id": image_dict.get("image_id", image_id),
            "width": image_dict["width"],
            "height": image_dict["height"],
            "file_name": image_dict["file_name"],

        anns_per_image = image_dict.get("annotations", [])
        for annotation in anns_per_image:
            # create a new dict with only COCO fields
            coco_annotation = {}

            # COCO requirement: XYWH box format
            bbox = annotation["bbox"]
            bbox_mode = annotation["bbox_mode"]
            bbox = BoxMode.convert(bbox, bbox_mode, BoxMode.XYWH_ABS)

            # COCO requirement: instance area
            if "segmentation" in annotation:
                # Computing areas for instances by counting the pixels
                segmentation = annotation["segmentation"]
                # TODO: check segmentation type: RLE, BinaryMask or Polygon
                if isinstance(segmentation, list):
                    polygons = PolygonMasks([segmentation])
                    area = polygons.area()[0].item()
                elif isinstance(segmentation, dict):  # RLE
                    area = mask_util.area(segmentation).item()
                    raise TypeError(f"Unknown segmentation type {type(segmentation)}!")
                # Computing areas using bounding boxes
                bbox_xy = BoxMode.convert(bbox, BoxMode.XYWH_ABS, BoxMode.XYXY_ABS)
                area = Boxes([bbox_xy]).area()[0].item()

            if "keypoints" in annotation:
                keypoints = annotation["keypoints"]  # list[int]
                for idx, v in enumerate(keypoints):
                    if idx % 3 != 2:
                        # COCO's segmentation coordinates are floating points in [0, H or W],
                        # but keypoint coordinates are integers in [0, H-1 or W-1]
                        # For COCO format consistency we substract 0.5
                        keypoints[idx] = v - 0.5
                if "num_keypoints" in annotation:
                    num_keypoints = annotation["num_keypoints"]
                    num_keypoints = sum(kp > 0 for kp in keypoints[2::3])

            # COCO requirement:
            #   linking annotations to images
            #   "id" field must start with 1
            coco_annotation["id"] = len(coco_annotations) + 1
            coco_annotation["image_id"] = coco_image["id"]
            coco_annotation["bbox"] = [round(float(x), 3) for x in bbox]
            coco_annotation["area"] = float(area)
            coco_annotation["iscrowd"] = annotation.get("iscrowd", 0)
            coco_annotation["category_id"] = reverse_id_mapper(annotation["category_id"])

            # Add optional fields
            if "keypoints" in annotation:
                coco_annotation["keypoints"] = keypoints
                coco_annotation["num_keypoints"] = num_keypoints

            if "segmentation" in annotation:
                seg = coco_annotation["segmentation"] = annotation["segmentation"]
                if isinstance(seg, dict):  # RLE
                    counts = seg["counts"]
                    if not isinstance(counts, str):
                        seg["counts"] = counts.decode("ascii")

        "Conversion finished, "
        f"#images: {len(coco_images)}, #annotations: {len(coco_annotations)}"

    info = {
        "date_created": str(,
        "description": "Automatically generated COCO json file for Detectron2.",
    coco_dict = {"info": info, "images": coco_images, "categories": categories, "licenses": None}
    if len(coco_annotations) > 0:
        coco_dict["annotations"] = coco_annotations
    return coco_dict

def convert_to_coco_json(dataset_name, output_file, allow_cached=True):
    将数据集转换为 COCO 格式,并保存为 json 文件.
    dataset_name must be registered in DatasetCatalog and in detectron2's standard format.

            reference from the config file to the catalogs
            must be registered in DatasetCatalog and in detectron2's standard format
        output_file: path of json file that will be saved to
        allow_cached: if json file is already present then skip conversion

    # TODO: The dataset or the conversion script *may* change,
    # a checksum would be useful for validating the cached data

    with file_lock(output_file):
        if PathManager.exists(output_file) and allow_cached:
                f"Using previously cached COCO format annotations at '{output_file}'. "
                "You need to clear the cache file if your dataset has been modified."
  "Converting annotations of dataset '{dataset_name}' to COCO format ...)")
            coco_dict = convert_to_coco_dict(dataset_name)

  "Caching COCO format annotations at '{output_file}' ...")
            with, "w") as f:
                json.dump(coco_dict, f)

3. detectron2 加载语义分割数据集

导入 Python 包代码段与上共用.

def load_sem_seg(gt_root, image_root, gt_ext="png", image_ext="jpg"):
    加载 semantic segmentation 数据集.
    "gt_root"路径中 "gt_ext" 扩展名的文件加载为 GT 标注;
    "image_root"路径中 "image_ext" 扩展名的文件为输入图片.
    GT 标注文件名和输入图片文件名除了扩展名,具有相同的 basename.

    This works for COCO as well as some other datasets.

        gt_root (str): GT 标注文件的完整路径. GT标注文件被保存为图像,
        image_root (str): the directory where the input images are.
        gt_ext (str): file extension for ground truth annotations.
        image_ext (str): file extension for input images.

            a list of dicts in detectron2 standard format without instance-level

        该函数不读取图片和 GT 文件. results 不包含 "image" 和 "sem_seg" 字段.

    # We match input images with ground truth based on their relative filepaths (without file
    # extensions) starting from 'image_root' and 'gt_root' respectively.
    def file2id(folder_path, file_path):
        # extract relative path starting from `folder_path`
        image_id = os.path.normpath(os.path.relpath(file_path, start=folder_path))
        # remove file extension
        image_id = os.path.splitext(image_id)[0]
        return image_id

    input_files = sorted(
        (os.path.join(image_root, f) for f in if f.endswith(image_ext)),
        key=lambda file_path: file2id(image_root, file_path),
    gt_files = sorted(
        (os.path.join(gt_root, f) for f in if f.endswith(gt_ext)),
        key=lambda file_path: file2id(gt_root, file_path),

    assert len(gt_files) > 0, "No annotations found in {}.".format(gt_root)

    # Use the intersection, so that val2017_100 annotations can run smoothly with val2017 images
    if len(input_files) != len(gt_files):
            "Directory {} and {} has {} and {} files, respectively.".format(
                image_root, gt_root, len(input_files), len(gt_files)
        input_basenames = [os.path.basename(f)[: -len(image_ext)] for f in input_files]
        gt_basenames = [os.path.basename(f)[: -len(gt_ext)] for f in gt_files]
        intersect = list(set(input_basenames) & set(gt_basenames))
        # sort, otherwise each worker may obtain a list[dict] in different order
        intersect = sorted(intersect)
        logger.warn("Will use their intersection of {} files.".format(len(intersect)))
        input_files = [os.path.join(image_root, f + image_ext) for f in intersect]
        gt_files = [os.path.join(gt_root, f + gt_ext) for f in intersect]
        "Loaded {} images with semantic segmentation from {}".format(len(input_files), image_root)

    dataset_dicts = []
    for (img_path, gt_path) in zip(input_files, gt_files):
        record = {}
        record["file_name"] = img_path
        record["sem_seg_file_name"] = gt_path

    return dataset_dicts
Last modification:July 4th, 2020 at 02:33 pm