use crate::bitmap::Bitmap;
use crate::pixmap::Pixmap;
#[derive(Debug, Default, Clone, Copy, PartialEq)]
pub enum Binarization {
#[default]
Fixed,
Sauvola { window: u32, k: f32 },
}
#[derive(Debug, Clone, Copy)]
pub struct SegmentOptions {
pub threshold: u8,
pub bg_subsample: u32,
pub adaptive_bg_subsample: bool,
pub binarization: Binarization,
pub bg_inpaint: bool,
pub bg_diffuse: bool,
pub deskew: bool,
pub block_classify: bool,
}
impl Default for SegmentOptions {
fn default() -> Self {
Self {
threshold: 128,
bg_subsample: 12,
adaptive_bg_subsample: false,
binarization: Binarization::Fixed,
bg_inpaint: false,
bg_diffuse: false,
deskew: false,
block_classify: false,
}
}
}
impl SegmentOptions {
pub fn archival() -> Self {
Self {
bg_subsample: 6,
adaptive_bg_subsample: false,
..Self::default()
}
}
}
pub struct SegmentedPage {
pub mask: Bitmap,
pub bg: Pixmap,
}
#[derive(Debug, Clone, Copy, Default)]
struct ColorAccum {
r: u64,
g: u64,
b: u64,
n: u64,
}
impl ColorAccum {
fn add(&mut self, r: u8, g: u8, b: u8) {
self.r += u64::from(r);
self.g += u64::from(g);
self.b += u64::from(b);
self.n += 1;
}
fn color(self) -> Option<(u8, u8, u8)> {
if self.n == 0 {
return None;
}
Some((
(self.r / self.n) as u8,
(self.g / self.n) as u8,
(self.b / self.n) as u8,
))
}
}
fn choose_bg_subsample(rgba: &Pixmap, mask: &Bitmap, ceiling: u32) -> u32 {
let (w, h) = (rgba.width as usize, rgba.height as usize);
const CELL: usize = 12;
let cols = w.div_ceil(CELL);
let rows = h.div_ceil(CELL);
let mut spread_hi = 0usize;
let mut cells = 0usize;
let mut cy = 0usize;
while cy < rows {
for cx in 0..cols {
let x0 = cx * CELL;
let y0 = cy * CELL;
let x1 = (x0 + CELL).min(w);
let y1 = (y0 + CELL).min(h);
let mut lo = 255u8;
let mut hi = 0u8;
let mut seen = false;
for y in y0..y1 {
let row = &rgba.data[y * w * 4..(y + 1) * w * 4];
for x in x0..x1 {
if mask.get(x as u32, y as u32) {
continue;
}
let l = luminance(row[x * 4], row[x * 4 + 1], row[x * 4 + 2]);
lo = lo.min(l);
hi = hi.max(l);
seen = true;
}
}
if seen {
cells += 1;
if hi - lo > 24 {
spread_hi += 1;
}
}
}
cy += 2;
}
if cells == 0 {
return ceiling;
}
let detail_pct = spread_hi * 100 / cells;
if detail_pct < 5 {
ceiling
} else if detail_pct < 30 {
6.min(ceiling)
} else {
3.min(ceiling)
}
}
#[inline]
fn luminance(r: u8, g: u8, b: u8) -> u8 {
(((r as u32) * 306 + (g as u32) * 601 + (b as u32) * 117) >> 10) as u8
}
pub fn estimate_skew(rgba: &Pixmap) -> f32 {
let (w, h) = (rgba.width as usize, rgba.height as usize);
if w < 64 || h < 64 {
return 0.0;
}
let stride = 2usize;
let mut ink: Vec<(u32, u32)> = Vec::new();
for y in 0..h {
let row = &rgba.data[y * w * 4..(y * w + w) * 4];
for x in (0..w).step_by(stride) {
let p = &row[x * 4..x * 4 + 3];
if luminance(p[0], p[1], p[2]) < 128 {
ink.push((x as u32, y as u32));
}
}
}
if ink.len() < 256 {
return 0.0; }
let score = |deg: f32| -> f64 {
let t = deg.to_radians().tan();
let bins = h + (w as f32 * t.abs()) as usize + 2;
let mut hist = vec![0u32; bins];
let last = bins - 1;
let off = if t < 0.0 { w as f32 * -t } else { 0.0 };
for &(x, y) in &ink {
let r = (y as f32 + x as f32 * t + off) as usize;
hist[r.min(last)] += 1;
}
hist.windows(2)
.map(|p| {
let d = p[1] as f64 - p[0] as f64;
d * d
})
.sum()
};
let sweep = |centre: f32, half: f32, step: f32| -> f32 {
let mut best = f64::MIN;
let (mut first, mut last_a) = (centre, centre);
let mut a = centre - half;
while a <= centre + half + 1e-6 {
let sc = score(a);
if sc > best {
best = sc;
first = a;
last_a = a;
} else if sc == best {
last_a = a;
}
a += step;
}
(first + last_a) / 2.0
};
let c = sweep(0.0, 5.0, 0.5);
let c = sweep(c, 0.5, 0.1);
let c = sweep(c, 0.1, 0.02);
let step = 0.02f32;
let (sl, sc, sr) = (score(c - step), score(c), score(c + step));
let denom = sl - 2.0 * sc + sr;
if denom.abs() > f64::EPSILON {
let shift = 0.5 * (sl - sr) / denom;
c + step * (shift as f32).clamp(-1.0, 1.0)
} else {
c
}
}
fn rotate_small(src: &Pixmap, deg: f32) -> Pixmap {
let rad = deg.to_radians();
let (sn, cs) = rad.sin_cos();
let (w, h) = (src.width as i32, src.height as i32);
let (cx, cy) = (w as f32 / 2.0, h as f32 / 2.0);
let mut out = Pixmap::white(src.width, src.height);
for y in 0..h {
for x in 0..w {
let dx = x as f32 - cx;
let dy = y as f32 - cy;
let sx = cs * dx + sn * dy + cx;
let sy = -sn * dx + cs * dy + cy;
let x0f = sx.floor();
let y0f = sy.floor();
if x0f < 0.0 || y0f < 0.0 || x0f as i32 + 1 >= w || y0f as i32 + 1 >= h {
continue; }
let (fx, fy) = (sx - x0f, sy - y0f);
let (x0, y0) = (x0f as usize, y0f as usize);
let idx = |xx: usize, yy: usize| (yy * src.width as usize + xx) * 4;
let di = (y as usize * src.width as usize + x as usize) * 4;
for ch in 0..3 {
let p00 = src.data[idx(x0, y0) + ch] as f32;
let p10 = src.data[idx(x0 + 1, y0) + ch] as f32;
let p01 = src.data[idx(x0, y0 + 1) + ch] as f32;
let p11 = src.data[idx(x0 + 1, y0 + 1) + ch] as f32;
let v = p00 * (1.0 - fx) * (1.0 - fy)
+ p10 * fx * (1.0 - fy)
+ p01 * (1.0 - fx) * fy
+ p11 * fx * fy;
out.data[di + ch] = v.round().clamp(0.0, 255.0) as u8;
}
}
}
out
}
const CLASSIFY_BLOCK: usize = 32;
fn clear_photo_blocks(mask: &mut Bitmap, rgba: &Pixmap) {
let (w, h) = (rgba.width as usize, rgba.height as usize);
let bw = w.div_ceil(CLASSIFY_BLOCK);
let bh = h.div_ceil(CLASSIFY_BLOCK);
if bw == 0 || bh == 0 {
return;
}
let mut photo = vec![false; bw * bh];
for by in 0..bh {
for bx in 0..bw {
let x0 = bx * CLASSIFY_BLOCK;
let y0 = by * CLASSIFY_BLOCK;
let x1 = (x0 + CLASSIFY_BLOCK).min(w);
let y1 = (y0 + CLASSIFY_BLOCK).min(h);
let mut midtone = 0u32;
let mut sharp = 0u32;
let mut flips = 0u32;
let mut pairs = 0u32;
let mut n = 0u32;
for y in y0..y1 {
let row = &rgba.data[y * w * 4..(y * w + w) * 4];
let mut prev: Option<u32> = None;
let mut prev_ink: Option<bool> = None;
for x in x0..x1 {
let p = &row[x * 4..x * 4 + 3];
let l = luminance(p[0], p[1], p[2]) as u32;
if l <= 223 {
midtone += 1;
}
let ink = mask.get(x as u32, y as u32);
if let Some(pl) = prev {
pairs += 1;
if pl.abs_diff(l) > 64 {
sharp += 1;
}
if prev_ink == Some(!ink) {
flips += 1;
}
}
prev = Some(l);
prev_ink = Some(ink);
n += 1;
}
}
let n = n.max(1);
let pairs = pairs.max(1);
let continuous = midtone * 5 > n * 3 && sharp * 200 < pairs * 3;
let halftone = flips * 4 > pairs;
photo[by * bw + bx] = continuous || halftone;
}
}
let smoothed: Vec<bool> = (0..bw * bh)
.map(|i| {
let (bx, by) = (i % bw, i / bw);
let mut yes = 0u32;
let mut total = 0u32;
for dy in -1i32..=1 {
for dx in -1i32..=1 {
let (nx, ny) = (bx as i32 + dx, by as i32 + dy);
if nx >= 0 && ny >= 0 && (nx as usize) < bw && (ny as usize) < bh {
total += 1;
if photo[ny as usize * bw + nx as usize] {
yes += 1;
}
}
}
}
yes * 2 > total
})
.collect();
for by in 0..bh {
for bx in 0..bw {
if !smoothed[by * bw + bx] {
continue;
}
let x1 = ((bx + 1) * CLASSIFY_BLOCK).min(w);
let y1 = ((by + 1) * CLASSIFY_BLOCK).min(h);
for y in (by * CLASSIFY_BLOCK)..y1 {
for x in (bx * CLASSIFY_BLOCK)..x1 {
mask.set(x as u32, y as u32, false);
}
}
}
}
}
const DESKEW_MIN_DEG: f32 = 0.15;
pub fn segment_page(rgba: &Pixmap, opts: &SegmentOptions) -> SegmentedPage {
if opts.deskew {
let correction = estimate_skew(rgba);
if correction.abs() >= DESKEW_MIN_DEG {
let upright = rotate_small(rgba, correction);
let mut inner = *opts;
inner.deskew = false; return segment_page(&upright, &inner);
}
}
let w = rgba.width;
let h = rgba.height;
let mut mask = Bitmap::new(w, h);
if w == 0 || h == 0 {
return SegmentedPage {
mask,
bg: Pixmap::default(),
};
}
match opts.binarization {
Binarization::Fixed => fill_fixed_mask(&mut mask, rgba, opts.threshold),
Binarization::Sauvola { window, k } => {
let luma = luminance_plane(rgba);
fill_sauvola_mask(&mut mask, &luma, w, h, window, k);
}
}
if opts.block_classify {
clear_photo_blocks(&mut mask, rgba);
}
let sub = if opts.adaptive_bg_subsample {
choose_bg_subsample(rgba, &mask, opts.bg_subsample.max(1))
} else {
opts.bg_subsample.max(1)
};
let bw = w.div_ceil(sub);
let bh = h.div_ceil(sub);
let mut bg = Pixmap::white(bw, bh);
#[cfg(feature = "parallel")]
{
use rayon::prelude::*;
let bwu = bw as usize;
bg.data
.par_chunks_mut(bwu * 4)
.enumerate()
.for_each(|(by, bg_row)| {
let by = by as u32;
for bx in 0..bw {
let (r, g, b) = bg_cell_color(rgba, &mask, opts, sub, w, h, bw, bh, bx, by);
let o = bx as usize * 4;
bg_row[o] = r;
bg_row[o + 1] = g;
bg_row[o + 2] = b;
}
});
}
#[cfg(not(feature = "parallel"))]
for by in 0..bh {
for bx in 0..bw {
let (r, g, b) = bg_cell_color(rgba, &mask, opts, sub, w, h, bw, bh, bx, by);
bg.set_rgb(bx, by, r, g, b);
}
}
if opts.bg_diffuse {
diffuse_masked_cells(&mut bg, rgba, &mask, sub, w, h);
}
SegmentedPage { mask, bg }
}
fn diffuse_masked_cells(bg: &mut Pixmap, rgba: &Pixmap, mask: &Bitmap, sub: u32, w: u32, h: u32) {
let bw = bg.width as usize;
let bh = bg.height as usize;
if bw == 0 || bh == 0 {
return;
}
let mut confident = vec![false; bw * bh];
let mut any_masked = false;
for by in 0..bh {
for bx in 0..bw {
let x0 = bx as u32 * sub;
let x1 = (x0 + sub).min(w);
let y0 = by as u32 * sub;
let y1 = (y0 + sub).min(h);
let c = block_mean(rgba, mask, x0, x1, y0, y1, true).is_some();
confident[by * bw + bx] = c;
any_masked |= !c;
}
}
if !any_masked {
return;
}
let mut plane = [
vec![0f32; bw * bh],
vec![0f32; bw * bh],
vec![0f32; bw * bh],
];
let mut seed = [0f64; 3];
let mut nconf = 0u64;
for i in 0..bw * bh {
let px = &bg.data[i * 4..i * 4 + 3];
for c in 0..3 {
plane[c][i] = px[c] as f32;
}
if confident[i] {
nconf += 1;
for c in 0..3 {
seed[c] += px[c] as f64;
}
}
}
if nconf == 0 {
return; }
let seed = [
(seed[0] / nconf as f64) as f32,
(seed[1] / nconf as f64) as f32,
(seed[2] / nconf as f64) as f32,
];
for i in 0..bw * bh {
if !confident[i] {
for c in 0..3 {
plane[c][i] = seed[c];
}
}
}
let max_iters = bw.max(bh).clamp(16, 512);
for _ in 0..max_iters {
let mut max_delta = 0f32;
for by in 0..bh {
for bx in 0..bw {
let i = by * bw + bx;
if confident[i] {
continue;
}
for p in plane.iter_mut() {
let mut sum = 0f32;
let mut n = 0f32;
if bx > 0 {
sum += p[i - 1];
n += 1.0;
}
if bx + 1 < bw {
sum += p[i + 1];
n += 1.0;
}
if by > 0 {
sum += p[i - bw];
n += 1.0;
}
if by + 1 < bh {
sum += p[i + bw];
n += 1.0;
}
let nv = sum / n;
let d = (nv - p[i]).abs();
if d > max_delta {
max_delta = d;
}
p[i] = nv;
}
}
}
if max_delta < 0.5 {
break;
}
}
for i in 0..bw * bh {
if confident[i] {
continue;
}
for (c, p) in plane.iter().enumerate() {
bg.data[i * 4 + c] = p[i].round().clamp(0.0, 255.0) as u8;
}
}
}
#[allow(clippy::too_many_arguments)]
fn bg_cell_color(
rgba: &Pixmap,
mask: &Bitmap,
opts: &SegmentOptions,
sub: u32,
w: u32,
h: u32,
bw: u32,
bh: u32,
bx: u32,
by: u32,
) -> (u8, u8, u8) {
let x0 = bx * sub;
let x1 = (x0 + sub).min(w);
let y0 = by * sub;
let y1 = (y0 + sub).min(h);
block_mean(rgba, mask, x0, x1, y0, y1, true)
.or_else(|| {
opts.bg_inpaint
.then(|| inpaint_block_mean(rgba, mask, bx, by, sub, bw, bh))
.flatten()
})
.or_else(|| block_mean(rgba, mask, x0, x1, y0, y1, false))
.unwrap_or((255, 255, 255))
}
fn luminance_plane(rgba: &Pixmap) -> Vec<u8> {
let mut luma = Vec::with_capacity((rgba.width * rgba.height) as usize);
for y in 0..rgba.height {
for x in 0..rgba.width {
let (r, g, b) = rgba.get_rgb(x, y);
luma.push(luminance(r, g, b));
}
}
luma
}
fn fill_fixed_mask(mask: &mut Bitmap, rgba: &Pixmap, threshold: u8) {
let threshold = u32::from(threshold);
let w = rgba.width as usize;
let mstride = mask.row_stride();
for y in 0..mask.height as usize {
let src = &rgba.data[y * w * 4..(y + 1) * w * 4];
let mrow = &mut mask.data[y * mstride..(y + 1) * mstride];
for (x, px) in src.chunks_exact(4).enumerate() {
if u32::from(luminance(px[0], px[1], px[2])) < threshold {
mrow[x >> 3] |= 0x80 >> (x & 7);
}
}
}
}
fn fill_sauvola_mask(mask: &mut Bitmap, luma: &[u8], w: u32, h: u32, window: u32, k: f32) {
let window = window.max(3);
let radius = window / 2;
let k = if k.is_finite() { k } else { 0.34 };
let k = k.clamp(0.0, 1.0);
let (sum, sum_sq) = integral_luma(luma, w, h);
let stride = w as usize + 1;
#[cfg(feature = "parallel")]
{
use rayon::prelude::*;
let mask_stride = mask.row_stride();
mask.data
.par_chunks_mut(mask_stride)
.enumerate()
.for_each(|(y, row)| {
fill_sauvola_row(row, luma, &sum, &sum_sq, stride, w, h, radius, k, y as u32);
});
}
#[cfg(not(feature = "parallel"))]
fill_sauvola_mask_sequential(mask, luma, &sum, &sum_sq, stride, w, h, radius, k);
}
#[allow(clippy::too_many_arguments)]
#[cfg_attr(feature = "parallel", allow(dead_code))]
fn fill_sauvola_mask_sequential(
mask: &mut Bitmap,
luma: &[u8],
sum: &[u64],
sum_sq: &[u64],
stride: usize,
w: u32,
h: u32,
radius: u32,
k: f32,
) {
let mask_stride = mask.row_stride();
for (y, row) in mask.data.chunks_mut(mask_stride).enumerate() {
fill_sauvola_row(row, luma, sum, sum_sq, stride, w, h, radius, k, y as u32);
}
}
#[allow(clippy::too_many_arguments)]
fn fill_sauvola_row(
mask_row: &mut [u8],
luma: &[u8],
sum: &[u64],
sum_sq: &[u64],
stride: usize,
w: u32,
h: u32,
radius: u32,
k: f32,
y: u32,
) {
let y0 = y.saturating_sub(radius);
let y1 = (y + radius + 1).min(h);
for x in 0..w {
let x0 = x.saturating_sub(radius);
let x1 = (x + radius + 1).min(w);
let area = f64::from((x1 - x0) * (y1 - y0));
let s = rect_sum(sum, stride, x0, y0, x1, y1) as f64;
let ss = rect_sum(sum_sq, stride, x0, y0, x1, y1) as f64;
let mean = s / area;
let variance = (ss / area - mean * mean).max(0.0);
let stddev = variance.sqrt();
let threshold = mean * (1.0 + f64::from(k) * (stddev / 128.0 - 1.0));
let idx = (y * w + x) as usize;
if f64::from(luma[idx]) < threshold {
mask_row[x as usize >> 3] |= 0x80 >> (x & 7);
}
}
}
fn integral_luma(luma: &[u8], w: u32, h: u32) -> (Vec<u64>, Vec<u64>) {
let stride = w as usize + 1;
let len = stride * (h as usize + 1);
let mut sum = vec![0u64; len];
let mut sum_sq = vec![0u64; len];
for y in 0..h as usize {
let mut row_sum = 0u64;
let mut row_sum_sq = 0u64;
for x in 0..w as usize {
let v = u64::from(luma[y * w as usize + x]);
row_sum += v;
row_sum_sq += v * v;
let dst = (y + 1) * stride + x + 1;
sum[dst] = sum[dst - stride] + row_sum;
sum_sq[dst] = sum_sq[dst - stride] + row_sum_sq;
}
}
(sum, sum_sq)
}
fn rect_sum(integral: &[u64], stride: usize, x0: u32, y0: u32, x1: u32, y1: u32) -> u64 {
let (x0, y0, x1, y1) = (x0 as usize, y0 as usize, x1 as usize, y1 as usize);
integral[y1 * stride + x1] + integral[y0 * stride + x0]
- integral[y0 * stride + x1]
- integral[y1 * stride + x0]
}
fn block_mean(
rgba: &Pixmap,
mask: &Bitmap,
x0: u32,
x1: u32,
y0: u32,
y1: u32,
unmasked_only: bool,
) -> Option<(u8, u8, u8)> {
let mut acc = ColorAccum::default();
let w = rgba.width as usize;
let mstride = mask.row_stride();
for y in y0..y1 {
let ry = y as usize;
let row = &rgba.data[(ry * w + x0 as usize) * 4..(ry * w + x1 as usize) * 4];
if unmasked_only {
let mrow = &mask.data[ry * mstride..(ry + 1) * mstride];
for (i, px) in row.chunks_exact(4).enumerate() {
let x = x0 as usize + i;
if (mrow[x >> 3] >> (7 - (x & 7))) & 1 != 0 {
continue;
}
acc.add(px[0], px[1], px[2]);
}
} else {
for px in row.chunks_exact(4) {
acc.add(px[0], px[1], px[2]);
}
}
}
acc.color()
}
fn inpaint_block_mean(
rgba: &Pixmap,
mask: &Bitmap,
bx: u32,
by: u32,
sub: u32,
bw: u32,
bh: u32,
) -> Option<(u8, u8, u8)> {
let max_radius = bw.max(bh);
for radius in 1..=max_radius {
let bx0 = bx.saturating_sub(radius);
let by0 = by.saturating_sub(radius);
let bx1 = (bx + radius + 1).min(bw);
let by1 = (by + radius + 1).min(bh);
let mut acc = ColorAccum::default();
for ny in by0..by1 {
for nx in bx0..bx1 {
let dx = nx.abs_diff(bx);
let dy = ny.abs_diff(by);
if dx.max(dy) != radius {
continue;
}
let x0 = nx * sub;
let x1 = (x0 + sub).min(rgba.width);
let y0 = ny * sub;
let y1 = (y0 + sub).min(rgba.height);
for y in y0..y1 {
for x in x0..x1 {
if !mask.get(x, y) {
let (r, g, b) = rgba.get_rgb(x, y);
acc.add(r, g, b);
}
}
}
}
}
if let Some(color) = acc.color() {
return Some(color);
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
fn striped_page(deg: f32) -> Pixmap {
let mut pm = Pixmap::white(512, 512);
let mut rng = 0x2545_F491u32;
let mut next = move |m: u32| {
rng = rng.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
(rng >> 16) % m
};
let mut y = 24u32;
while y + 10 < 488 {
let mut x = 24 + next(20);
while x + 12 < 488 {
let wlen = 12 + next(40);
for yy in y..y + 6 {
for xx in x..(x + wlen).min(488) {
pm.set_rgb(xx, yy, 0, 0, 0);
}
}
x += wlen + 6 + next(12);
}
y += 20 + next(9);
}
if deg == 0.0 {
pm
} else {
rotate_small(&pm, deg)
}
}
#[test]
fn estimate_skew_finds_synthetic_rotation() {
assert!(estimate_skew(&striped_page(0.0)).abs() <= 0.06);
let est = estimate_skew(&striped_page(1.0));
assert!(
(est + 1.0).abs() <= 0.06,
"1° skew must estimate a ≈−1° correction, got {est}"
);
}
#[test]
fn deskew_option_corrects_synthetic_skew() {
let upright = striped_page(0.0);
let skewed = striped_page(1.0);
let opts_plain = SegmentOptions::default();
let opts_deskew = SegmentOptions {
deskew: true,
..SegmentOptions::default()
};
let ref_mask = segment_page(&upright, &opts_plain).mask;
let diff = |m: &Bitmap| -> u64 {
let mut d = 0u64;
for y in 0..m.height {
for x in 0..m.width {
if m.get(x, y) != ref_mask.get(x, y) {
d += 1;
}
}
}
d
};
let skewed_diff = diff(&segment_page(&skewed, &opts_plain).mask);
let fixed_diff = diff(&segment_page(&skewed, &opts_deskew).mask);
assert!(
fixed_diff * 2 < skewed_diff,
"deskew must recover most of the skew: skewed {skewed_diff}, deskewed {fixed_diff}"
);
let a = segment_page(&upright, &opts_plain).mask;
let b = segment_page(&upright, &opts_deskew).mask;
assert_eq!(a.data, b.data, "deskew must be a no-op on an upright page");
}
#[test]
fn block_classify_clears_photo_keeps_text() {
let mut pm = Pixmap::white(256, 256);
for y in (16..112).step_by(16) {
for yy in y..y + 5 {
for x in 16..240 {
pm.set_rgb(x, yy, 0, 0, 0);
}
}
}
for y in 128..256 {
for x in 0..256 {
let v = (60 + (x + y) / 4) as u8;
pm.set_rgb(x, y, v, v, v);
}
}
let plain = segment_page(&pm, &SegmentOptions::default()).mask;
let classified = segment_page(
&pm,
&SegmentOptions {
block_classify: true,
..SegmentOptions::default()
},
)
.mask;
let ink = |m: &Bitmap, y0: u32, y1: u32| -> u64 {
(y0..y1)
.map(|y| (0..256).filter(|&x| m.get(x, y)).count() as u64)
.sum()
};
for y in 0..112 {
for x in 0..256 {
assert_eq!(
plain.get(x, y),
classified.get(x, y),
"text mask must be untouched at ({x},{y})"
);
}
}
assert!(ink(&plain, 128, 256) > 0, "gradient must binarize to ink");
assert_eq!(ink(&classified, 128, 256), 0, "photo half must be cleared");
}
fn fill(pm: &mut Pixmap, r: u8, g: u8, b: u8) {
for y in 0..pm.height {
for x in 0..pm.width {
pm.set_rgb(x, y, r, g, b);
}
}
}
#[test]
fn all_white_page_yields_empty_mask() {
let pm = Pixmap::white(24, 24);
let seg = segment_page(&pm, &SegmentOptions::default());
assert_eq!(seg.mask.width, 24);
assert_eq!(seg.mask.height, 24);
for y in 0..24 {
for x in 0..24 {
assert!(
!seg.mask.get(x, y),
"white pixel at ({x},{y}) should not be mask"
);
}
}
assert_eq!(seg.bg.width, 2);
assert_eq!(seg.bg.height, 2);
for chunk in seg.bg.data.chunks_exact(4) {
assert_eq!(&chunk[..3], &[255, 255, 255]);
}
}
#[test]
fn all_black_page_yields_full_mask_and_black_bg_fallback() {
let mut pm = Pixmap::white(12, 12);
fill(&mut pm, 0, 0, 0);
let seg = segment_page(&pm, &SegmentOptions::default());
for y in 0..12 {
for x in 0..12 {
assert!(seg.mask.get(x, y));
}
}
assert_eq!(seg.bg.width, 1);
assert_eq!(seg.bg.height, 1);
assert_eq!(&seg.bg.data[..3], &[0, 0, 0]);
}
#[test]
fn threshold_boundary_is_strict() {
let mut pm = Pixmap::white(4, 1);
pm.set_rgb(0, 0, 0, 0, 0);
pm.set_rgb(1, 0, 127, 127, 127);
pm.set_rgb(2, 0, 128, 128, 128);
pm.set_rgb(3, 0, 255, 255, 255);
let seg = segment_page(
&pm,
&SegmentOptions {
threshold: 128,
bg_subsample: 1,
adaptive_bg_subsample: false,
..SegmentOptions::default()
},
);
assert!(seg.mask.get(0, 0));
assert!(seg.mask.get(1, 0));
assert!(!seg.mask.get(2, 0));
assert!(!seg.mask.get(3, 0));
}
#[test]
fn bg_excludes_mask_pixels() {
let mut pm = Pixmap::white(4, 4);
fill(&mut pm, 240, 230, 100);
pm.set_rgb(1, 1, 0, 0, 0);
let seg = segment_page(
&pm,
&SegmentOptions {
threshold: 128,
bg_subsample: 4,
..SegmentOptions::default()
},
);
assert!(seg.mask.get(1, 1));
assert!(!seg.mask.get(0, 0));
assert_eq!(seg.bg.width, 1);
assert_eq!(seg.bg.height, 1);
let (r, g, b) = (seg.bg.data[0], seg.bg.data[1], seg.bg.data[2]);
assert_eq!(
(r, g, b),
(240, 230, 100),
"ink pixel should not contaminate BG mean"
);
}
#[test]
fn sauvola_handles_dark_background_and_light_ink() {
let mut pm = Pixmap::white(16, 8);
for y in 0..8 {
for x in 0..16 {
let v = if x < 8 { 80 } else { 220 };
pm.set_rgb(x, y, v, v, v);
}
}
pm.set_rgb(3, 3, 40, 40, 40);
pm.set_rgb(11, 3, 140, 140, 140);
let fixed = segment_page(&pm, &SegmentOptions::default());
let adaptive = segment_page(
&pm,
&SegmentOptions {
binarization: Binarization::Sauvola { window: 7, k: 0.34 },
..SegmentOptions::default()
},
);
let fixed_count = count_mask(&fixed.mask);
let adaptive_count = count_mask(&adaptive.mask);
assert!(fixed_count > 50, "fixed threshold masks the dark paper");
assert!(
adaptive_count < fixed_count / 2,
"adaptive mask should be much sparser than fixed ({adaptive_count} vs {fixed_count})"
);
assert!(adaptive.mask.get(3, 3), "dark ink on dark paper");
assert!(adaptive.mask.get(11, 3), "light ink on light paper");
assert!(!adaptive.mask.get(1, 1), "dark paper is background");
assert!(!adaptive.mask.get(9, 1), "bright paper is background");
}
#[cfg(feature = "parallel")]
#[test]
fn parallel_sauvola_mask_is_byte_identical_to_sequential() {
let (w, h) = (131, 97);
let luma: Vec<u8> = (0..w * h)
.map(|i| ((i * 73 + (i / w) * 29) & 0xff) as u8)
.collect();
let mut parallel = Bitmap::new(w, h);
fill_sauvola_mask(&mut parallel, &luma, w, h, 31, 0.34);
let window = 31_u32;
let radius = window / 2;
let (sum, sum_sq) = integral_luma(&luma, w, h);
let mut sequential = Bitmap::new(w, h);
fill_sauvola_mask_sequential(
&mut sequential,
&luma,
&sum,
&sum_sq,
w as usize + 1,
w,
h,
radius,
0.34,
);
assert_eq!(parallel.data, sequential.data);
}
#[test]
fn adaptive_bg_subsample_tracks_background_detail() {
let flat = Pixmap::white(120, 120);
let empty = Bitmap::new(120, 120);
assert_eq!(choose_bg_subsample(&flat, &empty, 12), 12);
let mut noisy = Pixmap::white(120, 120);
for y in 0..120 {
for x in 0..120 {
let v = ((x * 37 + y * 91) % 256) as u8;
noisy.set_rgb(x, y, v, v, v);
}
}
assert_eq!(choose_bg_subsample(&noisy, &empty, 12), 3);
let mut inked = Bitmap::new(120, 120);
for y in 0..120 {
for x in 0..120 {
inked.set_black(x, y);
}
}
assert_eq!(choose_bg_subsample(&noisy, &inked, 12), 12);
}
fn count_mask(mask: &Bitmap) -> u32 {
let mut n = 0;
for y in 0..mask.height {
for x in 0..mask.width {
n += u32::from(mask.get(x, y));
}
}
n
}
#[test]
fn inpaint_fully_masked_bg_block_from_neighbors() {
let mut pm = Pixmap::white(8, 4);
for y in 0..4 {
for x in 0..4 {
pm.set_rgb(x, y, 0, 0, 0);
}
for x in 4..8 {
pm.set_rgb(x, y, 210, 200, 160);
}
}
let opts = SegmentOptions {
threshold: 128,
bg_subsample: 4,
bg_inpaint: true,
..SegmentOptions::default()
};
let seg = segment_page(&pm, &opts);
assert_eq!(seg.bg.width, 2);
assert_eq!(seg.bg.height, 1);
assert_eq!(seg.bg.get_rgb(0, 0), (210, 200, 160));
assert_eq!(seg.bg.get_rgb(1, 0), (210, 200, 160));
}
#[test]
fn empty_input_returns_empty_outputs() {
let pm = Pixmap::default();
let seg = segment_page(&pm, &SegmentOptions::default());
assert_eq!(seg.mask.width, 0);
assert_eq!(seg.mask.height, 0);
assert_eq!(seg.bg.width, 0);
assert_eq!(seg.bg.height, 0);
}
#[test]
fn bg_dims_round_up() {
let pm = Pixmap::white(13, 7);
let seg = segment_page(
&pm,
&SegmentOptions {
threshold: 128,
bg_subsample: 12,
..SegmentOptions::default()
},
);
assert_eq!(seg.bg.width, 2);
assert_eq!(seg.bg.height, 1);
}
#[test]
fn inpaint_all_masked_single_block_falls_back_to_white() {
let mut pm = Pixmap::white(1, 1);
pm.set_rgb(0, 0, 0, 0, 0);
let opts = SegmentOptions {
threshold: 128,
bg_subsample: 1,
bg_inpaint: true,
..SegmentOptions::default()
};
let seg = segment_page(&pm, &opts);
assert_eq!(seg.bg.get_rgb(0, 0), (0, 0, 0));
}
#[test]
fn bg_subsample_zero_is_clamped_to_one() {
let pm = Pixmap::white(3, 3);
let seg = segment_page(
&pm,
&SegmentOptions {
threshold: 128,
bg_subsample: 0,
..SegmentOptions::default()
},
);
assert_eq!(seg.bg.width, 3);
assert_eq!(seg.bg.height, 3);
}
#[test]
fn bg_diffuse_smooths_masked_cells_and_keeps_confident_cells() {
let mut pm = Pixmap::white(4, 1);
pm.set_rgb(0, 0, 230, 150, 150);
pm.set_rgb(1, 0, 0, 0, 0); pm.set_rgb(2, 0, 0, 0, 0); pm.set_rgb(3, 0, 230, 150, 150);
let opts = SegmentOptions {
threshold: 128,
bg_subsample: 1,
bg_diffuse: true,
..SegmentOptions::default()
};
let seg = segment_page(&pm, &opts);
assert_eq!(seg.bg.width, 4);
assert_eq!(seg.bg.get_rgb(0, 0), (230, 150, 150));
assert_eq!(seg.bg.get_rgb(3, 0), (230, 150, 150));
for x in 1..=2 {
let (r, g, b) = seg.bg.get_rgb(x, 0);
assert_eq!(
(r, g, b),
(230, 150, 150),
"masked cell {x} should diffuse to the boundary, got {:?}",
(r, g, b)
);
}
let plain = SegmentOptions {
threshold: 128,
bg_subsample: 1,
..SegmentOptions::default()
};
let seg_plain = segment_page(&pm, &plain);
assert_eq!(seg_plain.bg.get_rgb(1, 0), (0, 0, 0));
}
}