1use std::sync::Arc;
2
3use crate::filter::{self, ColorSpec, Kernel, Matrix, Range};
4
5use crate::media::{Frame, Rgba};
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
8pub struct Color {
9 pub matrix: u8,
10 pub full_range: bool,
11}
12
13fn spec(color: Color, height: u32) -> ColorSpec {
14 let matrix = match color.matrix {
15 1 => Matrix::Bt709,
16 5 | 6 => Matrix::Bt601,
17 9 | 10 => Matrix::Bt2020,
18 _ if height >= 720 => Matrix::Bt709,
19 _ => Matrix::Bt601,
20 };
21 let range = if color.full_range {
22 Range::Full
23 } else {
24 Range::Limited
25 };
26 ColorSpec::new(matrix, range)
27}
28
29fn rgba_of(frame: &Frame, spec: ColorSpec) -> Rgba {
30 let mut out = Rgba::new(frame.width, frame.height);
31 if filter::yuv_to_rgba(frame, spec, &mut out).is_err() {
32 return Rgba::new(frame.width, frame.height);
33 }
34 Rgba {
35 width: out.width,
36 height: out.height,
37 data: out.data,
38 }
39}
40
41pub fn to_rgba(frame: &Frame) -> Rgba {
42 rgba_of(frame, spec(Color::default(), frame.height))
43}
44
45pub fn to_rgba_height(frame: &Frame, height: u32) -> Rgba {
46 to_rgba_color(frame, height, Color::default())
47}
48
49pub fn to_rgba_color(frame: &Frame, height: u32, color: Color) -> Rgba {
50 let c = spec(color, frame.height);
51 if height == 0 || height >= frame.height || frame.height == 0 {
52 return rgba_of(frame, c);
53 }
54 let w = fit_width(frame.width, frame.height, height);
55 rgba_of(&resize(frame, w, height), c)
56}
57
58pub fn resize(frame: &Frame, width: u32, height: u32) -> Frame {
59 if (frame.width, frame.height) == (width, height) {
60 return frame.clone();
61 }
62 filter::scale(frame, width, height, Kernel::Bilinear).unwrap_or_else(|_| frame.clone())
63}
64
65pub fn scale_to_height(img: &Rgba, height: u32) -> Rgba {
66 if height == 0 || img.height == 0 || img.height == height {
67 return img.clone();
68 }
69 scale(img, fit_width(img.width, img.height, height), height)
70}
71
72pub(super) fn even(v: u32) -> u32 {
73 v.max(2) & !1
74}
75
76pub(super) fn fit_width(width: u32, height: u32, to: u32) -> u32 {
77 even((width as u64 * to as u64 / height as u64) as u32)
78}
79
80pub fn scale(img: &Rgba, width: u32, height: u32) -> Rgba {
81 let mut out = Rgba::new(width, height);
82 let (sw, sh) = (img.width as usize, img.height as usize);
83 if sw == 0 || sh == 0 || width == 0 || height == 0 || img.data.len() < sw * sh * 4 {
84 return out;
85 }
86 let fx = sw as f32 / width as f32;
87 let fy = sh as f32 / height as f32;
88 let span = |o: usize, f: f32, max: usize| {
89 let a = ((o as f32 * f).floor() as usize).min(max - 1);
90 let b = (((o + 1) as f32 * f).ceil() as usize).clamp(a + 1, max);
91 (a, b)
92 };
93 for oy in 0..height as usize {
94 let (y0, y1) = span(oy, fy, sh);
95 for ox in 0..width as usize {
96 let (x0, x1) = span(ox, fx, sw);
97 let mut acc = [0u64; 4];
98 for y in y0..y1 {
99 let row = &img.data[(y * sw + x0) * 4..(y * sw + x1) * 4];
100 for px in row.chunks_exact(4) {
101 for (a, &v) in acc.iter_mut().zip(px) {
102 *a += v as u64;
103 }
104 }
105 }
106 let n = ((y1 - y0) * (x1 - x0)) as u64;
107 let o = (oy * width as usize + ox) * 4;
108 for (d, a) in out.data[o..o + 4].iter_mut().zip(acc) {
109 *d = ((a + n / 2) / n) as u8;
110 }
111 }
112 }
113 out
114}
115
116pub fn rotate(img: &Rgba, clockwise: u16) -> Rgba {
117 let (w, h) = (img.width as usize, img.height as usize);
118 let turns = (clockwise / 90) % 4;
119 if turns == 0 || img.data.len() < w * h * 4 {
120 return img.clone();
121 }
122 let (ow, oh) = if turns % 2 == 1 { (h, w) } else { (w, h) };
123 let mut out = Rgba::new(ow as u32, oh as u32);
124 for y in 0..h {
125 for x in 0..w {
126 let (nx, ny) = match turns {
127 1 => (h - 1 - y, x),
128 2 => (w - 1 - x, h - 1 - y),
129 _ => (y, w - 1 - x),
130 };
131 let s = (y * w + x) * 4;
132 let d = (ny * ow + nx) * 4;
133 out.data[d..d + 4].copy_from_slice(&img.data[s..s + 4]);
134 }
135 }
136 out
137}
138
139pub struct Blender {
140 layer: Option<(Arc<Rgba>, filter::Overlay)>,
141}
142
143impl Default for Blender {
144 fn default() -> Self {
145 Self::new()
146 }
147}
148
149impl Blender {
150 pub fn new() -> Self {
151 Blender { layer: None }
152 }
153
154 pub fn draw(&mut self, frame: &mut Frame, layer: &Arc<Rgba>) {
155 let fresh = match &self.layer {
156 Some((held, o)) => !Arc::ptr_eq(held, layer) || o.size() != (frame.width, frame.height),
157 None => true,
158 };
159 if fresh {
160 let fitted;
161 let image = if (layer.width, layer.height) == (frame.width, frame.height) {
162 layer.as_ref()
163 } else {
164 fitted = scale(layer, frame.width, frame.height);
165 &fitted
166 };
167 self.layer =
168 filter::Overlay::new(image, spec(Color::default(), frame.height), frame.format)
169 .ok()
170 .map(|o| (layer.clone(), o));
171 }
172 if let Some((_, o)) = &self.layer {
173 let _ = o.draw(frame, 0, 0);
174 }
175 }
176}
177
178pub fn blend(frame: &mut Frame, layer: &Rgba) {
179 Blender::new().draw(frame, &Arc::new(layer.clone()));
180}
181
182#[cfg(test)]
183mod tests {
184 use super::*;
185 use crate::media::PixelFormat;
186 use crate::media::TimeBase;
187
188 #[test]
189 fn an_opaque_layer_paints_its_colour_and_a_clear_one_keeps_the_frame() {
190 let mut f = Frame::new(PixelFormat::Yuv420p8, 4, 4, 0, TimeBase::new(1, 30));
191 for p in &mut f.planes {
192 p.data.fill(128);
193 }
194 f.planes[0].data.fill(100);
195 let mut layer = Rgba::new(4, 4);
196 for px in layer.data.chunks_exact_mut(4).take(8) {
197 px.copy_from_slice(&[255, 0, 0, 255]);
198 }
199 blend(&mut f, &layer);
200 let rgb = to_rgba(&f);
201 assert!(rgb.data[0] > 245 && rgb.data[1] < 10 && rgb.data[2] < 10);
202 let bottom = &rgb.data[(3 * 4) * 4..(3 * 4) * 4 + 3];
203 assert!(bottom.iter().all(|c| c.abs_diff(bottom[0]) <= 1));
204 assert_eq!(f.planes[0].data[15], 100);
205 }
206
207 #[test]
208 fn a_kept_blender_paints_what_a_fresh_blend_paints() {
209 let mut layer = Rgba::new(64, 36);
210 for (i, px) in layer.data.chunks_exact_mut(4).enumerate() {
211 px.copy_from_slice(&[(i * 7) as u8, 200, (i * 3) as u8, (i % 256) as u8]);
212 }
213 let layer = Arc::new(layer);
214 let mut kept = Blender::default();
215 for shade in [40u8, 170, 230] {
216 let mut a = Frame::new(PixelFormat::Yuv420p8, 64, 36, 0, TimeBase::new(1, 30));
217 for (k, p) in a.planes.iter_mut().enumerate() {
218 p.data.fill(shade.wrapping_add(k as u8 * 30));
219 }
220 let mut b = a.clone();
221 kept.draw(&mut a, &layer);
222 blend(&mut b, &layer);
223 assert_eq!(a, b, "{shade}");
224 }
225 }
226}