def register_tiny_imagenet(name, metadata: dict, data_root: str, anno_dir: str,
is_split: bool):
DatasetCatalog.register(
name,
lambda: load_tiny_annotations(data_root, anno_dir, name, is_split))
MetadataCatalog.get(name).set(data_root=data_root,
annotation_dir=anno_dir,
is_split=is_split,
evaluator_type="tiny",
**metadata)
def __init__(self, cfg):
self._logger = setup_logger(__name__, all_rank=True)
if dist.is_main_process():
self._logger.debug(f'Config File : \n{cfg}')
if cfg.VISUALIZE_DIR and not os.path.isdir(cfg.VISUALIZE_DIR) : os.makedirs(cfg.VISUALIZE_DIR)
self.visualize_dir = cfg.VISUALIZE_DIR
dist.synchronize()
self.test_loader = build_test_loader(cfg)
self.model = build_model(cfg)
self.model.eval()
if dist.is_main_process():
self._logger.debug(f"Model Structure\n{self.model}")
if dist.get_world_size() > 1:
self.model = DistributedDataParallel(self.model, device_ids=[dist.get_local_rank()], broadcast_buffers=False)
self.checkpointer = Checkpointer(
self.model,
cfg.OUTPUT_DIR,
)
self.checkpointer.load(cfg.WEIGHTS)
self.meta_data = MetadataCatalog.get(cfg.LOADER.TEST_DATASET)
self.class_color = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]
def __init__(self, cfg, distributed=True):
self._distributed = distributed
self._cpu_device = torch.device("cpu")
self._logger = setup_logger(__name__)
self._dataset_name = cfg.LOADER.TEST_DATASET
self._metadata = MetadataCatalog.get(self._dataset_name)
self._category = self._metadata.get("category_names")
def load_tiny_annotations(data_root: str, anno_dir: str, dataset_name: str,
is_split: bool):
_logger = setup_logger(__name__, all_rank=True)
meta = MetadataCatalog.get(dataset_name)
class_names = meta.category_names
dataset_dicts = []
if len(anno_dir):
annotation_dirname = os.path.join(data_root, anno_dir)
img_root = os.path.join(data_root, 'images')
f = open(annotation_dirname, 'r')
while True:
line = f.readline()
if not len(line): break
line = line.split('\t')
file_name, cls_name = line[:2]
img_id = int(file_name.split('.')[0].split('_')[1])
record = {
"file_name": os.path.join(img_root, file_name),
"image_id": img_id,
"annotations": class_names.index(cls_name),
}
dataset_dicts.append(record)
else:
annos = class_names if is_split else ['']
for anno in annos:
img_root = os.path.join(data_root, anno, 'images')
img_list = os.listdir(img_root)
for file_name in img_list:
img_id = int(file_name.split('.')[0].split('_')[1])
record = {
"file_name": os.path.join(img_root, file_name),
"image_id": img_id,
"annotations":
class_names.index(anno) if len(anno) else -1,
}
dataset_dicts.append(record)
_logger.info(
f"Loaded {len(dataset_dicts)} images in Tiny ImageNet from {dataset_name}"
)
return dataset_dicts
def __init__(self, cfg, distributed=True):
self._distributed = distributed
self._output_dir = cfg.OUTPUT_DIR
if self._output_dir and not os.path.isdir(self._output_dir):
os.makedirs(self._output_dir)
self._cpu_device = torch.device("cpu")
self._logger = setup_logger(__name__)
dataset_name = cfg.LOADER.TEST_DATASET
self._metadata = MetadataCatalog.get(dataset_name)
self._category = self._metadata.get("category_names")
with contextlib.redirect_stdout(io.StringIO()):
self._coco_api = COCO(self._metadata.json_file)
super().__init__(cfg)
def __init__(self, cfg, distributed=True):
self._distributed = distributed
self._cpu_device = torch.device("cpu")
self._logger = setup_logger(__name__)
self._dataset_name = cfg.LOADER.TEST_DATASET
self._metadata = MetadataCatalog.get(self._dataset_name)
self._category = self._metadata.get("category_names")
data_root = self._metadata.get('data_root')
self._anno_file_template = os.path.join(data_root, "Annotations",
"{}.xml")
self._image_set_path = os.path.join(
data_root, "ImageSets", "Main",
self._metadata.get('split') + ".txt")
year = self._metadata.get('year')
assert year in [2007, 2012], year
self._is_2007 = year == 2007
super().__init__(cfg)
def load_imagenet_annotations(data_root: str, dataset_name: str):
_logger = setup_logger(__name__, all_rank=True)
meta = MetadataCatalog.get(dataset_name)
class_names = meta.category_names
dataset_dicts = []
for anno in class_names:
img_root = os.path.join(data_root, anno)
img_list = os.listdir(img_root)
for file_name in img_list:
img_id = int(file_name.split('.')[0].split('_')[-1])
record = {
"file_name": os.path.join(img_root, file_name),
"image_id": img_id,
"annotations" : class_names.index(anno),
}
dataset_dicts.append(record)
_logger.info(f"Loaded {len(dataset_dicts)} images in ImageNet from {dataset_name}")
return dataset_dicts
def load_voc_instances(dirname: str, split: str, dataset_name, filter=False):
_logger = setup_logger(__name__, all_rank=True)
meta = MetadataCatalog.get(dataset_name)
class_names = meta.category_names
with open(os.path.join(dirname, "ImageSets", "Main", split + ".txt")) as f:
fileids = np.loadtxt(f, dtype=np.str)
annotation_dirname = os.path.join(dirname, "Annotations/")
dataset_dicts = []
for fileid in fileids:
anno_file = os.path.join(annotation_dirname, fileid + ".xml")
jpeg_file = os.path.join(dirname, "JPEGImages", fileid + ".jpg")
with open(anno_file) as f:
tree = ET.parse(f)
record = {
"file_name": jpeg_file,
"image_id": fileid,
"height": int(tree.findall("./size/height")[0].text),
"width": int(tree.findall("./size/width")[0].text),
}
objs = []
for obj in tree.findall("object"):
cls = obj.find("name").text
difficult = int(obj.find("difficult").text)
bbox = obj.find("bndbox")
bbox = [
float(bbox.find(x).text)
for x in ["xmin", "ymin", "xmax", "ymax"]
]
# Original annotations are integers in the range [1, W or H]
# Assuming they mean 1-based pixel indices (inclusive),
# a box with annotation (xmin=1, xmax=W) covers the whole image.
# In coordinate space this is represented by (xmin=0, xmax=W)
bbox[0] -= 1.0
bbox[1] -= 1.0
objs.append({
"category_id": class_names.index(cls),
"bbox": bbox,
"bbox_mode": BoxMode.XYXY_ABS,
'difficult': difficult
})
record["instances"] = objs
dataset_dicts.append(record)
_logger.info(
f"Loaded {len(dataset_dicts)} images in PASCAL VOC from {dirname}_{split}"
)
if filter: dataset_dicts = filter_images_with_difficult(dataset_dicts)
return dataset_dicts
def register_coco(name, metadata, json_file, image_root, filter):
DatasetCatalog.register(name, lambda: load_coco_json(json_file, image_root, name, filter=filter))
MetadataCatalog.get(name).set(
json_file=json_file, image_root=image_root, evaluator_type="coco", **metadata
)
def load_coco_json(json_file,
image_root,
dataset_name,
filter=True,
extra_annotation_keys=None):
_logger = setup_logger(__name__, all_rank=True)
start_time = time.time()
with contextlib.redirect_stdout(io.StringIO()):
coco_api = COCO(json_file)
end_time = time.time()
if end_time - start_time > 1:
_logger.info(
f"Loading {json_file} takes {end_time - start_time:.2f} seconds.")
meta = MetadataCatalog.get(dataset_name)
cat_ids = sorted(coco_api.getCatIds())
cats = coco_api.loadCats(cat_ids)
category_names = [c["name"] for c in sorted(cats, key=lambda x: x["id"])]
meta.category_names = category_names
id_map = {v: i for i, v in enumerate(cat_ids)}
meta.dataset_id_to_contiguous_id = id_map
img_ids = sorted(list(coco_api.imgs.keys()))
# list : 'license', 'url', 'file_name', 'height', 'width', 'id', 'date_captured'
imgs = coco_api.loadImgs(img_ids)
# list : 'segmentation', 'area', 'iscrowd', 'image_id', 'bbox', 'category_id', 'id'
anns = [coco_api.imgToAnns[img_id] for img_id in img_ids]
imgs_anns = list(zip(imgs, anns))
_logger.info(
f"Loaded {len(imgs_anns)} images in COCO format from {json_file}")
dataset_dicts = []
ann_keys = ["iscrowd", "bbox", "category_id"] + (extra_annotation_keys
or [])
for (img_dict, anno_dict_list) in imgs_anns:
record = {}
record["file_name"] = os.path.join(image_root, img_dict["file_name"])
record["height"] = img_dict["height"]
record["width"] = img_dict["width"]
image_id = record["image_id"] = img_dict["id"]
objs = []
for anno in anno_dict_list:
assert anno["image_id"] == image_id
assert anno.get("ignore", 0) == 0
obj = {key: anno[key] for key in ann_keys if key in anno}
obj["bbox_mode"] = BoxMode.XYWH_ABS
obj["category_id"] = id_map[obj["category_id"]]
objs.append(obj)
record["instances"] = objs
dataset_dicts.append(record)
if filter:
dataset_dicts = filter_images_with_only_crowd_annotations(
dataset_dicts)
return dataset_dicts
def register_pascal(name, metadata, data_root, split, year, filter):
DatasetCatalog.register(name, lambda: load_voc_instances(data_root, split, name, filter))
MetadataCatalog.get(name).set(
data_root=data_root, split=split, year=year, evaluator_type="voc", **metadata
)
def register_imagenet(name, metadata: dict, data_root: str):
DatasetCatalog.register(name,
lambda: load_imagenet_annotations(data_root, name))
MetadataCatalog.get(name).set(data_root=data_root,
evaluator_type="imagenet",
**metadata)