avism/data/datasets/avis.py

import contextlib
import io
import logging
import numpy as np
import os
import pycocotools.mask as mask_util
from fvcore.common.file_io import PathManager
from fvcore.common.timer import Timer

from detectron2.structures import Boxes, BoxMode, PolygonMasks
from detectron2.data import DatasetCatalog, MetadataCatalog

from .avis_api.avos import AVOS


"""
This file contains functions to parse AVIS dataset of
COCO-format annotations into dicts in "Detectron2 format".
"""

logger = logging.getLogger(__name__)

__all__ = ["load_avis_json", "register_avis_instances"]


AVIS_CATEGORIES = [
    {"color": [220, 20, 60], "isthing": 1, "id": 1, "name": "person"},
    {"color": [0, 82, 0], "isthing": 1, "id": 2, "name": "violin"},
    {"color": [119, 11, 32], "isthing": 1, "id": 3, "name": "guitar"},
    {"color": [165, 42, 42], "isthing": 1, "id": 4, "name": "cello"},
    {"color": [134, 134, 103], "isthing": 1, "id": 5, "name": "flute"},
    {"color": [0, 0, 142], "isthing": 1, "id": 6, "name": "piano"},
    {"color": [255, 109, 65], "isthing": 1, "id": 7, "name": "ukulele"},
    {"color": [0, 226, 252], "isthing": 1, "id": 8, "name": "accordion"},
    {"color": [5, 121, 0], "isthing": 1, "id": 9, "name": "guzheng"},
    {"color": [0, 60, 100], "isthing": 1, "id": 10, "name": "clarinet"},
    {"color": [250, 170, 30], "isthing": 1, "id": 11, "name": "cat"},
    {"color": [100, 170, 30], "isthing": 1, "id": 12, "name": "car"},
    {"color": [179, 0, 194], "isthing": 1, "id": 13, "name": "saxophone"},
    {"color": [255, 77, 255], "isthing": 1, "id": 14, "name": "dog"},
    {"color": [120, 166, 157], "isthing": 1, "id": 15, "name": "lawn_mover"},
    {"color": [73, 77, 174], "isthing": 1, "id": 16, "name": "tuba"},
    {"color": [0, 80, 100], "isthing": 1, "id": 17, "name": "banjo"},
    {"color": [182, 182, 255], "isthing": 1, "id": 18, "name": "pipa"},
    {"color": [0, 143, 149], "isthing": 1, "id": 19, "name": "bassoon"},
    {"color": [174, 57, 255], "isthing": 1, "id": 20, "name": "airplane"},
    {"color": [0, 0, 230], "isthing": 1, "id": 21, "name": "tree_harvester"},
    {"color": [72, 0, 118], "isthing": 1, "id": 22, "name": "trumpet"},
    {"color": [255, 179, 240], "isthing": 1, "id": 23, "name": "lion"},
    {"color": [0, 125, 92], "isthing": 1, "id": 24, "name": "bass"},
    {"color": [209, 0, 151], "isthing": 1, "id": 25, "name": "erhu"},
    {"color": [188, 208, 182], "isthing": 1, "id": 26, "name": "horse"}]


def _get_avis_instances_meta():
    thing_ids = [k["id"] for k in AVIS_CATEGORIES if k["isthing"] == 1]
    thing_colors = [k["color"] for k in AVIS_CATEGORIES if k["isthing"] == 1]
    assert len(thing_ids) == 26, len(thing_ids)
    # Mapping from the incontiguous AVIS category id to an id in [0, 25]
    thing_dataset_id_to_contiguous_id = {k: i for i, k in enumerate(thing_ids)}
    thing_classes = [k["name"] for k in AVIS_CATEGORIES if k["isthing"] == 1]
    ret = {
        "thing_dataset_id_to_contiguous_id": thing_dataset_id_to_contiguous_id,
        "thing_classes": thing_classes,
        "thing_colors": thing_colors,
    }
    return ret


def load_avis_json(json_file, image_root, dataset_name=None, extra_annotation_keys=None):

    timer = Timer()
    json_file = PathManager.get_local_path(json_file)
    with contextlib.redirect_stdout(io.StringIO()):
        avis_api = AVOS(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(avis_api.getCatIds())
        cats = avis_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

        # 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
    vid_ids = sorted(avis_api.vids.keys())
    vids = avis_api.loadVids(vid_ids)

    anns = [avis_api.vidToAnns[vid_id] for vid_id in vid_ids]
    total_num_valid_anns = sum([len(x) for x in anns])
    total_num_anns = len(avis_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."
        )

    vids_anns = list(zip(vids, anns))
    logger.info("Loaded {} videos in AVIS format from {}".format(len(vids_anns), json_file))

    dataset_dicts = []

    ann_keys = ["iscrowd", "category_id", "id"] + (extra_annotation_keys or [])

    num_instances_without_valid_segmentation = 0

    for (vid_dict, anno_dict_list) in vids_anns:
        record = {}
        record["file_names"] = [os.path.join(image_root, vid_dict["file_names"][i]) for i in range(vid_dict["length"])]
        record["height"] = vid_dict["height"]
        record["width"] = vid_dict["width"]
        record["length"] = vid_dict["length"]
        video_id = record["video_id"] = vid_dict["id"]

        video_objs = []
        for frame_idx in range(record["length"]):
            frame_objs = []
            for anno in anno_dict_list:
                assert anno["video_id"] == video_id

                obj = {key: anno[key] for key in ann_keys if key in anno}

                _bboxes = anno.get("bboxes", None)
                _segm = anno.get("segmentations", None)

                if not (_bboxes and _segm and _bboxes[frame_idx] and _segm[frame_idx]):
                    continue

                bbox = _bboxes[frame_idx]
                segm = _segm[frame_idx]

                obj["bbox"] = bbox
                obj["bbox_mode"] = BoxMode.XYWH_ABS

                if isinstance(segm, dict):
                    if isinstance(segm["counts"], list):
                        # convert to compressed RLE
                        segm = mask_util.frPyObjects(segm, *segm["size"])
                elif segm:
                    # 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

                if id_map:
                    obj["category_id"] = id_map[obj["category_id"]]
                frame_objs.append(obj)
            video_objs.append(frame_objs)
        record["annotations"] = video_objs

        # audio:
        audio_feats_pth = os.path.join(image_root[:-10], "FEATAudios", vid_dict['file_names'][0].split("/")[0] + '.npy')
        record["audio"] = np.load(audio_feats_pth)

        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. "
            "A valid polygon should be a list[float] with even length >= 6."
        )
    return dataset_dicts


def register_avis_instances(name, metadata, json_file, image_root):
    """
    Register a dataset in AVIS's json annotation format for
    instance tracking.

    Args:
        name (str): the name that identifies a dataset, e.g. "avis_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_avis_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="avis", **metadata
    )