1use super::color::{ColorSpec, Range};
2use super::plane::{Sample, check, plane_dims, reshape, samples, samples_mut, sized};
3use crate::media::{Error, Frame, PixelFormat, Result};
4use rayon::prelude::*;
5
6const TILE: usize = 64;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9pub enum Rotation {
10 R90,
11 R180,
12 R270,
13}
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
16pub struct Fill {
17 pub y: u16,
18 pub u: u16,
19 pub v: u16,
20}
21
22impl Fill {
23 pub fn black(range: Range, format: PixelFormat) -> Self {
24 let d = format.bit_depth() as u32;
25 match range {
26 Range::Limited => Fill {
27 y: 16 << (d - 8),
28 u: 128 << (d - 8),
29 v: 128 << (d - 8),
30 },
31 Range::Full => Fill {
32 y: 0,
33 u: 1 << (d - 1),
34 v: 1 << (d - 1),
35 },
36 }
37 }
38
39 pub fn rgb(r: u8, g: u8, b: u8, spec: ColorSpec, format: PixelFormat) -> Self {
40 let l = spec.levels(format.bit_depth());
41 let m = spec.matrix.forward();
42 let c = [r as f64 / 255.0, g as f64 / 255.0, b as f64 / 255.0];
43 let dot = |i: usize| m[i][0] * c[0] + m[i][1] * c[1] + m[i][2] * c[2];
44 let top = ((1u32 << format.bit_depth()) - 1) as f64;
45 let code = |v: f64| v.round().clamp(0.0, top) as u16;
46 Fill {
47 y: code(l.y_off + l.y_range * dot(0)),
48 u: code(l.c_off + l.c_range * dot(1)),
49 v: code(l.c_off + l.c_range * dot(2)),
50 }
51 }
52}
53
54pub fn rotate(src: &Frame, rotation: Rotation, out: &mut Frame) -> Result<()> {
55 check(src)?;
56 let (w, h) = match rotation {
57 Rotation::R180 => (src.width, src.height),
58 _ => (src.height, src.width),
59 };
60 reshape(out, src.format, w, h, src.pts, src.time_base);
61 for i in 0..3 {
62 match src.format {
63 PixelFormat::Yuv420p8 => rotate_plane::<u8>(src, out, i, rotation)?,
64 PixelFormat::Yuv420p10 => rotate_plane::<u16>(src, out, i, rotation)?,
65 }
66 }
67 Ok(())
68}
69
70fn rotate_plane<S: Sample>(
71 src: &Frame,
72 out: &mut Frame,
73 i: usize,
74 rotation: Rotation,
75) -> Result<()> {
76 let (sw, sh) = plane_dims(src.width, src.height, i);
77 let (dw, dh) = plane_dims(out.width, out.height, i);
78 let ss = src.planes[i].stride / S::BYTES;
79 let ds = out.planes[i].stride / S::BYTES;
80 let s = samples::<S>(&src.planes[i])?;
81 let d = samples_mut::<S>(&mut out.planes[i])?;
82 if dw == 0 || dh == 0 {
83 return Ok(());
84 }
85 match rotation {
86 Rotation::R180 => {
87 d.par_chunks_mut(ds).enumerate().for_each(|(y, row)| {
88 let from = &s[(sh - 1 - y) * ss..(sh - 1 - y) * ss + sw];
89 for (o, &v) in row[..dw].iter_mut().zip(from.iter().rev()) {
90 *o = v;
91 }
92 });
93 }
94 Rotation::R90 | Rotation::R270 => {
95 let n = 16 / S::BYTES;
96 d.par_chunks_mut(ds * TILE)
97 .enumerate()
98 .for_each(|(band, chunk)| {
99 let rows = chunk.len().div_ceil(ds).min(dh - band * TILE);
100 let cell = |y: usize, x: usize| {
101 if rotation == Rotation::R90 {
102 s[(sh - 1 - x) * ss + y]
103 } else {
104 s[x * ss + sw - 1 - y]
105 }
106 };
107 for x0 in (0..dw).step_by(n) {
108 for g in (0..rows).step_by(n) {
109 let y0 = band * TILE + g;
110 if x0 + n <= dw
111 && g + n <= rows
112 && transpose::<S>(s, ss, chunk, ds, sh, sw, rotation, g, y0, x0)
113 {
114 continue;
115 }
116 for r in g..(g + n).min(rows) {
117 let end = (x0 + n).min(dw);
118 for (x, o) in
119 chunk[r * ds + x0..r * ds + end].iter_mut().enumerate()
120 {
121 *o = cell(band * TILE + r, x0 + x);
122 }
123 }
124 }
125 }
126 });
127 }
128 }
129 Ok(())
130}
131
132#[allow(clippy::too_many_arguments)]
133#[inline(always)]
134fn transpose<S: Sample>(
135 s: &[S],
136 ss: usize,
137 chunk: &mut [S],
138 ds: usize,
139 sh: usize,
140 sw: usize,
141 rotation: Rotation,
142 g: usize,
143 y0: usize,
144 x0: usize,
145) -> bool {
146 #[cfg(target_arch = "x86_64")]
147 {
148 let n = 16 / S::BYTES;
149 let mut inp = [std::ptr::null::<u8>(); 16];
150 let mut out = [std::ptr::null_mut::<u8>(); 16];
151 let base = chunk.as_mut_ptr();
152 for k in 0..n {
153 let at = if rotation == Rotation::R90 {
154 (sh - 1 - x0 - k) * ss + y0
155 } else {
156 (x0 + k) * ss + sw - n - y0
157 };
158 inp[k] = s[at..at + n].as_ptr() as *const u8;
159 let r = if rotation == Rotation::R90 {
160 g + k
161 } else {
162 g + n - 1 - k
163 };
164 let at = r * ds + x0;
165 assert!(at + n <= chunk.len());
166 out[k] = unsafe { base.add(at) } as *mut u8;
167 }
168 unsafe {
169 if S::BYTES == 1 {
170 t16(&inp, &out);
171 } else {
172 t8(&inp, &out);
173 }
174 }
175 true
176 }
177 #[cfg(not(target_arch = "x86_64"))]
178 {
179 let _ = (s, ss, chunk, ds, sh, sw, rotation, g, y0, x0);
180 false
181 }
182}
183
184#[cfg(target_arch = "x86_64")]
185#[target_feature(enable = "sse2")]
186unsafe fn t16(inp: &[*const u8; 16], out: &[*mut u8; 16]) {
187 use std::arch::x86_64::*;
188 let mut x = [_mm_setzero_si128(); 16];
189 for (v, p) in x.iter_mut().zip(inp) {
190 *v = unsafe { _mm_loadu_si128(*p as *const __m128i) };
191 }
192 for _ in 0..4 {
193 let mut y = x;
194 for i in 0..8 {
195 y[2 * i] = _mm_unpacklo_epi8(x[i], x[i + 8]);
196 y[2 * i + 1] = _mm_unpackhi_epi8(x[i], x[i + 8]);
197 }
198 x = y;
199 }
200 for (v, p) in x.iter().zip(out) {
201 unsafe { _mm_storeu_si128(*p as *mut __m128i, *v) };
202 }
203}
204
205#[cfg(target_arch = "x86_64")]
206#[target_feature(enable = "sse2")]
207unsafe fn t8(inp: &[*const u8; 16], out: &[*mut u8; 16]) {
208 use std::arch::x86_64::*;
209 let mut x = [_mm_setzero_si128(); 8];
210 for (v, p) in x.iter_mut().zip(inp) {
211 *v = unsafe { _mm_loadu_si128(*p as *const __m128i) };
212 }
213 for _ in 0..3 {
214 let mut y = x;
215 for i in 0..4 {
216 y[2 * i] = _mm_unpacklo_epi16(x[i], x[i + 4]);
217 y[2 * i + 1] = _mm_unpackhi_epi16(x[i], x[i + 4]);
218 }
219 x = y;
220 }
221 for (v, p) in x.iter().zip(out) {
222 unsafe { _mm_storeu_si128(*p as *mut __m128i, *v) };
223 }
224}
225
226pub fn flip_horizontal(frame: &mut Frame) -> Result<()> {
227 check(frame)?;
228 let bytes = frame.format.bytes_per_sample();
229 for i in 0..3 {
230 let (w, h) = plane_dims(frame.width, frame.height, i);
231 if w == 0 || h == 0 {
232 continue;
233 }
234 let p = &mut frame.planes[i];
235 let stride = p.stride / bytes;
236 let end = stride * (h - 1) + w;
237 if bytes == 1 {
238 p.data[..end]
239 .par_chunks_mut(stride)
240 .for_each(|row| row[..w].reverse());
241 } else {
242 samples_mut::<u16>(p)?[..end]
243 .par_chunks_mut(stride)
244 .for_each(|row| row[..w].reverse());
245 }
246 }
247 Ok(())
248}
249
250pub fn flip_vertical(frame: &mut Frame) -> Result<()> {
251 check(frame)?;
252 let bytes = frame.format.bytes_per_sample();
253 for i in 0..3 {
254 let (w, h) = plane_dims(frame.width, frame.height, i);
255 if w == 0 || h < 2 {
256 continue;
257 }
258 let p = &mut frame.planes[i];
259 let stride = p.stride;
260 let (top, bottom) = p.data.split_at_mut(stride * (h / 2));
261 let bottom = &mut bottom[stride * (h % 2)..];
262 top.par_chunks_mut(stride)
263 .zip(bottom.par_chunks_mut(stride).rev())
264 .for_each(|(a, b)| a[..w * bytes].swap_with_slice(&mut b[..w * bytes]));
265 }
266 Ok(())
267}
268
269pub fn crop(src: &Frame, x: u32, y: u32, width: u32, height: u32, out: &mut Frame) -> Result<()> {
270 check(src)?;
271 if width == 0
272 || height == 0
273 || x as u64 + width as u64 > src.width as u64
274 || y as u64 + height as u64 > src.height as u64
275 {
276 return Err(Error::Invalid(format!(
277 "crop {width}x{height} at {x},{y} is outside {}x{}",
278 src.width, src.height
279 )));
280 }
281 reshape(out, src.format, width, height, src.pts, src.time_base);
282 let bytes = src.format.bytes_per_sample();
283 for i in 0..3 {
284 let (ox, oy) = if i == 0 { (x, y) } else { (x / 2, y / 2) };
285 let (w, _) = plane_dims(width, height, i);
286 let s = &src.planes[i];
287 let d = &mut out.planes[i];
288 let ds = d.stride;
289 d.data.par_chunks_mut(ds).enumerate().for_each(|(r, row)| {
290 let at = (oy as usize + r) * s.stride + ox as usize * bytes;
291 row[..w * bytes].copy_from_slice(&s.data[at..at + w * bytes]);
292 });
293 }
294 Ok(())
295}
296
297pub fn pad(
298 src: &Frame,
299 width: u32,
300 height: u32,
301 x: u32,
302 y: u32,
303 fill: Fill,
304 out: &mut Frame,
305) -> Result<()> {
306 check(src)?;
307 sized(width, height)?;
308 if x as u64 + src.width as u64 > width as u64 || y as u64 + src.height as u64 > height as u64 {
309 return Err(Error::Invalid(format!(
310 "{}x{} at {x},{y} does not fit in {width}x{height}",
311 src.width, src.height
312 )));
313 }
314 reshape(out, src.format, width, height, src.pts, src.time_base);
315 for i in 0..3 {
316 let value = [fill.y, fill.u, fill.v][i];
317 match src.format {
318 PixelFormat::Yuv420p8 => pad_plane::<u8>(src, out, i, x, y, value as u8)?,
319 PixelFormat::Yuv420p10 => pad_plane::<u16>(src, out, i, x, y, value)?,
320 }
321 }
322 Ok(())
323}
324
325fn pad_plane<S: Sample>(
326 src: &Frame,
327 out: &mut Frame,
328 i: usize,
329 x: u32,
330 y: u32,
331 value: S,
332) -> Result<()> {
333 let (ox, oy) = if i == 0 {
334 (x as usize, y as usize)
335 } else {
336 (x as usize / 2, y as usize / 2)
337 };
338 let (sw, sh) = plane_dims(src.width, src.height, i);
339 let (dw, _) = plane_dims(out.width, out.height, i);
340 let ss = src.planes[i].stride / S::BYTES;
341 let ds = out.planes[i].stride / S::BYTES;
342 let s = samples::<S>(&src.planes[i])?;
343 let d = samples_mut::<S>(&mut out.planes[i])?;
344 if dw == 0 {
345 return Ok(());
346 }
347 d.par_chunks_mut(ds).enumerate().for_each(|(r, row)| {
348 let row = &mut row[..dw];
349 if r < oy || r >= oy + sh {
350 row.fill(value);
351 return;
352 }
353 let end = (ox + sw).min(dw);
354 row[..ox].fill(value);
355 row[ox..end].copy_from_slice(&s[(r - oy) * ss..(r - oy) * ss + end - ox]);
356 row[end..].fill(value);
357 });
358 Ok(())
359}