align3d/io/dataset/
slamtb.rs

1use std::path::{Path, PathBuf};
2
3use super::core::{DatasetError, RgbdDataset};
4use crate::{
5    camera::CameraIntrinsics,
6    image::{IntoArray3, RgbdFrame, RgbdImage},
7    trajectory::Trajectory,
8    transform::Transform,
9};
10
11use nshare::ToNdarray2;
12
13pub struct SlamTbDataset {
14    cameras: Vec<CameraIntrinsics>,
15    extrinsic_cameras: Vec<Transform>,
16    rgb_images: Vec<String>,
17    depth_images: Vec<String>,
18    depth_scales: Vec<f64>,
19    base_dir: PathBuf,
20}
21
22mod json {
23    use serde_derive::Deserialize;
24    #[derive(Deserialize, Debug)]
25    pub struct KCam {
26        pub matrix: Vec<Vec<f64>>,
27        #[serde(rename = "undist_coeff")]
28        pub _undist_coeff: Vec<f32>,
29        pub image_size: (i32, i32),
30    }
31
32    #[derive(Deserialize, Debug)]
33    pub struct RTCam {
34        pub matrix: Vec<Vec<f32>>,
35    }
36
37    #[derive(Deserialize, Debug)]
38    pub struct Info {
39        pub kcam: KCam,
40        pub depth_scale: f64,
41        pub depth_bias: f64,
42        pub depth_max: f64,
43        pub rt_cam: RTCam,
44        pub timestamp: f64,
45    }
46
47    #[derive(Deserialize, Debug)]
48    pub struct Frame {
49        pub info: Info,
50        pub depth_image: String,
51        pub rgb_image: String,
52    }
53
54    #[derive(Deserialize, Debug)]
55    pub struct Document {
56        pub root: Vec<Frame>,
57    }
58}
59
60impl SlamTbDataset {
61    pub fn load(base_dir: &str) -> Result<Self, DatasetError> {
62        let buffer = std::io::BufReader::new(std::fs::File::open(
63            Path::new(base_dir).join("frames.json"),
64        )?);
65        serde_json::from_reader(buffer)
66            .map(|doc: json::Document| {
67                let mut cameras = Vec::new();
68                let mut extrinsic_cameras = Vec::new();
69                let mut rgb_images = Vec::new();
70                let mut depth_images = Vec::new();
71                let mut depth_scales = Vec::new();
72
73                for frame in doc.root.iter() {
74                    let info = &frame.info;
75
76                    let fx = info.kcam.matrix[0][0];
77                    let fy = info.kcam.matrix[1][1];
78                    let cx = info.kcam.matrix[0][2];
79                    let cy = info.kcam.matrix[1][2];
80
81                    let extrinsics = if info.rt_cam.matrix.len() == 4 {
82                        Transform::from_matrix4(&nalgebra::Matrix4::<f32>::from_fn(|r, c| {
83                            info.rt_cam.matrix[r][c]
84                        }))
85                    } else {
86                        Transform::eye()
87                    };
88
89                    cameras.push(CameraIntrinsics::from_simple_intrinsic(
90                        fx,
91                        fy,
92                        cx,
93                        cy,
94                        info.kcam.image_size.0 as usize,
95                        info.kcam.image_size.1 as usize,
96                    ));
97                    extrinsic_cameras.push(extrinsics);
98                    rgb_images.push(frame.rgb_image.clone());
99                    depth_images.push(frame.depth_image.clone());
100                    depth_scales.push(info.depth_scale);
101                }
102                Self {
103                    cameras,
104                    extrinsic_cameras,
105                    rgb_images,
106                    depth_images,
107                    depth_scales,
108                    base_dir: PathBuf::from(base_dir),
109                }
110            })
111            .map_err(|err| DatasetError::Parser(err.to_string()))
112    }
113}
114
115impl RgbdDataset for SlamTbDataset {
116    fn len(&self) -> usize {
117        self.rgb_images.len().min(self.depth_images.len())
118    }
119
120    fn is_empty(&self) -> bool {
121        self.len() == 0
122    }
123
124    fn get(&self, index: usize) -> Result<RgbdFrame, DatasetError> {
125        let rgb_image = image::open(self.base_dir.join(&self.rgb_images[index]))?
126            .into_rgb8()
127            .into_array3();
128
129        let rgb_image = rgb_image.as_standard_layout().into_owned();
130        let depth_image = image::open(self.base_dir.join(&self.depth_images[index]))?
131            .into_luma16()
132            .into_ndarray2();
133        Ok(RgbdFrame::new(
134            self.cameras[index].clone(),
135            RgbdImage::with_depth_scale(rgb_image, depth_image, self.depth_scales[index]),
136            Some(self.extrinsic_cameras[index].clone()),
137        ))
138    }
139
140    fn trajectory(&self) -> Option<Trajectory> {
141        let mut trajectory = Trajectory::default();
142        for (i, cam) in self.extrinsic_cameras.iter().enumerate() {
143            trajectory.push(cam.clone(), i as f32);
144        }
145        Some(trajectory)
146    }
147
148    fn camera(&self, index: usize) -> (CameraIntrinsics, Option<Transform>) {
149        (
150            self.cameras[index].clone(),
151            Some(self.extrinsic_cameras[index].clone()),
152        )
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159
160    #[test]
161    fn test_load() {
162        let rgbd_dataset = SlamTbDataset::load("tests/data/rgbd/sample1").unwrap();
163
164        let (camera, image, _) = rgbd_dataset.get(0).unwrap().into_parts();
165
166        assert_eq!(camera.fx, 544.4732666015625);
167        assert_eq!(camera.fy, 544.4732666015625);
168        assert_eq!(camera.cx, 320.0);
169        assert_eq!(camera.cy, 240.0);
170
171        assert_eq!(image.height(), 480);
172        assert_eq!(image.width(), 640);
173    }
174}
align3d/io/dataset/slamtb.rs

align3d/io/dataset/
slamtb.rs
