jawahar-konathala/Tryon2
0
1# Copyright (c) Facebook, Inc. and its affiliates.2import logging3import os4 5from detectron2.data import DatasetCatalog, MetadataCatalog6from detectron2.structures import BoxMode7from detectron2.utils.file_io import PathManager8from fvcore.common.timer import Timer9 10from .builtin_meta import _get_coco_instances_meta11from .lvis_v0_5_categories import LVIS_CATEGORIES as LVIS_V0_5_CATEGORIES12from .lvis_v1_categories import LVIS_CATEGORIES as LVIS_V1_CATEGORIES13from .lvis_v1_category_image_count import (14 LVIS_CATEGORY_IMAGE_COUNT as LVIS_V1_CATEGORY_IMAGE_COUNT,15)16 17"""18This file contains functions to parse LVIS-format annotations into dicts in the19"Detectron2 format".20"""21 22logger = logging.getLogger(__name__)23 24__all__ = ["load_lvis_json", "register_lvis_instances", "get_lvis_instances_meta"]25 26 27def register_lvis_instances(name, metadata, json_file, image_root):28 """29 Register a dataset in LVIS's json annotation format for instance detection and segmentation.30 31 Args:32 name (str): a name that identifies the dataset, e.g. "lvis_v0.5_train".33 metadata (dict): extra metadata associated with this dataset. It can be an empty dict.34 json_file (str): path to the json instance annotation file.35 image_root (str or path-like): directory which contains all the images.36 """37 DatasetCatalog.register(name, lambda: load_lvis_json(json_file, image_root, name))38 MetadataCatalog.get(name).set(39 json_file=json_file, image_root=image_root, evaluator_type="lvis", **metadata40 )41 42 43def load_lvis_json(44 json_file, image_root, dataset_name=None, extra_annotation_keys=None45):46 """47 Load a json file in LVIS's annotation format.48 49 Args:50 json_file (str): full path to the LVIS json annotation file.51 image_root (str): the directory where the images in this json file exists.52 dataset_name (str): the name of the dataset (e.g., "lvis_v0.5_train").53 If provided, this function will put "thing_classes" into the metadata54 associated with this dataset.55 extra_annotation_keys (list[str]): list of per-annotation keys that should also be56 loaded into the dataset dict (besides "bbox", "bbox_mode", "category_id",57 "segmentation"). The values for these keys will be returned as-is.58 59 Returns:60 list[dict]: a list of dicts in Detectron2 standard format. (See61 `Using Custom Datasets </tutorials/datasets.html>`_ )62 63 Notes:64 1. This function does not read the image files.65 The results do not have the "image" field.66 """67 from lvis import LVIS68 69 json_file = PathManager.get_local_path(json_file)70 71 timer = Timer()72 lvis_api = LVIS(json_file)73 if timer.seconds() > 1:74 logger.info(75 "Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds())76 )77 78 if dataset_name is not None:79 meta = get_lvis_instances_meta(dataset_name)80 MetadataCatalog.get(dataset_name).set(**meta)81 82 # sort indices for reproducible results83 img_ids = sorted(lvis_api.imgs.keys())84 # imgs is a list of dicts, each looks something like:85 # {'license': 4,86 # 'url': 'http://farm6.staticflickr.com/5454/9413846304_881d5e5c3b_z.jpg',87 # 'file_name': 'COCO_val2014_000000001268.jpg',88 # 'height': 427,89 # 'width': 640,90 # 'date_captured': '2013-11-17 05:57:24',91 # 'id': 1268}92 imgs = lvis_api.load_imgs(img_ids)93 # anns is a list[list[dict]], where each dict is an annotation94 # record for an object. The inner list enumerates the objects in an image95 # and the outer list enumerates over images. Example of anns[0]:96 # [{'segmentation': [[192.81,97 # 247.09,98 # ...99 # 219.03,100 # 249.06]],101 # 'area': 1035.749,102 # 'image_id': 1268,103 # 'bbox': [192.81, 224.8, 74.73, 33.43],104 # 'category_id': 16,105 # 'id': 42986},106 # ...]107 anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids]108 109 # Sanity check that each annotation has a unique id110 ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]111 assert len(set(ann_ids)) == len(112 ann_ids113 ), "Annotation ids in '{}' are not unique".format(json_file)114 115 imgs_anns = list(zip(imgs, anns))116 117 logger.info(118 "Loaded {} images in the LVIS format from {}".format(len(imgs_anns), json_file)119 )120 121 if extra_annotation_keys:122 logger.info(123 "The following extra annotation keys will be loaded: {} ".format(124 extra_annotation_keys125 )126 )127 else:128 extra_annotation_keys = []129 130 def get_file_name(img_root, img_dict):131 # Determine the path including the split folder ("train2017", "val2017", "test2017") from132 # the coco_url field. Example:133 # 'coco_url': 'http://images.cocodataset.org/train2017/000000155379.jpg'134 split_folder, file_name = img_dict["coco_url"].split("/")[-2:]135 return os.path.join(img_root + split_folder, file_name)136 137 dataset_dicts = []138 139 for (img_dict, anno_dict_list) in imgs_anns:140 record = {}141 record["file_name"] = get_file_name(image_root, img_dict)142 record["height"] = img_dict["height"]143 record["width"] = img_dict["width"]144 record["not_exhaustive_category_ids"] = img_dict.get(145 "not_exhaustive_category_ids", []146 )147 record["neg_category_ids"] = img_dict.get("neg_category_ids", [])148 image_id = record["image_id"] = img_dict["id"]149 150 objs = []151 for anno in anno_dict_list:152 # Check that the image_id in this annotation is the same as153 # the image_id we're looking at.154 # This fails only when the data parsing logic or the annotation file is buggy.155 assert anno["image_id"] == image_id156 obj = {"bbox": anno["bbox"], "bbox_mode": BoxMode.XYWH_ABS}157 # LVIS data loader can be used to load COCO dataset categories. In this case `meta`158 # variable will have a field with COCO-specific category mapping.159 if dataset_name is not None and "thing_dataset_id_to_contiguous_id" in meta:160 obj["category_id"] = meta["thing_dataset_id_to_contiguous_id"][161 anno["category_id"]162 ]163 else:164 obj["category_id"] = (165 anno["category_id"] - 1166 ) # Convert 1-indexed to 0-indexed167 segm = anno["segmentation"] # list[list[float]]168 # filter out invalid polygons (< 3 points)169 valid_segm = [170 poly for poly in segm if len(poly) % 2 == 0 and len(poly) >= 6171 ]172 assert len(segm) == len(173 valid_segm174 ), "Annotation contains an invalid polygon with < 3 points"175 assert len(segm) > 0176 obj["segmentation"] = segm177 for extra_ann_key in extra_annotation_keys:178 obj[extra_ann_key] = anno[extra_ann_key]179 objs.append(obj)180 record["annotations"] = objs181 dataset_dicts.append(record)182 183 return dataset_dicts184 185 186def get_lvis_instances_meta(dataset_name):187 """188 Load LVIS metadata.189 190 Args:191 dataset_name (str): LVIS dataset name without the split name (e.g., "lvis_v0.5").192 193 Returns:194 dict: LVIS metadata with keys: thing_classes195 """196 if "cocofied" in dataset_name:197 return _get_coco_instances_meta()198 if "v0.5" in dataset_name:199 return _get_lvis_instances_meta_v0_5()200 elif "v1" in dataset_name:201 return _get_lvis_instances_meta_v1()202 raise ValueError("No built-in metadata for dataset {}".format(dataset_name))203 204 205def _get_lvis_instances_meta_v0_5():206 assert len(LVIS_V0_5_CATEGORIES) == 1230207 cat_ids = [k["id"] for k in LVIS_V0_5_CATEGORIES]208 assert min(cat_ids) == 1 and max(cat_ids) == len(209 cat_ids210 ), "Category ids are not in [1, #categories], as expected"211 # Ensure that the category list is sorted by id212 lvis_categories = sorted(LVIS_V0_5_CATEGORIES, key=lambda x: x["id"])213 thing_classes = [k["synonyms"][0] for k in lvis_categories]214 meta = {"thing_classes": thing_classes}215 return meta216 217 218def _get_lvis_instances_meta_v1():219 assert len(LVIS_V1_CATEGORIES) == 1203220 cat_ids = [k["id"] for k in LVIS_V1_CATEGORIES]221 assert min(cat_ids) == 1 and max(cat_ids) == len(222 cat_ids223 ), "Category ids are not in [1, #categories], as expected"224 # Ensure that the category list is sorted by id225 lvis_categories = sorted(LVIS_V1_CATEGORIES, key=lambda x: x["id"])226 thing_classes = [k["synonyms"][0] for k in lvis_categories]227 meta = {228 "thing_classes": thing_classes,229 "class_image_count": LVIS_V1_CATEGORY_IMAGE_COUNT,230 }231 return meta232 233 234def main() -> None:235 global logger236 """237 Test the LVIS json dataset loader.238 239 Usage:240 python -m detectron2.data.datasets.lvis \241 path/to/json path/to/image_root dataset_name vis_limit242 """243 import sys244 245 import detectron2.data.datasets # noqa # add pre-defined metadata246 import numpy as np247 from detectron2.utils.logger import setup_logger248 from detectron2.utils.visualizer import Visualizer249 from PIL import Image250 251 logger = setup_logger(name=__name__)252 meta = MetadataCatalog.get(sys.argv[3])253 254 dicts = load_lvis_json(sys.argv[1], sys.argv[2], sys.argv[3])255 logger.info("Done loading {} samples.".format(len(dicts)))256 257 dirname = "lvis-data-vis"258 os.makedirs(dirname, exist_ok=True)259 for d in dicts[: int(sys.argv[4])]:260 img = np.array(Image.open(d["file_name"]))261 visualizer = Visualizer(img, metadata=meta)262 vis = visualizer.draw_dataset_dict(d)263 fpath = os.path.join(dirname, os.path.basename(d["file_name"]))264 vis.save(fpath)265 266 267if __name__ == "__main__":268 main() # pragma: no cover269