def _download(self):
_fpath = os.path.join(Path.db_root_dir(), self.FILE)
if os.path.isfile(_fpath):
print('Files already downloaded')
return
else:
print('Downloading ' + self.URL + ' to ' + _fpath)
def _progress(count, block_size, total_size):
sys.stdout.write('\r>> %s %.1f%%' %
(_fpath, float(count * block_size) /
float(total_size) * 100.0))
sys.stdout.flush()
urllib.request.urlretrieve(self.URL, _fpath, _progress)
# extract file
cwd = os.getcwd()
print('\nExtracting tar file')
tar = tarfile.open(_fpath)
os.chdir(Path.db_root_dir())
tar.extractall()
tar.close()
os.chdir(cwd)
print('Done!')
def __init__(self,
split,
area_range=[],
only_pascal_categories=False,
mask_per_class=True,
db_root=Path.db_root_dir('COCO'),
n_samples=-1,
transform=None,
retname=True,
overfit=False):
self.split = split
self.root = os.path.join(db_root, 'images', split)
annFile = os.path.join(db_root, 'annotations',
'instances_' + split + '.json')
self.coco = COCO(annFile)
self.pascal_cat_name = [
'person', 'bird', 'cat', 'cow', 'dog', 'horse', 'sheep',
'airplane', 'bicycle', 'boat', 'bus', 'car', 'motorcycle', 'train',
'bottle', 'chair', 'dining table', 'potted plant', 'couch', 'tv'
]
self.only_pascal_categories = only_pascal_categories
if self.only_pascal_categories:
cat_ids = self.coco.getCatIds(catNms=self.pascal_cat_name)
else:
cat_ids = self.coco.getCatIds()
self.img_ids = list(self.coco.imgs.keys())
self.ids = self.coco.getAnnIds(imgIds=self.img_ids,
areaRng=area_range,
catIds=cat_ids)
self.transform = transform
self.area_range = area_range
self.cat_ids = cat_ids
self.mask_per_class = mask_per_class
self.retname = retname
if self.mask_per_class:
self._select_imgs()
if n_samples > 0:
if self.mask_per_class:
self.img_ids = list(self.img_ids)[:n_samples]
else:
self.ids = self.ids[:n_samples]
if overfit:
n_of = 64
self.img_ids = list(self.img_ids)[:n_of]
# Display stats
if self.mask_per_class:
print("Number of images: {:d}".format(len(self.img_ids)))
else:
print('Number of images: {:d}\nNumber of objects: {:d}'.format(
len(self.coco.imgs), len(self.ids)))
def __init__(self,
root=Path.db_root_dir('MSRA10K'),
download=True,
split='trainval',
transform=None,
retname=True,
overfit=False):
if download:
self._download()
self.transform = transform
self.retname = retname
self.root = root
self.gt_dir = os.path.join(self.root, 'gt')
self.image_dir = os.path.join(self.root, 'Imgs')
_splits_dir = os.path.join(self.root, 'gt_sets')
self.split = split
if isinstance(self.split, str):
self.split = [self.split]
self.images = []
self.gts = []
self.im_ids = []
for sp in self.split:
with open(os.path.join(os.path.join(_splits_dir, sp + '.txt')),
"r") as f:
lines = f.readlines()
for line in lines:
line = line.strip()
_image = os.path.join(self.image_dir, line + ".jpg")
_gt = os.path.join(self.gt_dir, line + ".png")
assert os.path.isfile(_image)
assert os.path.isfile(_gt)
self.im_ids.append(line)
self.images.append(_image)
self.gts.append(_gt)
assert (len(self.images) == len(self.gts) == len(self.im_ids))
if overfit:
n_of = 64
self.images = self.images[:n_of]
self.im_ids = self.im_ids[:n_of]
# Display stats
print('Number of images: {:d}'.format(len(self.im_ids)))
def __init__(self, root=Path.db_root_dir('MNIST'), train=True, transform=None, target_transform=None, download=False,
multitask=False):
self.root = os.path.expanduser(root)
self.transform = transform
self.target_transform = target_transform
self.train = train # training set or test set
self.multitask = multitask
if download:
self.download()
if not self._check_exists():
raise RuntimeError('Dataset not found.' +
' You can use download=True to download it')
if self.train:
self.train_data, self.train_labels = torch.load(
os.path.join(self.root, self.processed_folder, self.training_file))
else:
self.test_data, self.test_labels = torch.load(
os.path.join(self.root, self.processed_folder, self.test_file))
if multitask:
self._process_labels()
def __init__(self,
root=Path.db_root_dir('BSDS500'),
download=True,
split=['train', 'val'],
transform=None,
retname=True,
n_votes=1,
overfit=False):
if download:
self._download()
self.transform = transform
self.retname = retname
self.n_votes = n_votes
self.root = root
self.gt_dir = os.path.join(self.root, 'data', 'groundTruth')
self.image_dir = os.path.join(self.root, 'data', 'images')
_splits_dir = os.path.join(self.root, 'lists')
if not os.path.exists(os.path.join(_splits_dir)):
os.mkdir(os.path.join(_splits_dir))
self.split = split
self._get_images_trainval()
if isinstance(self.split, str):
self.split = [self.split]
self.images = []
self.gts = []
self.im_ids = []
for sp in self.split:
with open(os.path.join(os.path.join(_splits_dir, sp + '.txt')), "r") as f:
lines = f.readlines()
for line in lines:
line = line.strip()
_image = os.path.join(self.image_dir, sp, line + ".jpg")
_gt = os.path.join(self.gt_dir, sp, line + ".mat")
assert os.path.isfile(_image)
assert os.path.isfile(_gt)
self.im_ids.append(line)
self.images.append(_image)
self.gts.append(_gt)
assert (len(self.images) == len(self.gts) == len(self.im_ids))
if overfit:
n_of = 16
self.images = self.images[:n_of]
self.gts = self.gts[:n_of]
self.im_ids = self.im_ids[:n_of]
# Display stats
print('Number of images: {:d}'.format(len(self.im_ids)))
def parse_args():
def str2bool(v):
return v.lower() in ("yes", "true", "t", "1")
parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
parser.add_argument('--data',
default=Path.db_root_dir('Imagenet'),
help='path to dataset')
parser.add_argument('--arch', '-a', metavar='ARCH', default='x50')
parser.add_argument('-j',
'--workers',
default=16,
type=int,
metavar='N',
help='number of data loading workers (default: 16)')
parser.add_argument('--epochs',
default=100,
type=int,
metavar='N',
help='number of total epochs to run')
parser.add_argument('--start-epoch',
default=0,
type=int,
metavar='N',
help='manual epoch number (useful on restarts)')
parser.add_argument('-b',
'--batch-size',
default=256,
type=int,
metavar='N',
help='mini-batch size (default: 256)')
parser.add_argument('--lr',
'--learning-rate',
default=0.1,
type=float,
metavar='LR',
help='initial learning rate')
parser.add_argument('--momentum',
default=0.9,
type=float,
metavar='M',
help='momentum')
parser.add_argument('--weight-decay',
'--wd',
default=1e-4,
type=float,
metavar='W',
help='weight decay (default: 1e-4)')
parser.add_argument('--print-freq',
'-p',
default=200,
type=int,
metavar='N',
help='print frequency (default: 100)')
parser.add_argument('--resume',
default='',
type=str,
metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('-e',
'--evaluate',
dest='evaluate',
action='store_true',
help='evaluate model on validation set')
parser.add_argument('--n_gpu', type=int, default=8, help='number of GPUs')
parser.add_argument('--group_norm',
type=str2bool,
default=False,
help='Group Normalization')
args = parser.parse_args()
args.prefix = time_file_str()
return args
def __init__(self, root=Path.db_root_dir('MNIST'), train=True, transform=None, target_transform=None, download=False):
super(MultiMNIST, self).__init__(root, train, transform, target_transform, download, multitask=False)
def __init__(
self,
root=Path.db_root_dir('PASCAL'),
download=True,
split='val',
transform=None,
retname=True,
do_semseg=True,
overfit=False,
):
self.root = root
_sbd_root = os.path.join(self.root, self.BASE_DIR)
_inst_dir = os.path.join(_sbd_root, 'inst')
_cat_dir = os.path.join(_sbd_root, 'cls')
_image_dir = os.path.join(_sbd_root, 'img')
if download:
self._download()
self.transform = transform
if isinstance(split, str):
self.split = [split]
else:
split.sort()
self.split = split
self.retname = retname
self.do_semseg = do_semseg
if self.do_semseg:
self.semsegs = []
# train/val/test splits are pre-cut
_splits_dir = os.path.join(_sbd_root)
self.im_ids = []
self.images = []
print("Initializing dataloader for SBD {} set".format(''.join(
self.split)))
for splt in self.split:
with open(os.path.join(os.path.join(_splits_dir, splt + '.txt')),
"r") as f:
lines = f.read().splitlines()
for ii, line in enumerate(lines):
# Images
_image = os.path.join(_image_dir, line + ".jpg")
assert os.path.isfile(_image)
self.images.append(_image)
self.im_ids.append(line.rstrip('\n'))
# Semantic Segmentation
if self.do_semseg:
_semseg = os.path.join(_cat_dir, line + '.mat')
assert os.path.isfile(_semseg)
self.semsegs.append(_semseg)
if self.do_semseg:
assert (len(self.images) == len(self.semsegs))
# Uncomment to overfit to one image
if overfit:
n_of = 32
self.im_ids = self.im_ids[:n_of]
self.images = self.images[:n_of]
if self.do_semseg:
self.semsegs = self.semsegs[:n_of]
# Display stats
print('Number of dataset images: {:d}'.format(len(self.images)))
def __init__(
self,
root=Path.db_root_dir('PASCAL_MT'),
download=True,
split='val',
transform=None,
area_thres=0,
retname=True,
overfit=False,
do_edge=True,
do_human_parts=False,
do_semseg=False,
do_normals=False,
do_sal=False,
num_human_parts=6,
):
self.root = root
if download:
self._download()
image_dir = os.path.join(self.root, 'JPEGImages')
self.transform = transform
if isinstance(split, str):
self.split = [split]
else:
split.sort()
self.split = split
self.area_thres = area_thres
self.retname = retname
# Edge Detection
self.do_edge = do_edge
self.edges = []
edge_gt_dir = os.path.join(self.root, 'pascal-context', 'trainval')
# Semantic Segmentation
self.do_semseg = do_semseg
self.semsegs = []
# Human Part Segmentation
self.do_human_parts = do_human_parts
part_gt_dir = os.path.join(self.root, 'human_parts')
self.parts = []
self.human_parts_category = 15
self.cat_part = json.load(
open(
os.path.join(os.path.dirname(__file__),
'../util/db_info/pascal_part.json'), 'r'))
self.cat_part["15"] = self.HUMAN_PART[num_human_parts]
self.parts_file = os.path.join(
os.path.join(self.root, 'ImageSets', 'Parts'),
''.join(self.split) + '.txt')
# Surface Normal Estimation
self.do_normals = do_normals
_normal_gt_dir = os.path.join(self.root, 'normals_distill')
self.normals = []
if self.do_normals:
with open(
os.path.join(PROJECT_ROOT_DIR,
'util/db_info/nyu_classes.json')) as f:
cls_nyu = json.load(f)
with open(
os.path.join(PROJECT_ROOT_DIR,
'util/db_info/context_classes.json')) as f:
cls_context = json.load(f)
self.normals_valid_classes = []
for cl_nyu in cls_nyu:
if cl_nyu in cls_context and cl_nyu != 'unknown':
self.normals_valid_classes.append(cls_context[cl_nyu])
# Custom additions due to incompatibilities
self.normals_valid_classes.append(cls_context['tvmonitor'])
# Saliency
self.do_sal = do_sal
_sal_gt_dir = os.path.join(self.root, 'sal_distill')
self.sals = []
# train/val/test splits are pre-cut
_splits_dir = os.path.join(self.root, 'ImageSets', 'Context')
self.im_ids = []
self.images = []
print("Initializing dataloader for PASCAL {} set".format(''.join(
self.split)))
for splt in self.split:
with open(os.path.join(os.path.join(_splits_dir, splt + '.txt')),
"r") as f:
lines = f.read().splitlines()
for ii, line in enumerate(lines):
# Images
_image = os.path.join(image_dir, line + ".jpg")
assert os.path.isfile(_image)
self.images.append(_image)
self.im_ids.append(line.rstrip('\n'))
# Edges
_edge = os.path.join(edge_gt_dir, line + ".mat")
assert os.path.isfile(_edge)
self.edges.append(_edge)
# Semantic Segmentation
_semseg = self._get_semseg_fname(line)
assert os.path.isfile(_semseg)
self.semsegs.append(_semseg)
# Human Parts
_human_part = os.path.join(self.root, part_gt_dir,
line + ".mat")
assert os.path.isfile(_human_part)
self.parts.append(_human_part)
_normal = os.path.join(self.root, _normal_gt_dir,
line + ".png")
assert os.path.isfile(_normal)
self.normals.append(_normal)
_sal = os.path.join(self.root, _sal_gt_dir, line + ".png")
assert os.path.isfile(_sal)
self.sals.append(_sal)
if self.do_edge:
assert (len(self.images) == len(self.edges))
if self.do_human_parts:
assert (len(self.images) == len(self.parts))
if self.do_semseg:
assert (len(self.images) == len(self.semsegs))
if self.do_normals:
assert (len(self.images) == len(self.normals))
if self.do_sal:
assert (len(self.images) == len(self.sals))
if not self._check_preprocess_parts():
print(
'Pre-processing PASCAL dataset for human parts, this will take long, but will be done only once.'
)
self._preprocess_parts()
if self.do_human_parts:
# Find images which have human parts
self.has_human_parts = []
for ii in range(len(self.im_ids)):
if self.human_parts_category in self.part_obj_dict[
self.im_ids[ii]]:
self.has_human_parts.append(1)
else:
self.has_human_parts.append(0)
# If the other tasks are disabled, select only the images that contain human parts, to allow batching
if not self.do_edge and not self.do_semseg and not self.do_sal and not self.do_normals:
print('Ignoring images that do not contain human parts')
for i in range(len(self.parts) - 1, -1, -1):
if self.has_human_parts[i] == 0:
del self.im_ids[i]
del self.images[i]
del self.parts[i]
del self.has_human_parts[i]
print('Number of images with human parts: {:d}'.format(
np.sum(self.has_human_parts)))
# Overfit to n_of images
if overfit:
n_of = 64
self.images = self.images[:n_of]
self.im_ids = self.im_ids[:n_of]
if self.do_edge:
self.edges = self.edges[:n_of]
if self.do_semseg:
self.semsegs = self.semsegs[:n_of]
if self.do_human_parts:
self.parts = self.parts[:n_of]
if self.do_normals:
self.normals = self.normals[:n_of]
if self.do_sal:
self.sals = self.sals[:n_of]
# Display stats
print('Number of dataset images: {:d}'.format(len(self.images)))
def __init__(
self,
root=Path.db_root_dir('FSV'),
split='test',
mini=True,
transform=None,
retname=True,
overfit=False,
do_semseg=False,
do_albedo=False,
do_depth=False,
prune_rare_classes=True,
):
self.root = root
self.transform = transform
self.prune = []
if prune_rare_classes:
self.prune = [1, 4, 5, 6, 7]
self.split = split
self.retname = retname
# Original Images
self.im_ids = []
self.images = []
_image_dir = os.path.join(root, 'gta_' + split)
# Semantic segmentation
self.do_semseg = do_semseg
self.semsegs = []
# Albedo
self.do_albedo = do_albedo
self.albedos = []
# Depth
self.do_depth = do_depth
self.depths = []
# train/val/test splits are pre-cut
_splits_dir = os.path.join(root, 'gt_sets')
print("Initializing dataloader for FSV GTA {} set".format(self.split))
with open(
os.path.join(
os.path.join(_splits_dir, 'gta_' + self.split + '.txt')),
"r") as f:
lines = f.read().splitlines()
if split == 'test' and mini:
lines = lines[0:len(lines):int(len(lines) / 5000)]
for ii, line in enumerate(lines):
# Images
_image = os.path.join(_image_dir, line + "_final.webp")
# assert os.path.isfile(_image)
self.images.append(_image)
self.im_ids.append(line.rstrip('\n'))
# Semantic Segmentation
_semseg = os.path.join(_image_dir, line + "_object_id.png")
# assert os.path.isfile(_semseg)
self.semsegs.append(_semseg)
# Albedo
_albedo = os.path.join(_image_dir, line + "_albedo.webp")
# assert os.path.isfile(_albedo)
self.albedos.append(_albedo)
# Depth Estimation
_depth = os.path.join(_image_dir, line + "_disparity.webp")
# assert os.path.isfile(_depth)
self.depths.append(_depth)
if self.do_semseg:
assert (len(self.images) == len(self.semsegs))
if self.do_albedo:
assert (len(self.images) == len(self.albedos))
if self.do_depth:
assert (len(self.images) == len(self.depths))
# Uncomment to overfit to one image
if overfit:
n_of = 64
self.images = self.images[:n_of]
self.im_ids = self.im_ids[:n_of]
# Display stats
print('Number of dataset images: {:d}'.format(len(self.images)))
def __init__(
self,
root=Path.db_root_dir('NYUD_MT'),
download=True,
split='val',
transform=None,
retname=True,
overfit=False,
do_edge=True,
do_semseg=False,
do_normals=False,
do_depth=False,
):
self.root = root
if download:
self._download()
self.transform = transform
if isinstance(split, str):
self.split = [split]
else:
split.sort()
self.split = split
self.retname = retname
# Original Images
self.im_ids = []
self.images = []
_image_dir = os.path.join(root, 'images')
# Edge Detection
self.do_edge = do_edge
self.edges = []
_edge_gt_dir = os.path.join(root, 'edge')
# Semantic segmentation
self.do_semseg = do_semseg
self.semsegs = []
_semseg_gt_dir = os.path.join(root, 'segmentation')
# Surface Normals
self.do_normals = do_normals
self.normals = []
_normal_gt_dir = os.path.join(root, 'normals')
# Depth
self.do_depth = do_depth
self.depths = []
_depth_gt_dir = os.path.join(root, 'depth')
# train/val/test splits are pre-cut
_splits_dir = os.path.join(root, 'gt_sets')
print('Initializing dataloader for NYUD {} set'.format(''.join(
self.split)))
for splt in self.split:
with open(os.path.join(os.path.join(_splits_dir, splt + '.txt')),
'r') as f:
lines = f.read().splitlines()
for ii, line in enumerate(lines):
# Images
_image = os.path.join(_image_dir, line + '.jpg')
assert os.path.isfile(_image)
self.images.append(_image)
self.im_ids.append(line.rstrip('\n'))
# Edges
_edge = os.path.join(self.root, _edge_gt_dir, line + '.png')
assert os.path.isfile(_edge)
self.edges.append(_edge)
# Semantic Segmentation
_semseg = os.path.join(self.root, _semseg_gt_dir,
line + '.mat')
assert os.path.isfile(_semseg)
self.semsegs.append(_semseg)
_normal = os.path.join(self.root, _normal_gt_dir,
line + '.jpg')
assert os.path.isfile(_normal)
self.normals.append(_normal)
_depth = os.path.join(self.root, _depth_gt_dir, line + '.mat')
assert os.path.isfile(_depth)
self.depths.append(_depth)
if self.do_edge:
assert (len(self.images) == len(self.edges))
if self.do_semseg:
assert (len(self.images) == len(self.semsegs))
if self.do_normals:
assert (len(self.images) == len(self.normals))
if self.do_depth:
assert (len(self.images) == len(self.depths))
# Uncomment to overfit to one image
if overfit:
n_of = 64
self.images = self.images[:n_of]
self.im_ids = self.im_ids[:n_of]
# Display stats
print('Number of dataset images: {:d}'.format(len(self.images)))
def __init__(
self,
root=Path.db_root_dir('NYUD_raw'),
split='train',
transform=None,
do_normals=True,
do_depth=False,
retname=True,
overfit=False,
):
self.root = root
self.transform = transform
self.split = split
self.retname = retname
self.do_normals = do_normals
self.do_depth = do_depth
# Original Images
self.im_ids = []
self.images = []
_image_dir = os.path.join(root, self.split, 'images')
_mask_gt_dir = os.path.join(root, self.split, 'masks')
# Surface Normals
self.normals = []
nrm_ext = '.png' if self.split == 'train' else '.jpg'
self.masks = []
_normal_gt_dir = os.path.join(root, self.split, 'normals')
# Monocular depth
self.depths = []
_depth_gt_dir = os.path.join(root, self.split, 'depth')
# train/val/test splits are pre-cut
_splits_dir = os.path.join(root, 'gt_sets')
print('Initializing dataloader for NYUD Raw, {} set'.format(
self.split))
with open(os.path.join(os.path.join(_splits_dir, self.split + '.txt')),
'r') as f:
lines = f.read().splitlines()
for ii, line in enumerate(lines):
# Images
_image = os.path.join(_image_dir, line + '.jpg')
assert os.path.isfile(_image)
self.images.append(_image)
self.im_ids.append(line.rstrip('\n'))
if self.do_normals:
# Normals
_normal = os.path.join(self.root, _normal_gt_dir,
line + nrm_ext)
assert os.path.isfile(_normal)
self.normals.append(_normal)
if self.do_depth:
# Depth
_depth = os.path.join(self.root, _depth_gt_dir, line + '.mat')
assert os.path.isfile(_depth)
self.depths.append(_depth)
if self.split == 'train':
# Masks (only available for train data)
_mask = os.path.join(self.root, _mask_gt_dir, line + '.png')
assert os.path.isfile(_mask)
self.masks.append(_mask)
if self.do_normals:
assert (len(self.images) == len(self.normals))
if self.do_depth:
assert (len(self.images) == len(self.depths))
if self.split == 'train':
assert (len(self.images) == len(self.masks))
# uncomment to overfit to one image
if overfit:
n_of = 64
self.images = self.images[:n_of]
self.im_ids = self.im_ids[:n_of]
# display stats
print('number of dataset images: {:d}'.format(len(self.images)))
