115 lines
3.2 KiB
Python
Executable File
115 lines
3.2 KiB
Python
Executable File
import os
|
|
|
|
|
|
class DatasetCatalog:
|
|
DATA_DIR = ''
|
|
DATASETS = {
|
|
'voc_ICCV_train': {
|
|
"data_dir": "VOCiccvtrain",
|
|
"split": "trainyolo"
|
|
},
|
|
'voc_ICCV_test': {
|
|
"data_dir": "VOCiccvtest",
|
|
"split": "testyolo"
|
|
},
|
|
'voc_2007_train': {
|
|
"data_dir": "VOC2007",
|
|
"split": "train"
|
|
},
|
|
'voc_3_train': {
|
|
"data_dir": "VOC3",
|
|
"split": "train"
|
|
},
|
|
'voc_3_test': {
|
|
"data_dir": "VOC3",
|
|
"split": "test"
|
|
},
|
|
'voc_2007_val': {
|
|
"data_dir": "VOC2007",
|
|
"split": "val"
|
|
},
|
|
'voc_2007_trainval': {
|
|
"data_dir": "VOC2007",
|
|
"split": "trainval"
|
|
},
|
|
'voc_2007_test': {
|
|
"data_dir": "VOC2007",
|
|
"split": "test"
|
|
},
|
|
'vocbird_2019_test': {
|
|
"data_dir": "HatcherW",
|
|
"split": "test"
|
|
},
|
|
'voc_2012_train': {
|
|
"data_dir": "VOC2012",
|
|
"split": "train"
|
|
},
|
|
'voc_2012_val': {
|
|
"data_dir": "VOC2012",
|
|
"split": "val"
|
|
},
|
|
'voc_2012_trainval': {
|
|
"data_dir": "VOC2012",
|
|
"split": "trainval"
|
|
},
|
|
'voc_2012_test': {
|
|
"data_dir": "VOC2012",
|
|
"split": "test"
|
|
},
|
|
'coco_2014_valminusminival': {
|
|
"data_dir": "val2014",
|
|
"ann_file": "annotations/instances_valminusminival2014.json"
|
|
},
|
|
'coco_2014_minival': {
|
|
"data_dir": "val2014",
|
|
"ann_file": "annotations/instances_minival2014.json"
|
|
},
|
|
'coco_2014_train': {
|
|
"data_dir": "train2014",
|
|
"ann_file": "annotations/instances_train2014.json"
|
|
},
|
|
'coco_2014_val': {
|
|
"data_dir": "val2014",
|
|
"ann_file": "annotations/instances_val2014.json"
|
|
},
|
|
"voc_2007_birdtest": {
|
|
"data_dir": "VOC2007",
|
|
"split": "bird_test"
|
|
},
|
|
"voc_2007_birdval": {
|
|
"data_dir": "VOC2007",
|
|
"split": "bird_trainval"
|
|
},
|
|
"voc_2007_birdtrain": {
|
|
"data_dir": "VOC2007",
|
|
"split": "bird_train"
|
|
},
|
|
}
|
|
|
|
@staticmethod
|
|
def get(name):
|
|
if "voc" in name:
|
|
voc_root = DatasetCatalog.DATA_DIR
|
|
if 'VOC_ROOT' in os.environ:
|
|
voc_root = os.environ['VOC_ROOT']
|
|
|
|
attrs = DatasetCatalog.DATASETS[name]
|
|
args = dict(
|
|
data_dir=os.path.join(voc_root, attrs["data_dir"]),
|
|
split=attrs["split"],
|
|
)
|
|
return dict(factory="VOCDataset", args=args)
|
|
elif "coco" in name:
|
|
coco_root = DatasetCatalog.DATA_DIR
|
|
if 'COCO_ROOT' in os.environ:
|
|
coco_root = os.environ['COCO_ROOT']
|
|
|
|
attrs = DatasetCatalog.DATASETS[name]
|
|
args = dict(
|
|
data_dir=os.path.join(coco_root, attrs["data_dir"]),
|
|
ann_file=os.path.join(coco_root, attrs["ann_file"]),
|
|
)
|
|
return dict(factory="COCODataset", args=args)
|
|
|
|
raise RuntimeError("Dataset not available: {}".format(name))
|