use serde::{Deserialize, Serialize};
use tokio::io::AsyncReadExt;
use tokio::process::Command;
use viser_ffmpeg::{ProbeCache, ffmpeg_path};
use crate::PooledStats;
#[derive(Debug, Clone, Default)]
pub struct NoRefOpts {
pub stride: usize,
pub probe_cache: Option<ProbeCache>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(default)]
pub struct NoRefResult {
pub sharpness: f64,
pub blockiness: f64,
pub noise: f64,
pub sharpness_pooled: PooledStats,
pub blockiness_pooled: PooledStats,
pub noise_pooled: PooledStats,
pub frames: usize,
}
pub fn variance_of_laplacian(frame: &[u8], w: usize, h: usize) -> f64 {
if w < 3 || h < 3 {
return 0.0;
}
let at = |x: usize, y: usize| frame[y * w + x] as f64;
let mut sum = 0.0;
let mut sum_sq = 0.0;
let mut n = 0.0;
for y in 1..h - 1 {
for x in 1..w - 1 {
let lap = 4.0 * at(x, y) - at(x - 1, y) - at(x + 1, y) - at(x, y - 1) - at(x, y + 1);
sum += lap;
sum_sq += lap * lap;
n += 1.0;
}
}
if n == 0.0 {
return 0.0;
}
let mean = sum / n;
(sum_sq / n) - mean * mean
}
pub fn blockiness(frame: &[u8], w: usize, h: usize) -> f64 {
if w < 9 || h < 9 {
return 0.0;
}
let at = |x: usize, y: usize| frame[y * w + x] as f64;
let (mut bnd, mut bnd_n, mut int, mut int_n) = (0.0, 0.0, 0.0, 0.0);
for y in 0..h {
for x in 1..w {
let d = (at(x, y) - at(x - 1, y)).abs();
if x % 8 == 0 {
bnd += d;
bnd_n += 1.0;
} else {
int += d;
int_n += 1.0;
}
}
}
for y in 1..h {
for x in 0..w {
let d = (at(x, y) - at(x, y - 1)).abs();
if y % 8 == 0 {
bnd += d;
bnd_n += 1.0;
} else {
int += d;
int_n += 1.0;
}
}
}
if bnd_n == 0.0 || int_n == 0.0 {
return 0.0;
}
(bnd / bnd_n - int / int_n).max(0.0)
}
pub fn noise_sigma(frame: &[u8], w: usize, h: usize) -> f64 {
if w < 3 || h < 3 {
return 0.0;
}
let at = |x: usize, y: usize| frame[y * w + x] as f64;
let mut sum_abs = 0.0;
for y in 1..h - 1 {
for x in 1..w - 1 {
let v = at(x - 1, y - 1) - 2.0 * at(x, y - 1) + at(x + 1, y - 1) - 2.0 * at(x - 1, y)
+ 4.0 * at(x, y)
- 2.0 * at(x + 1, y)
+ at(x - 1, y + 1)
- 2.0 * at(x, y + 1)
+ at(x + 1, y + 1);
sum_abs += v.abs();
}
}
let count = ((w - 2) * (h - 2)) as f64;
(std::f64::consts::PI / 2.0).sqrt() / (6.0 * count) * sum_abs
}
async fn luma_dims(input: &str, opts: &NoRefOpts) -> anyhow::Result<(usize, usize)> {
let info = if let Some(ref cache) = opts.probe_cache {
cache.probe(input).await?
} else {
viser_ffmpeg::probe(input).await?
};
let v = info.video_stream().ok_or_else(|| anyhow::anyhow!("no video stream in {input}"))?;
if v.width <= 0 || v.height <= 0 {
anyhow::bail!("invalid dimensions for {input}");
}
Ok((v.width as usize, v.height as usize))
}
pub async fn measure_noref(input: &str, opts: &NoRefOpts) -> anyhow::Result<NoRefResult> {
let (w, h) = luma_dims(input, opts).await?;
let frame_size = w * h;
let stride = opts.stride.max(1);
let mut child = Command::new(ffmpeg_path())
.args(["-i", input, "-vf", "format=gray", "-f", "rawvideo", "-"])
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::null())
.spawn()?;
let mut stdout = child.stdout.take().ok_or_else(|| anyhow::anyhow!("no ffmpeg stdout"))?;
let mut buf = vec![0u8; frame_size];
let (mut sharp, mut block, mut noise) = (Vec::new(), Vec::new(), Vec::new());
let mut idx = 0usize;
loop {
match stdout.read_exact(&mut buf).await {
Ok(_) => {
if idx.is_multiple_of(stride) {
sharp.push(variance_of_laplacian(&buf, w, h));
block.push(blockiness(&buf, w, h));
noise.push(noise_sigma(&buf, w, h));
}
idx += 1;
}
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => break,
Err(e) => return Err(e.into()),
}
}
let status = child.wait().await?;
if !status.success() && sharp.is_empty() {
anyhow::bail!("ffmpeg failed to decode {input}");
}
let sharpness_pooled = PooledStats::from_values(&sharp);
let blockiness_pooled = PooledStats::from_values(&block);
let noise_pooled = PooledStats::from_values(&noise);
Ok(NoRefResult {
sharpness: sharpness_pooled.mean,
blockiness: blockiness_pooled.mean,
noise: noise_pooled.mean,
sharpness_pooled,
blockiness_pooled,
noise_pooled,
frames: sharp.len(),
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn flat_frame_is_quiet() {
let frame = vec![128u8; 16 * 16];
assert_eq!(variance_of_laplacian(&frame, 16, 16), 0.0);
assert_eq!(blockiness(&frame, 16, 16), 0.0);
assert_eq!(noise_sigma(&frame, 16, 16), 0.0);
}
#[test]
fn block_grid_reads_as_blockiness() {
let (w, h) = (16, 16);
let mut frame = vec![0u8; w * h];
for y in 0..h {
for x in 0..w {
frame[y * w + x] = (40 * (x / 8)) as u8;
}
}
assert!(blockiness(&frame, w, h) > 0.0);
}
#[test]
fn edge_is_sharper_than_ramp() {
let (w, h) = (32, 32);
let mut edge = vec![0u8; w * h];
let mut ramp = vec![0u8; w * h];
for y in 0..h {
for x in 0..w {
edge[y * w + x] = if x < w / 2 { 0 } else { 255 };
ramp[y * w + x] = ((x * 255) / (w - 1)) as u8;
}
}
assert!(variance_of_laplacian(&edge, w, h) > variance_of_laplacian(&ramp, w, h));
}
#[test]
fn noise_estimate_separates_clean_from_noisy() {
let (w, h) = (32, 32);
let mut clean = vec![0u8; w * h];
let mut noisy = vec![0u8; w * h];
for y in 0..h {
for x in 0..w {
let base = ((x * 255) / (w - 1)) as u8;
clean[y * w + x] = base;
noisy[y * w + x] = if (x + y) % 2 == 0 {
base.saturating_add(30)
} else {
base.saturating_sub(30)
};
}
}
assert!(noise_sigma(&noisy, w, h) > noise_sigma(&clean, w, h));
}
}