densepose/data/datasets/lvis.py

# Copyright (c) Facebook, Inc. and its affiliates.
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)