def read_png_depth(file):
"""Reads a .png depth map."""
depth_png = np.array(load_image(file), dtype=int)
assert (np.max(depth_png) > 255), 'Wrong .png depth file'
depth = depth_png.astype(np.float) / 256.
depth[depth_png == 0] = -1.
return np.expand_dims(depth, axis=2)
def load_depth(file):
"""
Load a depth map from file
Parameters
----------
file : str
Depth map filename (.npz or .png)
Returns
-------
depth : np.array [H,W]
Depth map (invalid pixels are 0)
"""
if file.endswith('npz'):
return np.load(file)['depth']
elif file.endswith('png'):
depth_png = np.array(load_image(file), dtype=int)
assert (np.max(depth_png) > 255), 'Wrong .png depth file'
return depth_png.astype(np.float) / 256.
else:
raise NotImplementedError('Depth extension not supported.')
def __getitem__(self, idx):
"""Get dataset sample given an index."""
# Add image information
sample = {
'idx': idx,
'filename': '%s_%010d' % (self.split, idx),
'rgb': load_image(self.paths[idx]),
}
# Add intrinsics
parent_folder = self._get_parent_folder(self.paths[idx])
if parent_folder in self.calibration_cache:
c_data = self.calibration_cache[parent_folder]
else:
c_data = self._read_raw_calib_file(parent_folder)
self.calibration_cache[parent_folder] = c_data
sample.update({
'intrinsics':
self._get_intrinsics(self.paths[idx], c_data),
})
# Add pose information if requested
if self.with_pose:
sample.update({
'pose': self._get_pose(self.paths[idx]),
})
# Add depth information if requested
if self.with_depth:
sample.update({
'depth':
self._read_depth(self._get_depth_file(self.paths[idx])),
})
# Add context information if requested
if self.with_context:
# Add context images
all_context_idxs = self.backward_context_paths[idx] + \
self.forward_context_paths[idx]
image_context_paths, _ = \
self._get_context_files(self.paths[idx], all_context_idxs)
image_context = [load_image(f) for f in image_context_paths]
sample.update({'rgb_context': image_context})
# Add context poses
if self.with_pose:
first_pose = sample['pose']
image_context_pose = [
self._get_pose(f) for f in image_context_paths
]
image_context_pose = [
invert_pose_numpy(context_pose) @ first_pose
for context_pose in image_context_pose
]
sample.update({'pose_context': image_context_pose})
# Apply transformations
if self.data_transform:
sample = self.data_transform(sample)
# Return sample
return sample
def _read_rgb_context_files(self, session, filename):
context_paths = self._get_context_file_paths(filename)
return [
load_image(os.path.join(self.root_dir, session, filename))
for filename in context_paths
]
def _read_rgb_file(self, session, filename):
return load_image(os.path.join(self.root_dir, session, filename))