def __init__(self, data_dir, num_views=1000000, sidelength=None):
self.sidelength = sidelength
color_dir = os.path.join(data_dir, "rgb")
pose_dir = os.path.join(data_dir, "pose")
if not os.path.isdir(color_dir) or not os.path.isdir(pose_dir):
print("Error! root dir %s is wrong" % data_dir)
return
self.rgb_paths = sorted(data_util.glob_imgs(color_dir))[:num_views]
self.pose_paths = sorted(glob(os.path.join(pose_dir,
"*.txt")))[:num_views]
def __init__(self, data_dir, img_gamma=1.0):
super().__init__()
self.data_dir = data_dir
self.img_gamma = img_gamma
if not os.path.isdir(data_dir):
raise ValueError("Error! data dir is wrong")
# get path for all light probes
self.lp_fp_all = sorted(data_util.glob_imgs(self.data_dir))
self.lp_all = [None] * len(self.lp_fp_all)
def __init__(self, root_dir, img_size, sampling_pattern):
super().__init__()
self.img_size = img_size
self.color_dir = os.path.join(root_dir, 'rgb')
self.pose_dir = os.path.join(root_dir, 'pose')
if not os.path.isdir(self.color_dir):
print("Error! root dir is wrong")
return
self.all_color = sorted(data_util.glob_imgs(self.color_dir))
self.all_poses = sorted(glob(os.path.join(self.pose_dir, '*.txt')))
# Subsample the trajectory for training / test set split as well as the result matrix
file_lists = [self.all_color, self.all_poses]
if sampling_pattern != 'all':
if sampling_pattern.split('_')[0] == 'skip':
skip_val = int(sampling_pattern.split('_')[-1])
for i in range(len(file_lists)):
dummy_list = deepcopy(file_lists[i])
file_lists[i].clear()
file_lists[i].extend(dummy_list[::skip_val + 1])
else:
print("Unknown sampling pattern!")
return None
# Buffer files
print("Buffering files...")
self.all_views = []
self.all_views_nearest = []
for i in range(self.__len__()):
if not i % 10:
print(i)
self.all_views.append(self.read_view_tuple(i))
self.all_views_nearest.append(self.read_view_tuple_nearest(i))
# Calculate the ranking of nearest neigbors
self.nn_idcs, _ = data_util.get_nn_ranking(
[data_util.load_pose(pose) for pose in self.all_poses])
print("*" * 100)
print("Sampling pattern ", sampling_pattern)
print("Image size ", self.img_size)
print("*" * 100)
def __init__(
self,
instance_idx,
instance_dir,
specific_observation_idcs=None, # For few-shot case: Can pick specific observations only
img_sidelength=None,
num_images=-1):
self.instance_idx = instance_idx
self.img_sidelength = img_sidelength
self.instance_dir = instance_dir
color_dir = os.path.join(instance_dir, "rgb")
pose_dir = os.path.join(instance_dir, "pose")
param_dir = os.path.join(instance_dir, "params")
if not os.path.isdir(color_dir):
print("Error! root dir %s is wrong" % instance_dir)
return
self.has_params = os.path.isdir(param_dir)
self.color_paths = sorted(data_util.glob_imgs(color_dir))
self.pose_paths = sorted(glob(os.path.join(pose_dir, "*.txt")))
if self.has_params:
self.param_paths = sorted(glob(os.path.join(param_dir, "*.txt")))
else:
self.param_paths = []
if specific_observation_idcs is not None:
self.color_paths = pick(self.color_paths,
specific_observation_idcs)
self.pose_paths = pick(self.pose_paths, specific_observation_idcs)
self.param_paths = pick(self.param_paths,
specific_observation_idcs)
elif num_images != -1:
idcs = np.linspace(0,
stop=len(self.color_paths),
num=num_images,
endpoint=False,
dtype=int)
self.color_paths = pick(self.color_paths, idcs)
self.pose_paths = pick(self.pose_paths, idcs)
self.param_paths = pick(self.param_paths, idcs)
def __init__(self,
data_dir,
num_pairs_per_scene=-1,
num_scenes=-1,
sidelength=None):
"""
Args:
data_dir (string): Path to directory containing all instances
num_pairs_per_scene: Number of image pairs per scene instance.
if -1 use all images for scene instance.
num_scenes: if -1 use all scene in dir
"""
self.sidelength = sidelength
self.instance_dirs = sorted(glob(os.path.join(data_dir, "*/")))
assert (len(self.instance_dirs) !=
0), "No scene instances in the data directory"
if num_scenes != -1:
self.instance_dirs = self.instance_dirs[:num_scenes]
self.rgb_paths = []
self.pose_paths = []
for instance_dir in self.instance_dirs:
color_dir = os.path.join(instance_dir, "rgb")
pose_dir = os.path.join(instance_dir, "pose")
if not os.path.isdir(color_dir) or not os.path.isdir(pose_dir):
print("Error! root dir %s is wrong" % data_dir)
return
scene_rgb_paths = sorted(data_util.glob_imgs(color_dir))
scene_pose_paths = sorted(glob(os.path.join(pose_dir, "*.txt")))
if num_pairs_per_scene == -1:
self.rgb_paths += scene_rgb_paths
self.pose_paths += scene_pose_paths
else:
self.rgb_paths += scene_rgb_paths[0:num_pairs_per_scene * 2]
self.pose_paths += scene_pose_paths[0:num_pairs_per_scene * 2]
def __init__(self,
root_dir,
calib_path,
calib_format,
img_size,
sampling_pattern,
load_img=True,
img_dir=None,
ignore_dist_coeffs=True,
load_precompute=False,
precomp_high_dir=None,
precomp_low_dir=None,
img_gamma=1.0):
super().__init__()
self.root_dir = root_dir
self.calib_format = calib_format
self.img_size = img_size
self.ignore_dist_coeffs = ignore_dist_coeffs
self.load_img = load_img
self.load_precompute = load_precompute
self.precomp_high_dir = precomp_high_dir
self.precomp_low_dir = precomp_low_dir
self.img_gamma = img_gamma
if not os.path.isdir(root_dir):
raise ValueError("Error! root dir is wrong")
self.img_dir = img_dir
if self.load_img and not os.path.isdir(self.img_dir):
raise ValueError("Error! image dir is wrong")
# load calibration data
if calib_format == 'convert':
if not os.path.isfile(calib_path):
raise ValueError("Error! calib path is wrong")
self.calib = scipy.io.loadmat(calib_path)
self.global_RT = self.calib['global_RT']
num_view = self.calib['poses'].shape[0]
else:
raise ValueError('Unknown calib format')
self.global_RT_inv = np.linalg.inv(self.global_RT)
# get path for all input images
if self.load_img:
self.img_fp_all = sorted(data_util.glob_imgs(self.img_dir))
else:
self.img_fp_all = ['x.x'] * num_view
# get intrinsic/extrinsic of all input images
self.poses_all = []
img_fp_all_new = []
for idx in range(len(self.img_fp_all)):
img_fn = os.path.split(self.img_fp_all[idx])[-1]
self.poses_all.append(self.calib['poses'][idx, :, :])
img_fp_all_new.append(self.img_fp_all[idx])
# remove views without calibration result
self.img_fp_all = img_fp_all_new
# Subsample data
keep_idx = []
if sampling_pattern == 'all':
keep_idx = list(range(len(self.img_fp_all)))
else:
if sampling_pattern == 'filter':
img_fp_all_new = []
poses_all_new = []
for idx in self.calib['keep_id'][0, :]:
img_fp_all_new.append(self.img_fp_all[idx])
poses_all_new.append(self.poses_all[idx])
keep_idx.append(idx)
self.img_fp_all = img_fp_all_new
self.poses_all = poses_all_new
elif sampling_pattern.split('_')[0] == 'first':
first_val = int(sampling_pattern.split('_')[-1])
self.img_fp_all = self.img_fp_all[:first_val]
self.poses_all = self.poses_all[:first_val]
keep_idx = list(range(first_val))
elif sampling_pattern.split('_')[0] == 'after':
after_val = int(sampling_pattern.split('_')[-1])
keep_idx = list(range(after_val, len(self.img_fp_all)))
self.img_fp_all = self.img_fp_all[after_val:]
self.poses_all = self.poses_all[after_val:]
elif sampling_pattern.split('_')[0] == 'skip':
skip_val = int(sampling_pattern.split('_')[-1])
img_fp_all_new = []
poses_all_new = []
for idx in range(0, len(self.img_fp_all), skip_val):
img_fp_all_new.append(self.img_fp_all[idx])
poses_all_new.append(self.poses_all[idx])
keep_idx.append(idx)
self.img_fp_all = img_fp_all_new
self.poses_all = poses_all_new
elif sampling_pattern.split('_')[0] == 'skipinv':
skip_val = int(sampling_pattern.split('_')[-1])
img_fp_all_new = []
poses_all_new = []
for idx in range(0, len(self.img_fp_all)):
if idx % skip_val == 0:
continue
img_fp_all_new.append(self.img_fp_all[idx])
poses_all_new.append(self.poses_all[idx])
keep_idx.append(idx)
self.img_fp_all = img_fp_all_new
self.poses_all = poses_all_new
elif sampling_pattern.split('_')[0] == 'only':
choose_idx = int(sampling_pattern.split('_')[-1])
self.img_fp_all = [self.img_fp_all[choose_idx]]
self.poses_all = [self.poses_all[choose_idx]]
keep_idx.append(choose_idx)
else:
raise ValueError("Unknown sampling pattern!")
if self.calib_format == 'convert':
self.calib['img_hws'] = self.calib['img_hws'][keep_idx, ...]
self.calib['projs'] = self.calib['projs'][keep_idx, ...]
self.calib['poses'] = self.calib['poses'][keep_idx, ...]
self.calib['dist_coeffs'] = self.calib['dist_coeffs'][keep_idx,
...]
# get mapping from img_fn to idx and vice versa
self.img_fn2idx = {}
self.img_idx2fn = []
for idx in range(len(self.img_fp_all)):
img_fn = os.path.split(self.img_fp_all[idx])[-1]
self.img_fn2idx[img_fn] = idx
self.img_idx2fn.append(img_fn)
print("*" * 100)
print("Sampling pattern ", sampling_pattern)
print("Image size ", self.img_size)
print("*" * 100)