jawahar-konathala/Tryon2

1# Copyright (c) Facebook, Inc. and its affiliates.2import copy3import json4import os5 6from detectron2.data import DatasetCatalog, MetadataCatalog7from detectron2.utils.file_io import PathManager8 9from .coco import load_coco_json, load_sem_seg10 11__all__ = ["register_coco_panoptic", "register_coco_panoptic_separated"]12 13 14def load_coco_panoptic_json(json_file, image_dir, gt_dir, meta):15    """16    Args:17        image_dir (str): path to the raw dataset. e.g., "~/coco/train2017".18        gt_dir (str): path to the raw annotations. e.g., "~/coco/panoptic_train2017".19        json_file (str): path to the json file. e.g., "~/coco/annotations/panoptic_train2017.json".20 21    Returns:22        list[dict]: a list of dicts in Detectron2 standard format. (See23        `Using Custom Datasets </tutorials/datasets.html>`_ )24    """25 26    def _convert_category_id(segment_info, meta):27        if segment_info["category_id"] in meta["thing_dataset_id_to_contiguous_id"]:28            segment_info["category_id"] = meta["thing_dataset_id_to_contiguous_id"][29                segment_info["category_id"]30            ]31            segment_info["isthing"] = True32        else:33            segment_info["category_id"] = meta["stuff_dataset_id_to_contiguous_id"][34                segment_info["category_id"]35            ]36            segment_info["isthing"] = False37        return segment_info38 39    with PathManager.open(json_file) as f:40        json_info = json.load(f)41 42    ret = []43    for ann in json_info["annotations"]:44        image_id = int(ann["image_id"])45        # TODO: currently we assume image and label has the same filename but46        # different extension, and images have extension ".jpg" for COCO. Need47        # to make image extension a user-provided argument if we extend this48        # function to support other COCO-like datasets.49        image_file = os.path.join(image_dir, os.path.splitext(ann["file_name"])[0] + ".jpg")50        label_file = os.path.join(gt_dir, ann["file_name"])51        segments_info = [_convert_category_id(x, meta) for x in ann["segments_info"]]52        ret.append(53            {54                "file_name": image_file,55                "image_id": image_id,56                "pan_seg_file_name": label_file,57                "segments_info": segments_info,58            }59        )60    assert len(ret), f"No images found in {image_dir}!"61    assert PathManager.isfile(ret[0]["file_name"]), ret[0]["file_name"]62    assert PathManager.isfile(ret[0]["pan_seg_file_name"]), ret[0]["pan_seg_file_name"]63    return ret64 65 66def register_coco_panoptic(67    name, metadata, image_root, panoptic_root, panoptic_json, instances_json=None68):69    """70    Register a "standard" version of COCO panoptic segmentation dataset named `name`.71    The dictionaries in this registered dataset follows detectron2's standard format.72    Hence it's called "standard".73 74    Args:75        name (str): the name that identifies a dataset,76            e.g. "coco_2017_train_panoptic"77        metadata (dict): extra metadata associated with this dataset.78        image_root (str): directory which contains all the images79        panoptic_root (str): directory which contains panoptic annotation images in COCO format80        panoptic_json (str): path to the json panoptic annotation file in COCO format81        sem_seg_root (none): not used, to be consistent with82            `register_coco_panoptic_separated`.83        instances_json (str): path to the json instance annotation file84    """85    panoptic_name = name86    DatasetCatalog.register(87        panoptic_name,88        lambda: load_coco_panoptic_json(panoptic_json, image_root, panoptic_root, metadata),89    )90    MetadataCatalog.get(panoptic_name).set(91        panoptic_root=panoptic_root,92        image_root=image_root,93        panoptic_json=panoptic_json,94        json_file=instances_json,95        evaluator_type="coco_panoptic_seg",96        ignore_label=255,97        label_divisor=1000,98        **metadata,99    )100 101 102def register_coco_panoptic_separated(103    name, metadata, image_root, panoptic_root, panoptic_json, sem_seg_root, instances_json104):105    """106    Register a "separated" version of COCO panoptic segmentation dataset named `name`.107    The annotations in this registered dataset will contain both instance annotations and108    semantic annotations, each with its own contiguous ids. Hence it's called "separated".109 110    It follows the setting used by the PanopticFPN paper:111 112    1. The instance annotations directly come from polygons in the COCO113       instances annotation task, rather than from the masks in the COCO panoptic annotations.114 115       The two format have small differences:116       Polygons in the instance annotations may have overlaps.117       The mask annotations are produced by labeling the overlapped polygons118       with depth ordering.119 120    2. The semantic annotations are converted from panoptic annotations, where121       all "things" are assigned a semantic id of 0.122       All semantic categories will therefore have ids in contiguous123       range [1, #stuff_categories].124 125    This function will also register a pure semantic segmentation dataset126    named ``name + '_stuffonly'``.127 128    Args:129        name (str): the name that identifies a dataset,130            e.g. "coco_2017_train_panoptic"131        metadata (dict): extra metadata associated with this dataset.132        image_root (str): directory which contains all the images133        panoptic_root (str): directory which contains panoptic annotation images134        panoptic_json (str): path to the json panoptic annotation file135        sem_seg_root (str): directory which contains all the ground truth segmentation annotations.136        instances_json (str): path to the json instance annotation file137    """138    panoptic_name = name + "_separated"139    DatasetCatalog.register(140        panoptic_name,141        lambda: merge_to_panoptic(142            load_coco_json(instances_json, image_root, panoptic_name),143            load_sem_seg(sem_seg_root, image_root),144        ),145    )146    MetadataCatalog.get(panoptic_name).set(147        panoptic_root=panoptic_root,148        image_root=image_root,149        panoptic_json=panoptic_json,150        sem_seg_root=sem_seg_root,151        json_file=instances_json,  # TODO rename152        evaluator_type="coco_panoptic_seg",153        ignore_label=255,154        **metadata,155    )156 157    semantic_name = name + "_stuffonly"158    DatasetCatalog.register(semantic_name, lambda: load_sem_seg(sem_seg_root, image_root))159    MetadataCatalog.get(semantic_name).set(160        sem_seg_root=sem_seg_root,161        image_root=image_root,162        evaluator_type="sem_seg",163        ignore_label=255,164        **metadata,165    )166 167 168def merge_to_panoptic(detection_dicts, sem_seg_dicts):169    """170    Create dataset dicts for panoptic segmentation, by171    merging two dicts using "file_name" field to match their entries.172 173    Args:174        detection_dicts (list[dict]): lists of dicts for object detection or instance segmentation.175        sem_seg_dicts (list[dict]): lists of dicts for semantic segmentation.176 177    Returns:178        list[dict] (one per input image): Each dict contains all (key, value) pairs from dicts in179            both detection_dicts and sem_seg_dicts that correspond to the same image.180            The function assumes that the same key in different dicts has the same value.181    """182    results = []183    sem_seg_file_to_entry = {x["file_name"]: x for x in sem_seg_dicts}184    assert len(sem_seg_file_to_entry) > 0185 186    for det_dict in detection_dicts:187        dic = copy.copy(det_dict)188        dic.update(sem_seg_file_to_entry[dic["file_name"]])189        results.append(dic)190    return results191 192 193if __name__ == "__main__":194    """195    Test the COCO panoptic dataset loader.196 197    Usage:198        python -m detectron2.data.datasets.coco_panoptic \199            path/to/image_root path/to/panoptic_root path/to/panoptic_json dataset_name 10200 201        "dataset_name" can be "coco_2017_train_panoptic", or other202        pre-registered ones203    """204    from detectron2.utils.logger import setup_logger205    from detectron2.utils.visualizer import Visualizer206    import detectron2.data.datasets  # noqa # add pre-defined metadata207    import sys208    from PIL import Image209    import numpy as np210 211    logger = setup_logger(name=__name__)212    assert sys.argv[4] in DatasetCatalog.list()213    meta = MetadataCatalog.get(sys.argv[4])214 215    dicts = load_coco_panoptic_json(sys.argv[3], sys.argv[1], sys.argv[2], meta.as_dict())216    logger.info("Done loading {} samples.".format(len(dicts)))217 218    dirname = "coco-data-vis"219    os.makedirs(dirname, exist_ok=True)220    num_imgs_to_vis = int(sys.argv[5])221    for i, d in enumerate(dicts):222        img = np.array(Image.open(d["file_name"]))223        visualizer = Visualizer(img, metadata=meta)224        vis = visualizer.draw_dataset_dict(d)225        fpath = os.path.join(dirname, os.path.basename(d["file_name"]))226        vis.save(fpath)227        if i + 1 >= num_imgs_to_vis:228            break229