Spaces:
Sleeping
Sleeping
# Copyright (c) Facebook, Inc. and its affiliates. | |
import contextlib | |
import io | |
import logging | |
import os | |
from collections import defaultdict | |
from dataclasses import dataclass | |
from typing import Any, Dict, Iterable, List, Optional | |
from fvcore.common.timer import Timer | |
from detectron2.data import DatasetCatalog, MetadataCatalog | |
from detectron2.structures import BoxMode | |
from detectron2.utils.file_io import PathManager | |
from ..utils import maybe_prepend_base_path | |
DENSEPOSE_MASK_KEY = "dp_masks" | |
DENSEPOSE_IUV_KEYS_WITHOUT_MASK = ["dp_x", "dp_y", "dp_I", "dp_U", "dp_V"] | |
DENSEPOSE_CSE_KEYS_WITHOUT_MASK = ["dp_x", "dp_y", "dp_vertex", "ref_model"] | |
DENSEPOSE_ALL_POSSIBLE_KEYS = set( | |
DENSEPOSE_IUV_KEYS_WITHOUT_MASK + DENSEPOSE_CSE_KEYS_WITHOUT_MASK + [DENSEPOSE_MASK_KEY] | |
) | |
DENSEPOSE_METADATA_URL_PREFIX = "https://dl.fbaipublicfiles.com/densepose/data/" | |
class CocoDatasetInfo: | |
name: str | |
images_root: str | |
annotations_fpath: str | |
DATASETS = [ | |
CocoDatasetInfo( | |
name="densepose_coco_2014_train", | |
images_root="coco/train2014", | |
annotations_fpath="coco/annotations/densepose_train2014.json", | |
), | |
CocoDatasetInfo( | |
name="densepose_coco_2014_minival", | |
images_root="coco/val2014", | |
annotations_fpath="coco/annotations/densepose_minival2014.json", | |
), | |
CocoDatasetInfo( | |
name="densepose_coco_2014_minival_100", | |
images_root="coco/val2014", | |
annotations_fpath="coco/annotations/densepose_minival2014_100.json", | |
), | |
CocoDatasetInfo( | |
name="densepose_coco_2014_valminusminival", | |
images_root="coco/val2014", | |
annotations_fpath="coco/annotations/densepose_valminusminival2014.json", | |
), | |
CocoDatasetInfo( | |
name="densepose_coco_2014_train_cse", | |
images_root="coco/train2014", | |
annotations_fpath="coco_cse/densepose_train2014_cse.json", | |
), | |
CocoDatasetInfo( | |
name="densepose_coco_2014_minival_cse", | |
images_root="coco/val2014", | |
annotations_fpath="coco_cse/densepose_minival2014_cse.json", | |
), | |
CocoDatasetInfo( | |
name="densepose_coco_2014_minival_100_cse", | |
images_root="coco/val2014", | |
annotations_fpath="coco_cse/densepose_minival2014_100_cse.json", | |
), | |
CocoDatasetInfo( | |
name="densepose_coco_2014_valminusminival_cse", | |
images_root="coco/val2014", | |
annotations_fpath="coco_cse/densepose_valminusminival2014_cse.json", | |
), | |
CocoDatasetInfo( | |
name="densepose_chimps", | |
images_root="densepose_chimps/images", | |
annotations_fpath="densepose_chimps/densepose_chimps_densepose.json", | |
), | |
CocoDatasetInfo( | |
name="densepose_chimps_cse_train", | |
images_root="densepose_chimps/images", | |
annotations_fpath="densepose_chimps/densepose_chimps_cse_train.json", | |
), | |
CocoDatasetInfo( | |
name="densepose_chimps_cse_val", | |
images_root="densepose_chimps/images", | |
annotations_fpath="densepose_chimps/densepose_chimps_cse_val.json", | |
), | |
CocoDatasetInfo( | |
name="posetrack2017_train", | |
images_root="posetrack2017/posetrack_data_2017", | |
annotations_fpath="posetrack2017/densepose_posetrack_train2017.json", | |
), | |
CocoDatasetInfo( | |
name="posetrack2017_val", | |
images_root="posetrack2017/posetrack_data_2017", | |
annotations_fpath="posetrack2017/densepose_posetrack_val2017.json", | |
), | |
CocoDatasetInfo( | |
name="lvis_v05_train", | |
images_root="coco/train2017", | |
annotations_fpath="lvis/lvis_v0.5_plus_dp_train.json", | |
), | |
CocoDatasetInfo( | |
name="lvis_v05_val", | |
images_root="coco/val2017", | |
annotations_fpath="lvis/lvis_v0.5_plus_dp_val.json", | |
), | |
] | |
BASE_DATASETS = [ | |
CocoDatasetInfo( | |
name="base_coco_2017_train", | |
images_root="coco/train2017", | |
annotations_fpath="coco/annotations/instances_train2017.json", | |
), | |
CocoDatasetInfo( | |
name="base_coco_2017_val", | |
images_root="coco/val2017", | |
annotations_fpath="coco/annotations/instances_val2017.json", | |
), | |
CocoDatasetInfo( | |
name="base_coco_2017_val_100", | |
images_root="coco/val2017", | |
annotations_fpath="coco/annotations/instances_val2017_100.json", | |
), | |
] | |
def get_metadata(base_path: Optional[str]) -> Dict[str, Any]: | |
""" | |
Returns metadata associated with COCO DensePose datasets | |
Args: | |
base_path: Optional[str] | |
Base path used to load metadata from | |
Returns: | |
Dict[str, Any] | |
Metadata in the form of a dictionary | |
""" | |
meta = { | |
"densepose_transform_src": maybe_prepend_base_path(base_path, "UV_symmetry_transforms.mat"), | |
"densepose_smpl_subdiv": maybe_prepend_base_path(base_path, "SMPL_subdiv.mat"), | |
"densepose_smpl_subdiv_transform": maybe_prepend_base_path( | |
base_path, | |
"SMPL_SUBDIV_TRANSFORM.mat", | |
), | |
} | |
return meta | |
def _load_coco_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 pycocotools.coco import COCO | |
logger = logging.getLogger(__name__) | |
timer = Timer() | |
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())) | |
return coco_api | |
def _add_categories_metadata(dataset_name: str, categories: List[Dict[str, Any]]): | |
meta = MetadataCatalog.get(dataset_name) | |
meta.categories = {c["id"]: c["name"] for c in categories} | |
logger = logging.getLogger(__name__) | |
logger.info("Dataset {} categories: {}".format(dataset_name, meta.categories)) | |
def _verify_annotations_have_unique_ids(json_file: str, anns: List[List[Dict[str, Any]]]): | |
if "minival" in json_file: | |
# Skip validation on COCO2014 valminusminival and minival annotations | |
# The ratio of buggy annotations there is tiny and does not affect accuracy | |
# Therefore we explicitly white-list them | |
return | |
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]): | |
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]): | |
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]): | |
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]): | |
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]]], | |
): | |
ann_keys = ["iscrowd", "category_id"] | |
dataset_dicts = [] | |
contains_video_frame_info = False | |
for img_dict, ann_dicts in zip(img_datas, ann_datas): | |
record = {} | |
record["file_name"] = os.path.join(image_root, img_dict["file_name"]) | |
record["height"] = img_dict["height"] | |
record["width"] = img_dict["width"] | |
record["image_id"] = img_dict["id"] | |
record["dataset"] = dataset_name | |
if "frame_id" in img_dict: | |
record["frame_id"] = img_dict["frame_id"] | |
record["video_id"] = img_dict.get("vid_id", None) | |
contains_video_frame_info = True | |
objs = [] | |
for ann_dict in ann_dicts: | |
assert ann_dict["image_id"] == record["image_id"] | |
assert ann_dict.get("ignore", 0) == 0 | |
obj = {key: ann_dict[key] for key in ann_keys if key in ann_dict} | |
_maybe_add_bbox(obj, ann_dict) | |
_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) | |
if contains_video_frame_info: | |
create_video_frame_mapping(dataset_name, dataset_dicts) | |
return dataset_dicts | |
def get_contiguous_id_to_category_id_map(metadata): | |
cat_id_2_cont_id = metadata.thing_dataset_id_to_contiguous_id | |
cont_id_2_cat_id = {} | |
for cat_id, cont_id in cat_id_2_cont_id.items(): | |
if cont_id in cont_id_2_cat_id: | |
continue | |
cont_id_2_cat_id[cont_id] = cat_id | |
return cont_id_2_cat_id | |
def maybe_filter_categories_cocoapi(dataset_name, coco_api): | |
meta = MetadataCatalog.get(dataset_name) | |
cont_id_2_cat_id = get_contiguous_id_to_category_id_map(meta) | |
cat_id_2_cont_id = meta.thing_dataset_id_to_contiguous_id | |
# filter categories | |
cats = [] | |
for cat in coco_api.dataset["categories"]: | |
cat_id = cat["id"] | |
if cat_id not in cat_id_2_cont_id: | |
continue | |
cont_id = cat_id_2_cont_id[cat_id] | |
if (cont_id in cont_id_2_cat_id) and (cont_id_2_cat_id[cont_id] == cat_id): | |
cats.append(cat) | |
coco_api.dataset["categories"] = cats | |
# filter annotations, if multiple categories are mapped to a single | |
# contiguous ID, use only one category ID and map all annotations to that category ID | |
anns = [] | |
for ann in coco_api.dataset["annotations"]: | |
cat_id = ann["category_id"] | |
if cat_id not in cat_id_2_cont_id: | |
continue | |
cont_id = cat_id_2_cont_id[cat_id] | |
ann["category_id"] = cont_id_2_cat_id[cont_id] | |
anns.append(ann) | |
coco_api.dataset["annotations"] = anns | |
# recreate index | |
coco_api.createIndex() | |
def maybe_filter_and_map_categories_cocoapi(dataset_name, coco_api): | |
meta = MetadataCatalog.get(dataset_name) | |
category_id_map = meta.thing_dataset_id_to_contiguous_id | |
# map categories | |
cats = [] | |
for cat in coco_api.dataset["categories"]: | |
cat_id = cat["id"] | |
if cat_id not in category_id_map: | |
continue | |
cat["id"] = category_id_map[cat_id] | |
cats.append(cat) | |
coco_api.dataset["categories"] = cats | |
# map annotation categories | |
anns = [] | |
for ann in coco_api.dataset["annotations"]: | |
cat_id = ann["category_id"] | |
if cat_id not in category_id_map: | |
continue | |
ann["category_id"] = category_id_map[cat_id] | |
anns.append(ann) | |
coco_api.dataset["annotations"] = anns | |
# recreate index | |
coco_api.createIndex() | |
def create_video_frame_mapping(dataset_name, dataset_dicts): | |
mapping = defaultdict(dict) | |
for d in dataset_dicts: | |
video_id = d.get("video_id") | |
if video_id is None: | |
continue | |
mapping[video_id].update({d["frame_id"]: d["file_name"]}) | |
MetadataCatalog.get(dataset_name).set(video_frame_mapping=mapping) | |
def load_coco_json(annotations_json_file: str, image_root: str, dataset_name: str): | |
""" | |
Loads a JSON file with annotations in COCO instances format. | |
Replaces `detectron2.data.datasets.coco.load_coco_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. | |
""" | |
coco_api = _load_coco_annotations(PathManager.get_local_path(annotations_json_file)) | |
_add_categories_metadata(dataset_name, coco_api.loadCats(coco_api.getCatIds())) | |
# 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) | |
logger = logging.getLogger(__name__) | |
logger.info("Loaded {} images in COCO 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 = [coco_api.imgToAnns[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): | |
""" | |
Registers provided COCO 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_coco_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, | |
**get_metadata(DENSEPOSE_METADATA_URL_PREFIX) | |
) | |
def register_datasets( | |
datasets_data: Iterable[CocoDatasetInfo], datasets_root: Optional[str] = None | |
): | |
""" | |
Registers provided COCO 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) | |