kornia_3d/io/colmap/
text.rs

1use std::{
2    fs::File,
3    io::{BufRead, BufReader},
4    path::Path,
5};
6
7use super::{CameraModelId, ColmapCamera, ColmapImage, ColmapPoint3d};
8
9/// Error types for the COLMAP module.
10#[derive(Debug, thiserror::Error)]
11pub enum ColmapError {
12    /// Error reading or writing file
13    #[error("error reading or writing file")]
14    IoError(#[from] std::io::Error),
15
16    /// Invalid number of camera parameters
17    #[error("Invalid number of camera parameters")]
18    InvalidNumCameraParams(usize),
19
20    /// Parse error
21    #[error("Parse error {0}")]
22    ParseError(String),
23}
24
25/// Read the cameras.txt file and return a vector of ColmapCamera structs.
26///
27/// # Arguments
28///
29/// * `path` - The path to the cameras.txt file.
30///
31/// # Returns
32///
33/// A vector of ColmapCamera structs.
34pub fn read_cameras_txt(path: impl AsRef<Path>) -> Result<Vec<ColmapCamera>, ColmapError> {
35    // open the file and create a buffered reader
36    let file = File::open(path)?;
37    let reader = BufReader::new(file);
38
39    // skip the first 3 lines containing the header and parse the rest
40    let cameras = reader
41        .lines()
42        .skip(3)
43        .map(|line| -> Result<ColmapCamera, ColmapError> {
44            let line = line.map_err(ColmapError::from)?;
45            parse_camera_line(&line)
46        })
47        .collect::<Result<Vec<_>, _>>()?;
48
49    Ok(cameras)
50}
51
52/// Read the points3D.txt file and return a vector of ColmapPoint3d structs.
53///
54/// # Arguments
55///
56/// * `path` - The path to the points3D.txt file.
57///
58/// # Returns
59///
60/// A vector of ColmapPoint3d structs.
61pub fn read_points3d_txt(path: impl AsRef<Path>) -> Result<Vec<ColmapPoint3d>, ColmapError> {
62    // open the file and create a buffered reader
63    let file = File::open(path)?;
64    let reader = BufReader::new(file);
65
66    // skip the first 3 lines containing the header and parse the rest
67
68    let points = reader
69        .lines()
70        .skip(3)
71        .map(|line| -> Result<ColmapPoint3d, ColmapError> {
72            let line = line.map_err(ColmapError::from)?;
73            parse_point3d_line(&line)
74        })
75        .collect::<Result<Vec<_>, _>>()?;
76
77    Ok(points)
78}
79
80/// Read the images.txt file and return a vector of ColmapImage structs.
81///
82/// # Arguments
83///
84/// * `path` - The path to the images.txt file.
85///
86/// # Returns
87///
88/// A vector of ColmapImage structs.
89pub fn read_images_txt(path: impl AsRef<Path>) -> Result<Vec<ColmapImage>, ColmapError> {
90    // open the file and create a buffered reader
91    let file = File::open(path)?;
92    let reader = BufReader::new(file);
93
94    let images = reader
95        .lines()
96        .skip(4)
97        .collect::<Result<Vec<_>, _>>()?
98        .chunks(2)
99        .map(|chunk| match chunk {
100            [line1, line2] => parse_image_line(line1, line2),
101            _ => Err(ColmapError::ParseError(
102                "Invalid number of lines".to_string(),
103            )),
104        })
105        .collect::<Result<Vec<_>, _>>()?;
106
107    Ok(images)
108}
109
110/// Utility functions for parsing COLMAP text files
111fn parse_part<T: std::str::FromStr>(s: &str) -> Result<T, ColmapError>
112where
113    T::Err: std::fmt::Display,
114{
115    s.parse::<T>()
116        .map_err(|e| ColmapError::ParseError(format!("{s}: {e}")))
117}
118
119/// Parse a camera line and return a ColmapCamera struct.
120/// NOTE: The number of parameters depends on the camera model.
121///       CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[0], PARAMS[1], ...
122fn parse_camera_line(line: &str) -> Result<ColmapCamera, ColmapError> {
123    // split the line into parts by whitespace
124    let parts = line.split_whitespace().collect::<Vec<_>>();
125
126    if parts.len() < 5 {
127        return Err(ColmapError::ParseError(format!(
128            "Invalid number of parts: {parts_len}",
129            parts_len = parts.len()
130        )));
131    }
132
133    Ok(ColmapCamera {
134        camera_id: parse_part(parts[0])?,
135        model_id: parse_camera_model_id(parts[1])?,
136        width: parse_part(parts[2])?,
137        height: parse_part(parts[3])?,
138        params: parts[4..]
139            .iter()
140            .map(|s| parse_part(s))
141            .collect::<Result<Vec<_>, _>>()?,
142    })
143}
144
145fn parse_camera_model_id(model_id: &str) -> Result<CameraModelId, ColmapError> {
146    match model_id {
147        "SIMPLE_PINHOLE" => Ok(CameraModelId::CameraModelSimplePinhole),
148        "PINHOLE" => Ok(CameraModelId::CameraModelPinhole),
149        "SIMPLE_RADIAL" => Ok(CameraModelId::CameraModelSimplifiedRadial),
150        "RADIAL" => Ok(CameraModelId::CameraModelRadial),
151        "OPENCV" => Ok(CameraModelId::CameraModelOpenCV),
152        "OPENCV_FISHEYE" => Ok(CameraModelId::CameraModelOpenCVFisheye),
153        "FULL_OPENCV" => Ok(CameraModelId::CameraModelFullOpenCV),
154        "FOV" => Ok(CameraModelId::CameraModelFOV),
155        "SIMPLE_RADIAL_FISHEYE" => Ok(CameraModelId::CameraModelSimpleRadialFisheye),
156        "RADIAL_FISHEYE" => Ok(CameraModelId::CameraModelRadialFisheye),
157        "THIN_PRISM_FISHEYE" => Ok(CameraModelId::CameraModelThinPrismFisheye),
158        _ => Err(ColmapError::ParseError(format!(
159            "Invalid camera model id: {model_id}"
160        ))),
161    }
162}
163
164/// Parse a point3d line and return a ColmapPoint3d struct.
165/// NOTE: The number of parameters depends on the camera model.
166///       POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[0], TRACK[1], ...
167fn parse_point3d_line(line: &str) -> Result<ColmapPoint3d, ColmapError> {
168    // split the line into parts by whitespace
169    let parts = line.split_whitespace().collect::<Vec<_>>();
170
171    // check if the number of parts is correct
172    if parts.len() < 8 {
173        return Err(ColmapError::ParseError(format!(
174            "Invalid number of parts: {}",
175            parts.len()
176        )));
177    }
178
179    Ok(ColmapPoint3d {
180        point3d_id: parse_part(parts[0])?,
181        xyz: parts[1..4]
182            .iter()
183            .map(|s| parse_part(s))
184            .collect::<Result<Vec<_>, _>>()?
185            .try_into()
186            .map_err(|_| {
187                ColmapError::ParseError("Invalid number of xyz coordinates".to_string())
188            })?,
189        rgb: parts[4..7]
190            .iter()
191            .map(|s| parse_part(s))
192            .collect::<Result<Vec<_>, _>>()?
193            .try_into()
194            .map_err(|_| {
195                ColmapError::ParseError("Invalid number of rgb coordinates".to_string())
196            })?,
197        error: parse_part(parts[7])?,
198        track: parts[8..]
199            .chunks_exact(2)
200            .map(|chunk| -> Result<(u32, u32), ColmapError> {
201                Ok((parse_part(chunk[0])?, parse_part(chunk[1])?))
202            })
203            .collect::<Result<Vec<_>, _>>()?,
204    })
205}
206
207/// Parse an image line and return a ColmapImage struct.
208/// #   IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME
209/// #   POINTS2D[] as (X, Y, POINT3D_ID)
210fn parse_image_line(line1: &str, line2: &str) -> Result<ColmapImage, ColmapError> {
211    // split the line into parts by whitespace
212    let parts1 = line1.split_whitespace().collect::<Vec<_>>();
213    let parts2 = line2.split_whitespace().collect::<Vec<_>>();
214
215    Ok(ColmapImage {
216        image_id: parse_part(parts1[0])?,
217        rotation: parts1[1..5]
218            .iter()
219            .map(|s| parse_part(s))
220            .collect::<Result<Vec<_>, _>>()?
221            .try_into()
222            .map_err(|_| {
223                ColmapError::ParseError("Invalid number of rotation coordinates".to_string())
224            })?,
225        translation: parts1[5..8]
226            .iter()
227            .map(|s| parse_part(s))
228            .collect::<Result<Vec<_>, _>>()?
229            .try_into()
230            .map_err(|_| {
231                ColmapError::ParseError("Invalid number of translation coordinates".to_string())
232            })?,
233        camera_id: parse_part(parts1[8])?,
234        name: parts1[9].to_string(),
235        points2d: parts2
236            .chunks_exact(3)
237            .map(|chunk| -> Result<(f64, f64, i64), ColmapError> {
238                Ok((
239                    parse_part(chunk[0])?,
240                    parse_part(chunk[1])?,
241                    parse_part(chunk[2])?,
242                ))
243            })
244            .collect::<Result<Vec<_>, _>>()?,
245    })
246}
kornia_3d/io/colmap/text.rs

kornia_3d/io/colmap/
text.rs
