1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use ndarray::{Array2, Array3};

use crate::{
    bilateral::BilateralFilter,
    camera::{CameraIntrinsics, PinholeCamera},
    sampling::Downsample,
    transform::Transform,
};

use super::{py_scale_down, IntoImageRgb8};

/// A convinence struct that holds a color image, a depth image and its depth scale.
pub struct RgbdImage {
    pub color: Array3<u8>,
    pub depth: Array2<u16>,
    pub depth_scale: Option<f64>,
}

impl RgbdImage {
    pub fn new(color: Array3<u8>, depth: Array2<u16>) -> Self {
        Self {
            color,
            depth,
            depth_scale: None,
        }
    }

    pub fn with_depth_scale(color: Array3<u8>, depth: Array2<u16>, depth_scale: f64) -> Self {
        Self {
            color,
            depth,
            depth_scale: Some(depth_scale),
        }
    }

    pub fn width(&self) -> usize {
        self.color.shape()[1]
    }

    pub fn height(&self) -> usize {
        self.color.shape()[0]
    }
}

impl Downsample for RgbdImage {
    type Output = RgbdImage;
    fn downsample(&self, sigma: f32) -> RgbdImage {
        let resized_color = py_scale_down(&self.color.clone().into_image_rgb8(), sigma);
        let depth_filter = BilateralFilter::default();

        let resized_depth = depth_filter.scale_down(&self.depth);

        RgbdImage {
            color: resized_color,
            depth: resized_depth,
            depth_scale: self.depth_scale,
        }
    }
}

/// A struct that holds a camera intrinsics, a camera pose and an RGB-D image. Used by RGBD dataset readers.
pub struct RgbdFrame {
    /// The camera intrinsics.
    pub camera: CameraIntrinsics,
    /// The camera pose in the world coordinate system. None if the dataset has no ground truth.
    pub camera_to_world: Option<Transform>,
    /// The RGB-D image.
    pub image: RgbdImage,
}

impl RgbdFrame {
    pub fn new(
        camera: CameraIntrinsics,
        image: RgbdImage,
        camera_to_world: Option<Transform>,
    ) -> Self {
        Self {
            camera,
            image,
            camera_to_world,
        }
    }

    pub fn into_parts(self) -> (CameraIntrinsics, RgbdImage, Option<Transform>) {
        (self.camera, self.image, self.camera_to_world)
    }

    pub fn get_pinhole_camera(&self) -> Option<PinholeCamera> {
        self.camera_to_world
            .as_ref()
            .map(|camera_to_world| PinholeCamera::new(self.camera.clone(), camera_to_world.clone()))
    }
}

impl Downsample for RgbdFrame {
    type Output = RgbdFrame;

    fn downsample(&self, scale: f32) -> RgbdFrame {
        RgbdFrame {
            camera: self.camera.scale(0.5),
            image: self.image.downsample(scale),
            camera_to_world: self.camera_to_world.clone(),
        }
    }
}

#[cfg(test)]
mod tests {
    use rstest::rstest;

    use crate::{
        image::IntoImageRgb8, io::dataset::RgbdDataset, sampling::Downsample,
        unit_test::sample_rgbd_dataset1,
    };

    #[rstest]
    fn test_downsample(sample_rgbd_dataset1: impl RgbdDataset) {
        let image = sample_rgbd_dataset1.get(0).unwrap().image;
        let scale_05 = image.downsample(0.5);
        assert_eq!([240, 320, 3], scale_05.color.shape());
        assert_eq!([240, 320], scale_05.depth.shape());
        assert_eq!(320, scale_05.width());
        assert_eq!(240, scale_05.height());
        scale_05
            .color
            .into_image_rgb8()
            .save("scale_05_color.png")
            .unwrap();
    }
}