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)))