av1-grain 0.5.0

Helpers for generating and parsing AV1 film grain data
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

use std::f64::consts::PI;

use v_frame::{pixel::Pixel, plane::Plane};

/// Estimates the amount of noise within a plane.
/// Returns `None` if a reliable estimate cannot be obtained
/// due to too few smooth pixels.
///
/// Ranges seem to be approximately:
///
/// - `0.0..0.5` = no noticeable noise
/// - `0.5..1.0` = light noise, probably photon-noise-esque
///
/// # Panics
/// - If called with a `bit_depth` not between `8..=16`
#[must_use]
#[inline]
pub fn estimate_plane_noise<T: Pixel>(plane: &Plane<T>, bit_depth: usize) -> Option<f64> {
    const EDGE_THRESHOLD: u16 = 50;

    if size_of::<T>() == 1 {
        assert_eq!(bit_depth, 8);
    } else if size_of::<T>() == 2 {
        assert!(bit_depth > 8 && bit_depth <= 16);
    } else {
        unimplemented!("Bit depths greater than 16 are not currently supported");
    }

    let width = plane.width();
    let height = plane.height();
    let data_origin = plane.geometry().data_origin();
    let stride = plane.geometry().stride();
    let data = plane.data();

    let mut accum = 0u64;
    let mut count = 0u64;
    for i in 1..(height - 1) {
        for j in 1..(width - 1) {
            // Setup a small 3x3 matrix.
            let mut mat = [[0i16; 3]; 3];
            for ii in -1isize..=1isize {
                for jj in -1isize..=1isize {
                    unsafe {
                        let idx = (i * stride + j) as isize + ii * stride as isize + jj;
                        let pix: u16 = (*data.get_unchecked(data_origin + idx as usize)).into();

                        *mat.get_unchecked_mut((ii + 1) as usize)
                            .get_unchecked_mut((jj + 1) as usize) = if size_of::<T>() == 1 {
                            pix as i16
                        } else {
                            (pix >> (bit_depth - 8)) as i16
                        };
                    }
                }
            }

            // Compute sobel gradients
            let g_x =
                (mat[0][0] - mat[0][2]) + (mat[2][0] - mat[2][2]) + 2 * (mat[1][0] - mat[1][2]);
            let g_y =
                (mat[0][0] - mat[2][0]) + (mat[0][2] - mat[2][2]) + 2 * (mat[0][1] - mat[2][1]);
            let g_a = (g_x.abs() + g_y.abs()) as u16;
            // Accumulate Laplacian
            if g_a < EDGE_THRESHOLD {
                // Only count smooth pixels
                let v = 4 * mat[1][1] - 2 * (mat[0][1] + mat[2][1] + mat[1][0] + mat[1][2])
                    + (mat[0][0] + mat[0][2] + mat[2][0] + mat[2][2]);
                accum += u64::from(v.unsigned_abs());
                count += 1;
            }
        }
    }

    (count >= 16).then(|| accum as f64 / (6u64 * count) as f64 * (PI / 2f64).sqrt())
}

#[cfg(test)]
mod tests {
    use v_frame::{chroma::ChromaSubsampling, frame::FrameBuilder};

    use super::estimate_plane_noise;

    #[test]
    fn estimate_uses_visible_pixels_when_plane_has_padding() {
        let visible = vec![10u8; 6 * 6];
        let mut frame = FrameBuilder::new(6, 6, ChromaSubsampling::Yuv444, 8)
            .luma_padding_left(2)
            .luma_padding_right(2)
            .luma_padding_top(2)
            .luma_padding_bottom(2)
            .build::<u8>()
            .expect("valid padded frame");

        frame.y_plane.data_mut().fill(255);
        frame
            .y_plane
            .copy_from_slice(&visible)
            .expect("visible plane data matches dimensions");

        assert_eq!(estimate_plane_noise(&frame.y_plane, 8), Some(0.0));
    }
}