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)}