blazediff_milo/
lib.rs

1//! MILO: a learned perceptual image quality metric, run on the CPU.
2//!
3//! MILO (Çoğalan, Bemana, Myszkowski, Seidel and Groth, ACM TOG 2025) is a
4//! full-reference metric that models *visual masking*: an error hidden in
5//! texture counts for less than the same error on a flat surface. Where SSIM
6//! is a formula, MILO is a small convolutional network, 44,930 parameters,
7//! whose trained weights ship inside this crate. Given a reference and a
8//! distorted image it predicts a per-pixel visibility mask, weights the
9//! absolute error by it, and pools that into one number.
10//!
11//! This is a port of the authors' PyTorch implementation, held to its outputs
12//! by `tests/parity.rs`. Same pyramid, same network, same pooling; only the
13//! execution differs: a line-buffer pipeline so memory stays flat whatever
14//! the image size, lane-generic SIMD kernels picked at compile time, and row
15//! bands spread over threads with results that don't depend on the thread
16//! count.
17//!
18//! ```
19//! use blazediff_milo::{milo, MiloOptions, Rgba8};
20//!
21//! # fn main() -> Result<(), blazediff_milo::MiloError> {
22//! # let (width, height) = (64, 64);
23//! # let (rgba1, rgba2) = (vec![0u8; width * height * 4], vec![0u8; width * height * 4]);
24//! let outcome = milo(
25//!     Rgba8::new(&rgba1, width, height),
26//!     Rgba8::new(&rgba2, width, height),
27//!     &MiloOptions::default(),
28//! )?;
29//! println!("raw error {:.6}, MOS {:.3}", outcome.raw_error, outcome.mos);
30//! # Ok(())
31//! # }
32//! ```
33//!
34//! Two numbers come out. `raw_error` is the metric proper: the masked mean
35//! absolute error, exactly zero for identical images and growing with visible
36//! damage (the fixtures in this repo land between 0.0002 and 0.01). `mos` maps
37//! it onto KADID-10k's five-point mean opinion score through a learned
38//! calibration; note the calibration tops out around 4.35 rather than 5 for
39//! identical input, so threshold on `raw_error`, not `mos`.
40
41mod conv;
42mod error;
43mod map;
44mod metric;
45#[cfg(feature = "napi")]
46mod napi;
47mod network;
48mod pyramid;
49#[cfg(feature = "python")]
50mod python;
51mod simd;
52#[cfg(all(feature = "wasm", target_arch = "wasm32"))]
53mod wasm;
54mod weights;
55
56pub use error::MiloError;
57pub use map::render_map;
58
59/// The smallest side the reference implementation accepts: three halvings
60/// must leave at least two pixels for the seed mask.
61pub const MIN_SIDE: usize = 16;
62
63/// A borrowed RGBA8 image, 4 bytes per pixel, row-major. Alpha is ignored,
64/// as the reference converts to RGB.
65#[derive(Clone, Copy, Debug)]
66pub struct Rgba8<'a> {
67    pub data: &'a [u8],
68    pub width: usize,
69    pub height: usize,
70}
71
72impl<'a> Rgba8<'a> {
73    pub fn new(data: &'a [u8], width: usize, height: usize) -> Self {
74        Self {
75            data,
76            width,
77            height,
78        }
79    }
80
81    /// Reject a buffer that cannot hold `width * height` RGBA pixels.
82    fn validate(&self) -> Result<(), MiloError> {
83        let pixels = self.width * self.height;
84        if self.data.len() < pixels * 4 {
85            return Err(MiloError::Options(format!(
86                "image data is {} bytes, need {} for {}x{} RGBA",
87                self.data.len(),
88                pixels * 4,
89                self.width,
90                self.height
91            )));
92        }
93        Ok(())
94    }
95
96    /// Reject a pair the metric would otherwise index out of bounds or the
97    /// reference would refuse.
98    fn validate_pair(&self, other: Rgba8<'_>) -> Result<(), MiloError> {
99        if self.width != other.width || self.height != other.height {
100            return Err(MiloError::SizeMismatch {
101                img1_width: self.width as u32,
102                img1_height: self.height as u32,
103                img2_width: other.width as u32,
104                img2_height: other.height as u32,
105            });
106        }
107        if self.width < MIN_SIDE || self.height < MIN_SIDE {
108            return Err(MiloError::InputTooSmall {
109                width: self.width as u32,
110                height: self.height as u32,
111                minimum: MIN_SIDE as u32,
112            });
113        }
114        self.validate()?;
115        other.validate()
116    }
117}
118
119/// Knobs for a comparison.
120#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
121pub struct MiloOptions {
122    /// Threads to spread row bands over. `None` uses every core the OS
123    /// reports; `Some(1)` runs inline. The result is bit-identical either way.
124    /// Ignored on wasm32, which has no threads.
125    pub threads: Option<usize>,
126}
127
128/// A score and the maps it was pooled from.
129#[derive(Clone, Debug)]
130pub struct MiloOutcome {
131    /// Mean over pixels and channels of `mask * |reference - distorted|`.
132    /// Zero for identical images. This is the metric.
133    pub raw_error: f64,
134    /// `raw_error` on KADID-10k's 1..5 mean-opinion-score scale, via the
135    /// learned scaler. About 4.35 for identical images, lower with damage.
136    pub mos: f64,
137    /// The visibility mask, one value per pixel, row-major. Each pyramid
138    /// level adds a sigmoid in `(0, 1)`, so the range is `(0, 4)`.
139    pub mask: Vec<f32>,
140    /// Per-pixel perceived error, row-major, in `0..1`: the scaler applied to
141    /// each pixel's masked error, minus its value at zero. The reference's
142    /// `MILO_map`.
143    pub error_map: Vec<f32>,
144    pub width: usize,
145    pub height: usize,
146}
147
148/// Compare `distorted` against `reference`.
149///
150/// Both must be the same size, at least [`MIN_SIDE`] on each side.
151pub fn milo(
152    reference: Rgba8<'_>,
153    distorted: Rgba8<'_>,
154    options: &MiloOptions,
155) -> Result<MiloOutcome, MiloError> {
156    metric::run(reference, distorted, options)
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    fn rgba(width: usize, height: usize, value: u8) -> Vec<u8> {
164        let mut data = vec![0u8; width * height * 4];
165        for pixel in data.chunks_exact_mut(4) {
166            pixel.copy_from_slice(&[value, value, value, 255]);
167        }
168        data
169    }
170
171    #[test]
172    fn identical_images_score_zero() {
173        let data = rgba(20, 17, 128);
174        let outcome = milo(
175            Rgba8::new(&data, 20, 17),
176            Rgba8::new(&data, 20, 17),
177            &MiloOptions::default(),
178        )
179        .unwrap();
180        assert_eq!(outcome.raw_error, 0.0);
181        assert!(outcome.error_map.iter().all(|v| *v == 0.0));
182        assert_eq!(outcome.mask.len(), 20 * 17);
183        assert!((outcome.mos - 4.346).abs() < 0.01);
184    }
185
186    #[test]
187    fn a_visible_change_scores_above_zero() {
188        let a = rgba(32, 32, 30);
189        let mut b = a.clone();
190        for pixel in b.chunks_exact_mut(4).skip(300).take(200) {
191            pixel[..3].copy_from_slice(&[220, 220, 220]);
192        }
193        let outcome = milo(
194            Rgba8::new(&a, 32, 32),
195            Rgba8::new(&b, 32, 32),
196            &MiloOptions::default(),
197        )
198        .unwrap();
199        assert!(outcome.raw_error > 0.0);
200        assert!(outcome.mos < 4.3);
201        assert!(outcome.error_map.iter().any(|v| *v > 0.0));
202    }
203
204    #[test]
205    fn the_thread_count_does_not_change_the_answer() {
206        let a = rgba(40, 70, 90);
207        let mut b = a.clone();
208        for (i, pixel) in b.chunks_exact_mut(4).enumerate() {
209            pixel[0] = ((i * 31) % 255) as u8;
210        }
211        let one = milo(
212            Rgba8::new(&a, 40, 70),
213            Rgba8::new(&b, 40, 70),
214            &MiloOptions { threads: Some(1) },
215        )
216        .unwrap();
217        let four = milo(
218            Rgba8::new(&a, 40, 70),
219            Rgba8::new(&b, 40, 70),
220            &MiloOptions { threads: Some(4) },
221        )
222        .unwrap();
223        assert_eq!(one.raw_error, four.raw_error);
224        assert_eq!(one.mask, four.mask);
225        assert_eq!(one.error_map, four.error_map);
226    }
227
228    #[test]
229    fn alpha_is_ignored() {
230        let a = rgba(16, 16, 100);
231        let mut b = a.clone();
232        for pixel in b.chunks_exact_mut(4) {
233            pixel[3] = 0;
234        }
235        let outcome = milo(
236            Rgba8::new(&a, 16, 16),
237            Rgba8::new(&b, 16, 16),
238            &MiloOptions::default(),
239        )
240        .unwrap();
241        assert_eq!(outcome.raw_error, 0.0);
242    }
243
244    #[test]
245    fn mismatched_sizes_are_rejected() {
246        let a = rgba(16, 16, 0);
247        let b = rgba(16, 17, 0);
248        assert!(matches!(
249            milo(
250                Rgba8::new(&a, 16, 16),
251                Rgba8::new(&b, 16, 17),
252                &MiloOptions::default()
253            ),
254            Err(MiloError::SizeMismatch { .. })
255        ));
256    }
257
258    #[test]
259    fn inputs_below_the_reference_floor_are_rejected() {
260        let a = rgba(15, 40, 0);
261        assert!(matches!(
262            milo(
263                Rgba8::new(&a, 15, 40),
264                Rgba8::new(&a, 15, 40),
265                &MiloOptions::default()
266            ),
267            Err(MiloError::InputTooSmall { minimum: 16, .. })
268        ));
269    }
270
271    #[test]
272    fn a_short_buffer_is_rejected() {
273        let a = rgba(16, 16, 0);
274        assert!(matches!(
275            milo(
276                Rgba8::new(&a, 32, 16),
277                Rgba8::new(&a, 32, 16),
278                &MiloOptions::default()
279            ),
280            Err(MiloError::Options(_))
281        ));
282    }
283}
blazediff_milo/lib.rs

blazediff_milo/
lib.rs
