use crate::color::srgb_channel_to_linear;
use crate::score::clamp01;
use crate::types::{CaptureQualityReport, MeasurementMode, QualityCheck};
use image::{DynamicImage, GenericImageView, ImageReader};
use std::fmt;
use std::io::Cursor;
use thiserror::Error;
#[derive(Debug, Error)]
pub enum ImageError {
#[error("image decode failed: {0}")]
DecodeFailed(String),
#[error("image exceeds decode pixel limit")]
DecodeLimitExceeded,
#[error("image has no usable opaque pixels")]
NoUsablePixels,
#[error("color profile transform failed")]
ColorProfileTransformFailed,
}
const MAX_DECODE_PIXELS: u64 = 24_000_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IccProfileState {
NotEmbedded,
EmbeddedUnsupported,
TransformSucceeded,
TransformFailed,
}
#[derive(Clone)]
pub struct NormalizedImage {
pub width: u32,
pub height: u32,
pub pixels: Vec<LinearPixel>,
pub measurement_mode: MeasurementMode,
pub orientation_applied: bool,
pub icc_status: String,
}
impl fmt::Debug for NormalizedImage {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("NormalizedImage")
.field("dimensions", &(self.width, self.height))
.field("pixel_count", &self.pixels.len())
.field("measurement_mode", &self.measurement_mode)
.field("orientation_applied", &self.orientation_applied)
.field("icc_status", &self.icc_status)
.field("metadata_retained", &false)
.finish()
}
}
#[derive(Clone, Copy)]
pub struct LinearPixel {
pub r: f32,
pub g: f32,
pub b: f32,
pub y: f32,
pub opaque: bool,
}
impl fmt::Debug for LinearPixel {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("LinearPixel")
.field("channels", &"[REDACTED]")
.field("opaque", &self.opaque)
.finish()
}
}
pub fn decode_image(
bytes: &[u8],
allow_apparent_color_fallback: bool,
) -> Result<NormalizedImage, ImageError> {
if bytes.is_empty() {
return Err(ImageError::NoUsablePixels);
}
let icc_profile_state = detect_encoded_icc_profile_state(bytes, allow_apparent_color_fallback);
if let Some((width, height)) = png_header_dimensions(bytes) {
reject_oversized_dimensions(width, height)?;
}
let reader = ImageReader::new(Cursor::new(bytes))
.with_guessed_format()
.map_err(|err| ImageError::DecodeFailed(err.to_string()))?;
let (width, height) = reader
.into_dimensions()
.map_err(|err| ImageError::DecodeFailed(err.to_string()))?;
reject_oversized_dimensions(width, height)?;
let reader = ImageReader::new(Cursor::new(bytes))
.with_guessed_format()
.map_err(|err| ImageError::DecodeFailed(err.to_string()))?;
let image = reader
.decode()
.map_err(|err| ImageError::DecodeFailed(err.to_string()))?;
normalize_dynamic_image_with_profile_state(
&image,
icc_profile_state,
allow_apparent_color_fallback,
)
}
fn png_header_dimensions(bytes: &[u8]) -> Option<(u32, u32)> {
if bytes.len() < 24 {
return None;
}
if &bytes[0..8] != b"\x89PNG\r\n\x1a\n" || &bytes[12..16] != b"IHDR" {
return None;
}
let width = u32::from_be_bytes(bytes[16..20].try_into().ok()?);
let height = u32::from_be_bytes(bytes[20..24].try_into().ok()?);
Some((width, height))
}
pub fn normalize_dynamic_image(
image: &DynamicImage,
allow_apparent_color_fallback: bool,
) -> Result<NormalizedImage, ImageError> {
normalize_dynamic_image_with_profile_state(
image,
IccProfileState::NotEmbedded,
allow_apparent_color_fallback,
)
}
pub fn normalize_dynamic_image_with_profile_state(
image: &DynamicImage,
icc_profile_state: IccProfileState,
allow_apparent_color_fallback: bool,
) -> Result<NormalizedImage, ImageError> {
let (width, height) = image.dimensions();
reject_oversized_dimensions(width, height)?;
let (measurement_mode, icc_status) =
resolve_measurement_mode(icc_profile_state, allow_apparent_color_fallback)?;
let rgba = image.to_rgba8();
let mut pixels = Vec::with_capacity((width * height) as usize);
let mut opaque_count = 0usize;
for pixel in rgba.pixels() {
let [r8, g8, b8, a8] = pixel.0;
let opaque = a8 > 0;
if opaque {
opaque_count += 1;
}
let r = f32::from(r8) / 255.0;
let g = f32::from(g8) / 255.0;
let b = f32::from(b8) / 255.0;
let y = linear_luminance(r, g, b);
pixels.push(LinearPixel { r, g, b, y, opaque });
}
if opaque_count == 0 {
return Err(ImageError::NoUsablePixels);
}
Ok(NormalizedImage {
width,
height,
pixels,
measurement_mode,
orientation_applied: false,
icc_status: icc_status.to_string(),
})
}
pub fn resolve_measurement_mode(
icc_profile_state: IccProfileState,
allow_apparent_color_fallback: bool,
) -> Result<(MeasurementMode, &'static str), ImageError> {
match icc_profile_state {
IccProfileState::NotEmbedded => Ok((MeasurementMode::SrgbAssumed, "srgb_assumed")),
IccProfileState::EmbeddedUnsupported => Ok((
MeasurementMode::ApparentColorProfileUnsupported,
"embedded_icc_profile_unsupported",
)),
IccProfileState::TransformSucceeded => {
#[cfg(feature = "icc-lcms2")]
{
Ok((MeasurementMode::IccNormalized, "icc_normalized"))
}
#[cfg(not(feature = "icc-lcms2"))]
{
Ok((
MeasurementMode::ApparentColorProfileUnsupported,
"embedded_icc_profile_unsupported",
))
}
}
IccProfileState::TransformFailed if allow_apparent_color_fallback => Ok((
MeasurementMode::AllowedApparentFallback,
"icc_transform_failed_apparent_fallback",
)),
IccProfileState::TransformFailed => Err(ImageError::ColorProfileTransformFailed),
}
}
pub fn detect_encoded_icc_profile_state(
bytes: &[u8],
allow_apparent_color_fallback: bool,
) -> IccProfileState {
if !has_encoded_icc_profile(bytes) {
return IccProfileState::NotEmbedded;
}
if allow_apparent_color_fallback {
return IccProfileState::TransformFailed;
}
IccProfileState::EmbeddedUnsupported
}
fn has_encoded_icc_profile(bytes: &[u8]) -> bool {
has_png_iccp_chunk(bytes) || has_jpeg_icc_profile_segment(bytes)
}
fn has_png_iccp_chunk(bytes: &[u8]) -> bool {
if bytes.len() < 16 || &bytes[0..8] != b"\x89PNG\r\n\x1a\n" {
return false;
}
let mut offset = 8usize;
while offset + 12 <= bytes.len() {
let length = u32::from_be_bytes(match bytes[offset..offset + 4].try_into() {
Ok(value) => value,
Err(_) => return false,
}) as usize;
let chunk_type = &bytes[offset + 4..offset + 8];
if chunk_type == b"iCCP" {
return true;
}
let Some(next_offset) = offset
.checked_add(12)
.and_then(|base| base.checked_add(length))
else {
return false;
};
if next_offset > bytes.len() {
return false;
}
offset = next_offset;
}
false
}
fn has_jpeg_icc_profile_segment(bytes: &[u8]) -> bool {
if bytes.len() < 4 || bytes[0] != 0xff || bytes[1] != 0xd8 {
return false;
}
let mut offset = 2usize;
while offset + 4 <= bytes.len() {
if bytes[offset] != 0xff {
return false;
}
let marker = bytes[offset + 1];
if marker == 0xda || marker == 0xd9 {
return false;
}
let segment_length = u16::from_be_bytes([bytes[offset + 2], bytes[offset + 3]]) as usize;
if segment_length < 2 || offset + 2 + segment_length > bytes.len() {
return false;
}
let segment_data = &bytes[offset + 4..offset + 2 + segment_length];
if marker == 0xe2 && segment_data.starts_with(b"ICC_PROFILE\0") {
return true;
}
offset += 2 + segment_length;
}
false
}
fn reject_oversized_dimensions(width: u32, height: u32) -> Result<(), ImageError> {
let pixel_count = u64::from(width) * u64::from(height);
if pixel_count <= MAX_DECODE_PIXELS {
return Ok(());
}
Err(ImageError::DecodeLimitExceeded)
}
#[must_use]
pub fn linear_luminance(r: f32, g: f32, b: f32) -> f32 {
0.2126 * srgb_channel_to_linear(r)
+ 0.7152 * srgb_channel_to_linear(g)
+ 0.0722 * srgb_channel_to_linear(b)
}
pub fn capture_quality_report(image: &NormalizedImage) -> Result<CaptureQualityReport, ImageError> {
let valid_pixels: Vec<_> = image
.pixels
.iter()
.copied()
.filter(|pixel| pixel.opaque)
.collect();
if valid_pixels.is_empty() {
return Err(ImageError::NoUsablePixels);
}
let valid_count = valid_pixels.len() as f32;
let over_clip_fraction = valid_pixels
.iter()
.filter(|pixel| pixel.r.max(pixel.g).max(pixel.b) >= 250.0 / 255.0 || pixel.y >= 0.98)
.count() as f32
/ valid_count;
let under_clip_fraction =
valid_pixels.iter().filter(|pixel| pixel.y <= 0.02).count() as f32 / valid_count;
let white_balance = white_balance_check(&valid_pixels);
let blur = blur_check(image);
let shadow = shadow_check(image, &valid_pixels);
Ok(CaptureQualityReport {
decode_status: "decoded".to_string(),
orientation_applied: image.orientation_applied,
dimensions: (image.width, image.height),
icc_status: image.icc_status.clone(),
metadata_retained: false,
over_clip_fraction,
under_clip_fraction,
white_balance,
blur,
shadow,
face_angle: QualityCheck::NotMeasured {
reason: "deferred".to_string(),
deduction: 0.0,
},
filters_or_makeup: QualityCheck::NotMeasured {
reason: "deferred".to_string(),
deduction: 0.0,
},
occlusion: QualityCheck::NotMeasured {
reason: "deferred".to_string(),
deduction: 0.0,
},
calibration_card: QualityCheck::NotMeasured {
reason: "deferred".to_string(),
deduction: 0.0,
},
})
}
fn white_balance_check(valid_pixels: &[LinearPixel]) -> QualityCheck<f32> {
let pixels: Vec<_> = valid_pixels
.iter()
.filter(|pixel| pixel.r.max(pixel.g).max(pixel.b) < 250.0 / 255.0 && pixel.y > 0.02)
.collect();
if pixels.len() < 64 {
return QualityCheck::NotMeasured {
reason: "insufficient_non_clipped_pixels".to_string(),
deduction: 0.0,
};
}
let (sum_r, sum_g, sum_b) = pixels.iter().fold((0.0, 0.0, 0.0), |acc, pixel| {
(acc.0 + pixel.r, acc.1 + pixel.g, acc.2 + pixel.b)
});
let count = pixels.len() as f32;
let (mr, mg, mb) = (sum_r / count, sum_g / count, sum_b / count);
let gray = (mr + mg + mb) / 3.0;
let imbalance = clamp01(
((mr - gray)
.abs()
.max((mg - gray).abs())
.max((mb - gray).abs()))
/ gray.max(1e-6),
);
QualityCheck::Measured {
value: imbalance,
deduction: crate::score::white_balance_deduction(imbalance),
}
}
fn blur_check(image: &NormalizedImage) -> QualityCheck<f32> {
if image.width < 3 || image.height < 3 {
return QualityCheck::NotMeasured {
reason: "image_too_small_for_sobel".to_string(),
deduction: 0.0,
};
}
let mut energies = Vec::new();
let idx = |x: u32, y: u32| -> usize { (y * image.width + x) as usize };
for y in 1..image.height - 1 {
for x in 1..image.width - 1 {
if !image.pixels[idx(x, y)].opaque {
continue;
}
let p = |dx: i32, dy: i32| {
image.pixels[idx((x as i32 + dx) as u32, (y as i32 + dy) as u32)].y
};
let gx = -p(-1, -1) + p(1, -1) - 2.0 * p(-1, 0) + 2.0 * p(1, 0) - p(-1, 1) + p(1, 1);
let gy = -p(-1, -1) - 2.0 * p(0, -1) - p(1, -1) + p(-1, 1) + 2.0 * p(0, 1) + p(1, 1);
energies.push(gx.mul_add(gx, gy * gy));
}
}
if energies.is_empty() {
return QualityCheck::NotMeasured {
reason: "insufficient_sobel_pixels".to_string(),
deduction: 0.0,
};
}
let mean_energy = energies.iter().sum::<f32>() / energies.len() as f32;
let blur_score = clamp01(mean_energy / 0.25);
QualityCheck::Measured {
value: blur_score,
deduction: crate::score::blur_deduction(blur_score),
}
}
fn shadow_check(image: &NormalizedImage, valid_pixels: &[LinearPixel]) -> QualityCheck<f32> {
let mut tile_medians = Vec::new();
let tile_width = image.width.div_ceil(4);
let tile_height = image.height.div_ceil(4);
let idx = |x: u32, y: u32| -> usize { (y * image.width + x) as usize };
for ty in 0..4 {
for tx in 0..4 {
let x0 = tx * tile_width;
let y0 = ty * tile_height;
let x1 = ((tx + 1) * tile_width).min(image.width);
let y1 = ((ty + 1) * tile_height).min(image.height);
if x0 >= x1 || y0 >= y1 {
continue;
}
let area = ((x1 - x0) * (y1 - y0)).max(1) as f32;
let mut luminance = Vec::new();
for y in y0..y1 {
for x in x0..x1 {
let pixel = image.pixels[idx(x, y)];
if pixel.opaque {
luminance.push(pixel.y);
}
}
}
if luminance.len() >= 64 && luminance.len() as f32 / area >= 0.10 {
tile_medians.push(percentile(&mut luminance, 0.5));
}
}
}
if tile_medians.len() < 4 {
return QualityCheck::NotMeasured {
reason: "insufficient_luminance_tiles".to_string(),
deduction: 0.0,
};
}
let mut global: Vec<_> = valid_pixels.iter().map(|pixel| pixel.y).collect();
let global_median = percentile(&mut global, 0.5);
let p90 = percentile(&mut tile_medians.clone(), 0.9);
let p10 = percentile(&mut tile_medians, 0.1);
let unevenness = clamp01((p90 - p10) / global_median.max(0.05));
QualityCheck::Measured {
value: unevenness,
deduction: crate::score::shadow_deduction(unevenness),
}
}
fn percentile(values: &mut [f32], p: f32) -> f32 {
values.sort_by(|a, b| a.total_cmp(b));
let index = ((values.len() - 1) as f32 * p).round() as usize;
values[index]
}