use crate::par::prelude::*;
#[must_use]
pub fn lanczos3(x: f32) -> f32 {
const A: f32 = 3.0;
if x == 0.0 {
return 1.0;
}
let ax = x.abs();
if ax >= A {
return 0.0;
}
let px = std::f32::consts::PI * x;
(px.sin() / px) * ((px / A).sin() / (px / A))
}
struct Taps {
starts: Vec<usize>,
weights: Vec<f32>,
width: usize,
}
fn build_taps(src: usize, dst: usize) -> Taps {
let scale = src as f32 / dst as f32;
let filter_scale = if scale < 1.0 { 1.0 } else { scale };
let support = 3.0 * filter_scale;
let width = (support.ceil() as usize) * 2 + 1;
let mut starts = Vec::with_capacity(dst);
let mut weights = vec![0.0f32; dst * width];
for i in 0..dst {
let center = (i as f32 + 0.5) * scale;
let xmin = ((center - support + 0.5).floor().max(0.0)) as usize;
let xmax = (((center + support + 0.5).floor()) as usize).min(src);
let n = xmax.saturating_sub(xmin);
let row = &mut weights[i * width..i * width + width];
let mut sum = 0.0f32;
for (k, slot) in row.iter_mut().take(n.min(width)).enumerate() {
let x = (xmin + k) as f32 - center + 0.5;
let w = lanczos3(x / filter_scale);
*slot = w;
sum += w;
}
if sum != 0.0 {
for w in row.iter_mut().take(n.min(width)) {
*w /= sum;
}
}
starts.push(xmin);
}
Taps {
starts,
weights,
width,
}
}
#[must_use]
pub fn resize_lanczos(
src: &[f32],
sw: usize,
sh: usize,
dw: usize,
dh: usize,
channels: usize,
) -> Vec<f32> {
if sw == dw && sh == dh {
return src.to_vec();
}
let hx = build_taps(sw, dw);
let mut mid = vec![0.0f32; sh * dw * channels];
for y in 0..sh {
for x in 0..dw {
let start = hx.starts[x];
let row = &hx.weights[x * hx.width..x * hx.width + hx.width];
for c in 0..channels {
let mut acc = 0.0f32;
for (k, &w) in row.iter().enumerate() {
if w == 0.0 {
continue;
}
let sx = (start + k).min(sw - 1);
acc += w * src[(y * sw + sx) * channels + c];
}
mid[(y * dw + x) * channels + c] = acc;
}
}
}
let vy = build_taps(sh, dh);
let mut out = vec![0.0f32; dh * dw * channels];
for y in 0..dh {
let start = vy.starts[y];
let row = &vy.weights[y * vy.width..y * vy.width + vy.width];
for x in 0..dw {
for c in 0..channels {
let mut acc = 0.0f32;
for (k, &w) in row.iter().enumerate() {
if w == 0.0 {
continue;
}
let sy = (start + k).min(sh - 1);
acc += w * mid[(sy * dw + x) * channels + c];
}
out[(y * dw + x) * channels + c] = acc;
}
}
}
out
}
#[must_use]
pub const fn tile_grid(width: usize, height: usize, max_edge: usize) -> (usize, usize) {
if width <= max_edge && height <= max_edge {
return (0, 0);
}
(height.div_ceil(max_edge), width.div_ceil(max_edge))
}
#[must_use]
pub fn fit_longest_edge(width: usize, height: usize, longest: usize) -> (usize, usize) {
if width.max(height) == longest {
return (width, height);
}
let scale = longest as f32 / width.max(height) as f32;
(
((width as f32 * scale).round() as usize).max(1),
((height as f32 * scale).round() as usize).max(1),
)
}
#[must_use]
pub fn normalize_u8(pixels: &[u8]) -> Vec<f32> {
pixels.iter().map(|&v| f32::from(v) / 127.5 - 1.0).collect()
}
const PRECISION_BITS: i32 = 32 - 8 - 2;
fn quantise(w: f64) -> i32 {
let scaled = w * f64::from(1 << PRECISION_BITS);
if w < 0.0 {
(scaled - 0.5) as i32
} else {
(scaled + 0.5) as i32
}
}
struct FixedTaps {
starts: Vec<usize>,
lens: Vec<usize>,
k: Vec<i32>,
ksize: usize,
}
fn build_fixed_taps(src: usize, dst: usize) -> FixedTaps {
let scale = src as f64 / dst as f64;
let filter_scale = if scale < 1.0 { 1.0 } else { scale };
let support = 3.0 * filter_scale;
let ksize = (support.ceil() as usize) * 2 + 1;
let mut starts = Vec::with_capacity(dst);
let mut lens = Vec::with_capacity(dst);
let mut k = vec![0i32; dst * ksize];
let inv = 1.0 / filter_scale;
for xx in 0..dst {
let center = (xx as f64 + 0.5) * scale;
let xmin = ((center - support + 0.5) as isize).max(0) as usize;
let xmax = (((center + support + 0.5) as isize).max(0) as usize).min(src);
let n = xmax.saturating_sub(xmin);
let mut w = vec![0.0f64; ksize];
let mut ww = 0.0f64;
for (x, slot) in w.iter_mut().enumerate().take(n) {
let v = lanczos3_f64(((x + xmin) as f64 - center + 0.5) * inv);
*slot = v;
ww += v;
}
if ww != 0.0 {
for slot in w.iter_mut().take(n) {
*slot /= ww;
}
}
for (x, &v) in w.iter().enumerate() {
k[xx * ksize + x] = quantise(v);
}
starts.push(xmin);
lens.push(n);
}
FixedTaps {
starts,
lens,
k,
ksize,
}
}
fn lanczos3_f64(x: f64) -> f64 {
const A: f64 = 3.0;
if x == 0.0 {
return 1.0;
}
if x.abs() >= A {
return 0.0;
}
let px = std::f64::consts::PI * x;
(px.sin() / px) * ((px / A).sin() / (px / A))
}
const fn clip8(acc: i32) -> u8 {
let v = acc >> PRECISION_BITS;
if v <= 0 {
0
} else if v >= 255 {
255
} else {
v as u8
}
}
#[must_use]
pub fn resize_lanczos_u8(
src: &[u8],
sw: usize,
sh: usize,
dw: usize,
dh: usize,
channels: usize,
) -> Vec<u8> {
if sw == dw && sh == dh {
return src.to_vec();
}
let mid: Vec<u8> = if sw == dw {
src.to_vec()
} else {
let t = build_fixed_taps(sw, dw);
let mut out = vec![0u8; sh * dw * channels];
out.par_chunks_mut(dw * channels)
.enumerate()
.for_each(|(y, orow)| {
for x in 0..dw {
let (start, n) = (t.starts[x], t.lens[x]);
let row = &t.k[x * t.ksize..x * t.ksize + t.ksize];
for c in 0..channels {
let mut acc: i32 = 1 << (PRECISION_BITS - 1);
for (kk, &w) in row.iter().enumerate().take(n) {
acc += i32::from(src[(y * sw + start + kk) * channels + c]) * w;
}
orow[x * channels + c] = clip8(acc);
}
}
});
out
};
if sh == dh {
return mid;
}
let t = build_fixed_taps(sh, dh);
let mut out = vec![0u8; dh * dw * channels];
out.par_chunks_mut(dw * channels)
.enumerate()
.for_each(|(y, orow)| {
let (start, n) = (t.starts[y], t.lens[y]);
let row = &t.k[y * t.ksize..y * t.ksize + t.ksize];
for x in 0..dw {
for c in 0..channels {
let mut acc: i32 = 1 << (PRECISION_BITS - 1);
for (kk, &w) in row.iter().enumerate().take(n) {
acc += i32::from(mid[((start + kk) * dw + x) * channels + c]) * w;
}
orow[x * channels + c] = clip8(acc);
}
}
});
out
}
#[must_use]
pub fn vision_encoder_size(width: usize, height: usize, tile: usize) -> (usize, usize) {
let aspect = width as f64 / height as f64;
if width >= height {
let w = width.div_ceil(tile) * tile;
let h = ((w as f64 / aspect) as usize).div_ceil(tile) * tile;
(w, h.max(tile))
} else {
let h = height.div_ceil(tile) * tile;
let w = ((h as f64 * aspect) as usize).div_ceil(tile) * tile;
(w.max(tile), h)
}
}
pub struct Preprocessed {
pub pixel_values: Vec<f32>,
pub tiles: usize,
pub rows: usize,
pub cols: usize,
pub tile: usize,
}
#[must_use]
pub fn resized_size(width: usize, height: usize) -> (usize, usize) {
let (aw, ah) = fit_longest_edge(width, height, 2048);
vision_encoder_size(aw, ah, 512)
}
#[must_use]
pub fn tile_geometry(width: usize, height: usize, split: bool) -> (usize, usize, usize) {
const LONGEST: usize = 2048;
const TILE: usize = 512;
let (aw, ah) = fit_longest_edge(width, height, LONGEST);
let (bw, bh) = vision_encoder_size(aw, ah, TILE);
let (rows, cols) = if split { tile_grid(bw, bh, TILE) } else { (0, 0) };
(rows * cols + 1, rows, cols)
}
#[must_use]
pub fn preprocess_rgb8(rgb: &[u8], width: usize, height: usize) -> Preprocessed {
preprocess_rgb8_opts(rgb, width, height, true)
}
#[must_use]
pub fn preprocess_rgb8_opts(
rgb: &[u8],
width: usize,
height: usize,
split: bool,
) -> Preprocessed {
const LONGEST: usize = 2048;
const TILE: usize = 512;
let (aw, ah) = fit_longest_edge(width, height, LONGEST);
let a = resize_lanczos_u8(rgb, width, height, aw, ah, 3);
let (bw, bh) = vision_encoder_size(aw, ah, TILE);
let b = resize_lanczos_u8(&a, aw, ah, bw, bh, 3);
let (rows, cols) = if split {
tile_grid(bw, bh, TILE)
} else {
(0, 0)
};
let per = TILE * TILE * 3;
let tiles = rows * cols + 1;
let mut pixel_values = vec![0.0f32; tiles * per];
let mut write_tile = |idx: usize, src: &[u8]| {
let norm = normalize_u8(src);
let base = idx * per;
for c in 0..3 {
for i in 0..TILE * TILE {
pixel_values[base + c * TILE * TILE + i] = norm[i * 3 + c];
}
}
};
let mut tile_buf = vec![0u8; per];
for r in 0..rows {
for c in 0..cols {
for y in 0..TILE {
let s = ((r * TILE + y) * bw + c * TILE) * 3;
tile_buf[y * TILE * 3..(y + 1) * TILE * 3].copy_from_slice(&b[s..s + TILE * 3]);
}
write_tile(r * cols + c, &tile_buf);
}
}
let thumb = resize_lanczos_u8(&b, bw, bh, TILE, TILE, 3);
write_tile(rows * cols, &thumb);
Preprocessed {
pixel_values,
tiles,
rows,
cols,
tile: TILE,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_kernel_is_one_at_zero_and_zero_at_the_integers() {
assert!((lanczos3(0.0) - 1.0).abs() < 1e-6);
for k in [1.0f32, 2.0] {
assert!(lanczos3(k).abs() < 1e-5, "lanczos3({k}) should vanish");
assert!(lanczos3(-k).abs() < 1e-5);
}
assert_eq!(lanczos3(3.0), 0.0);
assert_eq!(lanczos3(4.5), 0.0);
}
#[test]
fn weights_sum_to_one_at_every_output_position() {
for (src, dst) in [(512, 2048), (2048, 512), (300, 512), (512, 300)] {
let t = build_taps(src, dst);
for i in 0..dst {
let s: f32 = t.weights[i * t.width..(i + 1) * t.width].iter().sum();
assert!(
(s - 1.0).abs() < 1e-4,
"{src}->{dst} position {i} sums to {s}"
);
}
}
}
#[test]
fn a_constant_image_survives_resizing_unchanged() {
for (sw, sh, dw, dh) in [(16, 16, 64, 64), (64, 64, 16, 16), (10, 20, 33, 7)] {
let src = vec![0.37f32; sw * sh * 3];
let out = resize_lanczos(&src, sw, sh, dw, dh, 3);
assert_eq!(out.len(), dw * dh * 3);
let worst = out.iter().map(|v| (v - 0.37).abs()).fold(0.0f32, f32::max);
assert!(worst < 1e-4, "{sw}x{sh}->{dw}x{dh} drifted by {worst}");
}
}
#[test]
fn an_identity_resize_is_a_copy() {
let src: Vec<f32> = (0..48).map(|i| i as f32).collect();
assert_eq!(resize_lanczos(&src, 4, 4, 4, 4, 3), src);
}
#[test]
fn the_grid_is_what_produced_seventeen_tiles() {
assert_eq!(fit_longest_edge(512, 512, 2048), (2048, 2048));
assert_eq!(tile_grid(2048, 2048, 512), (4, 4));
assert_eq!(4 * 4 + 1, 17);
assert_eq!(tile_grid(400, 300, 512), (0, 0));
}
#[test]
fn normalisation_lands_on_minus_one_to_one() {
let v = normalize_u8(&[0, 255, 128]);
assert!((v[0] + 1.0).abs() < 1e-6);
assert!((v[1] - 1.0).abs() < 1e-6);
assert!(v[2].abs() < 0.01);
}
#[test]
fn the_fixed_point_kernel_preserves_a_constant_image() {
for (sw, sh, dw, dh) in [(16, 16, 64, 64), (64, 64, 16, 16), (10, 20, 33, 7)] {
let src = vec![97u8; sw * sh * 3];
let out = resize_lanczos_u8(&src, sw, sh, dw, dh, 3);
assert_eq!(out.len(), dw * dh * 3);
let worst = out.iter().map(|&v| i32::from(v) - 97).map(i32::abs).max();
assert_eq!(worst, Some(0), "{sw}x{sh}->{dw}x{dh} drifted off 97");
}
}
#[test]
fn quantisation_rounds_away_from_zero_in_both_directions() {
let unit = f64::from(1 << PRECISION_BITS);
assert_eq!(quantise(1.0), 1 << PRECISION_BITS);
assert_eq!(quantise(-1.0), -(1 << PRECISION_BITS));
assert_eq!(quantise(0.5 / unit), 1);
assert_eq!(quantise(-0.5 / unit), -1);
}
#[test]
fn clip8_clamps_rather_than_wrapping() {
assert_eq!(clip8(-5 << PRECISION_BITS), 0);
assert_eq!(clip8(300 << PRECISION_BITS), 255);
assert_eq!(clip8(128 << PRECISION_BITS), 128);
assert_eq!(clip8((1 << PRECISION_BITS) - 1 + (1 << (PRECISION_BITS - 1))), 1);
}
#[test]
fn an_identity_resize_is_a_copy_in_the_fixed_point_path_too() {
let src: Vec<u8> = (0..48).map(|i| i as u8).collect();
assert_eq!(resize_lanczos_u8(&src, 4, 4, 4, 4, 3), src);
}
#[test]
fn the_vision_encoder_size_rounds_the_long_edge_first() {
assert_eq!(vision_encoder_size(2048, 2048, 512), (2048, 2048));
assert_eq!(vision_encoder_size(2048, 1536, 512), (2048, 1536));
for (w, h) in [(2048, 1153), (1000, 700), (513, 511), (100, 3000)] {
let (a, b) = vision_encoder_size(w, h, 512);
assert_eq!(a % 512, 0, "{w}x{h} -> {a}x{b}: width not a tile multiple");
assert_eq!(b % 512, 0, "{w}x{h} -> {a}x{b}: height not a tile multiple");
assert!(a >= 512 && b >= 512, "{w}x{h} -> {a}x{b}: degenerate");
}
}
#[test]
fn preprocessing_fills_every_tile_it_promises() {
let (w, h) = (300usize, 700usize);
let px: Vec<u8> = (0..w * h * 3).map(|i| (i % 251) as u8).collect();
let out = preprocess_rgb8(&px, w, h);
assert_eq!(out.pixel_values.len(), out.tiles * 3 * out.tile * out.tile);
assert_eq!(out.tiles, out.rows * out.cols + 1);
let per = 3 * out.tile * out.tile;
for t in 0..out.tiles {
let tile = &out.pixel_values[t * per..(t + 1) * per];
assert!(
tile.iter().any(|&v| v.abs() > 1e-6),
"tile {t} is entirely zero — an unfilled buffer, not an image"
);
}
}
}