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)