Spaces:
Running
on
Zero
Running
on
Zero
# Copyright (c) Facebook, Inc. and its affiliates. | |
import contextlib | |
import datetime | |
import io | |
import json | |
import logging | |
import numpy as np | |
import os | |
import shutil | |
import pycocotools.mask as mask_util | |
from fvcore.common.timer import Timer | |
from iopath.common.file_io import file_lock | |
from PIL import Image | |
from detectron2.structures import Boxes, BoxMode, PolygonMasks, RotatedBoxes | |
from detectron2.utils.file_io import PathManager | |
from .. import DatasetCatalog, MetadataCatalog | |
""" | |
This file contains functions to parse COCO-format annotations into dicts in "Detectron2 format". | |
""" | |
logger = logging.getLogger(__name__) | |
__all__ = [ | |
"load_coco_json", | |
"load_sem_seg", | |
"convert_to_coco_json", | |
"register_coco_instances", | |
] | |
def load_coco_json(json_file, image_root, dataset_name=None, extra_annotation_keys=None): | |
""" | |
Load a json file with COCO's instances annotation format. | |
Currently supports instance detection, instance segmentation, | |
and person keypoints annotations. | |
Args: | |
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 or None): the name of the dataset (e.g., coco_2017_train). | |
When provided, this function will also do the following: | |
* Put "thing_classes" into the metadata associated with this dataset. | |
* Map the category ids into a contiguous range (needed by standard dataset format), | |
and add "thing_dataset_id_to_contiguous_id" to the metadata associated | |
with this dataset. | |
This option should usually be provided, unless users need to load | |
the original json content and apply more processing manually. | |
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. | |
Returns: | |
list[dict]: a list of dicts in Detectron2 standard dataset dicts format (See | |
`Using Custom Datasets </tutorials/datasets.html>`_ ) when `dataset_name` is not None. | |
If `dataset_name` is None, the returned `category_ids` may be | |
incontiguous and may not conform to the Detectron2 standard format. | |
Notes: | |
1. This function does not read the image files. | |
The results do not have the "image" field. | |
""" | |
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) | |
if timer.seconds() > 1: | |
logger.info("Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds())) | |
id_map = None | |
if dataset_name is not None: | |
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: | |
logger.warning( | |
""" | |
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 is a list of dicts, each looks something like: | |
# {'license': 4, | |
# 'url': 'http://farm6.staticflickr.com/5454/9413846304_881d5e5c3b_z.jpg', | |
# 'file_name': 'COCO_val2014_000000001268.jpg', | |
# 'height': 427, | |
# 'width': 640, | |
# 'date_captured': '2013-11-17 05:57:24', | |
# 'id': 1268} | |
imgs = coco_api.loadImgs(img_ids) | |
# anns is a list[list[dict]], where each dict is an annotation | |
# record for an object. The inner list enumerates the objects in an image | |
# and the outer list enumerates over images. Example of 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}, | |
# ...] | |
anns = [coco_api.imgToAnns[img_id] for img_id in img_ids] | |
total_num_valid_anns = sum([len(x) for x in anns]) | |
total_num_anns = len(coco_api.anns) | |
if total_num_valid_anns < total_num_anns: | |
logger.warning( | |
f"{json_file} contains {total_num_anns} annotations, but only " | |
f"{total_num_valid_anns} of them match to images in the file." | |
) | |
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)) | |
logger.info("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: | |
# Check that the image_id in this annotation is the same as | |
# the image_id we're looking at. | |
# This fails only when the data parsing logic or the annotation file is buggy. | |
# 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} | |
if "bbox" in obj and len(obj["bbox"]) == 0: | |
raise ValueError( | |
f"One annotation of image {image_id} contains empty 'bbox' value! " | |
"This json does not have valid COCO format." | |
) | |
segm = anno.get("segmentation", None) | |
if segm: # either list[list[float]] or dict(RLE) | |
if isinstance(segm, dict): | |
if isinstance(segm["counts"], list): | |
# convert to compressed RLE | |
segm = mask_util.frPyObjects(segm, *segm["size"]) | |
else: | |
# filter out invalid polygons (< 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 | |
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: | |
annotation_category_id = obj["category_id"] | |
try: | |
obj["category_id"] = id_map[annotation_category_id] | |
except KeyError as e: | |
raise KeyError( | |
f"Encountered category_id={annotation_category_id} " | |
"but this id does not exist in 'categories' of the json file." | |
) from e | |
objs.append(obj) | |
record["annotations"] = objs | |
dataset_dicts.append(record) | |
if num_instances_without_valid_segmentation > 0: | |
logger.warning( | |
"Filtered out {} instances without valid segmentation. ".format( | |
num_instances_without_valid_segmentation | |
) | |
+ "There might be issues in your dataset generation process. Please " | |
"check https://detectron2.readthedocs.io/en/latest/tutorials/datasets.html carefully" | |
) | |
return dataset_dicts | |
def load_sem_seg(gt_root, image_root, gt_ext="png", image_ext="jpg"): | |
""" | |
Load semantic segmentation datasets. All files under "gt_root" with "gt_ext" extension are | |
treated as ground truth annotations and all files under "image_root" with "image_ext" extension | |
as input images. Ground truth and input images are matched using file paths relative to | |
"gt_root" and "image_root" respectively without taking into account file extensions. | |
This works for COCO as well as some other datasets. | |
Args: | |
gt_root (str): full path to ground truth semantic segmentation files. Semantic segmentation | |
annotations are stored as images with integer values in pixels that represent | |
corresponding semantic labels. | |
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. | |
Returns: | |
list[dict]: | |
a list of dicts in detectron2 standard format without instance-level | |
annotation. | |
Notes: | |
1. This function does not read the image and ground truth files. | |
The results do not have the "image" and "sem_seg" fields. | |
""" | |
# 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 PathManager.ls(image_root) 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 PathManager.ls(gt_root) 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): | |
logger.warn( | |
"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] | |
logger.info( | |
"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 | |
dataset_dicts.append(record) | |
return dataset_dicts | |
def convert_to_coco_dict(dataset_name): | |
""" | |
Convert an instance detection/segmentation or keypoint detection dataset | |
in detectron2's standard format into COCO json format. | |
Generic dataset description can be found here: | |
https://detectron2.readthedocs.io/tutorials/datasets.html#register-a-dataset | |
COCO data format description can be found here: | |
http://cocodataset.org/#format-data | |
Args: | |
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" | |
Returns: | |
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 | |
else: | |
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) | |
] | |
logger.info("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": int(image_dict["width"]), | |
"height": int(image_dict["height"]), | |
"file_name": str(image_dict["file_name"]), | |
} | |
coco_images.append(coco_image) | |
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 for axis-align and XYWHA for rotated | |
bbox = annotation["bbox"] | |
if isinstance(bbox, np.ndarray): | |
if bbox.ndim != 1: | |
raise ValueError(f"bbox has to be 1-dimensional. Got shape={bbox.shape}.") | |
bbox = bbox.tolist() | |
if len(bbox) not in [4, 5]: | |
raise ValueError(f"bbox has to has length 4 or 5. Got {bbox}.") | |
from_bbox_mode = annotation["bbox_mode"] | |
to_bbox_mode = BoxMode.XYWH_ABS if len(bbox) == 4 else BoxMode.XYWHA_ABS | |
bbox = BoxMode.convert(bbox, from_bbox_mode, to_bbox_mode) | |
# 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() | |
else: | |
raise TypeError(f"Unknown segmentation type {type(segmentation)}!") | |
else: | |
# Computing areas using bounding boxes | |
if to_bbox_mode == BoxMode.XYWH_ABS: | |
bbox_xy = BoxMode.convert(bbox, to_bbox_mode, BoxMode.XYXY_ABS) | |
area = Boxes([bbox_xy]).area()[0].item() | |
else: | |
area = RotatedBoxes([bbox]).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 | |
# https://github.com/facebookresearch/detectron2/pull/175#issuecomment-551202163 | |
keypoints[idx] = v - 0.5 | |
if "num_keypoints" in annotation: | |
num_keypoints = annotation["num_keypoints"] | |
else: | |
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"] = int(annotation.get("iscrowd", 0)) | |
coco_annotation["category_id"] = int(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): | |
# make it json-serializable | |
seg["counts"] = counts.decode("ascii") | |
coco_annotations.append(coco_annotation) | |
logger.info( | |
"Conversion finished, " | |
f"#images: {len(coco_images)}, #annotations: {len(coco_annotations)}" | |
) | |
info = { | |
"date_created": str(datetime.datetime.now()), | |
"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): | |
""" | |
Converts dataset into COCO format and saves it to a json file. | |
dataset_name must be registered in DatasetCatalog and in detectron2's standard format. | |
Args: | |
dataset_name: | |
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 | |
PathManager.mkdirs(os.path.dirname(output_file)) | |
with file_lock(output_file): | |
if PathManager.exists(output_file) and allow_cached: | |
logger.warning( | |
f"Using previously cached COCO format annotations at '{output_file}'. " | |
"You need to clear the cache file if your dataset has been modified." | |
) | |
else: | |
logger.info(f"Converting annotations of dataset '{dataset_name}' to COCO format ...)") | |
coco_dict = convert_to_coco_dict(dataset_name) | |
logger.info(f"Caching COCO format annotations at '{output_file}' ...") | |
tmp_file = output_file + ".tmp" | |
with PathManager.open(tmp_file, "w") as f: | |
json.dump(coco_dict, f) | |
shutil.move(tmp_file, output_file) | |
def register_coco_instances(name, metadata, json_file, image_root): | |
""" | |
Register a dataset in COCO's json annotation format for | |
instance detection, instance segmentation and keypoint detection. | |
(i.e., Type 1 and 2 in http://cocodataset.org/#format-data. | |
`instances*.json` and `person_keypoints*.json` in the dataset). | |
This is an example of how to register a new dataset. | |
You can do something similar to this function, to register new datasets. | |
Args: | |
name (str): the name that identifies a dataset, e.g. "coco_2014_train". | |
metadata (dict): extra metadata associated with this dataset. You can | |
leave it as an empty dict. | |
json_file (str): path to the json instance annotation file. | |
image_root (str or path-like): directory which contains all the images. | |
""" | |
assert isinstance(name, str), name | |
assert isinstance(json_file, (str, os.PathLike)), json_file | |
assert isinstance(image_root, (str, os.PathLike)), image_root | |
# 1. register a function which returns dicts | |
DatasetCatalog.register(name, lambda: load_coco_json(json_file, image_root, name)) | |
# 2. Optionally, add metadata about this dataset, | |
# since they might be useful in evaluation, visualization or logging | |
MetadataCatalog.get(name).set( | |
json_file=json_file, image_root=image_root, evaluator_type="coco", **metadata | |
) | |
def main() -> None: | |
global logger | |
""" | |
Test the COCO json dataset loader. | |
Usage: | |
python -m detectron2.data.datasets.coco \ | |
path/to/json path/to/image_root dataset_name | |
"dataset_name" can be "coco_2014_minival_100", or other | |
pre-registered ones | |
""" | |
import sys | |
import detectron2.data.datasets # noqa # add pre-defined metadata | |
from detectron2.utils.logger import setup_logger | |
from detectron2.utils.visualizer import Visualizer | |
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]) | |
logger.info("Done loading {} samples.".format(len(dicts))) | |
dirname = "coco-data-vis" | |
os.makedirs(dirname, exist_ok=True) | |
for d in 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 __name__ == "__main__": | |
main() # pragma: no cover | |