1pub mod pred;
2
3pub mod rc;
4
5mod bitstream;
6mod bool;
7mod entropy;
8mod frame;
9mod mvref;
10mod search;
11mod tables;
12pub mod tx;
13
14#[cfg(test)]
15mod tests;
16
17use std::sync::{Arc, Mutex};
18
19use bitstream::{CompressedParams, SegParams, Uncompressed, write_compressed, write_uncompressed};
20use bool::BoolWriter;
21use entropy::{Costs, FrameSyntax, NullSink, TileWriter};
22use frame::{MiGrid, Picture, SbCoefs};
23use rc::{Aq, FrameKind, RateControl, RcConfig, RcMode, TileLayout, Wavefront};
24use search::{FrameEnv, SbEncoder, Shared, Speed, Workspace};
25use tables::*;
26use tx::Quant;
27
28use crate::media::{Error, Frame, Packet, PixelFormat, Result};
29use crate::vp9::header::{LoopFilterParams, Segmentation};
30use crate::vp9::loopfilter::{LfFrame, PlaneView, filter_superblock};
31use crate::vp9::probs::{Counts, FrameContext, adapt_coef_probs, adapt_noncoef_probs};
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum Rate {
35 Crf(u8),
36 Bitrate(u32),
37 Q(u8),
38}
39
40#[derive(Debug, Clone, PartialEq)]
41pub struct EncoderConfig {
42 pub width: u32,
43 pub height: u32,
44 pub fps: (u32, u32),
45 pub rate: Rate,
46 pub gop: u32,
47 pub threads: usize,
48 pub aq: bool,
49}
50
51impl Default for EncoderConfig {
52 fn default() -> Self {
53 EncoderConfig {
54 width: 1920,
55 height: 1080,
56 fps: (30, 1),
57 rate: Rate::Crf(32),
58 gop: 240,
59 threads: 0,
60 aq: false,
61 }
62 }
63}
64
65pub struct Encoder {
66 cfg: EncoderConfig,
67 index: u64,
68 contexts: [FrameContext; 4],
69 slots: [Option<Arc<Picture>>; 8],
70 prev_mi: Option<MiGrid>,
71 last_key: bool,
72 mi_rows: usize,
73 mi_cols: usize,
74 sb_cols: usize,
75 coefs: Vec<Mutex<SbCoefs>>,
76 rc: RateControl,
77 wavefront: Wavefront,
78 workers: Vec<Workspace>,
79 recon_out: Option<Arc<Picture>>,
80 broken: bool,
81 color: (u8, bool),
82}
83
84fn quant(qindex: u8) -> [Quant; 2] {
85 let q = qindex as usize;
86 let y = Quant {
87 dc: DC_QLOOKUP[q],
88 ac: AC_QLOOKUP[q],
89 };
90 [y, y]
91}
92
93impl Encoder {
94 pub fn new(cfg: EncoderConfig) -> Result<Encoder> {
95 if cfg.width == 0 || cfg.height == 0 || cfg.width > 16384 || cfg.height > 16384 {
96 return Err(Error::Invalid(format!(
97 "vp9 encoder: invalid size {}x{}",
98 cfg.width, cfg.height
99 )));
100 }
101 let mi_cols = cfg.width.div_ceil(8) as usize;
102 let mi_rows = cfg.height.div_ceil(8) as usize;
103 let sb_cols = mi_cols.div_ceil(8);
104 let sb_rows = mi_rows.div_ceil(8);
105 let threads = match cfg.threads {
106 0 => std::thread::available_parallelism().map_or(1, |n| n.get()),
107 n => n,
108 };
109 let mode = match cfg.rate {
110 Rate::Crf(c) => RcMode::Crf { crf: c.min(63) },
111 Rate::Q(q) => RcMode::ConstantQ { q_index: q },
112 Rate::Bitrate(kbps) => RcMode::Bitrate {
113 kbps,
114 buffer_ms: 1000,
115 },
116 };
117 let mut rcc = RcConfig::new(mode, cfg.width as usize, cfg.height as usize);
118 rcc.fps_num = cfg.fps.0;
119 rcc.fps_den = cfg.fps.1;
120 rcc.keyint = cfg.gop;
121 rcc.aq = if cfg.aq { Aq::Variance } else { Aq::Off };
122 let layout = TileLayout::new(mi_cols, mi_rows, threads);
123 let workers = (0..threads.min(layout.cols.len() * sb_rows).max(1))
124 .map(|_| Workspace::default())
125 .collect();
126 let fc = FrameContext::default();
127 Ok(Encoder {
128 cfg,
129 index: 0,
130 contexts: [fc.clone(), fc.clone(), fc.clone(), fc],
131 slots: Default::default(),
132 prev_mi: None,
133 last_key: true,
134 mi_rows,
135 mi_cols,
136 sb_cols,
137 coefs: (0..sb_rows * sb_cols)
138 .map(|_| Mutex::new(SbCoefs::default()))
139 .collect(),
140 rc: RateControl::new(rcc),
141 wavefront: Wavefront::new(layout),
142 workers,
143 recon_out: None,
144 broken: false,
145 color: (2, false),
146 })
147 }
148
149 pub fn set_color(&mut self, color_space: u8, full_range: bool) -> Result<()> {
150 if color_space > 6 {
151 return Err(Error::Unsupported(format!(
152 "vp9 encoder: color space {color_space} in profile 0"
153 )));
154 }
155 self.color = (color_space, full_range);
156 Ok(())
157 }
158
159 pub fn color(&self) -> (u8, bool) {
160 self.color
161 }
162
163 pub fn encode(&mut self, frame: &Frame) -> Result<Vec<Packet>> {
164 if frame.format != PixelFormat::Yuv420p8 {
165 return Err(Error::Unsupported(format!(
166 "vp9 encoder: {:?} input, 8-bit 4:2:0 only",
167 frame.format
168 )));
169 }
170 if (frame.width, frame.height) != (self.cfg.width, self.cfg.height) {
171 return Err(Error::Invalid(format!(
172 "vp9 encoder: {}x{} frame for a {}x{} stream",
173 frame.width, frame.height, self.cfg.width, self.cfg.height
174 )));
175 }
176 if self.broken {
177 return Err(Error::Backend(
178 "vp9 encoder: an earlier frame failed, start a new encoder".into(),
179 ));
180 }
181 let (w, h) = (self.cfg.width as usize, self.cfg.height as usize);
182 for (p, fp) in frame.planes.iter().enumerate() {
183 let (pw, ph) = if p == 0 {
184 (w, h)
185 } else {
186 (w.div_ceil(2), h.div_ceil(2))
187 };
188 if fp.stride < pw || fp.data.len() < fp.stride * (ph - 1) + pw {
189 return Err(Error::Invalid(format!(
190 "vp9 encoder: plane {p} holds {} bytes at stride {} for {pw}x{ph}",
191 fp.data.len(),
192 fp.stride
193 )));
194 }
195 }
196 let mut src = Picture::new(w, h);
197 for p in 0..3 {
198 let (pw, ph) = src.visible(p);
199 let fp = &frame.planes[p];
200 for y in 0..ph {
201 let o = src.idx(p, 0, y);
202 src.planes[p][o..o + pw]
203 .copy_from_slice(&fp.data[y * fp.stride..y * fp.stride + pw]);
204 }
205 }
206 src.extend_borders();
207 let key = self.index.is_multiple_of(self.cfg.gop.max(1) as u64);
208 let data = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
209 self.encode_picture(&src, key)
210 })) {
211 Ok(d) => d,
212 Err(e) => {
213 self.broken = true;
214 let why = e
215 .downcast_ref::<&str>()
216 .map(|s| s.to_string())
217 .or_else(|| e.downcast_ref::<String>().cloned())
218 .unwrap_or_default();
219 return Err(Error::Backend(format!(
220 "vp9 encoder: frame {} failed: {why}",
221 self.index
222 )));
223 }
224 };
225 let pts = frame.pts;
226 self.index += 1;
227 Ok(vec![Packet {
228 data,
229 pts,
230 dts: pts,
231 keyframe: key,
232 }])
233 }
234
235 pub fn flush(&mut self) -> Result<Vec<Packet>> {
236 Ok(Vec::new())
237 }
238
239 #[cfg(test)]
240 fn recon(&self) -> Option<&Picture> {
241 self.recon_out.as_deref()
242 }
243
244 fn encode_picture(&mut self, src: &Picture, key: bool) -> Vec<u8> {
245 let (mi_rows, mi_cols) = (self.mi_rows, self.mi_cols);
246 if key {
247 let fc = FrameContext::default();
248 self.contexts = [fc.clone(), fc.clone(), fc.clone(), fc];
249 }
250 let kind = if key {
251 FrameKind::Key
252 } else {
253 FrameKind::Inter
254 };
255 let analysis = self
256 .rc
257 .analyse(&src.planes[0][src.origin[0]..], src.stride[0]);
258 let plan = self.rc.plan(kind, &analysis);
259 let base_q_idx = plan.base_q_idx.max(1);
260 let rseg = self.rc.segmentation();
261 let seg_on = rseg.enabled && base_q_idx > 0;
262 let mut seg_q = [quant(base_q_idx); 8];
263 let mut seg_lambda = [plan.lambda; 8];
264 if seg_on {
265 for s in 0..8 {
266 let qi = rseg.q_index(base_q_idx, s as u8).max(1);
267 seg_q[s] = quant(qi);
268 seg_lambda[s] = rc::lambda(qi);
269 }
270 }
271 let mut seg = Segmentation::default();
272 if seg_on {
273 seg.enabled = true;
274 seg.update_map = true;
275 seg.tree_probs = rseg.tree_probs;
276 seg.pred_probs = [255; 3];
277 seg.abs_or_delta_update = rseg.abs_delta;
278 for s in 0..8 {
279 seg.feature_enabled[s][0] = rseg.q[s] != 0;
280 seg.feature_data[s][0] = rseg.q[s];
281 seg.feature_enabled[s][1] = rseg.lf[s] != 0;
282 seg.feature_data[s][1] = rseg.lf[s] as i16;
283 }
284 }
285 let seg_map = seg_on.then(|| rseg.map.clone());
286 let tx_mode = TX_MODE_SELECT;
287 let allow_hp = false;
288 let interp_filter = 0u8;
289 let fc = self.contexts[0].clone();
290 let costs = Costs::new(&fc, key);
291 let mv_cost = pred::mvcost::MvCost::new(pred::mvcost::MvProbs {
292 joints: fc.mv_joint,
293 sign: fc.mv_sign,
294 classes: fc.mv_class,
295 class0: fc.mv_class0_bit,
296 bits: fc.mv_bits,
297 class0_fr: fc.mv_class0_fr,
298 fr: fc.mv_fr,
299 class0_hp: fc.mv_class0_hp,
300 hp: fc.mv_hp,
301 });
302 let speed = Speed {
303 intra_modes: if key {
304 &pred::mode::ALL_INTRA
305 } else {
306 &pred::mode::FAST_INTRA
307 },
308 subpel: pred::motion::Subpel::Quarter,
309 range: 32,
310 intra_margin: 110,
311 tx_depth: 1,
312 trellis: true,
313 };
314 let mut recon = Picture::new(src.width, src.height);
315 let mut mi = MiGrid::new(mi_rows, mi_cols, Vec::new());
316 let mut uv = vec![0u8; mi_rows * mi_cols];
317 let layout = self.wavefront.layout().clone();
318 let refs: [Option<&Picture>; 3] = if key {
319 [None, None, None]
320 } else {
321 [
322 self.slots[0].as_deref(),
323 self.slots[1].as_deref(),
324 self.slots[2].as_deref(),
325 ]
326 };
327 let prev_mi = if key {
328 None
329 } else {
330 self.prev_mi.as_ref().map(|m| m.as_slice())
331 };
332 {
333 let env = FrameEnv {
334 src,
335 refs,
336 key,
337 mi_rows,
338 mi_cols,
339 q: seg_q,
340 lambda: seg_lambda,
341 seg_map: seg_map.as_deref(),
342 tx_mode,
343 allow_hp,
344 costs: &costs,
345 mv_cost: &mv_cost,
346 speed: &speed,
347 prev_mi,
348 sign_bias: [false; 4],
349 filter: interp_filter,
350 };
351 let shared = Shared::new(&mut recon, mi.as_mut_slice(), &mut uv, mi_cols);
352 let coefs = &self.coefs;
353 let sb_cols = self.sb_cols;
354 self.wavefront.run(&mut self.workers, |ws, row| {
355 let tile = layout.cols[row.tile_col];
356 for sb_col in row.sbs() {
357 let mut co = coefs[row.sb_row * sb_cols + sb_col]
358 .lock()
359 .unwrap_or_else(|e| e.into_inner());
360 SbEncoder::encode(&env, &shared, ws, &mut co, tile, row.sb_row, sb_col);
361 }
362 });
363 }
364 let lf_level = plan.filter_level;
365 loop_filter(&mut recon, &mi, lf_level, &seg);
366 let coefs: Vec<SbCoefs> = self
367 .coefs
368 .iter_mut()
369 .map(|m| std::mem::take(m.get_mut().unwrap_or_else(|e| e.into_inner())))
370 .collect();
371 let syntax = FrameSyntax {
372 intra: key,
373 lossless: false,
374 tx_mode,
375 interp_filter,
376 allow_hp,
377 sign_bias: [false; 4],
378 seg: &seg,
379 mi: &mi,
380 uv: &uv,
381 prev_mi,
382 coefs: &coefs,
383 sb_cols: self.sb_cols,
384 fc: &fc,
385 };
386 let counts = {
387 let parts: Vec<Counts> = std::thread::scope(|s| {
388 let hs: Vec<_> = layout
389 .cols
390 .iter()
391 .map(|&(c0, c1)| {
392 let syntax = &syntax;
393 let rows = &layout.rows;
394 s.spawn(move || {
395 let mut tw = TileWriter::new(syntax, NullSink, c0, c1);
396 for &(r0, r1) in rows {
397 tw.tile(r0, r1);
398 }
399 tw.counts
400 })
401 })
402 .collect();
403 hs.into_iter()
404 .map(|h| h.join().unwrap_or_else(|e| std::panic::resume_unwind(e)))
405 .collect()
406 });
407 let mut c = Counts::default();
408 for p in &parts {
409 c.add(p);
410 }
411 c
412 };
413 let mut fc_upd = fc.clone();
414 let compressed = write_compressed(
415 &CompressedParams {
416 intra: key,
417 lossless: false,
418 tx_mode,
419 interp_filter,
420 allow_hp,
421 },
422 &mut fc_upd,
423 &counts,
424 );
425 let syntax = FrameSyntax {
426 fc: &fc_upd,
427 ..syntax
428 };
429 let tiles: Vec<Vec<Vec<u8>>> = std::thread::scope(|s| {
430 let hs: Vec<_> = layout
431 .cols
432 .iter()
433 .map(|&(c0, c1)| {
434 let syntax = &syntax;
435 let rows = &layout.rows;
436 s.spawn(move || {
437 let mut tw = TileWriter::new(syntax, BoolWriter::new(Vec::new()), c0, c1);
438 let mut out = Vec::new();
439 for &(r0, r1) in rows {
440 tw.tile(r0, r1);
441 let w = std::mem::replace(&mut tw.w, BoolWriter::new(Vec::new()));
442 out.push(w.finish());
443 }
444 out
445 })
446 })
447 .collect();
448 hs.into_iter()
449 .map(|h| h.join().unwrap_or_else(|e| std::panic::resume_unwind(e)))
450 .collect()
451 });
452 let hdr = Uncompressed {
453 key,
454 width: src.width as u32,
455 height: src.height as u32,
456 color_space: self.color.0,
457 full_range: self.color.1,
458 refresh_frame_flags: if key { 0xff } else { 0x01 },
459 ref_frame_idx: [0, 1, 2],
460 allow_hp,
461 interp_filter,
462 refresh_frame_context: true,
463 frame_parallel: false,
464 frame_context_idx: 0,
465 lf_level,
466 sharpness: plan.sharpness,
467 base_q_idx,
468 seg: seg_on.then(|| SegParams {
469 tree_probs: seg.tree_probs,
470 abs_delta: seg.abs_or_delta_update,
471 q: rseg_q(&seg),
472 lf: rseg_lf(&seg),
473 }),
474 tile_cols_log2: layout.log2_cols as u32,
475 tile_rows_log2: layout.log2_rows as u32,
476 };
477 let mut out = write_uncompressed(&hdr, compressed.len());
478 out.extend_from_slice(&compressed);
479 let (nr, nc) = (layout.rows.len(), layout.cols.len());
480 for r in 0..nr {
481 for c in 0..nc {
482 let t = &tiles[c][r];
483 if r + 1 < nr || c + 1 < nc {
484 out.extend_from_slice(&(t.len() as u32).to_be_bytes());
485 }
486 out.extend_from_slice(t);
487 }
488 }
489 for (m, c) in self.coefs.iter_mut().zip(coefs) {
490 *m.get_mut().unwrap_or_else(|e| e.into_inner()) = c;
491 }
492 let saved = &self.contexts[0];
493 let mut a = saved.clone();
494 a.tx = fc_upd.tx;
495 a.skip = fc_upd.skip;
496 adapt_coef_probs(&mut a, &counts, key, self.last_key);
497 if !key {
498 a.tx = saved.tx;
499 a.skip = saved.skip;
500 adapt_noncoef_probs(
501 &mut a,
502 &counts,
503 interp_filter == SWITCHABLE,
504 tx_mode == TX_MODE_SELECT,
505 allow_hp,
506 );
507 }
508 self.contexts[0] = a;
509 self.last_key = key;
510 self.rc.update(kind, out.len() as u64 * 8, base_q_idx);
511 recon.extend_borders();
512 let recon = Arc::new(recon);
513 let flags = hdr.refresh_frame_flags;
514 for i in 0..8 {
515 if (flags >> i) & 1 == 1 {
516 self.slots[i] = Some(recon.clone());
517 }
518 }
519 self.recon_out = Some(recon);
520 self.prev_mi = Some(mi);
521 out
522 }
523}
524
525fn rseg_q(seg: &Segmentation) -> [Option<i16>; 8] {
526 std::array::from_fn(|s| seg.feature_enabled[s][0].then_some(seg.feature_data[s][0]))
527}
528
529fn rseg_lf(seg: &Segmentation) -> [Option<i16>; 8] {
530 std::array::from_fn(|s| seg.feature_enabled[s][1].then_some(seg.feature_data[s][1]))
531}
532
533fn loop_filter(recon: &mut Picture, mi: &MiGrid, level: u8, seg: &Segmentation) {
534 if level == 0 {
535 return;
536 }
537 let lf = LoopFilterParams {
538 level,
539 sharpness: 0,
540 delta_enabled: true,
541 ref_deltas: [1, 0, -1, -1],
542 mode_deltas: [0, 0],
543 };
544 let f = LfFrame::new(&lf, seg, mi.rows, mi.cols, true, true, 8);
545 let mut bufs: Vec<Vec<u16>> = (0..3)
546 .map(|p| {
547 let (aw, ah) = recon.aligned[p];
548 let mut v = vec![0u16; aw * ah];
549 for y in 0..ah {
550 let o = recon.idx(p, 0, y);
551 for (d, s) in v[y * aw..(y + 1) * aw]
552 .iter_mut()
553 .zip(&recon.planes[p][o..o + aw])
554 {
555 *d = *s as u16;
556 }
557 }
558 v
559 })
560 .collect();
561 {
562 let [a, b, c] = &mut bufs[..] else {
563 return;
564 };
565 let mut views = [
566 PlaneView {
567 data: a,
568 stride: recon.aligned[0].0,
569 ox: 0,
570 oy: 0,
571 },
572 PlaneView {
573 data: b,
574 stride: recon.aligned[1].0,
575 ox: 0,
576 oy: 0,
577 },
578 PlaneView {
579 data: c,
580 stride: recon.aligned[2].0,
581 ox: 0,
582 oy: 0,
583 },
584 ];
585 for r in 0..mi.rows.div_ceil(8) {
586 for c in 0..mi.cols.div_ceil(8) {
587 filter_superblock(&mut views, mi, &f, r * 8, c * 8);
588 }
589 }
590 }
591 for p in 0..3 {
592 let (aw, ah) = recon.aligned[p];
593 for y in 0..ah {
594 let o = recon.idx(p, 0, y);
595 for (d, s) in recon.planes[p][o..o + aw]
596 .iter_mut()
597 .zip(&bufs[p][y * aw..(y + 1) * aw])
598 {
599 *d = *s as u8;
600 }
601 }
602 }
603}