nerf_homemade/poses/pose_utils.py

171 lines
5.4 KiB
Python
Raw Normal View History

import os
import logging
import typing
import numpy as np
from colmap_wrapper import run_colmap
import colmap_read_model as read_model
2021-03-23 15:19:50 +01:00
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:
"""
Geneate or retreive camera poses.
Retrieve the cameras either by generating using COLMAP or reading them if
they already have been written.
Parameters
----------
basedir: str
The path of the directory either containing raw images or the
arborescence of a COLMAP directory if camera poses have already been
computed
match_type: {'exhaustive', 'sequential'}
The type of match COLMAP should perform if it needs to be run.
Defaults to 'exhaustive'.
"""
files_needed = ["cameras.bin", "images.bin", "points3D.bin"]
path_to_sparse = os.path.join(basedir, "sparse/0")
if os.path.exists(path_to_sparse):
existing_files = os.listdir(path_to_sparse)
else:
existing_files = []
if not all([f in existing_files for f in files_needed]):
logging.info("Running COLMAP")
run_colmap(basedir, match_type)
else:
logging.info("Files genreated by COLMAP found. Skipping running COLMAP.")
logging.debug("Loading COLMAP data")
poses, points_3d, perm = load_colmap_data(basedir)
logging.debug("Saving COLMAP data to npy")
save_poses(basedir, poses, points_3d, perm)
def load_colmap_data(basedir: str) -> (np.ndarray, dict, np.ndarray):
"""
Load data from a COLMAP arborescence.
Parameters
----------
basedir: str
The path of the directory which contains the COLMAP arborescence
Returns
-------
poses: numpy.ndarray
List of poses for each image.
pts3d:
List of 3D points for each image.
perm:
Index list ordered by name of images.
Should be [1, 2, ..., N] if well ordered.
"""
# append prefix
basedir = os.path.join(basedir, "sparse/0/")
# read cameras data
cameras_file = os.path.join(basedir, "cameras.bin")
cameras_data = read_model.read_cameras_binary(cameras_file)
logging.debug(f"Loading camera model from {cameras_file}")
# extract intrinsic values for the camera
# asumption is made that it is unique
camera_h = cameras_data[1].height
camera_w = cameras_data[1].width
camera_f = cameras_data[1].params[0]
hwf = np.array([camera_h, camera_w, camera_f]).reshape([3,1])
logging.debug(f"Number of cameras: {len(cameras_data)}")
# read images data
images_file = os.path.join(basedir, "images.bin")
images_data = read_model.read_images_binary(images_file)
w2c_mats = []
bottom = np.array([0., 0., 0., 1.]).reshape([1, 4])
# sort by name
names = [images_data[k].name for k in images_data]
perm = np.argsort(names)
# equivalent to a range since dict is index by number from 1 to N
# create camera matrix
for k in images_data:
im = images_data[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)
# read 3d points data
pts3d_file = os.path.join(basedir, "points3D.bin")
pts3d = read_model.read_points3d_binary(pts3d_file)
# 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
def save_poses(basedir, poses, pts3d, perm) -> None:
"""
Save the COLMAP data in a `.npy` format.
Parameters
----------
basedir: str
The path of the directory in which to save the data.
Data will be saved in `basedir/poses_bounds.npy`.
poses:
The list of poses
pts3d:
The list of 3d points
perm:
The sorted index of the array
"""
2021-03-23 15:19:50 +01:00
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)