Spaces:
Paused
Paused
| import numpy as np | |
| import os, imageio | |
| import torch | |
| import scipy | |
| from tqdm import tqdm | |
| ########## Slightly modified version of LLFF data loading code | |
| ########## see https://github.com/Fyusion/LLFF for original | |
| def _minify(basedir, factors=[], resolutions=[]): | |
| needtoload = False | |
| for r in factors: | |
| imgdir = os.path.join(basedir, 'images_{}'.format(r)) | |
| if not os.path.exists(imgdir): | |
| needtoload = True | |
| for r in resolutions: | |
| imgdir = os.path.join(basedir, 'images_{}x{}'.format(r[1], r[0])) | |
| if not os.path.exists(imgdir): | |
| needtoload = True | |
| if not needtoload: | |
| return | |
| from shutil import copy | |
| from subprocess import check_output | |
| imgdir = os.path.join(basedir, 'images') | |
| imgs = [os.path.join(imgdir, f) for f in sorted(os.listdir(imgdir))] | |
| imgs = [f for f in imgs if any([f.endswith(ex) for ex in ['JPG', 'jpg', 'png', 'jpeg', 'PNG']])] | |
| imgdir_orig = imgdir | |
| wd = os.getcwd() | |
| for r in factors + resolutions: | |
| if isinstance(r, int): | |
| name = 'images_{}'.format(r) | |
| resizearg = '{}%'.format(100./r) | |
| else: | |
| name = 'images_{}x{}'.format(r[1], r[0]) | |
| resizearg = '{}x{}'.format(r[1], r[0]) | |
| imgdir = os.path.join(basedir, name) | |
| if os.path.exists(imgdir): | |
| continue | |
| print('Minifying', r, basedir) | |
| os.makedirs(imgdir) | |
| check_output('cp {}/* {}'.format(imgdir_orig, imgdir), shell=True) | |
| ext = imgs[0].split('.')[-1] | |
| args = ' '.join(['mogrify', '-resize', resizearg, '-format', 'png', '*.{}'.format(ext)]) | |
| print(args) | |
| os.chdir(imgdir) | |
| check_output(args, shell=True) | |
| os.chdir(wd) | |
| if ext != 'png': | |
| check_output('rm {}/*.{}'.format(imgdir, ext), shell=True) | |
| print('Removed duplicates') | |
| print('Done') | |
| def load_spin_data(basedir, spin_basedir, factor=None): | |
| spin_annotation_paths = os.listdir(spin_basedir) | |
| spin_annotation_paths = [n for n in spin_annotation_paths if 'cutout' not in n and 'pseudo' not in n and 'png' in n] | |
| # spin_annotation_paths = [os.path.join(spin_basedir, n) for n in spin_annotation_paths] | |
| # spin_annotation_paths = sorted(spin_annotation_paths) | |
| if factor is None: | |
| sfx = '_4' | |
| else: | |
| sfx = '_'+str(factor) | |
| imgdir = os.path.join(basedir, 'images' + sfx) | |
| if not os.path.exists(imgdir) and 'Truck' in imgdir: | |
| imgdir = os.path.join(basedir, 'train', 'rgb') | |
| elif not os.path.exists(imgdir) and 'lego' in imgdir: | |
| imgdir = os.path.join(basedir, 'rgb') | |
| from skimage.transform import resize | |
| elif not os.path.exists(imgdir): | |
| print( imgdir, 'does not exist, returning' ) | |
| return | |
| sorted_image_names = [f for f in sorted(os.listdir(imgdir)) if f.endswith('JPG') or f.endswith('jpg') or f.endswith('png')] | |
| id_to_gt_mask = {} | |
| ref_id = None | |
| for spin_annotation_name in spin_annotation_paths: | |
| for i in range(len(sorted_image_names)): | |
| if sorted_image_names[i].split('.')[-2] in spin_annotation_name.split('.')[-2]: | |
| if 'lego' in imgdir: | |
| tmp = resize(imageio.imread(os.path.join(spin_basedir, spin_annotation_name)).astype(np.float32), (768, 1020)) | |
| tmp[tmp >= 0.5] = 1 | |
| tmp[tmp != 1] = 0 | |
| id_to_gt_mask[i] = tmp | |
| # id_to_gt_mask[i] = imageio.imread(os.path.join(spin_basedir, spin_annotation_name)) | |
| # print(np.unique(id_to_gt_mask[i]), "??") | |
| else: | |
| id_to_gt_mask[i] = imageio.imread(os.path.join(spin_basedir, spin_annotation_name)) | |
| if ref_id is None: | |
| ref_id = i | |
| break | |
| return ref_id, id_to_gt_mask | |
| # if __name__ == '__main__': | |
| # print(load_spin_data('/datasets/nerf_data/nerf_llff_data(NVOS)/room/', '/datasets/nerf_data/MVSeg_data/room/')) |