You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
139 lines
4.2 KiB
Python
139 lines
4.2 KiB
Python
import os
|
|
import logging
|
|
import typing
|
|
|
|
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
|
|
|
|
|
|
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
|
|
"""
|
|
|
|
# TODO
|