1use std::fs;
2use std::io::Cursor;
3use std::path::Path;
4
5use image::{DynamicImage, ImageBuffer, ImageFormat, Rgba, RgbaImage};
6
7use crate::error::Result;
8
9#[derive(Clone, Copy, Debug, PartialEq)]
10pub struct Color {
11 pub r: f32,
12 pub g: f32,
13 pub b: f32,
14}
15
16impl Color {
17 pub const fn rgb(r: f32, g: f32, b: f32) -> Self {
18 Self { r, g, b }
19 }
20
21 pub fn clamp01(self) -> Self {
22 Self {
23 r: self.r.clamp(0.0, 1.0),
24 g: self.g.clamp(0.0, 1.0),
25 b: self.b.clamp(0.0, 1.0),
26 }
27 }
28
29 pub fn lerp(self, other: Self, alpha: f32) -> Self {
30 let beta = 1.0 - alpha;
31 Self {
32 r: self.r * beta + other.r * alpha,
33 g: self.g * beta + other.g * alpha,
34 b: self.b * beta + other.b * alpha,
35 }
36 }
37
38 pub fn luma(self) -> f32 {
39 self.r * 0.2126 + self.g * 0.7152 + self.b * 0.0722
40 }
41
42 pub fn abs_diff(self, other: Self) -> f32 {
43 ((self.r - other.r).abs() + (self.g - other.g).abs() + (self.b - other.b).abs()) / 3.0
44 }
45}
46
47#[derive(Clone, Debug)]
48pub struct ImageFrame {
49 width: usize,
50 height: usize,
51 pixels: Vec<Color>,
52}
53
54impl ImageFrame {
55 pub fn new(width: usize, height: usize) -> Self {
56 Self {
57 width,
58 height,
59 pixels: vec![Color::rgb(0.0, 0.0, 0.0); width * height],
60 }
61 }
62
63 pub fn from_pixels(width: usize, height: usize, pixels: Vec<Color>) -> Self {
64 assert_eq!(pixels.len(), width * height);
65 Self {
66 width,
67 height,
68 pixels,
69 }
70 }
71
72 pub fn width(&self) -> usize {
73 self.width
74 }
75
76 pub fn height(&self) -> usize {
77 self.height
78 }
79
80 pub fn len(&self) -> usize {
81 self.pixels.len()
82 }
83
84 pub fn is_empty(&self) -> bool {
85 self.pixels.is_empty()
86 }
87
88 pub fn pixels(&self) -> &[Color] {
89 &self.pixels
90 }
91
92 pub fn get(&self, x: usize, y: usize) -> Color {
93 self.pixels[y * self.width + x]
94 }
95
96 pub fn set(&mut self, x: usize, y: usize, value: Color) {
97 self.pixels[y * self.width + x] = value;
98 }
99
100 pub fn sample_clamped(&self, x: i32, y: i32) -> Color {
101 let clamped_x = x.clamp(0, self.width as i32 - 1) as usize;
102 let clamped_y = y.clamp(0, self.height as i32 - 1) as usize;
103 self.get(clamped_x, clamped_y)
104 }
105
106 pub fn sample_bilinear_clamped(&self, x: f32, y: f32) -> Color {
107 let x0 = x.floor();
108 let y0 = y.floor();
109 let x1 = x0 + 1.0;
110 let y1 = y0 + 1.0;
111 let tx = (x - x0).clamp(0.0, 1.0);
112 let ty = (y - y0).clamp(0.0, 1.0);
113
114 let c00 = self.sample_clamped(x0 as i32, y0 as i32);
115 let c10 = self.sample_clamped(x1 as i32, y0 as i32);
116 let c01 = self.sample_clamped(x0 as i32, y1 as i32);
117 let c11 = self.sample_clamped(x1 as i32, y1 as i32);
118
119 let top = c00.lerp(c10, tx);
120 let bottom = c01.lerp(c11, tx);
121 top.lerp(bottom, ty)
122 }
123
124 pub fn to_rgba_image(&self) -> RgbaImage {
125 let width = self.width as u32;
126 let height = self.height as u32;
127 ImageBuffer::from_fn(width, height, |x, y| {
128 let color = self.get(x as usize, y as usize).clamp01();
129 Rgba([
130 (color.r * 255.0).round() as u8,
131 (color.g * 255.0).round() as u8,
132 (color.b * 255.0).round() as u8,
133 255,
134 ])
135 })
136 }
137
138 pub fn encode_png(&self) -> Result<Vec<u8>> {
139 let image = DynamicImage::ImageRgba8(self.to_rgba_image());
140 let mut cursor = Cursor::new(Vec::new());
141 image.write_to(&mut cursor, ImageFormat::Png)?;
142 Ok(cursor.into_inner())
143 }
144
145 pub fn save_png(&self, path: &Path) -> Result<()> {
146 if let Some(parent) = path.parent() {
147 fs::create_dir_all(parent)?;
148 }
149 self.to_rgba_image().save(path)?;
150 Ok(())
151 }
152
153 pub fn load_png(path: &Path) -> Result<Self> {
154 let image = image::open(path)?.to_rgba8();
155 let width = image.width() as usize;
156 let height = image.height() as usize;
157 let pixels = image
158 .pixels()
159 .map(|pixel| {
160 Color::rgb(
161 pixel[0] as f32 / 255.0,
162 pixel[1] as f32 / 255.0,
163 pixel[2] as f32 / 255.0,
164 )
165 })
166 .collect();
167 Ok(Self::from_pixels(width, height, pixels))
168 }
169
170 pub fn crop(&self, bbox: BoundingBox) -> Self {
171 let mut cropped = ImageFrame::new(bbox.width(), bbox.height());
172 for y in 0..bbox.height() {
173 for x in 0..bbox.width() {
174 cropped.set(x, y, self.get(bbox.min_x + x, bbox.min_y + y));
175 }
176 }
177 cropped
178 }
179}
180
181#[derive(Clone, Debug)]
182pub struct ScalarField {
183 width: usize,
184 height: usize,
185 values: Vec<f32>,
186}
187
188impl ScalarField {
189 pub fn new(width: usize, height: usize) -> Self {
190 Self {
191 width,
192 height,
193 values: vec![0.0; width * height],
194 }
195 }
196
197 pub fn from_values(width: usize, height: usize, values: Vec<f32>) -> Self {
198 assert_eq!(values.len(), width * height);
199 Self {
200 width,
201 height,
202 values,
203 }
204 }
205
206 pub fn width(&self) -> usize {
207 self.width
208 }
209
210 pub fn height(&self) -> usize {
211 self.height
212 }
213
214 pub fn values(&self) -> &[f32] {
215 &self.values
216 }
217
218 pub fn len(&self) -> usize {
219 self.values.len()
220 }
221
222 pub fn get(&self, x: usize, y: usize) -> f32 {
223 self.values[y * self.width + x]
224 }
225
226 pub fn set(&mut self, x: usize, y: usize, value: f32) {
227 self.values[y * self.width + x] = value;
228 }
229
230 pub fn mean(&self) -> f32 {
231 if self.values.is_empty() {
232 return 0.0;
233 }
234 self.values.iter().sum::<f32>() / self.values.len() as f32
235 }
236
237 pub fn mean_over_mask(&self, mask: &[bool]) -> f32 {
238 let mut sum = 0.0;
239 let mut count = 0usize;
240 for (value, include) in self.values.iter().zip(mask.iter().copied()) {
241 if include {
242 sum += *value;
243 count += 1;
244 }
245 }
246 if count == 0 {
247 0.0
248 } else {
249 sum / count as f32
250 }
251 }
252}
253
254#[derive(Clone, Copy, Debug, PartialEq, Eq)]
255pub struct BoundingBox {
256 pub min_x: usize,
257 pub min_y: usize,
258 pub max_x: usize,
259 pub max_y: usize,
260}
261
262impl BoundingBox {
263 pub fn width(self) -> usize {
264 self.max_x - self.min_x + 1
265 }
266
267 pub fn height(self) -> usize {
268 self.max_y - self.min_y + 1
269 }
270
271 pub fn expand(self, width: usize, height: usize, margin: usize) -> Self {
272 Self {
273 min_x: self.min_x.saturating_sub(margin),
274 min_y: self.min_y.saturating_sub(margin),
275 max_x: (self.max_x + margin).min(width.saturating_sub(1)),
276 max_y: (self.max_y + margin).min(height.saturating_sub(1)),
277 }
278 }
279}
280
281pub fn bounding_box_from_mask(mask: &[bool], width: usize, height: usize) -> Option<BoundingBox> {
282 let mut min_x = width;
283 let mut min_y = height;
284 let mut max_x = 0usize;
285 let mut max_y = 0usize;
286 let mut found = false;
287
288 for y in 0..height {
289 for x in 0..width {
290 if mask[y * width + x] {
291 min_x = min_x.min(x);
292 min_y = min_y.min(y);
293 max_x = max_x.max(x);
294 max_y = max_y.max(y);
295 found = true;
296 }
297 }
298 }
299
300 found.then_some(BoundingBox {
301 min_x,
302 min_y,
303 max_x,
304 max_y,
305 })
306}
307
308pub fn mean_abs_error(frame_a: &ImageFrame, frame_b: &ImageFrame) -> f32 {
309 let mut sum = 0.0;
310 for (pixel_a, pixel_b) in frame_a.pixels().iter().zip(frame_b.pixels()) {
311 sum += pixel_a.abs_diff(*pixel_b);
312 }
313 sum / frame_a.len() as f32
314}
315
316pub fn mean_abs_error_over_mask(frame_a: &ImageFrame, frame_b: &ImageFrame, mask: &[bool]) -> f32 {
317 let mut sum = 0.0;
318 let mut count = 0usize;
319 for ((pixel_a, pixel_b), include) in frame_a
320 .pixels()
321 .iter()
322 .zip(frame_b.pixels())
323 .zip(mask.iter().copied())
324 {
325 if include {
326 sum += pixel_a.abs_diff(*pixel_b);
327 count += 1;
328 }
329 }
330 if count == 0 {
331 0.0
332 } else {
333 sum / count as f32
334 }
335}
336
337pub fn save_scalar_field_png(
338 field: &ScalarField,
339 path: &Path,
340 mapper: impl Fn(f32) -> [u8; 4],
341) -> Result<()> {
342 if let Some(parent) = path.parent() {
343 fs::create_dir_all(parent)?;
344 }
345 let image = ImageBuffer::from_fn(field.width as u32, field.height as u32, |x, y| {
346 let rgba = mapper(field.get(x as usize, y as usize));
347 Rgba(rgba)
348 });
349 image.save(path)?;
350 Ok(())
351}