async_cuda_npp/
remap.rs

1use async_cuda_core::runtime::Future;
2use async_cuda_core::DeviceBuffer2D;
3
4use crate::ffi;
5use crate::stream::Stream;
6
7type Result<T> = std::result::Result<T, crate::error::Error>;
8
9/// Remaps an image using bilinear interpolation. This function expects a reference to a device
10/// buffer as inputs, and a mutable reference to a device buffer to store the output of the
11/// operation in.
12///
13/// This function assumes the following about the input and output images:
14/// * Images are in RGB format.
15/// * Images are in standard memory order, i.e. HWC.
16///
17/// # Stream ordered semantics
18///
19/// This function uses stream ordered semantics. It can only be guaranteed to complete sequentially
20/// relative to operations scheduled on the same stream or the default stream.
21///
22/// # Arguments
23///
24/// * `input` - The on-device input image.
25/// * `output` - The on-device output image.
26/// * `map_x` - On-device X pixel map.
27/// * `map_y` - On-device Y pixel map.
28/// * `stream` - Stream to use.
29pub async fn remap(
30    input: &DeviceBuffer2D<u8>,
31    output: &mut DeviceBuffer2D<u8>,
32    map_x: &DeviceBuffer2D<f32>,
33    map_y: &DeviceBuffer2D<f32>,
34    stream: &Stream,
35) -> Result<()> {
36    assert_eq!(input.num_channels(), 3, "input image must be in RGB format");
37    assert_eq!(
38        output.num_channels(),
39        3,
40        "output image must be in RGB format"
41    );
42    assert_eq!(map_x.num_channels(), 1, "map must have one channel");
43    assert_eq!(map_y.num_channels(), 1, "map must have one channel");
44    assert_eq!(
45        output.width(),
46        map_x.width(),
47        "map x must have same width as output image"
48    );
49    assert_eq!(
50        output.height(),
51        map_x.height(),
52        "map x must have same height as output image"
53    );
54    assert_eq!(
55        output.width(),
56        map_y.width(),
57        "map y must have same width as output image"
58    );
59    assert_eq!(
60        output.height(),
61        map_y.height(),
62        "map y must have same height as output image"
63    );
64
65    let context = stream.to_context();
66    Future::new(move || {
67        ffi::remap::remap(
68            input.inner(),
69            output.inner_mut(),
70            map_x.inner(),
71            map_y.inner(),
72            &context,
73        )
74    })
75    .await
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    use crate::stream::Stream;
83    use crate::tests::image::*;
84    use crate::tests::memory::*;
85
86    use async_cuda_core::DeviceBuffer2D;
87
88    #[tokio::test]
89    async fn test_remap() {
90        const MAP_X: &[f32; 16] = &[
91            0.0, 1.0, 2.0, 3.0, // No mapping at all
92            1.0, 1.0, 2.0, 2.0, // Ignore the red border
93            1.0, 1.0, 2.0, 2.0, // Ignore the red border
94            1.0, 1.0, 2.0, 2.0, // Ignore the red border
95        ];
96        const MAP_Y: &[f32; 16] = &[
97            0.0, 0.0, 0.0, 0.0, // No mapping at all
98            1.0, 1.0, 1.0, 1.0, // Take from green band
99            1.0, 1.0, 1.0, 1.0, // Take from green band
100            2.0, 2.0, 2.0, 2.0, // Take from blue band
101        ];
102        const OUTPUT: Image4x4 = [
103            [R, R, R, R], // Red band
104            [G, G, G, G], // Green band
105            [G, G, G, G], // Green band
106            [B, B, B, B], // Blue band
107        ];
108        const OUTPUT_FLAT: [u8; 4 * 4 * 3] = flatten!(OUTPUT, 4 * 4 * 3);
109
110        let stream = Stream::new().await.unwrap();
111
112        let image = to_device_2d!(&RGB_FLAG, 4, 4, 3, &stream);
113        let map_x = to_device_2d!(MAP_X, 4, 4, 1, &stream);
114        let map_y = to_device_2d!(MAP_Y, 4, 4, 1, &stream);
115        let mut output = DeviceBuffer2D::<u8>::new(4, 4, 3).await;
116        assert!(remap(&image, &mut output, &map_x, &map_y, &stream)
117            .await
118            .is_ok());
119
120        let output = to_host_2d!(output, &stream);
121        assert_eq!(&output, &OUTPUT_FLAT);
122    }
123
124    #[tokio::test]
125    #[should_panic]
126    async fn test_it_panics_when_input_num_channels_incorrect() {
127        let input = DeviceBuffer2D::<u8>::new(100, 100, 2).await;
128        let map_x = DeviceBuffer2D::<f32>::new(100, 100, 1).await;
129        let map_y = DeviceBuffer2D::<f32>::new(100, 100, 1).await;
130        let mut output = DeviceBuffer2D::<u8>::new(100, 100, 3).await;
131        remap(&input, &mut output, &map_x, &map_y, &Stream::null().await)
132            .await
133            .unwrap();
134    }
135
136    #[tokio::test]
137    #[should_panic]
138    async fn test_it_panics_when_output_num_channels_incorrect() {
139        let input = DeviceBuffer2D::<u8>::new(100, 100, 3).await;
140        let map_x = DeviceBuffer2D::<f32>::new(100, 100, 1).await;
141        let map_y = DeviceBuffer2D::<f32>::new(100, 100, 1).await;
142        let mut output = DeviceBuffer2D::<u8>::new(100, 100, 2).await;
143        remap(&input, &mut output, &map_x, &map_y, &Stream::null().await)
144            .await
145            .unwrap();
146    }
147
148    #[tokio::test]
149    #[should_panic]
150    async fn test_it_panics_when_map_num_channels_incorrect() {
151        let input = DeviceBuffer2D::<u8>::new(100, 100, 3).await;
152        let map_x = DeviceBuffer2D::<f32>::new(100, 100, 2).await;
153        let map_y = DeviceBuffer2D::<f32>::new(100, 100, 3).await;
154        let mut output = DeviceBuffer2D::<u8>::new(100, 100, 3).await;
155        remap(&input, &mut output, &map_x, &map_y, &Stream::null().await)
156            .await
157            .unwrap();
158    }
159
160    #[tokio::test]
161    #[should_panic]
162    async fn test_it_panics_when_map_width_incorrect() {
163        let input = DeviceBuffer2D::<u8>::new(100, 100, 3).await;
164        let map_x = DeviceBuffer2D::<f32>::new(120, 100, 1).await;
165        let map_y = DeviceBuffer2D::<f32>::new(120, 100, 1).await;
166        let mut output = DeviceBuffer2D::<u8>::new(100, 100, 3).await;
167        remap(&input, &mut output, &map_x, &map_y, &Stream::null().await)
168            .await
169            .unwrap();
170    }
171
172    #[tokio::test]
173    #[should_panic]
174    async fn test_it_panics_when_map_height_incorrect() {
175        let input = DeviceBuffer2D::<u8>::new(100, 100, 3).await;
176        let map_x = DeviceBuffer2D::<f32>::new(100, 120, 1).await;
177        let map_y = DeviceBuffer2D::<f32>::new(100, 120, 1).await;
178        let mut output = DeviceBuffer2D::<u8>::new(100, 100, 3).await;
179        remap(&input, &mut output, &map_x, &map_y, &Stream::null().await)
180            .await
181            .unwrap();
182    }
183}