fds_process_custom_data.py

import numpy as np
import pycolmap

import os
import sys
import shutil
import argparse
sys.path.append('.')

from util.read_write_fused_vis import read_fused, write_fused, MeshPoint

### This script is used to convert a COLMAP reconstruction to the Fast Dipole Sums format ###
### It has been adapted from the original script in the Fast Dipole Sums repository to accept input arguments ###
### The original script can be found at "misc/process_custom_data.py" in the Fast Dipole Sums repository ###
### https://github.com/cmu-ci-lab/fast_dipole_sums ###

def recover_P(intrinsics, pose):
    # Extract R and t from the pose matrix
    K = intrinsics[:3, :3]
    R = pose[:3, :3].transpose()  # Transpose back to get original R
    t = pose[:3, 3]

    P = np.eye(4)
    P[:3, :3] = K @ R
    P[:3, 3] = -K @ R @ t
    
    return P

def convert(colmap_dir, out_dir, bs_radius, bs_center):
    data_out = dict()

    # changes these manually so that scene is centered at origin
    # preferably bounded within a unit sphere
    scale = bs_radius
    dx = bs_center[0]
    dy = bs_center[1]
    dz = bs_center[2]

    S = np.eye(4)
    S[:3, :3] *= scale
    S[0, 3] = dx
    S[1, 3] = dy
    S[2, 3] = dz

    image_dir = f'{colmap_dir}/dense/images'
    reconstruction = pycolmap.Reconstruction(f'{colmap_dir}/dense/sparse')

    os.makedirs(f'{out_dir}/image', exist_ok=True)

    image_idx_map = {}
    for idx, image in reconstruction.images.items():
        image_idx_map[image.name] = idx

    for i, (im, idx) in enumerate(image_idx_map.items()):
        image = reconstruction.images[idx]
        camera = reconstruction.cameras[image.camera_id]

        # assume pinhole camera for now
        focal_x = camera.params[0]
        focal_y = camera.params[1]
        cx = camera.params[2]
        cy = camera.params[3]

        # simple radial model (colmap default)
        # focal_x = focal_y = camera.params[0]
        # cx = camera.params[1]
        # cy = camera.params[2]

        K = np.array([[focal_x, 0, cx],
                      [0, focal_y, cy],
                      [0, 0, 1]])

        pose = image.cam_from_world().inverse().matrix()
        P = recover_P(K, pose)

        data_out[f'world_mat_{i}'] = P
        data_out[f'scale_mat_{i}'] = S

        suffix = im.split('.')[-1]
        shutil.copy(f'{image_dir}/{im}', f'{out_dir}/image/{i:04}.{suffix}')

    np.savez(f'{out_dir}/cameras_sphere.npz', **data_out)

    mesh_points = read_fused(f'{colmap_dir}/dense/fused.ply', f'{colmap_dir}/dense/fused.ply.vis')
    points = np.array([m.position for m in mesh_points], dtype=np.float32)
    normals = np.array([m.normal for m in mesh_points], dtype=np.float32)
    colors = np.array([m.color for m in mesh_points], dtype=np.float32) / 255

    points = (points - np.array([dx, dy, dz])) / scale
    mesh_points = [MeshPoint(position=points[i], color=colors[i], normal=normals[i], num_visible_images=0, visible_image_idxs=[])
                   for i in range(len(points))]
    write_fused(mesh_points, f'{out_dir}/points.ply', f'{out_dir}/points.ply.vis')

def main():
    parser = argparse.ArgumentParser(description='Convert COLMAP reconstruction to Fast Dipole Sums format')
    
    parser.add_argument('--colmap_dir', type=str, required=True, help='Path to COLMAP reconstruction directory')
    parser.add_argument('--out_dir', type=str, required=True, help='Output path')
    parser.add_argument('--bs_center', nargs = 3, type = float, required = True, help = 'Bounding sphere center as "x y z"')
    parser.add_argument('--bs_radius', type=float, required=True, help='Bounding sphere radius')
    args = parser.parse_args()
    
    print(f"Converting COLMAP data from {args.colmap_dir} to {args.out_dir}")
    convert(args.colmap_dir, args.out_dir, args.bs_radius, args.bs_center)
    print("Conversion complete!")

if __name__ == '__main__':
    main()