1use spatialrust_image::{Image, ImageView};
4
5use crate::{
6 estimate_homography_ransac, PerspectiveTransform, PointCorrespondence2,
7 RobustEstimationOptions, VisionError, VisionResult,
8};
9
10pub 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#[derive(Clone, Copy, Debug, PartialEq)]
39pub struct ExposureFusionOptions {
40 pub well_exposed_sigma: f64,
42 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
52pub 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
102pub struct PanoramaOptions {
103 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#[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 pub const fn image(&self) -> &Image<u8, 3> {
124 &self.image
125 }
126 pub const fn origin_x(&self) -> i32 {
128 self.origin_x
129 }
130 pub const fn origin_y(&self) -> i32 {
132 self.origin_y
133 }
134}
135
136pub 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
153pub 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}