use heifer_hevc_dec::recon::Frame;
use crate::image::Image;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ColorParams {
pub matrix: u16,
pub full_range: bool,
}
impl ColorParams {
pub fn from_frame(frame: &Frame) -> Self {
Self {
matrix: u16::from(frame.matrix_coeffs),
full_range: frame.full_range,
}
}
}
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), }
}
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 {
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;
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 {
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 {
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]);
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]
);
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]
);
}
}