heifer 0.1.1

Pure-Rust HEIF/HEIC image decoder: no C dependencies, no unsafe code
Documentation
//! YCbCr to RGB conversion (ITU-T H.273 matrix coefficients).

use heifer_hevc_dec::recon::Frame;

use crate::image::Image;

/// How to interpret YCbCr samples.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ColorParams {
    /// `matrix_coefficients` (H.273): 0 = identity (GBR), 1 = BT.709, 5/6 = BT.601, 9 = BT.2020.
    pub matrix: u16,
    /// Whether samples use the full range (otherwise "video" range, e.g. 16..235 for 8-bit).
    pub full_range: bool,
}

impl ColorParams {
    /// Colour description of an HEVC stream's VUI. When absent, HEVC specifies limited range,
    /// and the matrix is unspecified (treated as BT.601).
    pub fn from_frame(frame: &Frame) -> Self {
        Self {
            matrix: u16::from(frame.matrix_coeffs),
            full_range: frame.full_range,
        }
    }
}

/// `(Kr, Kb)` for a matrix coefficients code point.
fn kr_kb(matrix: u16) -> (f64, f64) {
    match matrix {
        1 => (0.2126, 0.0722),
        4 => (0.30, 0.11),
        7 => (0.212, 0.087),
        9 | 10 => (0.2627, 0.0593),
        _ => (0.299, 0.114), // 5, 6 and unspecified: BT.601
    }
}

/// Converts a decoded frame to an RGBA image (alpha opaque). Chroma is upsampled by
/// nearest neighbour.
pub fn frame_to_rgba(frame: &Frame, params: ColorParams) -> Image {
    let (w, h) = (frame.widths[0], frame.heights[0]);
    let bd_y = frame.bit_depth[0];
    let bd_c = frame.bit_depth[1];
    let out_depth = bd_y.max(bd_c);
    let max_out = f64::from((1u32 << out_depth) - 1);
    let mut data = Vec::with_capacity(w as usize * h as usize * 4);
    let alpha = (1u32 << out_depth) as u16 - 1;

    if frame.widths[1] == 0 {
        // Monochrome.
        let scale = max_out / f64::from((1u32 << bd_y) - 1);
        for &v in &frame.planes[0] {
            let g = (f64::from(v) * scale).round() as u16;
            data.extend_from_slice(&[g, g, g, alpha]);
        }
        return Image {
            width: w,
            height: h,
            bit_depth: out_depth,
            has_alpha: false,
            data,
        };
    }

    let (cw, chh) = (frame.widths[1], frame.heights[1]);
    let (sx, sy) = (w / cw.max(1), h / chh.max(1));
    let (kr, kb) = kr_kb(params.matrix);
    let kg = 1.0 - kr - kb;

    // Normalize luma to 0..1 and chroma to -0.5..0.5.
    let (y_off, y_scale, c_scale) = {
        let (fy, fc) = (f64::from(1u32 << (bd_y - 8)), f64::from(1u32 << (bd_c - 8)));
        if params.full_range {
            (
                0.0,
                f64::from((1u32 << bd_y) - 1),
                f64::from((1u32 << bd_c) - 1),
            )
        } else {
            (16.0 * fy, 219.0 * fy, 224.0 * fc)
        }
    };
    let c_mid = f64::from(1u32 << (bd_c - 1));
    let to_out = |v: f64| (v * max_out).round().clamp(0.0, max_out) as u16;

    if params.matrix != 0 {
        // Same arithmetic as the generic loop below, with the per-sample terms looked up in
        // tables (one entry per possible sample value), so the result is identical.
        let lut = |bits: u8, f: &dyn Fn(f64) -> f64| -> Vec<f64> {
            (0..1u32 << bits).map(|v| f(f64::from(v))).collect()
        };
        let yn_lut = lut(bd_y, &|v| (v - y_off) / y_scale);
        let r_lut = lut(bd_c, &|v| 2.0 * (1.0 - kr) * ((v - c_mid) / c_scale));
        let b_lut = lut(bd_c, &|v| 2.0 * (1.0 - kb) * ((v - c_mid) / c_scale));
        let at = |t: &[f64], v: u16| t[usize::from(v).min(t.len() - 1)];
        let xcs: Vec<usize> = (0..w).map(|x| (x / sx).min(cw - 1) as usize).collect();
        data.resize(w as usize * h as usize * 4, alpha);
        for (y, out_row) in data.chunks_exact_mut(w as usize * 4).enumerate() {
            let yc = (y as u32 / sy).min(chh - 1) as usize;
            let luma = &frame.planes[0][y * w as usize..(y + 1) * w as usize];
            let cb_row = &frame.planes[1][yc * cw as usize..(yc + 1) * cw as usize];
            let cr_row = &frame.planes[2][yc * cw as usize..(yc + 1) * cw as usize];
            for ((px, &yv), &xc) in out_row.chunks_exact_mut(4).zip(luma).zip(&xcs) {
                let yn = at(&yn_lut, yv);
                let r = yn + at(&r_lut, cr_row[xc]);
                let b = yn + at(&b_lut, cb_row[xc]);
                let g = (yn - kr * r - kb * b) / kg;
                px[..3].copy_from_slice(&[to_out(r), to_out(g), to_out(b)]);
            }
        }
        return Image {
            width: w,
            height: h,
            bit_depth: out_depth,
            has_alpha: false,
            data,
        };
    }

    for y in 0..h {
        let yc = (y / sy).min(chh - 1) as usize;
        for x in 0..w {
            let xc = (x / sx).min(cw - 1) as usize;
            let yv = f64::from(frame.planes[0][(y * w + x) as usize]);
            let cb = f64::from(frame.planes[1][yc * cw as usize + xc]);
            let cr = f64::from(frame.planes[2][yc * cw as usize + xc]);
            let (r, g, b) = if params.matrix == 0 {
                // Identity: planes are G, B, R.
                let s = max_out / f64::from((1u32 << bd_y) - 1);
                ((cr * s) / max_out, (yv * s) / max_out, (cb * s) / max_out)
            } else {
                let yn = (yv - y_off) / y_scale;
                let cbn = (cb - c_mid) / c_scale;
                let crn = (cr - c_mid) / c_scale;
                let r = yn + 2.0 * (1.0 - kr) * crn;
                let b = yn + 2.0 * (1.0 - kb) * cbn;
                let g = (yn - kr * r - kb * b) / kg;
                (r, g, b)
            };
            data.extend_from_slice(&[to_out(r), to_out(g), to_out(b), alpha]);
        }
    }
    Image {
        width: w,
        height: h,
        bit_depth: out_depth,
        has_alpha: false,
        data,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn frame(y: u16, cb: u16, cr: u16) -> Frame {
        Frame {
            widths: [2, 1, 1],
            heights: [2, 1, 1],
            planes: [vec![y; 4], vec![cb], vec![cr]],
            bit_depth: [8, 8],
            chroma_format: 1,
            full_range: false,
            matrix_coeffs: 2,
        }
    }

    #[test]
    fn grey_and_primaries() {
        let full = ColorParams {
            matrix: 6,
            full_range: true,
        };
        let img = frame_to_rgba(&frame(128, 128, 128), full);
        assert_eq!(&img.data[..4], &[128, 128, 128, 255]);
        // Limited range black and white.
        let limited = ColorParams {
            matrix: 1,
            full_range: false,
        };
        assert_eq!(
            &frame_to_rgba(&frame(16, 128, 128), limited).data[..3],
            &[0, 0, 0]
        );
        assert_eq!(
            &frame_to_rgba(&frame(235, 128, 128), limited).data[..3],
            &[255, 255, 255]
        );
        // Pure red in BT.601 full range: Y=76, Cb=85, Cr=255.
        let red = frame_to_rgba(&frame(76, 85, 255), full);
        assert!(
            red.data[0] >= 253 && red.data[1] <= 2 && red.data[2] <= 2,
            "{:?}",
            &red.data[..3]
        );
    }
}