Add save function

poses/pose_utils.py

@@ -6,9 +6,9 @@ import numpy as np
 
 from colmap_wrapper import run_colmap
 import colmap_read_model as read_model
-#from nerf_homemade.poses.colmap_wrapper import run_colmap
-#import nerf_homemade.poses.colmap_read_model as read_model
 
+FORMAT = "%(asctime)s %(levelname)s \t %(message)s"
+logging.basicConfig(format=FORMAT, level=logging.DEBUG)
 
 def gen_poses(basedir: str, match_type: str='exhaustive') -> None:
     """
@@ -135,4 +135,36 @@ def save_poses(basedir, poses, pts3d, perm) -> None:
         The sorted index of the array
     """
 
-    # TODO
+    pts_arr = []
+    vis_arr = []
+    for k in pts3d:
+        pts_arr.append(pts3d[k].xyz)
+        cams = [0] * poses.shape[-1]
+        for ind in pts3d[k].image_ids:
+            if len(cams) < ind - 1:
+                logging.error("The correct camera poses for current points cannot be accessed")
+                return
+            cams[ind - 1] = 1
+        vis_arr.append(cams)
+
+    pts_arr = np.array(pts_arr)
+    vis_arr = np.array(vis_arr)
+    logging.info(f"Points {pts_arr.shape} Visibility {vis_arr.shape}")
+    
+    zvals = np.sum(-(pts_arr[:, np.newaxis, :].transpose([2,0,1]) - poses[:3, 3:4, :]) * poses[:3, 2:3, :], 0)
+    valid_z = zvals[vis_arr == 1]
+    logging.info(f"Depths stats - min: {valid_z.min()} max: {valid_z.max()}
+    mean: {valid_z.mean()}")
+    
+    save_arr = []
+    for i in perm:
+        vis = vis_arr[:, i]
+        zs = zvals[:, i]
+        zs = zs[vis==1]
+        close_depth, inf_depth = np.percentile(zs, .1), np.percentile(zs, 99.9)
+        
+        save_arr.append(np.concatenate([poses[..., i].ravel(), np.array([close_depth, inf_depth])], 0))
+    save_arr = np.array(save_arr)
+    
+    save_path = os.path.join(basedir, "poses_bounds.npy")
+    np.save(save_path, save_arr)