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
9pub 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, 1.0, 1.0, 2.0, 2.0, 1.0, 1.0, 2.0, 2.0, 1.0, 1.0, 2.0, 2.0, ];
96 const MAP_Y: &[f32; 16] = &[
97 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0, ];
102 const OUTPUT: Image4x4 = [
103 [R, R, R, R], [G, G, G, G], [G, G, G, G], [B, B, B, B], ];
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}