From f08026d13318b05cc15186ea5e25779b35db4d6a Mon Sep 17 00:00:00 2001
From: star <1170771394@qq.com>
Date: Thu, 2 Dec 2021 17:07:57 +0800
Subject: [PATCH] Update pose_utils.py

fix the error when existing invalid views in the scene.
Beside the pose_boundary.npy will be generated, it also will generate "view_imgs.txt" to mark which views are used.
---
 llff/poses/pose_utils.py | 103 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 96 insertions(+), 7 deletions(-)

diff --git a/llff/poses/pose_utils.py b/llff/poses/pose_utils.py
index 3cc4eb5..9e7de61 100644
--- a/llff/poses/pose_utils.py
+++ b/llff/poses/pose_utils.py
@@ -7,6 +7,96 @@
 from llff.poses.colmap_wrapper import run_colmap
 import llff.poses.colmap_read_model as read_model
 
+def save_views(realdir,names):
+    with open(os.path.join(realdir,'view_imgs.txt'), mode='w') as f:
+        f.writelines('\n'.join(names))
+    f.close()
+
+
+def load_save_pose(realdir):
+
+    # load colmap data 
+    camerasfile = os.path.join(realdir, 'sparse/0/cameras.bin')
+    camdata = read_model.read_cameras_binary(camerasfile)
+    
+    list_of_keys = list(camdata.keys())
+    cam = camdata[list_of_keys[0]]
+    print( 'Cameras', cam)
+
+    h, w, f = cam.height, cam.width, cam.params[0]
+    hwf = np.array([h,w,f]).reshape([3,1])
+    
+    imagesfile = os.path.join(realdir, 'sparse/0/images.bin')
+    imdata = read_model.read_images_binary(imagesfile)
+    
+    real_ids = [k for k in imdata]
+
+    w2c_mats = []
+    bottom = np.array([0,0,0,1.]).reshape([1,4])
+    
+    names = [imdata[k].name for k in imdata]
+    print( 'Images #', len(names))
+
+    # if (len(names)< 32):
+    #     raise ValueError(f'{realdir} only {len(names)} images register, need Re-run colmap or reset the threshold')
+
+
+    perm = np.argsort(names)
+    sort_names = [names[i] for i in perm]
+    save_views(realdir,sort_names)
+
+    for k in imdata:
+        im = imdata[k]
+        R = im.qvec2rotmat()
+        t = im.tvec.reshape([3,1])
+        m = np.concatenate([np.concatenate([R, t], 1), bottom], 0)
+        w2c_mats.append(m)
+    
+    w2c_mats = np.stack(w2c_mats, 0)
+    c2w_mats = np.linalg.inv(w2c_mats)
+    
+    poses = c2w_mats[:, :3, :4].transpose([1,2,0])
+    poses = np.concatenate([poses, np.tile(hwf[..., np.newaxis], [1,1,poses.shape[-1]])], 1)
+    
+    points3dfile = os.path.join(realdir, 'sparse/0/points3D.bin')
+    pts3d = read_model.read_points3d_binary(points3dfile)
+    
+    # must switch to [-u, r, -t] from [r, -u, t], NOT [r, u, -t]
+    poses = np.concatenate([poses[:, 1:2, :], poses[:, 0:1, :], -poses[:, 2:3, :], poses[:, 3:4, :], poses[:, 4:5, :]], 1)
+
+    # save pose
+    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) < real_ids.index(ind):
+                print('ERROR: the correct camera poses for current points cannot be accessed')
+                return
+            cams[real_ids.index(ind)] = 1
+        vis_arr.append(cams)
+
+    pts_arr = np.array(pts_arr)
+    vis_arr = np.array(vis_arr)
+    print( '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]
+    print( 'Depth stats', valid_z.min(), valid_z.max(), 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)
+    
+    np.save(os.path.join(realdir, 'poses_bounds.npy'), save_arr)
+    
+
 
 def load_colmap_data(realdir):
     
@@ -49,7 +139,6 @@ def load_colmap_data(realdir):
     
     # must switch to [-u, r, -t] from [r, -u, t], NOT [r, u, -t]
     poses = np.concatenate([poses[:, 1:2, :], poses[:, 0:1, :], -poses[:, 2:3, :], poses[:, 3:4, :], poses[:, 4:5, :]], 1)
-    
     return poses, pts3d, perm
 
 
@@ -68,8 +157,7 @@ def save_poses(basedir, poses, pts3d, perm):
 
     pts_arr = np.array(pts_arr)
     vis_arr = np.array(vis_arr)
-    print( 'Points', pts_arr.shape, 'Visibility', vis_arr.shape )
-    
+    print( '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]
     print( 'Depth stats', valid_z.min(), valid_z.max(), valid_z.mean() )
@@ -268,12 +356,13 @@ def gen_poses(basedir, match_type, factors=None):
         run_colmap(basedir, match_type)
     else:
         print('Don\'t need to run COLMAP')
-        
     print( 'Post-colmap')
+
+    load_save_pose(basedir)
+
+    # poses, pts3d, perm = load_colmap_data(basedir)
     
-    poses, pts3d, perm = load_colmap_data(basedir)
-    
-    save_poses(basedir, poses, pts3d, perm)
+    # save_poses(basedir, poses, pts3d, perm)
     
     if factors is not None:
         print( 'Factors:', factors)