def infer_and_save_depth(input_file, output_file, model_wrapper, image_shape, half, save): """ Process a single input file to produce and save visualization Parameters ---------- input_file : str Image file output_file : str Output file, or folder where the output will be saved model_wrapper : nn.Module Model wrapper used for inference image_shape : Image shape Input image shape half: bool use half precision (fp16) save: str Save format (npz or png) """ if not is_image(output_file): # If not an image, assume it's a folder and append the input name os.makedirs(output_file, exist_ok=True) output_file = os.path.join(output_file, os.path.basename(input_file)) # change to half precision for evaluation if requested dtype = torch.float16 if half else None # Load image image = load_image(input_file) # Resize and to tensor image = resize_image(image, image_shape) image = to_tensor(image).unsqueeze(0) # Send image to GPU if available if torch.cuda.is_available(): image = image.to('cuda:{}'.format(rank()), dtype=dtype) # Depth inference (returns predicted inverse depth) pred_inv_depth = model_wrapper.depth(image)[0] if save == 'npz' or save == 'png': # Get depth from predicted depth map and save to different formats filename = '{}.{}'.format(os.path.splitext(output_file)[0], save) print('Saving {} to {}'.format( pcolor(input_file, 'cyan', attrs=['bold']), pcolor(filename, 'magenta', attrs=['bold']))) write_depth(filename, depth=inv2depth(pred_inv_depth)) else: # Prepare RGB image rgb = image[0].permute(1, 2, 0).detach().cpu().numpy() * 255 # Prepare inverse depth viz_pred_inv_depth = viz_inv_depth(pred_inv_depth[0]) * 255 # Concatenate both vertically image = np.concatenate([rgb, viz_pred_inv_depth], 0) # Save visualization print('Saving {} to {}'.format( pcolor(input_file, 'cyan', attrs=['bold']), pcolor(output_file, 'magenta', attrs=['bold']))) imwrite(output_file, image[:, :, ::-1])
def save_depth(batch, output, args, dataset, save): """ Save depth predictions in various ways Parameters ---------- batch : dict Batch from dataloader output : dict Output from model args : tuple Step arguments dataset : CfgNode Dataset configuration save : CfgNode Save configuration """ # If there is no save folder, don't save if save.folder is '': return # If we want to save if save.depth.rgb or save.depth.viz or save.depth.npz or save.depth.png: # Retrieve useful tensors rgb = batch['rgb'] pred_inv_depth = output['inv_depth'] # Prepare path strings print("save:",batch['sensor_name']) # print("save:",batch['filename']) filename = [os.path.split(batch['filename'][i])[-1] for i in range(len(batch['filename']))] dataset_idx = 0 if len(args) == 1 else args[1] save_path = [os.path.join(save.folder, 'depth', prepare_dataset_prefix(dataset, dataset_idx), os.path.basename(save.pretrained).split('.')[0] ,batch['sensor_name'][i]) for i in range(len(batch['sensor_name']))] # Create folder for i in range(len(save_path)): os.makedirs(save_path[i], exist_ok=True) # For each image in the batch length = rgb.shape[0] for i in range(length): # Save numpy depth maps if save.depth.npz: write_depth('{}/{}_depth.npz'.format(save_path[i], filename[i]), depth=inv2depth(pred_inv_depth[i]), intrinsics=batch['intrinsics'][i] if 'intrinsics' in batch else None) # Save png depth maps if save.depth.png: write_depth('{}/{}_depth.png'.format(save_path[i], filename[i]), depth=inv2depth(pred_inv_depth[i])) # Save rgb images if save.depth.rgb: rgb_i = rgb[i].permute(1, 2, 0).detach().cpu().numpy() * 255 write_image('{}/{}_rgb.png'.format(save_path[i], filename[i]), rgb_i) # Save inverse depth visualizations if save.depth.viz: viz_i = viz_inv_depth(pred_inv_depth[i]) * 255 write_image('{}/{}_viz.png'.format(save_path[i], filename[i]), viz_i)