def stack_sample(sample):
"""Stack a sample from multiple sensors"""
# If there is only one sensor don't do anything
if len(sample) == 1:
return sample[0]
# Otherwise, stack sample
stacked_sample = {}
for key in sample[0]:
# Global keys (do not stack)
if key in ['idx', 'dataset_idx', 'sensor_name', 'filename']:
stacked_sample[key] = sample[0][key]
else:
# Stack torch tensors
if is_tensor(sample[0][key]):
stacked_sample[key] = torch.stack([s[key] for s in sample], 0)
# Stack numpy arrays
elif is_numpy(sample[0][key]):
stacked_sample[key] = np.stack([s[key] for s in sample], 0)
# Stack list
elif is_list(sample[0][key]):
stacked_sample[key] = []
# Stack list of torch tensors
if is_tensor(sample[0][key][0]):
for i in range(len(sample[0][key])):
stacked_sample[key].append(
torch.stack([s[key][i] for s in sample], 0))
# Stack list of numpy arrays
if is_numpy(sample[0][key][0]):
for i in range(len(sample[0][key])):
stacked_sample[key].append(
np.stack([s[key][i] for s in sample], 0))
# Return stacked sample
return stacked_sample
def write_depth(filename, depth, intrinsics=None):
"""
Write a depth map to file, and optionally its corresponding intrinsics.
Parameters
----------
filename : str
File where depth map will be saved (.npz or .png)
depth : np.array [H,W]
Depth map
intrinsics : np.array [3,3]
Optional camera intrinsics matrix
"""
# If depth is a tensor
if is_tensor(depth):
depth = depth.detach().squeeze().cpu()
# If intrinsics is a tensor
if is_tensor(intrinsics):
intrinsics = intrinsics.detach().cpu()
# If we are saving as a .npz
if filename.endswith('.npz'):
np.savez_compressed(filename, depth=depth, intrinsics=intrinsics)
# If we are saving as a .png
elif filename.endswith('.png'):
depth = transforms.ToPILImage()((depth * 256).int())
depth.save(filename)
# Something is wrong
else:
raise NotImplementedError('Depth filename not valid.')
def viz_inv_depth(inv_depth, normalizer=None, percentile=95,
colormap='plasma', filter_zeros=False):
"""
Converts an inverse depth map to a colormap for visualization.
Parameters
----------
inv_depth : torch.Tensor [B,1,H,W]
Inverse depth map to be converted
normalizer : float
Value for inverse depth map normalization
percentile : float
Percentile value for automatic normalization
colormap : str
Colormap to be used
filter_zeros : bool
If True, do not consider zero values during normalization
Returns
-------
colormap : np.array [H,W,3]
Colormap generated from the inverse depth map
"""
# If a tensor is provided, convert to numpy
if is_tensor(inv_depth):
# Squeeze if depth channel exists
if len(inv_depth.shape) == 3:
inv_depth = inv_depth.squeeze(0)
inv_depth = inv_depth.detach().cpu().numpy()
cm = get_cmap(colormap)
if normalizer is None:
normalizer = np.percentile(
inv_depth[inv_depth > 0] if filter_zeros else inv_depth, percentile)
inv_depth /= (normalizer + 1e-6)
return cm(np.clip(inv_depth, 0., 1.0))[:, :, :3]
def stack_batch(batch):
"""
Stack multi-camera batches (B,N,C,H,W becomes BN,C,H,W)
Parameters
----------
batch : dict
Batch
Returns
-------
batch : dict
Stacked batch
"""
# If there is multi-camera information
if len(batch['rgb'].shape) == 5:
assert batch['rgb'].shape[
0] == 1, 'Only batch size 1 is supported for multi-cameras'
# Loop over all keys
for key in batch.keys():
# If list, stack every item
if is_list(batch[key]):
if is_tensor(batch[key][0]) or is_numpy(batch[key][0]):
batch[key] = [sample[0] for sample in batch[key]]
# Else, stack single item
else:
batch[key] = batch[key][0]
return batch
def write_depth(filename, depth, intrinsics=None):
"""
Write a depth map to file, and optionally its corresponding intrinsics.
This code is modified to export compatible-format depth image to openVSLAM
Parameters
----------
filename : str
File where depth map will be saved (.npz or .png)
depth : np.array [H,W]
Depth map
intrinsics : np.array [3,3]
Optional camera intrinsics matrix
"""
# If depth is a tensor
if is_tensor(depth):
# torch.cuda.synchronize()
# start_time = time.perf_counter()
depth = depth.detach().squeeze().cpu(
) # This is the bottle neck of 300mss
# torch.cuda.synchronize()
# end_time = time.perf_counter()
# print("The download from gpu ", end_time - start_time)
depth2 = depth.numpy()
depth2 = cv2.resize(src=depth2, dsize=(1226, 370))
# print(depth2.shape)
depth2 = np.clip(depth2, 0, 100)
# If intrinsics is a tensor
if is_tensor(intrinsics):
intrinsics = intrinsics.detach().cpu()
# If we are saving as a .npz
if filename.endswith('.npz'):
np.savez_compressed(filename, depth=depth, intrinsics=intrinsics)
# If we are saving as a .png
elif filename.endswith('.png'):
# depth = transforms.ToPILImage()((depth * 256).int())
# depth.save(filename)
depth2 = np.uint16(depth2 * 256)
cv2.imwrite(filename, depth2)
# Something is wrong
else:
raise NotImplementedError('Depth filename not valid.')
def stack_sample(sample):
"""Stack a sample from multiple sensors"""
# If there is only one sensor don't do anything
if len(sample) == 1:
return sample[0]
# Otherwise, stack sample
stacked_sample = {}
for key in sample[0]:
# Global keys (do not stack)
if key in ['idx', 'dataset_idx', 'sensor_name', 'filename']:
stacked_sample[key] = sample[0][key]
else:
# Stack torch tensors
if is_tensor(sample[0][key]):
stacked_sample[key] = torch.cat(
[s[key].unsqueeze(0) for s in sample], 0)
# Return stacked sample
return stacked_sample
def write_depth(self, depth):
"""
Write a depth map to file, and optionally its corresponding intrinsics.
This code is modified to export compatible-format depth image to openVSLAM
Parameters
----------
depth : np.array [H,W]
Depth map
"""
# If depth is a tensor
if is_tensor(depth):
depth = depth.detach().squeeze().cpu()
depth = np.clip(depth, 0, 100)
# make depth image to 16 bit format following TUM RGBD dataset format
# it is also ROS standard(?)
depth = np.uint16(depth * 256)
return depth
def prep_image(prefix, key, image):
"""
Prepare image for wandb logging
Parameters
----------
prefix : str
Prefix added to the key for logging
key : str
Key from data containing the inverse depth map
image : torch.Tensor [3,H,W]
Image to be logged
Returns
-------
output : dict
Dictionary with key and value for logging
"""
if is_tensor(image):
image = image.detach().permute(1, 2, 0).cpu().numpy()
prefix_key = '{}-{}'.format(prefix, key)
return {prefix_key: wandb.Image(image, caption=key)}
