implement small colmap interop crate

[edgarriba]

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::path::Path;

#[derive(Debug, thiserror::Error)]
pub enum ColmapError {
    #[error("error reading or writing file")]
    IoError(#[from] std::io::Error),

    #[error("Invalid number of camera parameters")]
    InvalidNumCameraParams(usize),
}

/// Represents a 2D vector.
///
/// # Attributes
///
/// * `x` - x coordinate.
/// * `y` - y coordinate.
///
#[derive(Clone)]
pub struct Vector2d {
    pub x: f64,
    pub y: f64,
}

/// Represents a Colmap camera model id.
///
/// # Variants
///
/// * `CameraModelInvalid` - Invalid camera model.
/// * `CameraModelSimplePinhole` - Simple pinhole camera model.
/// * `CameraModelPinhole` - Pinhole camera model.
/// * `CameraModelSimplifiedRadial` - Simplified radial camera model.
/// * `CameraModelRadial` - Radial camera model.
/// * `CameraModelOpenCV` - OpenCV camera model.
/// * `CameraModelOpenCVFisheye` - OpenCV fisheye camera model.
/// * `CameraModelFullOpenCV` - Full OpenCV camera model.
/// * `CameraModelFOV` - Field of view camera model.
/// * `CameraModelSimpleRadialFisheye` - Simple radial fisheye camera model.
/// * `CameraModelRadialFisheye` - Radial fisheye camera model.
/// * `CameraModelThinPrismFisheye` - Thin prism fisheye camera model.
/// * `CameraModelCount` - Number of camera models.
///
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum CameraModelId {
    CameraModelInvalid = -1,
    CameraModelSimplePinhole = 0,
    CameraModelPinhole = 1,
    CameraModelSimplifiedRadial = 2,
    CameraModelRadial = 3,
    CameraModelOpenCV = 4,
    CameraModelOpenCVFisheye = 5,
    CameraModelFullOpenCV = 6,
    CameraModelFOV = 7,
    CameraModelSimpleRadialFisheye = 8,
    CameraModelRadialFisheye = 9,
    CameraModelThinPrismFisheye = 10,
    CameraModelCount = 11,
}

/// Represents a camera in the Colmap system.
///
/// # Attributes
///
/// * `camera_id` - Unique camera id.
/// * `model_id` - Camera model id.
/// * `width` - Image width.
/// * `height` - Image height.
/// * `params` - Camera parameters.
///
pub struct ColmapCamera {
    pub camera_id: u32,
    pub model_id: CameraModelId,
    pub width: usize,
    pub height: usize,
    pub params: Vec<f64>,
}

/// Represents an image in the Colmap system.
///
/// # Attributes
///
/// * `name` - Image name.
/// * `image_id` - Unique image id.
/// * `camera_id` - Camera id.
/// * `rotation` - Rotation quaternion (qw, qx, qy, qz).
/// * `translation` - Translation vector (x, y, z).
/// * `points2d` - 2D points.
/// * `point3d_ids` - 3D point ids.
///
pub struct ColmapImage {
    pub name: String,
    pub image_id: u32,
    pub camera_id: u32,
    pub rotation: [f64; 4],    // qw, qx, qy, qz
    pub translation: [f64; 3], // x, y, z
    pub points2d: Vec<Vector2d>,
    pub point3d_ids: Vec<u64>,
}

impl ColmapCamera {
    /// Create a new ColmapCamera.
    ///
    /// # Arguments
    ///
    /// * `camera_id` - Unique camera id.
    /// * `model_id` - Camera model id.
    /// * `width` - Image width.
    /// * `height` - Image height.
    /// * `params` - Camera parameters.
    ///
    /// # Returns
    ///
    /// * Result containing the new ColmapCamera.
    ///
    /// # Errors
    ///
    /// * If the number of camera parameters is invalid.
    pub fn new(
        camera_id: u32,
        model_id: CameraModelId,
        width: usize,
        height: usize,
        params: Vec<f64>,
    ) -> Result<Self, ColmapError> {
        if params.len() != camera_model_num_params(model_id) {
            return Err(ColmapError::InvalidNumCameraParams(params.len()));
        }
        Ok(ColmapCamera {
            camera_id,
            model_id,
            width,
            height,
            params,
        })
    }
}

fn camera_model_num_params(model_id: CameraModelId) -> usize {
    match model_id {
        CameraModelId::CameraModelSimplePinhole => 3,
        CameraModelId::CameraModelPinhole => 4,
        CameraModelId::CameraModelSimplifiedRadial => 4,
        CameraModelId::CameraModelRadial => 5,
        CameraModelId::CameraModelOpenCV => 8,
        CameraModelId::CameraModelOpenCVFisheye => 8,
        CameraModelId::CameraModelFullOpenCV => 12,
        CameraModelId::CameraModelFOV => 5,
        CameraModelId::CameraModelSimpleRadialFisheye => 4,
        CameraModelId::CameraModelRadialFisheye => 5,
        CameraModelId::CameraModelThinPrismFisheye => 12,
        _ => 0,
    }
}

/// Write colmap images to binary file.
///
/// # Arguments
///
/// * `path` - Path to the output file.
/// * `images` - Vector of colmap images.
///
pub fn write_images_binary(path: &Path, images: &Vec<ColmapImage>) -> Result<(), ColmapError> {
    let mut f = std::fs::File::create(path)?;

    f.write_u64::<LittleEndian>(images.len() as u64)?;

    for image in images {
        f.write_u32::<LittleEndian>(image.image_id)?;

        // write qw, qx, qy, qz
        for i in 0..4 {
            f.write_f64::<LittleEndian>(image.rotation[i])?;
        }

        // write x, y, z
        for i in 0..3 {
            f.write_f64::<LittleEndian>(image.translation[i])?;
        }

        f.write_u32::<LittleEndian>(image.camera_id)?;

        // image name
        for c in image.name.chars() {
            f.write_u8(c as u8)?;
        }
        f.write_u8('\0' as u8)?;

        f.write_u64::<LittleEndian>(image.points2d.len() as u64)?;

        for j in 0..image.points2d.len() {
            f.write_f64::<LittleEndian>(image.points2d[j].x)?;
            f.write_f64::<LittleEndian>(image.points2d[j].y)?;
            f.write_u64::<LittleEndian>(image.point3d_ids[j])?;
        }
    }

    Ok(())
}

/// Write colmap cameras to binary file.
///
/// # Arguments
///
/// * `path` - Path to the output file.
/// * `cameras` - Vector of colmap cameras.
///
pub fn write_cameras_binary(path: &Path, cameras: &Vec<ColmapCamera>) -> Result<(), ColmapError> {
    let mut f = std::fs::File::create(path)?;

    f.write_u64::<LittleEndian>(cameras.len() as u64)?;

    for camera in cameras {
        f.write_u32::<LittleEndian>(camera.camera_id)?;
        f.write_i32::<LittleEndian>(camera.model_id as i32)?; // enum to i32 for c++ compatibility
        f.write_u64::<LittleEndian>(camera.width as u64)?;
        f.write_u64::<LittleEndian>(camera.height as u64)?;
        for param in &camera.params {
            f.write_f64::<LittleEndian>(*param)?;
        }
    }

    Ok(())
}

/// Read colmap cameras from binary file.
///
/// # Arguments
///
/// * `path` - Path to the input file.
///
/// # Returns
///
/// * Vector of colmap cameras.
///
pub fn read_cameras_binary(path: &Path) -> Result<Vec<ColmapCamera>, ColmapError> {
    let mut f = std::fs::File::open(path)?;

    let mut cameras = Vec::new();

    let num_cameras = f.read_u64::<LittleEndian>()?;

    for _ in 0..num_cameras {
        let camera_id = f.read_u32::<LittleEndian>()?;
        let model_id = match f.read_i32::<LittleEndian>()? {
            0 => CameraModelId::CameraModelSimplePinhole,
            1 => CameraModelId::CameraModelPinhole,
            2 => CameraModelId::CameraModelSimplifiedRadial,
            3 => CameraModelId::CameraModelRadial,
            4 => CameraModelId::CameraModelOpenCV,
            5 => CameraModelId::CameraModelOpenCVFisheye,
            6 => CameraModelId::CameraModelFullOpenCV,
            7 => CameraModelId::CameraModelFOV,
            8 => CameraModelId::CameraModelSimpleRadialFisheye,
            9 => CameraModelId::CameraModelRadialFisheye,
            10 => CameraModelId::CameraModelThinPrismFisheye,
            11 => CameraModelId::CameraModelCount,
            _ => CameraModelId::CameraModelInvalid,
        };
        let width = f.read_u64::<LittleEndian>()? as usize;
        let height = f.read_u64::<LittleEndian>()? as usize;

        let mut params = Vec::new();
        for _ in 0..camera_model_num_params(model_id) {
            params.push(f.read_f64::<LittleEndian>()?);
        }

        cameras.push(ColmapCamera::new(
            camera_id, model_id, width, height, params,
        )?);
    }

    Ok(cameras)
}

/// Read colmap images from binary file.
///
/// # Arguments
///
/// * `path` - Path to the input file.
///
/// # Returns
///
/// * Vector of colmap images.
///
pub fn read_images_binary(path: &Path) -> Result<Vec<ColmapImage>, ColmapError> {
    let mut f = std::fs::File::open(path)?;

    let mut images = Vec::new();

    let num_images = f.read_u64::<LittleEndian>()?;

    for _ in 0..num_images {
        let image_id = f.read_u32::<LittleEndian>()?;

        let mut rotation = [0.0; 4];
        for i in 0..4 {
            rotation[i] = f.read_f64::<LittleEndian>()?;
        }

        let mut translation = [0.0; 3];
        for i in 0..3 {
            translation[i] = f.read_f64::<LittleEndian>()?;
        }

        let camera_id = f.read_u32::<LittleEndian>()?;

        let mut name = String::new();
        loop {
            let c = f.read_u8()?;
            if c == 0 {
                break;
            }
            name.push(c as char);
        }

        let num_points_2d = f.read_u64::<LittleEndian>()?;

        let mut points2d = Vec::with_capacity(num_points_2d as usize);
        let mut point3d_ids = Vec::with_capacity(num_points_2d as usize);

        for _ in 0..num_points_2d {
            let x = f.read_f64::<LittleEndian>()?;
            let y = f.read_f64::<LittleEndian>()?;
            let point3d_id = f.read_u64::<LittleEndian>()?;

            points2d.push(Vector2d { x, y });
            point3d_ids.push(point3d_id);
        }

        images.push(ColmapImage {
            name,
            image_id,
            camera_id,
            rotation,
            translation,
            points2d,
            point3d_ids,
        });
    }

    Ok(images)
}

#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Result;

    #[test]
    fn test_write_images_binary() -> Result<()> {
        let images = vec![
            ColmapImage {
                name: "image1".to_string(),
                image_id: 1,
                camera_id: 1,
                rotation: [1.0, 0.0, 0.0, 0.0],
                translation: [0.0, 0.0, 0.0],
                point3d_ids: vec![1, 2, 3],
                points2d: vec![
                    Vector2d { x: 1.0, y: 2.0 },
                    Vector2d { x: 3.0, y: 4.0 },
                    Vector2d { x: 5.0, y: 6.0 },
                ],
            },
            ColmapImage {
                name: "image2".to_string(),
                image_id: 2,
                camera_id: 1,
                rotation: [1.0, 0.0, 0.0, 0.0],
                translation: [0.0, 0.0, 0.0],
                point3d_ids: vec![4, 5, 6],
                points2d: vec![Vector2d { x: 2.0, y: 3.0 }, Vector2d { x: 4.0, y: 5.0 }],
            },
        ];

        // write to file
        let temp_dir = tempfile::TempDir::new()?;
        let path = temp_dir.path().join("test_images.bin");

        write_images_binary(&path, &images).unwrap();

        // read back from file
        let read_images = read_images_binary(&path).unwrap();

        assert_eq!(images.len(), read_images.len());

        for i in 0..images.len() {
            assert_eq!(images[i].name, read_images[i].name);
            assert_eq!(images[i].image_id, read_images[i].image_id);
            assert_eq!(images[i].camera_id, read_images[i].camera_id);
            for j in 0..4 {
                assert_eq!(images[i].rotation[j], read_images[i].rotation[j]);
            }
            for j in 0..3 {
                assert_eq!(images[i].translation[j], read_images[i].translation[j]);
            }
            for j in 0..images[i].points2d.len() {
                assert_eq!(images[i].points2d[j].x, read_images[i].points2d[j].x);
                assert_eq!(images[i].points2d[j].y, read_images[i].points2d[j].y);
                assert_eq!(images[i].point3d_ids[j], read_images[i].point3d_ids[j]);
            }
        }

        Ok(())
    }

    #[test]
    fn test_write_cameras_binary() -> Result<()> {
        let cameras = vec![
            ColmapCamera::new(
                1,
                CameraModelId::CameraModelSimplePinhole,
                1920,
                1080,
                vec![1.0, 2.0, 3.0],
            )?,
            ColmapCamera::new(
                2,
                CameraModelId::CameraModelOpenCV,
                1920,
                1080,
                vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
            )?,
        ];

        // write to file
        //let path = Path::new("test_cameras.bin");
        let temp_dir = tempfile::TempDir::new()?;
        let path = temp_dir.path().join("test_cameras.bin");

        write_cameras_binary(&path, &cameras)?;

        // read back from file
        let read_cameras = read_cameras_binary(&path)?;

        assert_eq!(cameras.len(), read_cameras.len());

        for i in 0..cameras.len() {
            println!("{:?}", cameras[i].model_id);
            println!("{:?}", read_cameras[i].model_id);
            assert_eq!(cameras[i].camera_id, read_cameras[i].camera_id);
            assert_eq!(cameras[i].model_id, read_cameras[i].model_id);
            assert_eq!(cameras[i].width, read_cameras[i].width);
            assert_eq!(cameras[i].height, read_cameras[i].height);
            for j in 0..cameras[i].params.len() {
                assert_eq!(cameras[i].params[j], read_cameras[i].params[j]);
            }
        }

        Ok(())
    }
}