def generate_depth_map(self, sample_idx, datum_idx, filename):
# Generate depth filename
filename = '{}/{}.npz'.format(
os.path.dirname(self.path), filename.format('depth/{}'.format(self.depth_type)))
# Load and return if exists
if os.path.exists(filename):
return np.load(filename)['depth']
# Otherwise, create, save and return
else:
# Get pointcloud
scene_idx, sample_idx_in_scene, _ = self.dataset.dataset_item_index[sample_idx]
pc_datum_idx_in_sample = self.dataset.get_datum_index_for_datum_name(
scene_idx, sample_idx_in_scene, self.depth_type)
pc_datum_data = self.dataset.get_point_cloud_from_datum(
scene_idx, sample_idx_in_scene, pc_datum_idx_in_sample)
# Create camera
camera_rgb = self.get_current('rgb', datum_idx)
camera_pose = self.get_current('pose', datum_idx)
camera_intrinsics = self.get_current('intrinsics', datum_idx)
camera = Camera(K=camera_intrinsics, p_cw=camera_pose.inverse())
# Generate depth map
world_points = pc_datum_data['pose'] * pc_datum_data['point_cloud']
depth = generate_depth_map(camera, world_points, camera_rgb.size[::-1])
# Save depth map
os.makedirs(os.path.dirname(filename), exist_ok=True)
np.savez_compressed(filename, depth=depth)
# Return depth map
return depth
def get_depth_from_point_cloud(dataset, scene_idx, sample_idx_in_scene,
cam_datum_idx_in_sample,
pc_datum_idx_in_sample):
"""Generate the depth map in the camera view using the provided point cloud
datum within the sample.
Parameters
----------
dataset: dgp.dataset.BaseDataset
Inherited base dataset to augment with depth data.
scene_idx: int
Index of the scene.
sample_idx_in_scene: int
Index of the sample within the scene at scene_idx.
cam_datum_idx_in_sample: int
Index of the camera datum within the sample.
pc_datum_idx_in_sample: int
Index of the point cloud datum within the sample.
Returns
-------
depth: np.ndarray
Depth map from the camera's viewpoint.
"""
# Get point cloud datum and load it
pc_datum = dataset.get_datum(scene_idx, sample_idx_in_scene,
pc_datum_idx_in_sample)
pc_datum_type = pc_datum.datum.WhichOneof('datum_oneof')
assert pc_datum_type == 'point_cloud', 'Depth cannot be generated from {} '.format(
pc_datum_type)
pc_datum_data = dataset.get_point_cloud_from_datum(scene_idx,
sample_idx_in_scene,
pc_datum_idx_in_sample)
X_W = pc_datum_data['pose'] * pc_datum_data['point_cloud']
# Get target camera datum for projection
cam_datum = dataset.get_datum(scene_idx, sample_idx_in_scene,
cam_datum_idx_in_sample)
cam_datum_type = cam_datum.datum.WhichOneof('datum_oneof')
assert cam_datum_type == 'image', 'Depth cannot be projected onto {} '.format(
cam_datum_type)
cam_datum_data = dataset.get_image_from_datum(scene_idx,
sample_idx_in_scene,
cam_datum_idx_in_sample)
p_WC = cam_datum_data['pose']
camera = Camera(K=cam_datum_data['intrinsics'], p_cw=p_WC.inverse())
(W, H) = cam_datum_data['rgb'].size[:2]
return generate_depth_map(camera, X_W, (H, W))
def generate_depth_map(self, sample_idx, datum_idx, filename):
"""
Generates the depth map for a camera by projecting LiDAR information.
It also caches the depth map following DGP folder structure, so it's not recalculated
Parameters
----------
sample_idx : int
sample index
datum_idx : int
Datum index
filename :
Filename used for loading / saving
Returns
-------
depth : np.array [H, W]
Depth map for that datum in that sample
"""
# Generate depth filename
filename = '{}/{}.npz'.format(
os.path.dirname(self.path),
filename.format('depth/{}'.format(self.depth_type)))
# Load and return if exists
if os.path.exists(filename):
return np.load(filename, allow_pickle=True)['depth']
# Otherwise, create, save and return
else:
# Get pointcloud
scene_idx, sample_idx_in_scene, _ = self.dataset.dataset_item_index[
sample_idx]
pc_datum_idx_in_sample = self.dataset.get_datum_index_for_datum_name(
scene_idx, sample_idx_in_scene, self.depth_type)
pc_datum_data = self.dataset.get_point_cloud_from_datum(
scene_idx, sample_idx_in_scene, pc_datum_idx_in_sample)
# Create camera
camera_rgb = self.get_current('rgb', datum_idx)
camera_pose = self.get_current('pose', datum_idx)
camera_intrinsics = self.get_current('intrinsics', datum_idx)
camera = Camera(K=camera_intrinsics, p_cw=camera_pose.inverse())
# Generate depth map
world_points = pc_datum_data['pose'] * pc_datum_data['point_cloud']
depth = generate_depth_map(camera, world_points,
camera_rgb.size[::-1])
# Save depth map
os.makedirs(os.path.dirname(filename), exist_ok=True)
np.savez_compressed(filename, depth=depth)
# Return depth map
return depth