support saving disparity prediction as .pfm file

Author: haofeixu

evaluate_depth.py

@@ -10,7 +10,7 @@
 from dataloader.depth.datasets import ScannetDataset, DemonDataset
 from loss.depth_loss import compute_errors
 from utils.utils import InputPadder
-from utils.visualization import viz_depth_tensor_from_monodepth2
+from utils.visualization import viz_depth_tensor
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
@@ -132,9 +132,9 @@ def validate_scannet(model,
 
         if save_vis_depth:
             filename = os.path.join(save_dir, '%04d_depth_pred.png' % valid_samples)
-            viz_inv_depth = viz_depth_tensor_from_monodepth2(1. / pred_depth.cpu(),
-                                                             return_numpy=True,
-                                                             colormap='plasma')  # [H, W, 3] uint8
+            viz_inv_depth = viz_depth_tensor(1. / pred_depth.cpu(),
+                                             return_numpy=True,
+                                             colormap='plasma')  # [H, W, 3] uint8
             Image.fromarray(viz_inv_depth).save(filename)
 
         gt_depth = gt_depth.cpu().numpy()
@@ -271,9 +271,9 @@ def validate_demon(model,
 
         if save_vis_depth:
             filename = os.path.join(save_dir, '%04d.png' % valid_samples)
-            viz_inv_depth = viz_depth_tensor_from_monodepth2(1. / pred_depth.cpu(),
-                                                             return_numpy=True,
-                                                             colormap='plasma')  # [H, W, 3] uint8
+            viz_inv_depth = viz_depth_tensor(1. / pred_depth.cpu(),
+                                             return_numpy=True,
+                                             colormap='plasma')  # [H, W, 3] uint8
             Image.fromarray(viz_inv_depth).save(filename)
 
         gt_depth = gt_depth.cpu().numpy()
@@ -402,8 +402,8 @@ def inference_depth(model,
         pr_depth = pred_depth[0]
 
         filename = os.path.join(output_path, os.path.basename(imgs[i])[:-4] + '.png')
-        viz_inv_depth = viz_depth_tensor_from_monodepth2(1. / pr_depth.cpu(),
-                                                         return_numpy=True)  # [H, W, 3] uint8
+        viz_inv_depth = viz_depth_tensor(1. / pr_depth.cpu(),
+                                         return_numpy=True)  # [H, W, 3] uint8
         Image.fromarray(viz_inv_depth).save(filename)
 
         if pred_bidir_depth:
@@ -412,8 +412,8 @@ def inference_depth(model,
             pr_depth_bwd = pred_depth[1]
 
             filename = os.path.join(output_path, os.path.basename(imgs[i])[:-4] + '_bwd.png')
-            viz_inv_depth = viz_depth_tensor_from_monodepth2(1. / pr_depth_bwd.cpu(),
-                                                             return_numpy=True)  # [H, W, 3] uint8
+            viz_inv_depth = viz_depth_tensor(1. / pr_depth_bwd.cpu(),
+                                             return_numpy=True)  # [H, W, 3] uint8
             Image.fromarray(viz_inv_depth).save(filename)
 
     print('Done!')

evaluate_flow.py

@@ -798,7 +798,6 @@ def inference_flow(model,
                 output_file_bwd = os.path.join(output_path, os.path.basename(filenames[test_id])[:-4] + '_pred_bwd.flo')
                 frame_utils.writeFlow(output_file_bwd, flow_bwd)
 
-
     if save_video:
         suffix = '_flow_img.mp4' if concat_flow_img else '_flow.mp4'
         output_file = os.path.join(output_path, os.path.basename(inference_video)[:-4] + suffix)

evaluate_stereo.py

@@ -723,6 +723,7 @@ def inference_stereo(model,
                      num_reg_refine=1,
                      pred_bidir_disp=False,
                      pred_right_disp=False,
+                     save_pfm_disp=False,
                      ):
     model.eval()
 
@@ -818,6 +819,10 @@ def inference_stereo(model,
 
         disp = pred_disp[0].cpu().numpy()
 
+        if save_pfm_disp:
+            save_name_pfm = save_name[:-4] + '.pfm'
+            write_pfm(save_name_pfm, disp)
+
         disp = vis_disparity(disp)
         cv2.imwrite(save_name, disp)
 
@@ -828,6 +833,10 @@ def inference_stereo(model,
             # flip back
             disp = hflip(pred_disp[1]).cpu().numpy()
 
+            if save_pfm_disp:
+                save_name_pfm = save_name[:-4] + '.pfm'
+                write_pfm(save_name_pfm, disp)
+
             disp = vis_disparity(disp)
             cv2.imwrite(save_name, disp)
 

main_stereo.py

@@ -114,6 +114,8 @@ def get_args_parser():
                         help='predict both left and right disparities')
     parser.add_argument('--pred_right_disp', action='store_true',
                         help='predict right disparity')
+    parser.add_argument('--save_pfm_disp', action='store_true',
+                        help='save predicted disparity as .pfm format')
 
     parser.add_argument('--debug', action='store_true')
 
@@ -340,6 +342,7 @@ def main(args):
                          num_reg_refine=args.num_reg_refine,
                          pred_bidir_disp=args.pred_bidir_disp,
                          pred_right_disp=args.pred_right_disp,
+                         save_pfm_disp=args.save_pfm_disp,
                          )
 
         return

utils/visualization.py

@@ -89,7 +89,7 @@ def tensor2numpy(var_dict):
     return var_dict
 
 
-def viz_depth_tensor_from_monodepth2(disp, return_numpy=False, colormap='plasma'):
+def viz_depth_tensor(disp, return_numpy=False, colormap='plasma'):
     # visualize inverse depth
     assert isinstance(disp, torch.Tensor)