1use crate::color::srgb_channel_to_linear;
2use crate::score::clamp01;
3use crate::types::{CaptureQualityReport, MeasurementMode, QualityCheck};
4use image::{DynamicImage, GenericImageView, ImageReader};
5use std::fmt;
6use std::io::Cursor;
7use thiserror::Error;
8
9#[derive(Debug, Error)]
10pub enum ImageError {
11 #[error("image decode failed: {0}")]
12 DecodeFailed(String),
13 #[error("image exceeds decode pixel limit")]
14 DecodeLimitExceeded,
15 #[error("image has no usable opaque pixels")]
16 NoUsablePixels,
17 #[error("color profile transform failed")]
18 ColorProfileTransformFailed,
19}
20
21const MAX_DECODE_PIXELS: u64 = 24_000_000;
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum IccProfileState {
25 NotEmbedded,
26 EmbeddedUnsupported,
27 TransformSucceeded,
28 TransformFailed,
29}
30
31#[derive(Clone)]
32pub struct NormalizedImage {
33 pub width: u32,
34 pub height: u32,
35 pub pixels: Vec<LinearPixel>,
36 pub measurement_mode: MeasurementMode,
37 pub orientation_applied: bool,
38 pub icc_status: String,
39}
40
41impl fmt::Debug for NormalizedImage {
42 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
43 formatter
44 .debug_struct("NormalizedImage")
45 .field("dimensions", &(self.width, self.height))
46 .field("pixel_count", &self.pixels.len())
47 .field("measurement_mode", &self.measurement_mode)
48 .field("orientation_applied", &self.orientation_applied)
49 .field("icc_status", &self.icc_status)
50 .field("metadata_retained", &false)
51 .finish()
52 }
53}
54
55#[derive(Clone, Copy)]
56pub struct LinearPixel {
57 pub r: f32,
58 pub g: f32,
59 pub b: f32,
60 pub y: f32,
61 pub opaque: bool,
62}
63
64impl fmt::Debug for LinearPixel {
65 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
66 formatter
67 .debug_struct("LinearPixel")
68 .field("channels", &"[REDACTED]")
69 .field("opaque", &self.opaque)
70 .finish()
71 }
72}
73
74pub fn decode_image(
75 bytes: &[u8],
76 allow_apparent_color_fallback: bool,
77) -> Result<NormalizedImage, ImageError> {
78 if bytes.is_empty() {
79 return Err(ImageError::NoUsablePixels);
80 }
81 let icc_profile_state = detect_encoded_icc_profile_state(bytes, allow_apparent_color_fallback);
82 if let Some((width, height)) = png_header_dimensions(bytes) {
83 reject_oversized_dimensions(width, height)?;
84 }
85 let reader = ImageReader::new(Cursor::new(bytes))
86 .with_guessed_format()
87 .map_err(|err| ImageError::DecodeFailed(err.to_string()))?;
88 let (width, height) = reader
89 .into_dimensions()
90 .map_err(|err| ImageError::DecodeFailed(err.to_string()))?;
91 reject_oversized_dimensions(width, height)?;
92 let reader = ImageReader::new(Cursor::new(bytes))
93 .with_guessed_format()
94 .map_err(|err| ImageError::DecodeFailed(err.to_string()))?;
95 let image = reader
96 .decode()
97 .map_err(|err| ImageError::DecodeFailed(err.to_string()))?;
98 normalize_dynamic_image_with_profile_state(
99 &image,
100 icc_profile_state,
101 allow_apparent_color_fallback,
102 )
103}
104
105fn png_header_dimensions(bytes: &[u8]) -> Option<(u32, u32)> {
106 if bytes.len() < 24 {
107 return None;
108 }
109 if &bytes[0..8] != b"\x89PNG\r\n\x1a\n" || &bytes[12..16] != b"IHDR" {
110 return None;
111 }
112 let width = u32::from_be_bytes(bytes[16..20].try_into().ok()?);
113 let height = u32::from_be_bytes(bytes[20..24].try_into().ok()?);
114 Some((width, height))
115}
116
117pub fn normalize_dynamic_image(
118 image: &DynamicImage,
119 allow_apparent_color_fallback: bool,
120) -> Result<NormalizedImage, ImageError> {
121 normalize_dynamic_image_with_profile_state(
122 image,
123 IccProfileState::NotEmbedded,
124 allow_apparent_color_fallback,
125 )
126}
127
128pub fn normalize_dynamic_image_with_profile_state(
129 image: &DynamicImage,
130 icc_profile_state: IccProfileState,
131 allow_apparent_color_fallback: bool,
132) -> Result<NormalizedImage, ImageError> {
133 let (width, height) = image.dimensions();
134 reject_oversized_dimensions(width, height)?;
135 let (measurement_mode, icc_status) =
136 resolve_measurement_mode(icc_profile_state, allow_apparent_color_fallback)?;
137 let rgba = image.to_rgba8();
138 let mut pixels = Vec::with_capacity((width * height) as usize);
139 let mut opaque_count = 0usize;
140 for pixel in rgba.pixels() {
141 let [r8, g8, b8, a8] = pixel.0;
142 let opaque = a8 > 0;
143 if opaque {
144 opaque_count += 1;
145 }
146 let r = f32::from(r8) / 255.0;
147 let g = f32::from(g8) / 255.0;
148 let b = f32::from(b8) / 255.0;
149 let y = linear_luminance(r, g, b);
150 pixels.push(LinearPixel { r, g, b, y, opaque });
151 }
152 if opaque_count == 0 {
153 return Err(ImageError::NoUsablePixels);
154 }
155 Ok(NormalizedImage {
156 width,
157 height,
158 pixels,
159 measurement_mode,
160 orientation_applied: false,
161 icc_status: icc_status.to_string(),
162 })
163}
164
165pub fn resolve_measurement_mode(
166 icc_profile_state: IccProfileState,
167 allow_apparent_color_fallback: bool,
168) -> Result<(MeasurementMode, &'static str), ImageError> {
169 match icc_profile_state {
170 IccProfileState::NotEmbedded => Ok((MeasurementMode::SrgbAssumed, "srgb_assumed")),
171 IccProfileState::EmbeddedUnsupported => Ok((
172 MeasurementMode::ApparentColorProfileUnsupported,
173 "embedded_icc_profile_unsupported",
174 )),
175 IccProfileState::TransformSucceeded => {
176 #[cfg(feature = "icc-lcms2")]
177 {
178 Ok((MeasurementMode::IccNormalized, "icc_normalized"))
179 }
180 #[cfg(not(feature = "icc-lcms2"))]
181 {
182 Ok((
183 MeasurementMode::ApparentColorProfileUnsupported,
184 "embedded_icc_profile_unsupported",
185 ))
186 }
187 }
188 IccProfileState::TransformFailed if allow_apparent_color_fallback => Ok((
189 MeasurementMode::AllowedApparentFallback,
190 "icc_transform_failed_apparent_fallback",
191 )),
192 IccProfileState::TransformFailed => Err(ImageError::ColorProfileTransformFailed),
193 }
194}
195
196pub fn detect_encoded_icc_profile_state(
197 bytes: &[u8],
198 allow_apparent_color_fallback: bool,
199) -> IccProfileState {
200 if !has_encoded_icc_profile(bytes) {
201 return IccProfileState::NotEmbedded;
202 }
203 if allow_apparent_color_fallback {
204 return IccProfileState::TransformFailed;
205 }
206 IccProfileState::EmbeddedUnsupported
207}
208
209fn has_encoded_icc_profile(bytes: &[u8]) -> bool {
210 has_png_iccp_chunk(bytes) || has_jpeg_icc_profile_segment(bytes)
211}
212
213fn has_png_iccp_chunk(bytes: &[u8]) -> bool {
214 if bytes.len() < 16 || &bytes[0..8] != b"\x89PNG\r\n\x1a\n" {
215 return false;
216 }
217 let mut offset = 8usize;
218 while offset + 12 <= bytes.len() {
219 let length = u32::from_be_bytes(match bytes[offset..offset + 4].try_into() {
220 Ok(value) => value,
221 Err(_) => return false,
222 }) as usize;
223 let chunk_type = &bytes[offset + 4..offset + 8];
224 if chunk_type == b"iCCP" {
225 return true;
226 }
227 let Some(next_offset) = offset
228 .checked_add(12)
229 .and_then(|base| base.checked_add(length))
230 else {
231 return false;
232 };
233 if next_offset > bytes.len() {
234 return false;
235 }
236 offset = next_offset;
237 }
238 false
239}
240
241fn has_jpeg_icc_profile_segment(bytes: &[u8]) -> bool {
242 if bytes.len() < 4 || bytes[0] != 0xff || bytes[1] != 0xd8 {
243 return false;
244 }
245 let mut offset = 2usize;
246 while offset + 4 <= bytes.len() {
247 if bytes[offset] != 0xff {
248 return false;
249 }
250 let marker = bytes[offset + 1];
251 if marker == 0xda || marker == 0xd9 {
252 return false;
253 }
254 let segment_length = u16::from_be_bytes([bytes[offset + 2], bytes[offset + 3]]) as usize;
255 if segment_length < 2 || offset + 2 + segment_length > bytes.len() {
256 return false;
257 }
258 let segment_data = &bytes[offset + 4..offset + 2 + segment_length];
259 if marker == 0xe2 && segment_data.starts_with(b"ICC_PROFILE\0") {
260 return true;
261 }
262 offset += 2 + segment_length;
263 }
264 false
265}
266
267fn reject_oversized_dimensions(width: u32, height: u32) -> Result<(), ImageError> {
268 let pixel_count = u64::from(width) * u64::from(height);
269 if pixel_count <= MAX_DECODE_PIXELS {
270 return Ok(());
271 }
272 Err(ImageError::DecodeLimitExceeded)
273}
274
275#[must_use]
276pub fn linear_luminance(r: f32, g: f32, b: f32) -> f32 {
277 0.2126 * srgb_channel_to_linear(r)
278 + 0.7152 * srgb_channel_to_linear(g)
279 + 0.0722 * srgb_channel_to_linear(b)
280}
281
282pub fn capture_quality_report(image: &NormalizedImage) -> Result<CaptureQualityReport, ImageError> {
283 let valid_pixels: Vec<_> = image
284 .pixels
285 .iter()
286 .copied()
287 .filter(|pixel| pixel.opaque)
288 .collect();
289 if valid_pixels.is_empty() {
290 return Err(ImageError::NoUsablePixels);
291 }
292 let valid_count = valid_pixels.len() as f32;
293 let over_clip_fraction = valid_pixels
294 .iter()
295 .filter(|pixel| pixel.r.max(pixel.g).max(pixel.b) >= 250.0 / 255.0 || pixel.y >= 0.98)
296 .count() as f32
297 / valid_count;
298 let under_clip_fraction =
299 valid_pixels.iter().filter(|pixel| pixel.y <= 0.02).count() as f32 / valid_count;
300 let white_balance = white_balance_check(&valid_pixels);
301 let blur = blur_check(image);
302 let shadow = shadow_check(image, &valid_pixels);
303 Ok(CaptureQualityReport {
304 decode_status: "decoded".to_string(),
305 orientation_applied: image.orientation_applied,
306 dimensions: (image.width, image.height),
307 icc_status: image.icc_status.clone(),
308 metadata_retained: false,
309 over_clip_fraction,
310 under_clip_fraction,
311 white_balance,
312 blur,
313 shadow,
314 face_angle: QualityCheck::NotMeasured {
315 reason: "deferred".to_string(),
316 deduction: 0.0,
317 },
318 filters_or_makeup: QualityCheck::NotMeasured {
319 reason: "deferred".to_string(),
320 deduction: 0.0,
321 },
322 occlusion: QualityCheck::NotMeasured {
323 reason: "deferred".to_string(),
324 deduction: 0.0,
325 },
326 calibration_card: QualityCheck::NotMeasured {
327 reason: "deferred".to_string(),
328 deduction: 0.0,
329 },
330 })
331}
332
333fn white_balance_check(valid_pixels: &[LinearPixel]) -> QualityCheck<f32> {
334 let pixels: Vec<_> = valid_pixels
335 .iter()
336 .filter(|pixel| pixel.r.max(pixel.g).max(pixel.b) < 250.0 / 255.0 && pixel.y > 0.02)
337 .collect();
338 if pixels.len() < 64 {
339 return QualityCheck::NotMeasured {
340 reason: "insufficient_non_clipped_pixels".to_string(),
341 deduction: 0.0,
342 };
343 }
344 let (sum_r, sum_g, sum_b) = pixels.iter().fold((0.0, 0.0, 0.0), |acc, pixel| {
345 (acc.0 + pixel.r, acc.1 + pixel.g, acc.2 + pixel.b)
346 });
347 let count = pixels.len() as f32;
348 let (mr, mg, mb) = (sum_r / count, sum_g / count, sum_b / count);
349 let gray = (mr + mg + mb) / 3.0;
350 let imbalance = clamp01(
351 ((mr - gray)
352 .abs()
353 .max((mg - gray).abs())
354 .max((mb - gray).abs()))
355 / gray.max(1e-6),
356 );
357 QualityCheck::Measured {
358 value: imbalance,
359 deduction: crate::score::white_balance_deduction(imbalance),
360 }
361}
362
363fn blur_check(image: &NormalizedImage) -> QualityCheck<f32> {
364 if image.width < 3 || image.height < 3 {
365 return QualityCheck::NotMeasured {
366 reason: "image_too_small_for_sobel".to_string(),
367 deduction: 0.0,
368 };
369 }
370 let mut energies = Vec::new();
371 let idx = |x: u32, y: u32| -> usize { (y * image.width + x) as usize };
372 for y in 1..image.height - 1 {
373 for x in 1..image.width - 1 {
374 if !image.pixels[idx(x, y)].opaque {
375 continue;
376 }
377 let p = |dx: i32, dy: i32| {
378 image.pixels[idx((x as i32 + dx) as u32, (y as i32 + dy) as u32)].y
379 };
380 let gx = -p(-1, -1) + p(1, -1) - 2.0 * p(-1, 0) + 2.0 * p(1, 0) - p(-1, 1) + p(1, 1);
381 let gy = -p(-1, -1) - 2.0 * p(0, -1) - p(1, -1) + p(-1, 1) + 2.0 * p(0, 1) + p(1, 1);
382 energies.push(gx.mul_add(gx, gy * gy));
383 }
384 }
385 if energies.is_empty() {
386 return QualityCheck::NotMeasured {
387 reason: "insufficient_sobel_pixels".to_string(),
388 deduction: 0.0,
389 };
390 }
391 let mean_energy = energies.iter().sum::<f32>() / energies.len() as f32;
392 let blur_score = clamp01(mean_energy / 0.25);
393 QualityCheck::Measured {
394 value: blur_score,
395 deduction: crate::score::blur_deduction(blur_score),
396 }
397}
398
399fn shadow_check(image: &NormalizedImage, valid_pixels: &[LinearPixel]) -> QualityCheck<f32> {
400 let mut tile_medians = Vec::new();
401 let tile_width = image.width.div_ceil(4);
402 let tile_height = image.height.div_ceil(4);
403 let idx = |x: u32, y: u32| -> usize { (y * image.width + x) as usize };
404 for ty in 0..4 {
405 for tx in 0..4 {
406 let x0 = tx * tile_width;
407 let y0 = ty * tile_height;
408 let x1 = ((tx + 1) * tile_width).min(image.width);
409 let y1 = ((ty + 1) * tile_height).min(image.height);
410 if x0 >= x1 || y0 >= y1 {
411 continue;
412 }
413 let area = ((x1 - x0) * (y1 - y0)).max(1) as f32;
414 let mut luminance = Vec::new();
415 for y in y0..y1 {
416 for x in x0..x1 {
417 let pixel = image.pixels[idx(x, y)];
418 if pixel.opaque {
419 luminance.push(pixel.y);
420 }
421 }
422 }
423 if luminance.len() >= 64 && luminance.len() as f32 / area >= 0.10 {
424 tile_medians.push(percentile(&mut luminance, 0.5));
425 }
426 }
427 }
428 if tile_medians.len() < 4 {
429 return QualityCheck::NotMeasured {
430 reason: "insufficient_luminance_tiles".to_string(),
431 deduction: 0.0,
432 };
433 }
434 let mut global: Vec<_> = valid_pixels.iter().map(|pixel| pixel.y).collect();
435 let global_median = percentile(&mut global, 0.5);
436 let p90 = percentile(&mut tile_medians.clone(), 0.9);
437 let p10 = percentile(&mut tile_medians, 0.1);
438 let unevenness = clamp01((p90 - p10) / global_median.max(0.05));
439 QualityCheck::Measured {
440 value: unevenness,
441 deduction: crate::score::shadow_deduction(unevenness),
442 }
443}
444
445fn percentile(values: &mut [f32], p: f32) -> f32 {
446 values.sort_by(|a, b| a.total_cmp(b));
447 let index = ((values.len() - 1) as f32 * p).round() as usize;
448 values[index]
449}