# Copyright (c) Facebook, Inc. and its affiliates.
import copy
import json
import os
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.utils.file_io import PathManager
from .coco import load_coco_json, load_sem_seg
__all__ = ["register_coco_panoptic", "register_coco_panoptic_separated"]
def load_coco_panoptic_json(json_file, image_dir, gt_dir, meta):
"""
Args:
image_dir (str): path to the raw dataset. e.g., "~/coco/train2017".
gt_dir (str): path to the raw annotations. e.g., "~/coco/panoptic_train2017".
json_file (str): path to the json file. e.g., "~/coco/annotations/panoptic_train2017.json".
Returns:
list[dict]: a list of dicts in Detectron2 standard format. (See
`Using Custom Datasets </tutorials/datasets.html>`_ )
"""
def _convert_category_id(segment_info, meta):
if segment_info["category_id"] in meta["thing_dataset_id_to_contiguous_id"]:
segment_info["category_id"] = meta["thing_dataset_id_to_contiguous_id"][
segment_info["category_id"]
]
segment_info["isthing"] = True
else:
segment_info["category_id"] = meta["stuff_dataset_id_to_contiguous_id"][
segment_info["category_id"]
]
segment_info["isthing"] = False
return segment_info
with PathManager.open(json_file) as f:
json_info = json.load(f)
ret = []
for ann in json_info["annotations"]:
image_id = int(ann["image_id"])
# TODO: currently we assume image and label has the same filename but
# different extension, and images have extension ".jpg" for COCO. Need
# to make image extension a user-provided argument if we extend this
# function to support other COCO-like datasets.
image_file = os.path.join(image_dir, os.path.splitext(ann["file_name"])[0] + ".jpg")
label_file = os.path.join(gt_dir, ann["file_name"])
segments_info = [_convert_category_id(x, meta) for x in ann["segments_info"]]
ret.append(
{
"file_name": image_file,
"image_id": image_id,
"pan_seg_file_name": label_file,
"segments_info": segments_info,
}
)
assert len(ret), f"No images found in {image_dir}!"
assert PathManager.isfile(ret[0]["file_name"]), ret[0]["file_name"]
assert PathManager.isfile(ret[0]["pan_seg_file_name"]), ret[0]["pan_seg_file_name"]
return ret
[docs]def register_coco_panoptic(
name, metadata, image_root, panoptic_root, panoptic_json, instances_json=None
):
"""
Register a "standard" version of COCO panoptic segmentation dataset named `name`.
The dictionaries in this registered dataset follows detectron2's standard format.
Hence it's called "standard".
Args:
name (str): the name that identifies a dataset,
e.g. "coco_2017_train_panoptic"
metadata (dict): extra metadata associated with this dataset.
image_root (str): directory which contains all the images
panoptic_root (str): directory which contains panoptic annotation images in COCO format
panoptic_json (str): path to the json panoptic annotation file in COCO format
sem_seg_root (none): not used, to be consistent with
`register_coco_panoptic_separated`.
instances_json (str): path to the json instance annotation file
"""
panoptic_name = name
DatasetCatalog.register(
panoptic_name,
lambda: load_coco_panoptic_json(panoptic_json, image_root, panoptic_root, metadata),
)
MetadataCatalog.get(panoptic_name).set(
panoptic_root=panoptic_root,
image_root=image_root,
panoptic_json=panoptic_json,
json_file=instances_json,
evaluator_type="coco_panoptic_seg",
ignore_label=255,
label_divisor=1000,
**metadata,
)
[docs]def register_coco_panoptic_separated(
name, metadata, image_root, panoptic_root, panoptic_json, sem_seg_root, instances_json
):
"""
Register a "separated" version of COCO panoptic segmentation dataset named `name`.
The annotations in this registered dataset will contain both instance annotations and
semantic annotations, each with its own contiguous ids. Hence it's called "separated".
It follows the setting used by the PanopticFPN paper:
1. The instance annotations directly come from polygons in the COCO
instances annotation task, rather than from the masks in the COCO panoptic annotations.
The two format have small differences:
Polygons in the instance annotations may have overlaps.
The mask annotations are produced by labeling the overlapped polygons
with depth ordering.
2. The semantic annotations are converted from panoptic annotations, where
all "things" are assigned a semantic id of 0.
All semantic categories will therefore have ids in contiguous
range [1, #stuff_categories].
This function will also register a pure semantic segmentation dataset
named ``name + '_stuffonly'``.
Args:
name (str): the name that identifies a dataset,
e.g. "coco_2017_train_panoptic"
metadata (dict): extra metadata associated with this dataset.
image_root (str): directory which contains all the images
panoptic_root (str): directory which contains panoptic annotation images
panoptic_json (str): path to the json panoptic annotation file
sem_seg_root (str): directory which contains all the ground truth segmentation annotations.
instances_json (str): path to the json instance annotation file
"""
panoptic_name = name + "_separated"
DatasetCatalog.register(
panoptic_name,
lambda: merge_to_panoptic(
load_coco_json(instances_json, image_root, panoptic_name),
load_sem_seg(sem_seg_root, image_root),
),
)
MetadataCatalog.get(panoptic_name).set(
panoptic_root=panoptic_root,
image_root=image_root,
panoptic_json=panoptic_json,
sem_seg_root=sem_seg_root,
json_file=instances_json, # TODO rename
evaluator_type="coco_panoptic_seg",
ignore_label=255,
**metadata,
)
semantic_name = name + "_stuffonly"
DatasetCatalog.register(semantic_name, lambda: load_sem_seg(sem_seg_root, image_root))
MetadataCatalog.get(semantic_name).set(
sem_seg_root=sem_seg_root,
image_root=image_root,
evaluator_type="sem_seg",
ignore_label=255,
**metadata,
)
def merge_to_panoptic(detection_dicts, sem_seg_dicts):
"""
Create dataset dicts for panoptic segmentation, by
merging two dicts using "file_name" field to match their entries.
Args:
detection_dicts (list[dict]): lists of dicts for object detection or instance segmentation.
sem_seg_dicts (list[dict]): lists of dicts for semantic segmentation.
Returns:
list[dict] (one per input image): Each dict contains all (key, value) pairs from dicts in
both detection_dicts and sem_seg_dicts that correspond to the same image.
The function assumes that the same key in different dicts has the same value.
"""
results = []
sem_seg_file_to_entry = {x["file_name"]: x for x in sem_seg_dicts}
assert len(sem_seg_file_to_entry) > 0
for det_dict in detection_dicts:
dic = copy.copy(det_dict)
dic.update(sem_seg_file_to_entry[dic["file_name"]])
results.append(dic)
return results
if __name__ == "__main__":
"""
Test the COCO panoptic dataset loader.
Usage:
python -m detectron2.data.datasets.coco_panoptic \
path/to/image_root path/to/panoptic_root path/to/panoptic_json dataset_name 10
"dataset_name" can be "coco_2017_train_panoptic", or other
pre-registered ones
"""
from detectron2.utils.logger import setup_logger
from detectron2.utils.visualizer import Visualizer
import detectron2.data.datasets # noqa # add pre-defined metadata
import sys
from PIL import Image
import numpy as np
logger = setup_logger(name=__name__)
assert sys.argv[4] in DatasetCatalog.list()
meta = MetadataCatalog.get(sys.argv[4])
dicts = load_coco_panoptic_json(sys.argv[3], sys.argv[1], sys.argv[2], meta.as_dict())
logger.info("Done loading {} samples.".format(len(dicts)))
dirname = "coco-data-vis"
os.makedirs(dirname, exist_ok=True)
num_imgs_to_vis = int(sys.argv[5])
for i, d in enumerate(dicts):
img = np.array(Image.open(d["file_name"]))
visualizer = Visualizer(img, metadata=meta)
vis = visualizer.draw_dataset_dict(d)
fpath = os.path.join(dirname, os.path.basename(d["file_name"]))
vis.save(fpath)
if i + 1 >= num_imgs_to_vis:
break