Skip to main content

spatialrust_vision/
photography.rs

1//! Runtime-free computational photography and pairwise panorama composition.
2
3use spatialrust_image::{Image, ImageView};
4
5use crate::{
6    estimate_homography_ransac, PerspectiveTransform, PointCorrespondence2,
7    RobustEstimationOptions, VisionError, VisionResult,
8};
9
10/// Applies deterministic gray-world white balance to an RGB image.
11pub fn gray_world_white_balance(input: ImageView<'_, u8, 3>) -> VisionResult<Image<u8, 3>> {
12    if input.width() == 0 || input.height() == 0 {
13        return Ok(Image::try_new_with_metadata(0, 0, Vec::new(), input.metadata())?);
14    }
15    let pixels = input.width().saturating_mul(input.height()) as f64;
16    let mut sums = [0.0; 3];
17    for y in 0..input.height() {
18        let row = input.row(y).expect("validated RGB row");
19        for pixel in row.chunks_exact(3) {
20            for channel in 0..3 {
21                sums[channel] += f64::from(pixel[channel]);
22            }
23        }
24    }
25    let means = sums.map(|sum| sum / pixels);
26    let target = means.iter().sum::<f64>() / 3.0;
27    let gains = means.map(|mean| if mean > 0.0 { target / mean } else { 1.0 });
28    let mut data = Vec::with_capacity(input.width() * input.height() * 3);
29    for y in 0..input.height() {
30        for (index, &value) in input.row(y).expect("validated RGB row").iter().enumerate() {
31            data.push((f64::from(value) * gains[index % 3]).round().clamp(0.0, 255.0) as u8);
32        }
33    }
34    Ok(Image::try_new_with_metadata(input.width(), input.height(), data, input.metadata())?)
35}
36
37/// Well-exposedness fusion controls.
38#[derive(Clone, Copy, Debug, PartialEq)]
39pub struct ExposureFusionOptions {
40    /// Standard deviation around middle gray in normalized intensity.
41    pub well_exposed_sigma: f64,
42    /// Positive floor preventing completely black weights.
43    pub weight_floor: f64,
44}
45
46impl Default for ExposureFusionOptions {
47    fn default() -> Self {
48        Self { well_exposed_sigma: 0.2, weight_floor: 1e-6 }
49    }
50}
51
52/// Fuses aligned RGB exposures with normalized per-pixel well-exposedness.
53pub fn fuse_exposures(
54    inputs: &[ImageView<'_, u8, 3>],
55    options: ExposureFusionOptions,
56) -> VisionResult<Image<u8, 3>> {
57    let first = *inputs.first().ok_or_else(|| {
58        VisionError::InvalidParameter("exposure fusion requires at least one image".into())
59    })?;
60    if !options.well_exposed_sigma.is_finite()
61        || options.well_exposed_sigma <= 0.0
62        || !options.weight_floor.is_finite()
63        || options.weight_floor <= 0.0
64    {
65        return Err(VisionError::InvalidParameter("invalid exposure fusion weights".into()));
66    }
67    if inputs.iter().any(|image| {
68        image.width() != first.width()
69            || image.height() != first.height()
70            || image.metadata() != first.metadata()
71    }) {
72        return Err(VisionError::ShapeMismatch(
73            "exposure inputs must share dimensions and metadata".into(),
74        ));
75    }
76    let mut output = Vec::with_capacity(first.width() * first.height() * 3);
77    let sigma2 = 2.0 * options.well_exposed_sigma.powi(2);
78    for y in 0..first.height() {
79        for x in 0..first.width() {
80            let mut weighted = [0.0; 3];
81            let mut total = 0.0;
82            for image in inputs {
83                let pixel = image.get(x, y).expect("validated coordinates");
84                let luminance = (0.2126 * f64::from(pixel[0])
85                    + 0.7152 * f64::from(pixel[1])
86                    + 0.0722 * f64::from(pixel[2]))
87                    / 255.0;
88                let weight = (-(luminance - 0.5).powi(2) / sigma2).exp() + options.weight_floor;
89                for channel in 0..3 {
90                    weighted[channel] += f64::from(pixel[channel]) * weight;
91                }
92                total += weight;
93            }
94            output.extend(weighted.map(|value| (value / total).round().clamp(0.0, 255.0) as u8));
95        }
96    }
97    Ok(Image::try_new_with_metadata(first.width(), first.height(), output, first.metadata())?)
98}
99
100/// Bounded pairwise panorama settings.
101#[derive(Clone, Copy, Debug, PartialEq, Eq)]
102pub struct PanoramaOptions {
103    /// Hard ceiling for output pixels before allocation.
104    pub max_output_pixels: usize,
105}
106
107impl Default for PanoramaOptions {
108    fn default() -> Self {
109        Self { max_output_pixels: 64 * 1024 * 1024 }
110    }
111}
112
113/// Pairwise panorama and its integer world-coordinate origin.
114#[derive(Clone, Debug, PartialEq)]
115pub struct Panorama {
116    image: Image<u8, 3>,
117    origin_x: i32,
118    origin_y: i32,
119}
120
121impl Panorama {
122    /// Returns the blended RGB canvas.
123    pub const fn image(&self) -> &Image<u8, 3> {
124        &self.image
125    }
126    /// Returns the right-image world x coordinate represented by canvas x=0.
127    pub const fn origin_x(&self) -> i32 {
128        self.origin_x
129    }
130    /// Returns the right-image world y coordinate represented by canvas y=0.
131    pub const fn origin_y(&self) -> i32 {
132        self.origin_y
133    }
134}
135
136/// Estimates a source-to-target homography and stitches the pair.
137pub fn estimate_and_stitch_panorama(
138    source: ImageView<'_, u8, 3>,
139    target: ImageView<'_, u8, 3>,
140    correspondences: &[PointCorrespondence2],
141    robust: RobustEstimationOptions,
142    options: PanoramaOptions,
143) -> VisionResult<Panorama> {
144    let estimate = estimate_homography_ransac(correspondences, robust)?;
145    stitch_panorama_pair(
146        source,
147        target,
148        PerspectiveTransform { matrix: estimate.model().matrix().m },
149        options,
150    )
151}
152
153/// Warps `source` into `target` coordinates and feather-blends their overlap.
154pub fn stitch_panorama_pair(
155    source: ImageView<'_, u8, 3>,
156    target: ImageView<'_, u8, 3>,
157    source_to_target: PerspectiveTransform,
158    options: PanoramaOptions,
159) -> VisionResult<Panorama> {
160    if source.metadata() != target.metadata() {
161        return Err(VisionError::ShapeMismatch("panorama images must share color metadata".into()));
162    }
163    if source.width() == 0 || source.height() == 0 || target.width() == 0 || target.height() == 0 {
164        return Err(VisionError::InvalidDimensions("panorama inputs must be non-empty".into()));
165    }
166    if options.max_output_pixels == 0 {
167        return Err(VisionError::InvalidParameter("panorama pixel budget must be positive".into()));
168    }
169    let mut corners = vec![
170        (0.0, 0.0),
171        ((target.width() - 1) as f64, 0.0),
172        (0.0, (target.height() - 1) as f64),
173        ((target.width() - 1) as f64, (target.height() - 1) as f64),
174    ];
175    for &(x, y) in &[
176        (0.0, 0.0),
177        ((source.width() - 1) as f64, 0.0),
178        (0.0, (source.height() - 1) as f64),
179        ((source.width() - 1) as f64, (source.height() - 1) as f64),
180    ] {
181        corners.push(source_to_target.map_point(x, y).ok_or(VisionError::SingularTransform)?);
182    }
183    let min_x = corners.iter().map(|p| p.0).fold(f64::INFINITY, f64::min).floor() as i32;
184    let min_y = corners.iter().map(|p| p.1).fold(f64::INFINITY, f64::min).floor() as i32;
185    let max_x = corners.iter().map(|p| p.0).fold(f64::NEG_INFINITY, f64::max).ceil() as i32;
186    let max_y = corners.iter().map(|p| p.1).fold(f64::NEG_INFINITY, f64::max).ceil() as i32;
187    let width = usize::try_from(max_x - min_x + 1).map_err(|_| {
188        VisionError::InvalidDimensions("panorama width is not representable".into())
189    })?;
190    let height = usize::try_from(max_y - min_y + 1).map_err(|_| {
191        VisionError::InvalidDimensions("panorama height is not representable".into())
192    })?;
193    if width.checked_mul(height).filter(|&pixels| pixels <= options.max_output_pixels).is_none() {
194        return Err(VisionError::InvalidDimensions("panorama exceeds output pixel budget".into()));
195    }
196    let inverse = source_to_target.inverse()?;
197    let mut data = Vec::with_capacity(width * height * 3);
198    for canvas_y in 0..height {
199        for canvas_x in 0..width {
200            let world_x = f64::from(min_x) + canvas_x as f64;
201            let world_y = f64::from(min_y) + canvas_y as f64;
202            let target_sample = sample_inside(target, world_x, world_y);
203            let source_sample =
204                inverse.map_point(world_x, world_y).and_then(|(x, y)| sample_inside(source, x, y));
205            let pixel = match (source_sample, target_sample) {
206                (Some((left, lw)), Some((right, rw))) => {
207                    let total = lw + rw;
208                    std::array::from_fn(|channel| {
209                        ((f64::from(left[channel]) * lw + f64::from(right[channel]) * rw) / total)
210                            .round()
211                            .clamp(0.0, 255.0) as u8
212                    })
213                }
214                (Some((pixel, _)), None) | (None, Some((pixel, _))) => pixel,
215                (None, None) => [0; 3],
216            };
217            data.extend_from_slice(&pixel);
218        }
219    }
220    Ok(Panorama {
221        image: Image::try_new_with_metadata(width, height, data, target.metadata())?,
222        origin_x: min_x,
223        origin_y: min_y,
224    })
225}
226
227fn sample_inside(image: ImageView<'_, u8, 3>, x: f64, y: f64) -> Option<([u8; 3], f64)> {
228    if x < 0.0 || y < 0.0 || x > (image.width() - 1) as f64 || y > (image.height() - 1) as f64 {
229        return None;
230    }
231    let x0 = x.floor() as usize;
232    let y0 = y.floor() as usize;
233    let x1 = (x0 + 1).min(image.width() - 1);
234    let y1 = (y0 + 1).min(image.height() - 1);
235    let wx = x - x0 as f64;
236    let wy = y - y0 as f64;
237    let pixel = std::array::from_fn(|channel| {
238        let p00 = f64::from(image.get(x0, y0).unwrap()[channel]);
239        let p10 = f64::from(image.get(x1, y0).unwrap()[channel]);
240        let p01 = f64::from(image.get(x0, y1).unwrap()[channel]);
241        let p11 = f64::from(image.get(x1, y1).unwrap()[channel]);
242        let top = p00 * (1.0 - wx) + p10 * wx;
243        let bottom = p01 * (1.0 - wx) + p11 * wx;
244        (top * (1.0 - wy) + bottom * wy).round().clamp(0.0, 255.0) as u8
245    });
246    let edge =
247        x.min(y).min((image.width() - 1) as f64 - x).min((image.height() - 1) as f64 - y) + 1.0;
248    Some((pixel, edge.max(1e-6)))
249}
250
251#[cfg(test)]
252mod tests {
253    use super::{
254        fuse_exposures, gray_world_white_balance, stitch_panorama_pair, ExposureFusionOptions,
255        PanoramaOptions,
256    };
257    use crate::PerspectiveTransform;
258    use spatialrust_image::Image;
259
260    #[test]
261    fn white_balance_equalizes_channel_means() {
262        let image = Image::try_new(2, 1, vec![40, 80, 120, 20, 40, 60]).unwrap();
263        let balanced = gray_world_white_balance(image.view()).unwrap();
264        let sums = [
265            balanced.as_slice()[0] as u16 + balanced.as_slice()[3] as u16,
266            balanced.as_slice()[1] as u16 + balanced.as_slice()[4] as u16,
267            balanced.as_slice()[2] as u16 + balanced.as_slice()[5] as u16,
268        ];
269        assert!(sums.iter().max().unwrap() - sums.iter().min().unwrap() <= 1);
270    }
271
272    #[test]
273    fn exposure_fusion_prefers_middle_gray() {
274        let dark = Image::from_pixel(1, 1, [10, 10, 10]).unwrap();
275        let middle = Image::from_pixel(1, 1, [128, 128, 128]).unwrap();
276        let fused = fuse_exposures(&[dark.view(), middle.view()], ExposureFusionOptions::default())
277            .unwrap();
278        assert!(fused.as_slice()[0] >= 120);
279    }
280
281    #[test]
282    fn translated_pair_expands_canvas_and_preserves_sides() {
283        let source = Image::from_pixel(3, 2, [200, 0, 0]).unwrap();
284        let target = Image::from_pixel(3, 2, [0, 0, 200]).unwrap();
285        let panorama = stitch_panorama_pair(
286            source.view(),
287            target.view(),
288            PerspectiveTransform { matrix: [[1.0, 0.0, 2.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]] },
289            PanoramaOptions::default(),
290        )
291        .unwrap();
292        assert_eq!((panorama.image().width(), panorama.image().height()), (5, 2));
293        assert_eq!(panorama.image().get(0, 0).unwrap(), &[0, 0, 200]);
294        assert_eq!(panorama.image().get(4, 0).unwrap(), &[200, 0, 0]);
295        assert!(stitch_panorama_pair(
296            source.view(),
297            target.view(),
298            PerspectiveTransform::identity(),
299            PanoramaOptions { max_output_pixels: 5 },
300        )
301        .is_err());
302    }
303}