def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
**kwargs):
root = root if root is not None else os.path.join(
dataset_dir(), 'NYUv2')
super(NYUv2, self).__init__(root, split, mode, transform, **kwargs)
_img_dir = os.path.join(root, 'images')
_mask_dir = os.path.join(root, 'labels40')
if split == 'train':
_split_f = os.path.join(root, 'train.txt')
elif split == 'val':
_split_f = os.path.join(root, 'val.txt')
else:
raise RuntimeError('Unknown dataset split: {}'.format(split))
self.images = []
self.masks = []
with open(os.path.join(_split_f), 'r') as lines:
for line in lines:
_image = os.path.join(_img_dir, line.strip() + '.jpg')
assert os.path.isfile(_image)
self.images.append(_image)
_mask = os.path.join(_mask_dir, line.strip() + '.png')
assert os.path.isfile(_mask)
self.masks.append(_mask)
assert len(self.images) == len(self.masks)
def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
**kwargs):
root = root if root is not None else os.path.join(
dataset_dir(), 'COCO')
super(MSCOCO, self).__init__(root, split, mode, transform, **kwargs)
try_import_pycocotools()
from pycocotools.coco import COCO
from pycocotools import mask
if split == 'train':
print('train set')
ann_file = os.path.join(root,
'annotations/instances_train2017.json')
ids_file = os.path.join(root, 'annotations/train_ids.mx')
self.root = os.path.join(root, 'train2017')
else:
print('val set')
ann_file = os.path.join(root, 'annotations/instances_val2017.json')
ids_file = os.path.join(root, 'annotations/val_ids.mx')
self.root = os.path.join(root, 'val2017')
self.coco = COCO(ann_file)
self.coco_mask = mask
if os.path.exists(ids_file):
with open(ids_file, 'rb') as f:
self.ids = pickle.load(f)
else:
ids = list(self.coco.imgs.keys())
self.ids = self._preprocess(ids, ids_file)
self.transform = transform
def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
**kwargs):
root = root if root is not None else os.path.join(
dataset_dir(), 'VOCdevkit', 'VOCAug')
super(PascalVOCAug, self).__init__(root, split, mode, transform,
**kwargs)
_img_dir = os.path.join(root, 'img_aug')
_mask_dir = os.path.join(root, 'cls_aug')
if split == 'train':
# _split_f = os.path.join(root, 'trainval_aug.txt')
_split_f = os.path.join(root, 'train_aug.txt')
elif split == 'val':
_split_f = os.path.join(root, 'val.txt')
else:
raise RuntimeError('Unknown dataset split: {}'.format(split))
self.images = []
self.masks = []
with open(os.path.join(_split_f), 'r') as lines:
for line in lines:
_image_name, _mask_name = line.strip().split(' ')
_image = os.path.join(_img_dir, _image_name)
assert os.path.isfile(_image)
self.images.append(_image)
_mask = os.path.join(_mask_dir, _mask_name)
assert os.path.isfile(_mask)
self.masks.append(_mask)
assert len(self.images) == len(self.masks)
def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
**kwargs):
root = root if root is not None else os.path.join(
dataset_dir(), 'WeizHorses')
super(WeizmannHorses, self).__init__(root, split, mode, transform,
**kwargs)
_img_dir = os.path.join(root, 'horse')
_mask_dir = os.path.join(root, 'mask')
if split == 'train':
_split_f = os.path.join(root, 'train.txt')
elif split == 'val':
_split_f = os.path.join(root, 'val.txt')
else:
raise RuntimeError(f'Unknown dataset split: {split}')
self.images = []
self.masks = []
with open(_split_f, 'r') as lines:
for line in lines:
_img = os.path.join(_img_dir, line.strip())
assert os.path.isfile(_img)
self.images.append(_img)
_mask = os.path.join(_mask_dir, line.strip())
assert os.path.isfile(_mask)
self.masks.append(_mask)
assert len(self.images) == len(self.masks)
def __init__(self, root=None, split='train', mode=None, transform=None, **kwargs):
root = root if root is not None else os.path.join(dataset_dir(), 'VOCdevkit', 'SBD')
super(SBD, self).__init__(root, split, mode, transform, **kwargs)
# train/val/test splits are pre-cut
_mask_dir = os.path.join(root, 'cls')
_image_dir = os.path.join(root, 'img')
if split == 'train':
_split_f = os.path.join(root, 'trainval.txt')
# _split_f = os.path.join(root, 'train.txt')
elif split == 'val':
_split_f = os.path.join(root, 'val.txt')
else:
raise RuntimeError('Unknown dataset split: {}'.format(split))
self.images = []
self.masks = []
with open(os.path.join(_split_f), "r") as lines:
for line in lines:
_image = os.path.join(_image_dir, line.rstrip('\n') + ".jpg")
assert os.path.isfile(_image)
self.images.append(_image)
_mask = os.path.join(_mask_dir, line.rstrip('\n') + ".mat")
assert os.path.isfile(_mask)
self.masks.append(_mask)
assert (len(self.images) == len(self.masks))
def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
**kwargs):
root = root if root is not None else os.path.join(
dataset_dir(), 'Aeroscapes')
super(Aeroscapes, self).__init__(root, split, mode, transform,
**kwargs)
_img_dir = os.path.join(root, 'JPEGImages')
_mask_dir = os.path.join(root, 'SegmentationClass')
_splits_dir = os.path.join(root, 'ImageSets')
if split == 'train':
_split_f = os.path.join(_splits_dir, 'trn.txt')
elif split == 'val':
_split_f = os.path.join(_splits_dir, 'val.txt')
else:
raise RuntimeError('Unknown dataset split.')
self.images = []
self.masks = []
with open(os.path.join(_split_f, ), 'r') as lines:
for line in lines:
_image = os.path.join(_img_dir, line.rstrip('\n') + '.jpg')
assert os.path.isfile(_image)
self.images.append(_image)
_mask = os.path.join(_mask_dir, line.rstrip('\n') + ".png")
assert os.path.isfile(_mask)
self.masks.append(_mask)
assert len(self.images) == len(self.masks)
def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
**kwargs):
root = root if root is not None else os.path.join(
dataset_dir(), 'Mapillary')
super(Mapillary, self).__init__(root, split, mode, transform, **kwargs)
self.images, self.masks = _get_mapi_pairs(self.root, self.split)
def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
**kwargs):
root = root if root is not None else os.path.join(
dataset_dir(), 'BDD', 'seg')
super(BDD100K, self).__init__(root, split, mode, transform, **kwargs)
self.images, self.masks = _get_bdd_pairs(root, self.split)
if split in ('train', 'val'):
assert (len(self.images) == len(self.masks))
def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
**kwargs):
root = root if root is not None else os.path.join(
dataset_dir(), 'Cityscapes')
super(CityCoarse, self).__init__(root, split, mode, transform,
**kwargs)
assert self.mode in ('train', 'val')
self.images, self.masks = _get_city_pairs(self.root, self.split)
assert (len(self.images) == len(self.masks))
def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
**kwargs):
root = root if root is not None else os.path.join(
dataset_dir(), 'ADE20K/ADEChallengeData2016')
super(ADE20K, self).__init__(root, split, mode, transform, **kwargs)
self.images, self.masks = _get_ade20k_pairs(root, split)
assert (len(self.images) == len(self.masks))
if len(self.images) == 0:
raise (RuntimeError("Found 0 images in sub-folders of: " + root +
"\n"))
def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
base_size=768,
**kwargs):
root = root if root is not None else os.path.join(
dataset_dir(), 'MHP', 'v1')
super(MHPV1, self).__init__(root, split, mode, transform, base_size,
**kwargs)
self.images, self.masks = _get_mhp_pairs(root, split)
assert (len(self.images) == len(self.masks))
if len(self.images) == 0:
raise (RuntimeError("Found 0 images in sub-folders of: " + root +
"\n"))
def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
**kwargs):
root = root if root is not None else os.path.join(
dataset_dir(), 'Cityscapes')
super(Cityscapes, self).__init__(root, split, mode, transform,
**kwargs)
self.images, self.mask_paths = _get_city_pairs(self.root, self.split)
assert (len(self.images) == len(self.mask_paths))
self.valid_classes = [
7, 8, 11, 12, 13, 17, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 31,
32, 33
]
self._key = np.array([
-1, -1, -1, -1, -1, -1, -1, -1, 0, 1, -1, -1, 2, 3, 4, -1, -1, -1,
5, -1, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, -1, -1, 16, 17, 18
])
self._mapping = np.array(range(-1, len(self._key) - 1)).astype('int32')
def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
full_resolution=True,
**kwargs):
base_dir = 'CamVidFull' if full_resolution else 'CamVid'
root = root if root is not None else os.path.join(
dataset_dir(), base_dir)
super(CamVid, self).__init__(root, split, mode, transform, **kwargs)
_img_dir = os.path.join(root, 'images')
_mask_dir = os.path.join(root, 'labels')
if split == 'train':
_split_f = os.path.join(root, 'trainval.txt')
print(f'Training list: {_split_f}')
elif split == 'val':
_split_f = os.path.join(root, 'test.txt')
print(f'Validation list: {_split_f}')
elif split == 'test':
_split_f = os.path.join(root, 'test.txt')
else:
raise RuntimeError(f'Unknown dataset split: {split}')
self.images = []
self.masks = []
with open(os.path.join(_split_f), 'r') as lines:
for line in lines:
_img_name, _mask_name = line.strip().split(' ')
_image = os.path.join(_img_dir, _img_name)
assert os.path.isfile(_image)
self.images.append(_image)
_mask = os.path.join(_mask_dir, _mask_name)
assert os.path.isfile(_mask)
self.masks.append(_mask)
assert len(self.images) == len(self.masks)
def __init__(self,
root=None,
split='train',
mode=None,
transform=None,
**kwargs):
root = root if root is not None else os.path.join(
dataset_dir(), 'VOCdevkit', 'VOC2012')
super(PascalVOC, self).__init__(root, split, mode, transform, **kwargs)
_mask_dir = os.path.join(root, 'SegmentationClass')
_image_dir = os.path.join(root, 'JPEGImages')
# train/val/test splits are pre-cut
_splits_dir = os.path.join(root, 'ImageSets/Segmentation')
if split == 'train':
_split_f = os.path.join(_splits_dir, 'trainval.txt')
# _split_f = os.path.join(_splits_dir, 'train.txt')
elif split == 'val':
_split_f = os.path.join(_splits_dir, 'val.txt')
elif split == 'test':
_split_f = os.path.join(_splits_dir, 'test.txt')
else:
raise RuntimeError('Unknown dataset split.')
self.images = []
self.masks = []
with open(os.path.join(_split_f), "r") as lines:
for line in lines:
_image = os.path.join(_image_dir, line.rstrip('\n') + ".jpg")
assert os.path.isfile(_image)
self.images.append(_image)
if split != 'test':
_mask = os.path.join(_mask_dir, line.rstrip('\n') + ".png")
assert os.path.isfile(_mask)
self.masks.append(_mask)
if split != 'test':
assert (len(self.images) == len(self.masks))