# Copyright (c) Facebook, Inc. and its affiliates. # pyre-unsafe import logging import os from typing import Any, Dict, Iterable, List, Optional from fvcore.common.timer import Timer from detectron2.data import DatasetCatalog, MetadataCatalog from detectron2.data.datasets.lvis import get_lvis_instances_meta from detectron2.structures import BoxMode from detectron2.utils.file_io import PathManager from ..utils import maybe_prepend_base_path from .coco import ( DENSEPOSE_ALL_POSSIBLE_KEYS, DENSEPOSE_METADATA_URL_PREFIX, CocoDatasetInfo, get_metadata, ) DATASETS = [ CocoDatasetInfo( name="densepose_lvis_v1_ds1_train_v1", images_root="coco_", annotations_fpath="lvis/densepose_lvis_v1_ds1_train_v1.json", ), CocoDatasetInfo( name="densepose_lvis_v1_ds1_val_v1", images_root="coco_", annotations_fpath="lvis/densepose_lvis_v1_ds1_val_v1.json", ), CocoDatasetInfo( name="densepose_lvis_v1_ds2_train_v1", images_root="coco_", annotations_fpath="lvis/densepose_lvis_v1_ds2_train_v1.json", ), CocoDatasetInfo( name="densepose_lvis_v1_ds2_val_v1", images_root="coco_", annotations_fpath="lvis/densepose_lvis_v1_ds2_val_v1.json", ), CocoDatasetInfo( name="densepose_lvis_v1_ds1_val_animals_100", images_root="coco_", annotations_fpath="lvis/densepose_lvis_v1_val_animals_100_v2.json", ), ] def _load_lvis_annotations(json_file: str): """ Load COCO annotations from a JSON file Args: json_file: str Path to the file to load annotations from Returns: Instance of `pycocotools.coco.COCO` that provides access to annotations data """ from lvis import LVIS json_file = PathManager.get_local_path(json_file) logger = logging.getLogger(__name__) timer = Timer() lvis_api = LVIS(json_file) if timer.seconds() > 1: logger.info("Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds())) return lvis_api def _add_categories_metadata(dataset_name: str) -> None: metadict = get_lvis_instances_meta(dataset_name) categories = metadict["thing_classes"] metadata = MetadataCatalog.get(dataset_name) metadata.categories = {i + 1: categories[i] for i in range(len(categories))} logger = logging.getLogger(__name__) logger.info(f"Dataset {dataset_name} has {len(categories)} categories") def _verify_annotations_have_unique_ids(json_file: str, anns: List[List[Dict[str, Any]]]) -> None: 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 ) def _maybe_add_bbox(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None: if "bbox" not in ann_dict: return obj["bbox"] = ann_dict["bbox"] obj["bbox_mode"] = BoxMode.XYWH_ABS def _maybe_add_segm(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None: if "segmentation" not in ann_dict: return segm = ann_dict["segmentation"] if not isinstance(segm, dict): # 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: return obj["segmentation"] = segm def _maybe_add_keypoints(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None: if "keypoints" not in ann_dict: return keypts = ann_dict["keypoints"] # 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 def _maybe_add_densepose(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None: for key in DENSEPOSE_ALL_POSSIBLE_KEYS: if key in ann_dict: obj[key] = ann_dict[key] def _combine_images_with_annotations( dataset_name: str, image_root: str, img_datas: Iterable[Dict[str, Any]], ann_datas: Iterable[Iterable[Dict[str, Any]]], ): dataset_dicts = [] def get_file_name(img_root, img_dict): # Determine the path including the split folder ("train2017", "val2017", "test2017") from # the coco_url field. Example: # 'coco_url': 'http://images.cocodataset.org/train2017/000000155379.jpg' split_folder, file_name = img_dict["coco_url"].split("/")[-2:] return os.path.join(img_root + split_folder, file_name) for img_dict, ann_dicts in zip(img_datas, ann_datas): record = {} record["file_name"] = get_file_name(image_root, img_dict) record["height"] = img_dict["height"] record["width"] = img_dict["width"] record["not_exhaustive_category_ids"] = img_dict.get("not_exhaustive_category_ids", []) record["neg_category_ids"] = img_dict.get("neg_category_ids", []) record["image_id"] = img_dict["id"] record["dataset"] = dataset_name objs = [] for ann_dict in ann_dicts: assert ann_dict["image_id"] == record["image_id"] obj = {} _maybe_add_bbox(obj, ann_dict) obj["iscrowd"] = ann_dict.get("iscrowd", 0) obj["category_id"] = ann_dict["category_id"] _maybe_add_segm(obj, ann_dict) _maybe_add_keypoints(obj, ann_dict) _maybe_add_densepose(obj, ann_dict) objs.append(obj) record["annotations"] = objs dataset_dicts.append(record) return dataset_dicts def load_lvis_json(annotations_json_file: str, image_root: str, dataset_name: str): """ Loads a JSON file with annotations in LVIS instances format. Replaces `detectron2.data.datasets.coco.load_lvis_json` to handle metadata in a more flexible way. Postpones category mapping to a later stage to be able to combine several datasets with different (but coherent) sets of categories. Args: annotations_json_file: str Path to the JSON file with annotations in COCO instances format. image_root: str directory that contains all the images dataset_name: str the name that identifies a dataset, e.g. "densepose_coco_2014_train" extra_annotation_keys: Optional[List[str]] If provided, these keys are used to extract additional data from the annotations. """ lvis_api = _load_lvis_annotations(PathManager.get_local_path(annotations_json_file)) _add_categories_metadata(dataset_name) # sort indices for reproducible results img_ids = sorted(lvis_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 = lvis_api.load_imgs(img_ids) logger = logging.getLogger(__name__) logger.info("Loaded {} images in LVIS format from {}".format(len(imgs), annotations_json_file)) # 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. anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids] _verify_annotations_have_unique_ids(annotations_json_file, anns) dataset_records = _combine_images_with_annotations(dataset_name, image_root, imgs, anns) return dataset_records def register_dataset(dataset_data: CocoDatasetInfo, datasets_root: Optional[str] = None) -> None: """ Registers provided LVIS DensePose dataset Args: dataset_data: CocoDatasetInfo Dataset data datasets_root: Optional[str] Datasets root folder (default: None) """ annotations_fpath = maybe_prepend_base_path(datasets_root, dataset_data.annotations_fpath) images_root = maybe_prepend_base_path(datasets_root, dataset_data.images_root) def load_annotations(): return load_lvis_json( annotations_json_file=annotations_fpath, image_root=images_root, dataset_name=dataset_data.name, ) DatasetCatalog.register(dataset_data.name, load_annotations) MetadataCatalog.get(dataset_data.name).set( json_file=annotations_fpath, image_root=images_root, evaluator_type="lvis", **get_metadata(DENSEPOSE_METADATA_URL_PREFIX), ) def register_datasets( datasets_data: Iterable[CocoDatasetInfo], datasets_root: Optional[str] = None ) -> None: """ Registers provided LVIS DensePose datasets Args: datasets_data: Iterable[CocoDatasetInfo] An iterable of dataset datas datasets_root: Optional[str] Datasets root folder (default: None) """ for dataset_data in datasets_data: register_dataset(dataset_data, datasets_root)