nerf_homemade/poses/pose_utils.py

import os
import logging
import typing

import numpy as np

from colmap_wrapper import run_colmap
import 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:
    """
    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
    """

    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)