densepose/data/datasets/coco.py

# 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/"


@dataclass
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)