Skip to main content

spatialrust_vision/
warp.rs

1//! Image remapping and geometric warp primitives.
2
3use spatialrust_image::{Image, ImageView};
4
5use crate::border::{constant_pixel, fetch};
6use crate::{BorderMode, Interpolation, PixelComponent, VisionError, VisionResult};
7
8/// A source-to-destination 2D affine transform.
9#[derive(Clone, Copy, Debug, PartialEq)]
10pub struct AffineTransform {
11    /// First two rows of a homogeneous 3x3 transform.
12    pub matrix: [[f64; 3]; 2],
13}
14
15impl AffineTransform {
16    /// Identity affine transform.
17    #[must_use]
18    pub const fn identity() -> Self {
19        Self { matrix: [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]] }
20    }
21
22    /// Applies the source-to-destination transform.
23    #[must_use]
24    pub fn map_point(self, x: f64, y: f64) -> (f64, f64) {
25        (
26            self.matrix[0][0].mul_add(x, self.matrix[0][1].mul_add(y, self.matrix[0][2])),
27            self.matrix[1][0].mul_add(x, self.matrix[1][1].mul_add(y, self.matrix[1][2])),
28        )
29    }
30
31    /// Computes the inverse affine transform.
32    pub fn inverse(self) -> VisionResult<Self> {
33        let a = self.matrix[0][0];
34        let b = self.matrix[0][1];
35        let c = self.matrix[0][2];
36        let d = self.matrix[1][0];
37        let e = self.matrix[1][1];
38        let f = self.matrix[1][2];
39        let determinant = a.mul_add(e, -(b * d));
40        if !determinant.is_finite() || determinant.abs() <= f64::EPSILON {
41            return Err(VisionError::SingularTransform);
42        }
43        let inv = 1.0 / determinant;
44        Ok(Self {
45            matrix: [
46                [e * inv, -b * inv, (b * f - e * c) * inv],
47                [-d * inv, a * inv, (d * c - a * f) * inv],
48            ],
49        })
50    }
51}
52
53/// A source-to-destination projective 2D transform.
54#[derive(Clone, Copy, Debug, PartialEq)]
55pub struct PerspectiveTransform {
56    /// Homogeneous 3x3 transform matrix.
57    pub matrix: [[f64; 3]; 3],
58}
59
60impl PerspectiveTransform {
61    /// Identity projective transform.
62    #[must_use]
63    pub const fn identity() -> Self {
64        Self { matrix: [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]] }
65    }
66
67    /// Applies the source-to-destination homogeneous transform.
68    #[must_use]
69    pub fn map_point(self, x: f64, y: f64) -> Option<(f64, f64)> {
70        let denominator =
71            self.matrix[2][0].mul_add(x, self.matrix[2][1].mul_add(y, self.matrix[2][2]));
72        if !denominator.is_finite() || denominator.abs() <= f64::EPSILON {
73            return None;
74        }
75        Some((
76            self.matrix[0][0].mul_add(x, self.matrix[0][1].mul_add(y, self.matrix[0][2]))
77                / denominator,
78            self.matrix[1][0].mul_add(x, self.matrix[1][1].mul_add(y, self.matrix[1][2]))
79                / denominator,
80        ))
81    }
82
83    /// Computes the inverse projective transform.
84    pub fn inverse(self) -> VisionResult<Self> {
85        let m = self.matrix;
86        let c00 = m[1][1].mul_add(m[2][2], -(m[1][2] * m[2][1]));
87        let c01 = -(m[1][0].mul_add(m[2][2], -(m[1][2] * m[2][0])));
88        let c02 = m[1][0].mul_add(m[2][1], -(m[1][1] * m[2][0]));
89        let c10 = -(m[0][1].mul_add(m[2][2], -(m[0][2] * m[2][1])));
90        let c11 = m[0][0].mul_add(m[2][2], -(m[0][2] * m[2][0]));
91        let c12 = -(m[0][0].mul_add(m[2][1], -(m[0][1] * m[2][0])));
92        let c20 = m[0][1].mul_add(m[1][2], -(m[0][2] * m[1][1]));
93        let c21 = -(m[0][0].mul_add(m[1][2], -(m[0][2] * m[1][0])));
94        let c22 = m[0][0].mul_add(m[1][1], -(m[0][1] * m[1][0]));
95        let determinant = m[0][0].mul_add(c00, m[0][1].mul_add(c01, m[0][2] * c02));
96        if !determinant.is_finite() || determinant.abs() <= f64::EPSILON {
97            return Err(VisionError::SingularTransform);
98        }
99        let inv = 1.0 / determinant;
100        Ok(Self {
101            matrix: [
102                [c00 * inv, c10 * inv, c20 * inv],
103                [c01 * inv, c11 * inv, c21 * inv],
104                [c02 * inv, c12 * inv, c22 * inv],
105            ],
106        })
107    }
108}
109
110/// Samples an image at absolute coordinates supplied by two single-channel maps.
111pub fn remap<T: PixelComponent, const CHANNELS: usize>(
112    input: ImageView<'_, T, CHANNELS>,
113    map_x: ImageView<'_, f32, 1>,
114    map_y: ImageView<'_, f32, 1>,
115    interpolation: Interpolation,
116    border: BorderMode<T, CHANNELS>,
117) -> VisionResult<Image<T, CHANNELS>> {
118    if map_x.width() != map_y.width() || map_x.height() != map_y.height() {
119        return Err(VisionError::ShapeMismatch("map_x and map_y dimensions must match".to_owned()));
120    }
121    if interpolation == Interpolation::Area {
122        return Err(VisionError::InvalidParameter(
123            "area interpolation is not defined for arbitrary remap".to_owned(),
124        ));
125    }
126    let mut output = Vec::with_capacity(map_x.width() * map_x.height() * CHANNELS);
127    for y in 0..map_x.height() {
128        for x in 0..map_x.width() {
129            let sx = f64::from(map_x.get(x, y).expect("map coordinate in bounds")[0]);
130            let sy = f64::from(map_y.get(x, y).expect("map coordinate in bounds")[0]);
131            let pixel = sample(input, sx, sy, interpolation, border);
132            output.extend_from_slice(&pixel);
133        }
134    }
135    Ok(Image::try_new_with_metadata(map_x.width(), map_x.height(), output, input.metadata())?)
136}
137
138/// Warps an image with a source-to-destination affine transform.
139pub fn warp_affine<T: PixelComponent, const CHANNELS: usize>(
140    input: ImageView<'_, T, CHANNELS>,
141    transform: AffineTransform,
142    output_width: usize,
143    output_height: usize,
144    interpolation: Interpolation,
145    border: BorderMode<T, CHANNELS>,
146) -> VisionResult<Image<T, CHANNELS>> {
147    let inverse = transform.inverse()?;
148    warp_with_mapping(input, output_width, output_height, interpolation, border, |x, y| {
149        Some(inverse.map_point(x, y))
150    })
151}
152
153/// Warps an image with a source-to-destination projective transform.
154pub fn warp_perspective<T: PixelComponent, const CHANNELS: usize>(
155    input: ImageView<'_, T, CHANNELS>,
156    transform: PerspectiveTransform,
157    output_width: usize,
158    output_height: usize,
159    interpolation: Interpolation,
160    border: BorderMode<T, CHANNELS>,
161) -> VisionResult<Image<T, CHANNELS>> {
162    let inverse = transform.inverse()?;
163    warp_with_mapping(input, output_width, output_height, interpolation, border, |x, y| {
164        inverse.map_point(x, y)
165    })
166}
167
168fn warp_with_mapping<T: PixelComponent, const CHANNELS: usize>(
169    input: ImageView<'_, T, CHANNELS>,
170    output_width: usize,
171    output_height: usize,
172    interpolation: Interpolation,
173    border: BorderMode<T, CHANNELS>,
174    mut mapping: impl FnMut(f64, f64) -> Option<(f64, f64)>,
175) -> VisionResult<Image<T, CHANNELS>> {
176    if interpolation == Interpolation::Area {
177        return Err(VisionError::InvalidParameter(
178            "area interpolation is not defined for geometric warp".to_owned(),
179        ));
180    }
181    let mut output = Vec::with_capacity(output_width * output_height * CHANNELS);
182    for y in 0..output_height {
183        for x in 0..output_width {
184            let pixel = mapping(x as f64, y as f64).map_or_else(
185                || constant_pixel(border),
186                |(sx, sy)| sample(input, sx, sy, interpolation, border),
187            );
188            output.extend_from_slice(&pixel);
189        }
190    }
191    Ok(Image::try_new_with_metadata(output_width, output_height, output, input.metadata())?)
192}
193
194fn sample<T: PixelComponent, const CHANNELS: usize>(
195    input: ImageView<'_, T, CHANNELS>,
196    x: f64,
197    y: f64,
198    interpolation: Interpolation,
199    border: BorderMode<T, CHANNELS>,
200) -> [T; CHANNELS] {
201    if !x.is_finite() || !y.is_finite() || input.width() == 0 || input.height() == 0 {
202        return constant_pixel(border);
203    }
204    match interpolation {
205        Interpolation::Nearest => fetch(input, x.round() as isize, y.round() as isize, border),
206        Interpolation::Bilinear => sample_bilinear(input, x, y, border),
207        Interpolation::Bicubic => sample_bicubic(input, x, y, border),
208        Interpolation::Area => unreachable!("area rejected before sampling"),
209    }
210}
211
212fn sample_bilinear<T: PixelComponent, const CHANNELS: usize>(
213    input: ImageView<'_, T, CHANNELS>,
214    x: f64,
215    y: f64,
216    border: BorderMode<T, CHANNELS>,
217) -> [T; CHANNELS] {
218    let x0 = x.floor() as isize;
219    let y0 = y.floor() as isize;
220    let wx = x - x.floor();
221    let wy = y - y.floor();
222    let p00 = fetch(input, x0, y0, border);
223    let p10 = fetch(input, x0 + 1, y0, border);
224    let p01 = fetch(input, x0, y0 + 1, border);
225    let p11 = fetch(input, x0 + 1, y0 + 1, border);
226    std::array::from_fn(|channel| {
227        let top = p00[channel].to_f64().mul_add(1.0 - wx, p10[channel].to_f64() * wx);
228        let bottom = p01[channel].to_f64().mul_add(1.0 - wx, p11[channel].to_f64() * wx);
229        T::from_f64(top.mul_add(1.0 - wy, bottom * wy))
230    })
231}
232
233fn cubic_weight(distance: f64) -> f64 {
234    let x = distance.abs();
235    const A: f64 = -0.75;
236    if x <= 1.0 {
237        (A + 2.0) * x * x * x - (A + 3.0) * x * x + 1.0
238    } else if x < 2.0 {
239        A * x * x * x - 5.0 * A * x * x + 8.0 * A * x - 4.0 * A
240    } else {
241        0.0
242    }
243}
244
245fn sample_bicubic<T: PixelComponent, const CHANNELS: usize>(
246    input: ImageView<'_, T, CHANNELS>,
247    x: f64,
248    y: f64,
249    border: BorderMode<T, CHANNELS>,
250) -> [T; CHANNELS] {
251    let base_x = x.floor() as isize;
252    let base_y = y.floor() as isize;
253    let mut sums = [0.0_f64; CHANNELS];
254    let mut total_weight = 0.0;
255    for dy in -1..=2 {
256        let wy = cubic_weight(y - (base_y + dy) as f64);
257        for dx in -1..=2 {
258            let weight = wy * cubic_weight(x - (base_x + dx) as f64);
259            let pixel = fetch(input, base_x + dx, base_y + dy, border);
260            for channel in 0..CHANNELS {
261                sums[channel] += pixel[channel].to_f64() * weight;
262            }
263            total_weight += weight;
264        }
265    }
266    std::array::from_fn(|channel| T::from_f64(sums[channel] / total_weight))
267}
268
269#[cfg(test)]
270mod tests {
271    use super::{remap, warp_affine, warp_perspective, AffineTransform, PerspectiveTransform};
272    use crate::{BorderMode, Interpolation};
273    use spatialrust_image::Image;
274
275    #[test]
276    fn identity_remap_is_exact() {
277        let input = Image::<u8, 1>::try_new(2, 2, vec![1, 2, 3, 4]).unwrap();
278        let mx = Image::<f32, 1>::try_new(2, 2, vec![0.0, 1.0, 0.0, 1.0]).unwrap();
279        let my = Image::<f32, 1>::try_new(2, 2, vec![0.0, 0.0, 1.0, 1.0]).unwrap();
280        let output = remap(
281            input.view(),
282            mx.view(),
283            my.view(),
284            Interpolation::Bilinear,
285            BorderMode::Constant([0]),
286        )
287        .unwrap();
288        assert_eq!(output, input);
289    }
290
291    #[test]
292    fn affine_translation_uses_constant_border() {
293        let input = Image::<u8, 1>::try_new(3, 1, vec![1, 2, 3]).unwrap();
294        let transform = AffineTransform { matrix: [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]] };
295        let output = warp_affine(
296            input.view(),
297            transform,
298            3,
299            1,
300            Interpolation::Nearest,
301            BorderMode::Constant([9]),
302        )
303        .unwrap();
304        assert_eq!(output.as_slice(), &[9, 1, 2]);
305    }
306
307    #[test]
308    fn perspective_identity_is_exact() {
309        let input = Image::<u8, 1>::try_new(2, 2, vec![1, 2, 3, 4]).unwrap();
310        let output = warp_perspective(
311            input.view(),
312            PerspectiveTransform::identity(),
313            2,
314            2,
315            Interpolation::Nearest,
316            BorderMode::Replicate,
317        )
318        .unwrap();
319        assert_eq!(output, input);
320    }
321
322    #[test]
323    fn border_modes_are_distinct() {
324        let input = Image::<u8, 1>::try_new(3, 1, vec![10, 20, 30]).unwrap();
325        let mx = Image::<f32, 1>::try_new(1, 1, vec![-1.0]).unwrap();
326        let my = Image::<f32, 1>::try_new(1, 1, vec![0.0]).unwrap();
327        let run = |border| {
328            remap(input.view(), mx.view(), my.view(), Interpolation::Nearest, border)
329                .unwrap()
330                .as_slice()[0]
331        };
332        assert_eq!(run(BorderMode::Constant([5])), 5);
333        assert_eq!(run(BorderMode::Replicate), 10);
334        assert_eq!(run(BorderMode::Reflect), 10);
335        assert_eq!(run(BorderMode::Reflect101), 20);
336        assert_eq!(run(BorderMode::Wrap), 30);
337    }
338}