1// Copyright (c) 2018-2022, The rav1e contributors. All rights reserved
2//
3// This source code is subject to the terms of the BSD 2 Clause License and
4// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
5// was not distributed with this source code in the LICENSE file, you can
6// obtain it at www.aomedia.org/license/software. If the Alliance for Open
7// Media Patent License 1.0 was not distributed with this source code in the
8// PATENTS file, you can obtain it at www.aomedia.org/license/patent.
9#![deny(missing_docs)]
10
11use crate::activity::ActivityMask;
12use crate::api::lookahead::*;
13use crate::api::{
14 EncoderConfig, EncoderStatus, FrameType, Opaque, Packet, T35,
15};
16use crate::color::ChromaSampling::Cs400;
17use crate::cpu_features::CpuFeatureLevel;
18use crate::dist::get_satd;
19use crate::encoder::*;
20use crate::frame::*;
21use crate::partition::*;
22use crate::rate::{
23 RCState, FRAME_NSUBTYPES, FRAME_SUBTYPE_I, FRAME_SUBTYPE_P,
24 FRAME_SUBTYPE_SEF,
25};
26use crate::scenechange::SceneChangeDetector;
27use crate::stats::EncoderStats;
28use crate::tiling::Area;
29use crate::util::Pixel;
30use arrayvec::ArrayVec;
31use std::cmp;
32use std::collections::{BTreeMap, BTreeSet};
33use std::env;
34use std::fs;
35use std::path::PathBuf;
36use std::sync::Arc;
37
38/// The set of options that controls frame re-ordering and reference picture
39/// selection.
40/// The options stored here are invariant over the whole encode.
41#[derive(Debug, Clone, Copy)]
42pub struct InterConfig {
43 /// Whether frame re-ordering is enabled.
44 reorder: bool,
45 /// Whether P-frames can use multiple references.
46 pub(crate) multiref: bool,
47 /// The depth of the re-ordering pyramid.
48 /// The current code cannot support values larger than 2.
49 pub(crate) pyramid_depth: u64,
50 /// Number of input frames in group.
51 pub(crate) group_input_len: u64,
52 /// Number of output frames in group.
53 /// This includes both hidden frames and "show existing frame" frames.
54 group_output_len: u64,
55 /// Interval between consecutive S-frames.
56 /// Keyframes reset this interval.
57 /// This MUST be a multiple of group_input_len.
58 pub(crate) switch_frame_interval: u64,
59}
60
61impl InterConfig {
62 pub(crate) fn new(enc_config: &EncoderConfig) -> InterConfig {
63 let reorder = !enc_config.low_latency;
64 // A group always starts with (group_output_len - group_input_len) hidden
65 // frames, followed by group_input_len shown frames.
66 // The shown frames iterate over the input frames in order, with frames
67 // already encoded as hidden frames now displayed with Show Existing
68 // Frame.
69 // For example, for a pyramid depth of 2, the group is as follows:
70 // |TU |TU |TU |TU
71 // idx_in_group_output: 0 1 2 3 4 5
72 // input_frameno: 4 2 1 SEF 3 SEF
73 // output_frameno: 1 2 3 4 5 6
74 // level: 0 1 2 1 2 0
75 // ^^^^^ ^^^^^^^^^^^^^
76 // hidden shown
77 // TODO: This only works for pyramid_depth <= 2 --- after that we need
78 // more hidden frames in the middle of the group.
79 let pyramid_depth = if reorder { 2 } else { 0 };
80 let group_input_len = 1 << pyramid_depth;
81 let group_output_len = group_input_len + pyramid_depth;
82 let switch_frame_interval = enc_config.switch_frame_interval;
83 assert!(switch_frame_interval % group_input_len == 0);
84 InterConfig {
85 reorder,
86 multiref: reorder || enc_config.speed_settings.multiref,
87 pyramid_depth,
88 group_input_len,
89 group_output_len,
90 switch_frame_interval,
91 }
92 }
93
94 /// Get the index of an output frame in its re-ordering group given the output
95 /// frame number of the frame in the current keyframe gop.
96 /// When re-ordering is disabled, this always returns 0.
97 pub(crate) fn get_idx_in_group_output(
98 &self, output_frameno_in_gop: u64,
99 ) -> u64 {
100 // The first frame in the GOP should be a keyframe and is not re-ordered,
101 // so we should not be calling this function on it.
102 debug_assert!(output_frameno_in_gop > 0);
103 (output_frameno_in_gop - 1) % self.group_output_len
104 }
105
106 /// Get the order-hint of an output frame given the output frame number of the
107 /// frame in the current keyframe gop and the index of that output frame
108 /// in its re-ordering gorup.
109 pub(crate) fn get_order_hint(
110 &self, output_frameno_in_gop: u64, idx_in_group_output: u64,
111 ) -> u32 {
112 // The first frame in the GOP should be a keyframe, but currently this
113 // function only handles inter frames.
114 // We could return 0 for keyframes if keyframe support is needed.
115 debug_assert!(output_frameno_in_gop > 0);
116 // Which P-frame group in the current gop is this output frame in?
117 // Subtract 1 because the first frame in the gop is always a keyframe.
118 let group_idx = (output_frameno_in_gop - 1) / self.group_output_len;
119 // Get the offset to the corresponding input frame.
120 // TODO: This only works with pyramid_depth <= 2.
121 let offset = if idx_in_group_output < self.pyramid_depth {
122 self.group_input_len >> idx_in_group_output
123 } else {
124 idx_in_group_output - self.pyramid_depth + 1
125 };
126 // Construct the final order hint relative to the start of the group.
127 (self.group_input_len * group_idx + offset) as u32
128 }
129
130 /// Get the level of the current frame in the pyramid.
131 pub(crate) const fn get_level(&self, idx_in_group_output: u64) -> u64 {
132 if !self.reorder {
133 0
134 } else if idx_in_group_output < self.pyramid_depth {
135 // Hidden frames are output first (to be shown in the future).
136 idx_in_group_output
137 } else {
138 // Shown frames
139 // TODO: This only works with pyramid_depth <= 2.
140 pos_to_lvl(
141 idx_in_group_output - self.pyramid_depth + 1,
142 self.pyramid_depth,
143 )
144 }
145 }
146
147 pub(crate) const fn get_slot_idx(&self, level: u64, order_hint: u32) -> u32 {
148 // Frames with level == 0 are stored in slots 0..4, and frames with higher
149 // values of level in slots 4..8
150 if level == 0 {
151 (order_hint >> self.pyramid_depth) & 3
152 } else {
153 // This only works with pyramid_depth <= 4.
154 3 + level as u32
155 }
156 }
157
158 pub(crate) const fn get_show_frame(&self, idx_in_group_output: u64) -> bool {
159 idx_in_group_output >= self.pyramid_depth
160 }
161
162 pub(crate) const fn get_show_existing_frame(
163 &self, idx_in_group_output: u64,
164 ) -> bool {
165 // The self.reorder test here is redundant, but short-circuits the rest,
166 // avoiding a bunch of work when it's false.
167 self.reorder
168 && self.get_show_frame(idx_in_group_output)
169 && (idx_in_group_output - self.pyramid_depth + 1).count_ones() == 1
170 && idx_in_group_output != self.pyramid_depth
171 }
172
173 pub(crate) fn get_input_frameno(
174 &self, output_frameno_in_gop: u64, gop_input_frameno_start: u64,
175 ) -> u64 {
176 if output_frameno_in_gop == 0 {
177 gop_input_frameno_start
178 } else {
179 let idx_in_group_output =
180 self.get_idx_in_group_output(output_frameno_in_gop);
181 let order_hint =
182 self.get_order_hint(output_frameno_in_gop, idx_in_group_output);
183 gop_input_frameno_start + order_hint as u64
184 }
185 }
186
187 const fn max_reordering_latency(&self) -> u64 {
188 self.group_input_len
189 }
190
191 pub(crate) const fn keyframe_lookahead_distance(&self) -> u64 {
192 self.max_reordering_latency() + 1
193 }
194
195 pub(crate) const fn allowed_ref_frames(&self) -> &[RefType] {
196 use crate::partition::RefType::*;
197 if self.reorder {
198 &ALL_INTER_REFS
199 } else if self.multiref {
200 &[LAST_FRAME, LAST2_FRAME, LAST3_FRAME, GOLDEN_FRAME]
201 } else {
202 &[LAST_FRAME]
203 }
204 }
205}
206
207// Thin wrapper for frame-related data
208// that gets cached and reused throughout the life of a frame.
209#[derive(Clone)]
210pub(crate) struct FrameData<T: Pixel> {
211 pub(crate) fi: FrameInvariants<T>,
212 pub(crate) fs: FrameState<T>,
213}
214
215impl<T: Pixel> FrameData<T> {
216 pub(crate) fn new(fi: FrameInvariants<T>, frame: Arc<Frame<T>>) -> Self {
217 let fs: FrameState = FrameState::new_with_frame(&fi, frame);
218 FrameData { fi, fs }
219 }
220}
221
222type FrameQueue<T> = BTreeMap<u64, Option<Arc<Frame<T>>>>;
223type FrameDataQueue<T> = BTreeMap<u64, Option<FrameData<T>>>;
224
225// the fields pub(super) are accessed only by the tests
226pub(crate) struct ContextInner<T: Pixel> {
227 pub(crate) frame_count: u64,
228 pub(crate) limit: Option<u64>,
229 pub(crate) output_frameno: u64,
230 pub(super) inter_cfg: InterConfig,
231 pub(super) frames_processed: u64,
232 /// Maps *input_frameno* to frames
233 pub(super) frame_q: FrameQueue<T>,
234 /// Maps *output_frameno* to frame data
235 pub(super) frame_data: FrameDataQueue<T>,
236 /// A list of the input_frameno for keyframes in this encode.
237 /// Needed so that we don't need to keep all of the frame_invariants in
238 /// memory for the whole life of the encode.
239 // TODO: Is this needed at all?
240 keyframes: BTreeSet<u64>,
241 // TODO: Is this needed at all?
242 keyframes_forced: BTreeSet<u64>,
243 /// A storage space for reordered frames.
244 packet_data: Vec<u8>,
245 /// Maps `output_frameno` to `gop_output_frameno_start`.
246 gop_output_frameno_start: BTreeMap<u64, u64>,
247 /// Maps `output_frameno` to `gop_input_frameno_start`.
248 pub(crate) gop_input_frameno_start: BTreeMap<u64, u64>,
249 keyframe_detector: SceneChangeDetector<T>,
250 pub(crate) config: Arc<EncoderConfig>,
251 seq: Arc<Sequence>,
252 pub(crate) rc_state: RCState,
253 maybe_prev_log_base_q: Option<i64>,
254 /// The next `input_frameno` to be processed by lookahead.
255 next_lookahead_frame: u64,
256 /// The next `output_frameno` to be computed by lookahead.
257 next_lookahead_output_frameno: u64,
258 /// Optional opaque to be sent back to the user
259 opaque_q: BTreeMap<u64, Opaque>,
260 /// Optional T35 metadata per frame
261 t35_q: BTreeMap<u64, Box<[T35]>>,
262}
263
264impl<T: Pixel> ContextInner<T> {
265 pub fn new(enc: &EncoderConfig) -> Self {
266 // initialize with temporal delimiter
267 let packet_data = TEMPORAL_DELIMITER.to_vec();
268 let mut keyframes = BTreeSet::new();
269 keyframes.insert(0);
270
271 let maybe_ac_qi_max =
272 if enc.quantizer < 255 { Some(enc.quantizer as u8) } else { None };
273
274 let seq = Arc::new(Sequence::new(enc));
275 let inter_cfg = InterConfig::new(enc);
276 let lookahead_distance = inter_cfg.keyframe_lookahead_distance() as usize;
277
278 ContextInner {
279 frame_count: 0,
280 limit: None,
281 inter_cfg,
282 output_frameno: 0,
283 frames_processed: 0,
284 frame_q: BTreeMap::new(),
285 frame_data: BTreeMap::new(),
286 keyframes,
287 keyframes_forced: BTreeSet::new(),
288 packet_data,
289 gop_output_frameno_start: BTreeMap::new(),
290 gop_input_frameno_start: BTreeMap::new(),
291 keyframe_detector: SceneChangeDetector::new(
292 enc.clone(),
293 CpuFeatureLevel::default(),
294 lookahead_distance,
295 seq.clone(),
296 ),
297 config: Arc::new(enc.clone()),
298 seq,
299 rc_state: RCState::new(
300 enc.width as i32,
301 enc.height as i32,
302 enc.time_base.den as i64,
303 enc.time_base.num as i64,
304 enc.bitrate,
305 maybe_ac_qi_max,
306 enc.min_quantizer,
307 enc.max_key_frame_interval as i32,
308 enc.reservoir_frame_delay,
309 ),
310 maybe_prev_log_base_q: None,
311 next_lookahead_frame: 1,
312 next_lookahead_output_frameno: 0,
313 opaque_q: BTreeMap::new(),
314 t35_q: BTreeMap::new(),
315 }
316 }
317
318 #[profiling::function]
319 pub fn send_frame(
320 &mut self, mut frame: Option<Arc<Frame<T>>>,
321 params: Option<FrameParameters>,
322 ) -> Result<(), EncoderStatus> {
323 if let Some(ref mut frame) = frame {
324 use crate::api::color::ChromaSampling;
325 let EncoderConfig { width, height, chroma_sampling, .. } = *self.config;
326 let planes =
327 if chroma_sampling == ChromaSampling::Cs400 { 1 } else { 3 };
328 // Try to add padding
329 if let Some(ref mut frame) = Arc::get_mut(frame) {
330 for plane in frame.planes[..planes].iter_mut() {
331 plane.pad(width, height);
332 }
333 }
334 // Enforce that padding is added
335 for (p, plane) in frame.planes[..planes].iter().enumerate() {
336 assert!(
337 plane.probe_padding(width, height),
338 "Plane {p} was not padded before passing Frame to send_frame()."
339 );
340 }
341 }
342
343 let input_frameno = self.frame_count;
344 let is_flushing = frame.is_none();
345 if !is_flushing {
346 self.frame_count += 1;
347 }
348 self.frame_q.insert(input_frameno, frame);
349
350 if let Some(params) = params {
351 if params.frame_type_override == FrameTypeOverride::Key {
352 self.keyframes_forced.insert(input_frameno);
353 }
354 if let Some(op) = params.opaque {
355 self.opaque_q.insert(input_frameno, op);
356 }
357 self.t35_q.insert(input_frameno, params.t35_metadata);
358 }
359
360 if !self.needs_more_frame_q_lookahead(self.next_lookahead_frame) {
361 let lookahead_frames = self
362 .frame_q
363 .range(self.next_lookahead_frame - 1..)
364 .filter_map(|(&_input_frameno, frame)| frame.as_ref())
365 .collect::<Vec<&Arc<Frame<T>>>>();
366
367 if is_flushing {
368 // This is the last time send_frame is called, process all the
369 // remaining frames.
370 for cur_lookahead_frames in
371 std::iter::successors(Some(&lookahead_frames[..]), |s| s.get(1..))
372 {
373 if cur_lookahead_frames.len() < 2 {
374 // All frames have been processed
375 break;
376 }
377
378 Self::compute_keyframe_placement(
379 cur_lookahead_frames,
380 &self.keyframes_forced,
381 &mut self.keyframe_detector,
382 &mut self.next_lookahead_frame,
383 &mut self.keyframes,
384 );
385 }
386 } else {
387 Self::compute_keyframe_placement(
388 &lookahead_frames,
389 &self.keyframes_forced,
390 &mut self.keyframe_detector,
391 &mut self.next_lookahead_frame,
392 &mut self.keyframes,
393 );
394 }
395 }
396
397 self.compute_frame_invariants();
398
399 Ok(())
400 }
401
402 /// Indicates whether more frames need to be read into the frame queue
403 /// in order for frame queue lookahead to be full.
404 fn needs_more_frame_q_lookahead(&self, input_frameno: u64) -> bool {
405 let lookahead_end = self.frame_q.keys().last().cloned().unwrap_or(0);
406 let frames_needed =
407 input_frameno + self.inter_cfg.keyframe_lookahead_distance() + 1;
408 lookahead_end < frames_needed && self.needs_more_frames(lookahead_end)
409 }
410
411 /// Indicates whether more frames need to be processed into `FrameInvariants`
412 /// in order for FI lookahead to be full.
413 pub fn needs_more_fi_lookahead(&self) -> bool {
414 let ready_frames = self.get_rdo_lookahead_frames().count();
415 ready_frames < self.config.speed_settings.rdo_lookahead_frames + 1
416 && self.needs_more_frames(self.next_lookahead_frame)
417 }
418
419 pub fn needs_more_frames(&self, frame_count: u64) -> bool {
420 self.limit.map(|limit| frame_count < limit).unwrap_or(true)
421 }
422
423 fn get_rdo_lookahead_frames(
424 &self,
425 ) -> impl Iterator<Item = (&u64, &FrameData<T>)> {
426 self
427 .frame_data
428 .iter()
429 .skip_while(move |(&output_frameno, _)| {
430 output_frameno < self.output_frameno
431 })
432 .filter_map(|(fno, data)| data.as_ref().map(|data| (fno, data)))
433 .filter(|(_, data)| !data.fi.is_show_existing_frame())
434 .take(self.config.speed_settings.rdo_lookahead_frames + 1)
435 }
436
437 fn next_keyframe_input_frameno(
438 &self, gop_input_frameno_start: u64, ignore_limit: bool,
439 ) -> u64 {
440 let next_detected = self
441 .keyframes
442 .iter()
443 .find(|&&input_frameno| input_frameno > gop_input_frameno_start)
444 .cloned();
445 let mut next_limit =
446 gop_input_frameno_start + self.config.max_key_frame_interval;
447 if !ignore_limit && self.limit.is_some() {
448 next_limit = next_limit.min(self.limit.unwrap());
449 }
450 if next_detected.is_none() {
451 return next_limit;
452 }
453 cmp::min(next_detected.unwrap(), next_limit)
454 }
455
456 fn set_frame_properties(
457 &mut self, output_frameno: u64,
458 ) -> Result<(), EncoderStatus> {
459 let fi = self.build_frame_properties(output_frameno)?;
460
461 self.frame_data.insert(
462 output_frameno,
463 fi.map(|fi| {
464 let frame = self
465 .frame_q
466 .get(&fi.input_frameno)
467 .as_ref()
468 .unwrap()
469 .as_ref()
470 .unwrap();
471 FrameData::new(fi, frame.clone())
472 }),
473 );
474
475 Ok(())
476 }
477
478 #[allow(unused)]
479 pub fn build_dump_properties() -> PathBuf {
480 let mut data_location = PathBuf::new();
481 if env::var_os("RAV1E_DATA_PATH").is_some() {
482 data_location.push(&env::var_os("RAV1E_DATA_PATH").unwrap());
483 } else {
484 data_location.push(&env::current_dir().unwrap());
485 data_location.push(".lookahead_data");
486 }
487 fs::create_dir_all(&data_location).unwrap();
488 data_location
489 }
490
491 fn build_frame_properties(
492 &mut self, output_frameno: u64,
493 ) -> Result<Option<FrameInvariants<T>>, EncoderStatus> {
494 let (prev_gop_output_frameno_start, prev_gop_input_frameno_start) =
495 if output_frameno == 0 {
496 (0, 0)
497 } else {
498 (
499 self.gop_output_frameno_start[&(output_frameno - 1)],
500 self.gop_input_frameno_start[&(output_frameno - 1)],
501 )
502 };
503
504 self
505 .gop_output_frameno_start
506 .insert(output_frameno, prev_gop_output_frameno_start);
507 self
508 .gop_input_frameno_start
509 .insert(output_frameno, prev_gop_input_frameno_start);
510
511 let output_frameno_in_gop =
512 output_frameno - self.gop_output_frameno_start[&output_frameno];
513 let mut input_frameno = self.inter_cfg.get_input_frameno(
514 output_frameno_in_gop,
515 self.gop_input_frameno_start[&output_frameno],
516 );
517
518 if self.needs_more_frame_q_lookahead(input_frameno) {
519 return Err(EncoderStatus::NeedMoreData);
520 }
521
522 let t35_metadata = if let Some(t35) = self.t35_q.remove(&input_frameno) {
523 t35
524 } else {
525 Box::new([])
526 };
527
528 if output_frameno_in_gop > 0 {
529 let next_keyframe_input_frameno = self.next_keyframe_input_frameno(
530 self.gop_input_frameno_start[&output_frameno],
531 false,
532 );
533 let prev_input_frameno =
534 self.get_previous_fi(output_frameno).input_frameno;
535 if input_frameno >= next_keyframe_input_frameno {
536 if !self.inter_cfg.reorder
537 || ((output_frameno_in_gop - 1) % self.inter_cfg.group_output_len
538 == 0
539 && prev_input_frameno == (next_keyframe_input_frameno - 1))
540 {
541 input_frameno = next_keyframe_input_frameno;
542
543 // If we'll return early, do it before modifying the state.
544 match self.frame_q.get(&input_frameno) {
545 Some(Some(_)) => {}
546 _ => {
547 return Err(EncoderStatus::NeedMoreData);
548 }
549 }
550
551 *self.gop_output_frameno_start.get_mut(&output_frameno).unwrap() =
552 output_frameno;
553 *self.gop_input_frameno_start.get_mut(&output_frameno).unwrap() =
554 next_keyframe_input_frameno;
555 } else {
556 let fi = FrameInvariants::new_inter_frame(
557 self.get_previous_coded_fi(output_frameno),
558 &self.inter_cfg,
559 self.gop_input_frameno_start[&output_frameno],
560 output_frameno_in_gop,
561 next_keyframe_input_frameno,
562 self.config.error_resilient,
563 t35_metadata,
564 );
565 assert!(fi.is_none());
566 return Ok(fi);
567 }
568 }
569 }
570
571 match self.frame_q.get(&input_frameno) {
572 Some(Some(_)) => {}
573 _ => {
574 return Err(EncoderStatus::NeedMoreData);
575 }
576 }
577
578 // Now that we know the input_frameno, look up the correct frame type
579 let frame_type = if self.keyframes.contains(&input_frameno) {
580 FrameType::KEY
581 } else {
582 FrameType::INTER
583 };
584 if frame_type == FrameType::KEY {
585 *self.gop_output_frameno_start.get_mut(&output_frameno).unwrap() =
586 output_frameno;
587 *self.gop_input_frameno_start.get_mut(&output_frameno).unwrap() =
588 input_frameno;
589 }
590
591 let output_frameno_in_gop =
592 output_frameno - self.gop_output_frameno_start[&output_frameno];
593 if output_frameno_in_gop == 0 {
594 let fi = FrameInvariants::new_key_frame(
595 self.config.clone(),
596 self.seq.clone(),
597 self.gop_input_frameno_start[&output_frameno],
598 t35_metadata,
599 );
600 Ok(Some(fi))
601 } else {
602 let next_keyframe_input_frameno = self.next_keyframe_input_frameno(
603 self.gop_input_frameno_start[&output_frameno],
604 false,
605 );
606 let fi = FrameInvariants::new_inter_frame(
607 self.get_previous_coded_fi(output_frameno),
608 &self.inter_cfg,
609 self.gop_input_frameno_start[&output_frameno],
610 output_frameno_in_gop,
611 next_keyframe_input_frameno,
612 self.config.error_resilient,
613 t35_metadata,
614 );
615 assert!(fi.is_some());
616 Ok(fi)
617 }
618 }
619
620 fn get_previous_fi(&self, output_frameno: u64) -> &FrameInvariants<T> {
621 let res = self
622 .frame_data
623 .iter()
624 .filter(|(fno, _)| **fno < output_frameno)
625 .rfind(|(_, fd)| fd.is_some())
626 .unwrap();
627 &res.1.as_ref().unwrap().fi
628 }
629
630 fn get_previous_coded_fi(&self, output_frameno: u64) -> &FrameInvariants<T> {
631 let res = self
632 .frame_data
633 .iter()
634 .filter(|(fno, _)| **fno < output_frameno)
635 .rfind(|(_, fd)| {
636 fd.as_ref().map(|fd| !fd.fi.is_show_existing_frame()).unwrap_or(false)
637 })
638 .unwrap();
639 &res.1.as_ref().unwrap().fi
640 }
641
642 pub(crate) fn done_processing(&self) -> bool {
643 self.limit.map(|limit| self.frames_processed == limit).unwrap_or(false)
644 }
645
646 /// Computes lookahead motion vectors and fills in `lookahead_mvs`,
647 /// `rec_buffer` and `lookahead_rec_buffer` on the `FrameInvariants`. This
648 /// function must be called after every new `FrameInvariants` is initially
649 /// computed.
650 #[profiling::function]
651 fn compute_lookahead_motion_vectors(&mut self, output_frameno: u64) {
652 let frame_data = self.frame_data.get(&output_frameno).unwrap();
653
654 // We're only interested in valid frames which are not show-existing-frame.
655 // Those two don't modify the rec_buffer so there's no need to do anything
656 // special about it either, it'll propagate on its own.
657 if frame_data
658 .as_ref()
659 .map(|fd| fd.fi.is_show_existing_frame())
660 .unwrap_or(true)
661 {
662 return;
663 }
664
665 let qps = {
666 let fti = frame_data.as_ref().unwrap().fi.get_frame_subtype();
667 self.rc_state.select_qi(
668 self,
669 output_frameno,
670 fti,
671 self.maybe_prev_log_base_q,
672 0,
673 )
674 };
675
676 let frame_data =
677 self.frame_data.get_mut(&output_frameno).unwrap().as_mut().unwrap();
678 let fs = &mut frame_data.fs;
679 let fi = &mut frame_data.fi;
680 let coded_data = fi.coded_frame_data.as_mut().unwrap();
681
682 #[cfg(feature = "dump_lookahead_data")]
683 {
684 let data_location = Self::build_dump_properties();
685 let plane = &fs.input_qres;
686 let mut file_name = format!("{:010}-qres", fi.input_frameno);
687 let buf: Vec<_> = plane.iter().map(|p| p.as_()).collect();
688 image::GrayImage::from_vec(
689 plane.cfg.width as u32,
690 plane.cfg.height as u32,
691 buf,
692 )
693 .unwrap()
694 .save(data_location.join(file_name).with_extension("png"))
695 .unwrap();
696 let plane = &fs.input_hres;
697 file_name = format!("{:010}-hres", fi.input_frameno);
698 let buf: Vec<_> = plane.iter().map(|p| p.as_()).collect();
699 image::GrayImage::from_vec(
700 plane.cfg.width as u32,
701 plane.cfg.height as u32,
702 buf,
703 )
704 .unwrap()
705 .save(data_location.join(file_name).with_extension("png"))
706 .unwrap();
707 }
708
709 // Do not modify the next output frame's FrameInvariants.
710 if self.output_frameno == output_frameno {
711 // We do want to propagate the lookahead_rec_buffer though.
712 let rfs = Arc::new(ReferenceFrame {
713 order_hint: fi.order_hint,
714 width: fi.width as u32,
715 height: fi.height as u32,
716 render_width: fi.render_width,
717 render_height: fi.render_height,
718 // Use the original frame contents.
719 frame: fs.input.clone(),
720 input_hres: fs.input_hres.clone(),
721 input_qres: fs.input_qres.clone(),
722 cdfs: fs.cdfs,
723 frame_me_stats: fs.frame_me_stats.clone(),
724 output_frameno,
725 segmentation: fs.segmentation,
726 });
727 for i in 0..REF_FRAMES {
728 if (fi.refresh_frame_flags & (1 << i)) != 0 {
729 coded_data.lookahead_rec_buffer.frames[i] = Some(Arc::clone(&rfs));
730 coded_data.lookahead_rec_buffer.deblock[i] = fs.deblock;
731 }
732 }
733
734 return;
735 }
736
737 // Our lookahead_rec_buffer should be filled with correct original frame
738 // data from the previous frames. Copy it into rec_buffer because that's
739 // what the MV search uses. During the actual encoding rec_buffer is
740 // overwritten with its correct values anyway.
741 fi.rec_buffer = coded_data.lookahead_rec_buffer.clone();
742
743 // Estimate lambda with rate-control dry-run
744 fi.set_quantizers(&qps);
745
746 // TODO: as in the encoding code, key frames will have no references.
747 // However, for block importance purposes we want key frames to act as
748 // P-frames in this instance.
749 //
750 // Compute the motion vectors.
751 compute_motion_vectors(fi, fs, &self.inter_cfg);
752
753 let coded_data = fi.coded_frame_data.as_mut().unwrap();
754
755 #[cfg(feature = "dump_lookahead_data")]
756 {
757 use crate::partition::RefType::*;
758 let data_location = Self::build_dump_properties();
759 let file_name = format!("{:010}-mvs", fi.input_frameno);
760 let second_ref_frame = if !self.inter_cfg.multiref {
761 LAST_FRAME // make second_ref_frame match first
762 } else if fi.idx_in_group_output == 0 {
763 LAST2_FRAME
764 } else {
765 ALTREF_FRAME
766 };
767
768 // Use the default index, it corresponds to the last P-frame or to the
769 // backwards lower reference (so the closest previous frame).
770 let index = if second_ref_frame.to_index() != 0 { 0 } else { 1 };
771
772 let me_stats = &fs.frame_me_stats.read().expect("poisoned lock")[index];
773 use byteorder::{NativeEndian, WriteBytesExt};
774 // dynamic allocation: debugging only
775 let mut buf = vec![];
776 buf.write_u64::<NativeEndian>(me_stats.rows as u64).unwrap();
777 buf.write_u64::<NativeEndian>(me_stats.cols as u64).unwrap();
778 for y in 0..me_stats.rows {
779 for x in 0..me_stats.cols {
780 let mv = me_stats[y][x].mv;
781 buf.write_i16::<NativeEndian>(mv.row).unwrap();
782 buf.write_i16::<NativeEndian>(mv.col).unwrap();
783 }
784 }
785 ::std::fs::write(
786 data_location.join(file_name).with_extension("bin"),
787 buf,
788 )
789 .unwrap();
790 }
791
792 // Set lookahead_rec_buffer on this FrameInvariants for future
793 // FrameInvariants to pick it up.
794 let rfs = Arc::new(ReferenceFrame {
795 order_hint: fi.order_hint,
796 width: fi.width as u32,
797 height: fi.height as u32,
798 render_width: fi.render_width,
799 render_height: fi.render_height,
800 // Use the original frame contents.
801 frame: fs.input.clone(),
802 input_hres: fs.input_hres.clone(),
803 input_qres: fs.input_qres.clone(),
804 cdfs: fs.cdfs,
805 frame_me_stats: fs.frame_me_stats.clone(),
806 output_frameno,
807 segmentation: fs.segmentation,
808 });
809 for i in 0..REF_FRAMES {
810 if (fi.refresh_frame_flags & (1 << i)) != 0 {
811 coded_data.lookahead_rec_buffer.frames[i] = Some(Arc::clone(&rfs));
812 coded_data.lookahead_rec_buffer.deblock[i] = fs.deblock;
813 }
814 }
815 }
816
817 /// Computes lookahead intra cost approximations and fills in
818 /// `lookahead_intra_costs` on the `FrameInvariants`.
819 fn compute_lookahead_intra_costs(&mut self, output_frameno: u64) {
820 let frame_data = self.frame_data.get(&output_frameno).unwrap();
821 let fd = &frame_data.as_ref();
822
823 // We're only interested in valid frames which are not show-existing-frame.
824 if fd.map(|fd| fd.fi.is_show_existing_frame()).unwrap_or(true) {
825 return;
826 }
827
828 let fi = &fd.unwrap().fi;
829
830 self
831 .frame_data
832 .get_mut(&output_frameno)
833 .unwrap()
834 .as_mut()
835 .unwrap()
836 .fi
837 .coded_frame_data
838 .as_mut()
839 .unwrap()
840 .lookahead_intra_costs = self
841 .keyframe_detector
842 .intra_costs
843 .remove(&fi.input_frameno)
844 .unwrap_or_else(|| {
845 let frame = self.frame_q[&fi.input_frameno].as_ref().unwrap();
846
847 let temp_plane = self
848 .keyframe_detector
849 .temp_plane
850 .get_or_insert_with(|| frame.planes[0].clone());
851
852 // We use the cached values from scenechange if available,
853 // otherwise we need to calculate them here.
854 estimate_intra_costs(
855 temp_plane,
856 &**frame,
857 fi.sequence.bit_depth,
858 fi.cpu_feature_level,
859 )
860 });
861 }
862
863 #[profiling::function]
864 pub fn compute_keyframe_placement(
865 lookahead_frames: &[&Arc<Frame<T>>], keyframes_forced: &BTreeSet<u64>,
866 keyframe_detector: &mut SceneChangeDetector<T>,
867 next_lookahead_frame: &mut u64, keyframes: &mut BTreeSet<u64>,
868 ) {
869 if keyframes_forced.contains(next_lookahead_frame)
870 || keyframe_detector.analyze_next_frame(
871 lookahead_frames,
872 *next_lookahead_frame,
873 *keyframes.iter().last().unwrap(),
874 )
875 {
876 keyframes.insert(*next_lookahead_frame);
877 }
878
879 *next_lookahead_frame += 1;
880 }
881
882 #[profiling::function]
883 pub fn compute_frame_invariants(&mut self) {
884 while self.set_frame_properties(self.next_lookahead_output_frameno).is_ok()
885 {
886 self
887 .compute_lookahead_motion_vectors(self.next_lookahead_output_frameno);
888 if self.config.temporal_rdo() {
889 self.compute_lookahead_intra_costs(self.next_lookahead_output_frameno);
890 }
891 self.next_lookahead_output_frameno += 1;
892 }
893 }
894
895 #[profiling::function]
896 fn update_block_importances(
897 fi: &FrameInvariants<T>, me_stats: &crate::me::FrameMEStats,
898 frame: &Frame<T>, reference_frame: &Frame<T>, bit_depth: usize,
899 bsize: BlockSize, len: usize,
900 reference_frame_block_importances: &mut [f32],
901 ) {
902 let coded_data = fi.coded_frame_data.as_ref().unwrap();
903 let plane_org = &frame.planes[0];
904 let plane_ref = &reference_frame.planes[0];
905 let lookahead_intra_costs_lines =
906 coded_data.lookahead_intra_costs.chunks_exact(coded_data.w_in_imp_b);
907 let block_importances_lines =
908 coded_data.block_importances.chunks_exact(coded_data.w_in_imp_b);
909
910 lookahead_intra_costs_lines
911 .zip(block_importances_lines)
912 .zip(me_stats.rows_iter().step_by(2))
913 .enumerate()
914 .flat_map(
915 |(y, ((lookahead_intra_costs, block_importances), me_stats_line))| {
916 lookahead_intra_costs
917 .iter()
918 .zip(block_importances.iter())
919 .zip(me_stats_line.iter().step_by(2))
920 .enumerate()
921 .map(move |(x, ((&intra_cost, &future_importance), &me_stat))| {
922 let mv = me_stat.mv;
923
924 // Coordinates of the top-left corner of the reference block, in MV
925 // units.
926 let reference_x =
927 x as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.col as i64;
928 let reference_y =
929 y as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.row as i64;
930
931 let region_org = plane_org.region(Area::Rect {
932 x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
933 y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
934 width: IMPORTANCE_BLOCK_SIZE,
935 height: IMPORTANCE_BLOCK_SIZE,
936 });
937
938 let region_ref = plane_ref.region(Area::Rect {
939 x: reference_x as isize
940 / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
941 y: reference_y as isize
942 / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
943 width: IMPORTANCE_BLOCK_SIZE,
944 height: IMPORTANCE_BLOCK_SIZE,
945 });
946
947 let inter_cost = get_satd(
948 &region_org,
949 &region_ref,
950 bsize.width(),
951 bsize.height(),
952 bit_depth,
953 fi.cpu_feature_level,
954 ) as f32;
955
956 let intra_cost = intra_cost as f32;
957 // let intra_cost = lookahead_intra_costs[x] as f32;
958 // let future_importance = block_importances[x];
959
960 let propagate_fraction = if intra_cost <= inter_cost {
961 0.
962 } else {
963 1. - inter_cost / intra_cost
964 };
965
966 let propagate_amount = (intra_cost + future_importance)
967 * propagate_fraction
968 / len as f32;
969 (propagate_amount, reference_x, reference_y)
970 })
971 },
972 )
973 .for_each(|(propagate_amount, reference_x, reference_y)| {
974 let mut propagate =
975 |block_x_in_mv_units, block_y_in_mv_units, fraction| {
976 let x = block_x_in_mv_units / IMP_BLOCK_SIZE_IN_MV_UNITS;
977 let y = block_y_in_mv_units / IMP_BLOCK_SIZE_IN_MV_UNITS;
978
979 // TODO: propagate partially if the block is partially off-frame
980 // (possible on right and bottom edges)?
981 if x >= 0
982 && y >= 0
983 && (x as usize) < coded_data.w_in_imp_b
984 && (y as usize) < coded_data.h_in_imp_b
985 {
986 reference_frame_block_importances
987 [y as usize * coded_data.w_in_imp_b + x as usize] +=
988 propagate_amount * fraction;
989 }
990 };
991
992 // Coordinates of the top-left corner of the block intersecting the
993 // reference block from the top-left.
994 let top_left_block_x = (reference_x
995 - if reference_x < 0 { IMP_BLOCK_SIZE_IN_MV_UNITS - 1 } else { 0 })
996 / IMP_BLOCK_SIZE_IN_MV_UNITS
997 * IMP_BLOCK_SIZE_IN_MV_UNITS;
998 let top_left_block_y = (reference_y
999 - if reference_y < 0 { IMP_BLOCK_SIZE_IN_MV_UNITS - 1 } else { 0 })
1000 / IMP_BLOCK_SIZE_IN_MV_UNITS
1001 * IMP_BLOCK_SIZE_IN_MV_UNITS;
1002
1003 debug_assert!(reference_x >= top_left_block_x);
1004 debug_assert!(reference_y >= top_left_block_y);
1005
1006 let top_right_block_x = top_left_block_x + IMP_BLOCK_SIZE_IN_MV_UNITS;
1007 let top_right_block_y = top_left_block_y;
1008 let bottom_left_block_x = top_left_block_x;
1009 let bottom_left_block_y =
1010 top_left_block_y + IMP_BLOCK_SIZE_IN_MV_UNITS;
1011 let bottom_right_block_x = top_right_block_x;
1012 let bottom_right_block_y = bottom_left_block_y;
1013
1014 let top_left_block_fraction = ((top_right_block_x - reference_x)
1015 * (bottom_left_block_y - reference_y))
1016 as f32
1017 / IMP_BLOCK_AREA_IN_MV_UNITS as f32;
1018
1019 propagate(top_left_block_x, top_left_block_y, top_left_block_fraction);
1020
1021 let top_right_block_fraction =
1022 ((reference_x + IMP_BLOCK_SIZE_IN_MV_UNITS - top_right_block_x)
1023 * (bottom_left_block_y - reference_y)) as f32
1024 / IMP_BLOCK_AREA_IN_MV_UNITS as f32;
1025
1026 propagate(
1027 top_right_block_x,
1028 top_right_block_y,
1029 top_right_block_fraction,
1030 );
1031
1032 let bottom_left_block_fraction = ((top_right_block_x - reference_x)
1033 * (reference_y + IMP_BLOCK_SIZE_IN_MV_UNITS - bottom_left_block_y))
1034 as f32
1035 / IMP_BLOCK_AREA_IN_MV_UNITS as f32;
1036
1037 propagate(
1038 bottom_left_block_x,
1039 bottom_left_block_y,
1040 bottom_left_block_fraction,
1041 );
1042
1043 let bottom_right_block_fraction =
1044 ((reference_x + IMP_BLOCK_SIZE_IN_MV_UNITS - top_right_block_x)
1045 * (reference_y + IMP_BLOCK_SIZE_IN_MV_UNITS - bottom_left_block_y))
1046 as f32
1047 / IMP_BLOCK_AREA_IN_MV_UNITS as f32;
1048
1049 propagate(
1050 bottom_right_block_x,
1051 bottom_right_block_y,
1052 bottom_right_block_fraction,
1053 );
1054 });
1055 }
1056
1057 /// Computes the block importances for the current output frame.
1058 #[profiling::function]
1059 fn compute_block_importances(&mut self) {
1060 // SEF don't need block importances.
1061 if self.frame_data[&self.output_frameno]
1062 .as_ref()
1063 .unwrap()
1064 .fi
1065 .is_show_existing_frame()
1066 {
1067 return;
1068 }
1069
1070 // Get a list of output_framenos that we want to propagate through.
1071 let output_framenos = self
1072 .get_rdo_lookahead_frames()
1073 .map(|(&output_frameno, _)| output_frameno)
1074 .collect::<Vec<_>>();
1075
1076 // The first one should be the current output frame.
1077 assert_eq!(output_framenos[0], self.output_frameno);
1078
1079 // First, initialize them all with zeros.
1080 for output_frameno in output_framenos.iter() {
1081 let fi = &mut self
1082 .frame_data
1083 .get_mut(output_frameno)
1084 .unwrap()
1085 .as_mut()
1086 .unwrap()
1087 .fi;
1088 for x in
1089 fi.coded_frame_data.as_mut().unwrap().block_importances.iter_mut()
1090 {
1091 *x = 0.;
1092 }
1093 }
1094
1095 // Now compute and propagate the block importances from the end. The
1096 // current output frame will get its block importances from the future
1097 // frames.
1098 let bsize = BlockSize::from_width_and_height(
1099 IMPORTANCE_BLOCK_SIZE,
1100 IMPORTANCE_BLOCK_SIZE,
1101 );
1102
1103 for &output_frameno in output_framenos.iter().skip(1).rev() {
1104 // TODO: see comment above about key frames not having references.
1105 if self
1106 .frame_data
1107 .get(&output_frameno)
1108 .unwrap()
1109 .as_ref()
1110 .unwrap()
1111 .fi
1112 .frame_type
1113 == FrameType::KEY
1114 {
1115 continue;
1116 }
1117
1118 // Remove fi from the map temporarily and put it back in in the end of
1119 // the iteration. This is required because we need to mutably borrow
1120 // referenced fis from the map, and that wouldn't be possible if this was
1121 // an active borrow.
1122 //
1123 // Performance note: Contrary to intuition,
1124 // removing the data and re-inserting it at the end
1125 // is more performant because it avoids a very expensive clone.
1126 let output_frame_data =
1127 self.frame_data.remove(&output_frameno).unwrap().unwrap();
1128 {
1129 let fi = &output_frame_data.fi;
1130 let fs = &output_frame_data.fs;
1131
1132 let frame = self.frame_q[&fi.input_frameno].as_ref().unwrap();
1133
1134 // There can be at most 3 of these.
1135 let mut unique_indices = ArrayVec::<_, 3>::new();
1136
1137 for (mv_index, &rec_index) in fi.ref_frames.iter().enumerate() {
1138 if !unique_indices.iter().any(|&(_, r)| r == rec_index) {
1139 unique_indices.push((mv_index, rec_index));
1140 }
1141 }
1142
1143 let bit_depth = self.config.bit_depth;
1144 let frame_data = &mut self.frame_data;
1145 let len = unique_indices.len();
1146
1147 let lookahead_me_stats =
1148 fs.frame_me_stats.read().expect("poisoned lock");
1149
1150 // Compute and propagate the importance, split evenly between the
1151 // referenced frames.
1152 unique_indices.iter().for_each(|&(mv_index, rec_index)| {
1153 // Use rec_buffer here rather than lookahead_rec_buffer because
1154 // rec_buffer still contains the reference frames for the current frame
1155 // (it's only overwritten when the frame is encoded), while
1156 // lookahead_rec_buffer already contains reference frames for the next
1157 // frame (for the reference propagation to work correctly).
1158 let reference =
1159 fi.rec_buffer.frames[rec_index as usize].as_ref().unwrap();
1160 let reference_frame = &reference.frame;
1161 let reference_output_frameno = reference.output_frameno;
1162 let me_stats = &lookahead_me_stats[mv_index];
1163
1164 // We should never use frame as its own reference.
1165 assert_ne!(reference_output_frameno, output_frameno);
1166
1167 if let Some(reference_frame_block_importances) =
1168 frame_data.get_mut(&reference_output_frameno).map(|data| {
1169 &mut data
1170 .as_mut()
1171 .unwrap()
1172 .fi
1173 .coded_frame_data
1174 .as_mut()
1175 .unwrap()
1176 .block_importances
1177 })
1178 {
1179 Self::update_block_importances(
1180 fi,
1181 me_stats,
1182 frame,
1183 reference_frame,
1184 bit_depth,
1185 bsize,
1186 len,
1187 reference_frame_block_importances,
1188 );
1189 }
1190 });
1191 }
1192 self.frame_data.insert(output_frameno, Some(output_frame_data));
1193 }
1194
1195 if !output_framenos.is_empty() {
1196 let fi = &mut self
1197 .frame_data
1198 .get_mut(&output_framenos[0])
1199 .unwrap()
1200 .as_mut()
1201 .unwrap()
1202 .fi;
1203 let coded_data = fi.coded_frame_data.as_mut().unwrap();
1204 let block_importances = coded_data.block_importances.iter();
1205 let lookahead_intra_costs = coded_data.lookahead_intra_costs.iter();
1206 let distortion_scales = coded_data.distortion_scales.iter_mut();
1207 for ((&propagate_cost, &intra_cost), distortion_scale) in
1208 block_importances.zip(lookahead_intra_costs).zip(distortion_scales)
1209 {
1210 *distortion_scale = crate::rdo::distortion_scale_for(
1211 propagate_cost as f64,
1212 intra_cost as f64,
1213 );
1214 }
1215 #[cfg(feature = "dump_lookahead_data")]
1216 {
1217 use byteorder::{NativeEndian, WriteBytesExt};
1218
1219 let coded_data = fi.coded_frame_data.as_ref().unwrap();
1220
1221 let mut buf = vec![];
1222 let data_location = Self::build_dump_properties();
1223 let file_name = format!("{:010}-imps", fi.input_frameno);
1224 buf.write_u64::<NativeEndian>(coded_data.h_in_imp_b as u64).unwrap();
1225 buf.write_u64::<NativeEndian>(coded_data.w_in_imp_b as u64).unwrap();
1226 buf.write_u64::<NativeEndian>(fi.get_frame_subtype() as u64).unwrap();
1227 for y in 0..coded_data.h_in_imp_b {
1228 for x in 0..coded_data.w_in_imp_b {
1229 buf
1230 .write_f32::<NativeEndian>(f64::from(
1231 coded_data.distortion_scales[y * coded_data.w_in_imp_b + x],
1232 ) as f32)
1233 .unwrap();
1234 }
1235 }
1236 ::std::fs::write(
1237 data_location.join(file_name).with_extension("bin"),
1238 buf,
1239 )
1240 .unwrap();
1241 }
1242 }
1243 }
1244
1245 pub(crate) fn encode_packet(
1246 &mut self, cur_output_frameno: u64,
1247 ) -> Result<Packet<T>, EncoderStatus> {
1248 if self
1249 .frame_data
1250 .get(&cur_output_frameno)
1251 .unwrap()
1252 .as_ref()
1253 .unwrap()
1254 .fi
1255 .is_show_existing_frame()
1256 {
1257 if !self.rc_state.ready() {
1258 return Err(EncoderStatus::NotReady);
1259 }
1260
1261 self.encode_show_existing_packet(cur_output_frameno)
1262 } else if let Some(Some(_)) = self.frame_q.get(
1263 &self
1264 .frame_data
1265 .get(&cur_output_frameno)
1266 .unwrap()
1267 .as_ref()
1268 .unwrap()
1269 .fi
1270 .input_frameno,
1271 ) {
1272 if !self.rc_state.ready() {
1273 return Err(EncoderStatus::NotReady);
1274 }
1275
1276 self.encode_normal_packet(cur_output_frameno)
1277 } else {
1278 Err(EncoderStatus::NeedMoreData)
1279 }
1280 }
1281
1282 #[profiling::function]
1283 pub fn encode_show_existing_packet(
1284 &mut self, cur_output_frameno: u64,
1285 ) -> Result<Packet<T>, EncoderStatus> {
1286 let frame_data =
1287 self.frame_data.get_mut(&cur_output_frameno).unwrap().as_mut().unwrap();
1288 let sef_data = encode_show_existing_frame(
1289 &frame_data.fi,
1290 &mut frame_data.fs,
1291 &self.inter_cfg,
1292 );
1293 let bits = (sef_data.len() * 8) as i64;
1294 self.packet_data.extend(sef_data);
1295 self.rc_state.update_state(
1296 bits,
1297 FRAME_SUBTYPE_SEF,
1298 frame_data.fi.show_frame,
1299 0,
1300 false,
1301 false,
1302 );
1303 let (rec, source) = if frame_data.fi.show_frame {
1304 (Some(frame_data.fs.rec.clone()), Some(frame_data.fs.input.clone()))
1305 } else {
1306 (None, None)
1307 };
1308
1309 self.output_frameno += 1;
1310
1311 let input_frameno = frame_data.fi.input_frameno;
1312 let frame_type = frame_data.fi.frame_type;
1313 let qp = frame_data.fi.base_q_idx;
1314 let enc_stats = frame_data.fs.enc_stats.clone();
1315 self.finalize_packet(rec, source, input_frameno, frame_type, qp, enc_stats)
1316 }
1317
1318 #[profiling::function]
1319 pub fn encode_normal_packet(
1320 &mut self, cur_output_frameno: u64,
1321 ) -> Result<Packet<T>, EncoderStatus> {
1322 let mut frame_data =
1323 self.frame_data.remove(&cur_output_frameno).unwrap().unwrap();
1324
1325 let mut log_isqrt_mean_scale = 0i64;
1326
1327 if let Some(coded_data) = frame_data.fi.coded_frame_data.as_mut() {
1328 if self.config.tune == Tune::Psychovisual {
1329 let frame =
1330 self.frame_q[&frame_data.fi.input_frameno].as_ref().unwrap();
1331 coded_data.activity_mask = ActivityMask::from_plane(&frame.planes[0]);
1332 coded_data.activity_mask.fill_scales(
1333 frame_data.fi.sequence.bit_depth,
1334 &mut coded_data.activity_scales,
1335 );
1336 log_isqrt_mean_scale = coded_data.compute_spatiotemporal_scores();
1337 } else {
1338 coded_data.activity_mask = ActivityMask::default();
1339 log_isqrt_mean_scale = coded_data.compute_temporal_scores();
1340 }
1341 #[cfg(feature = "dump_lookahead_data")]
1342 {
1343 use crate::encoder::Scales::*;
1344 let input_frameno = frame_data.fi.input_frameno;
1345 if self.config.tune == Tune::Psychovisual {
1346 coded_data.dump_scales(
1347 Self::build_dump_properties(),
1348 ActivityScales,
1349 input_frameno,
1350 );
1351 coded_data.dump_scales(
1352 Self::build_dump_properties(),
1353 SpatiotemporalScales,
1354 input_frameno,
1355 );
1356 }
1357 coded_data.dump_scales(
1358 Self::build_dump_properties(),
1359 DistortionScales,
1360 input_frameno,
1361 );
1362 }
1363 }
1364
1365 let fti = frame_data.fi.get_frame_subtype();
1366 let qps = self.rc_state.select_qi(
1367 self,
1368 cur_output_frameno,
1369 fti,
1370 self.maybe_prev_log_base_q,
1371 log_isqrt_mean_scale,
1372 );
1373 frame_data.fi.set_quantizers(&qps);
1374
1375 if self.rc_state.needs_trial_encode(fti) {
1376 let mut trial_fs = frame_data.fs.clone();
1377 let data = encode_frame(&frame_data.fi, &mut trial_fs, &self.inter_cfg);
1378 self.rc_state.update_state(
1379 (data.len() * 8) as i64,
1380 fti,
1381 frame_data.fi.show_frame,
1382 qps.log_target_q,
1383 true,
1384 false,
1385 );
1386 let qps = self.rc_state.select_qi(
1387 self,
1388 cur_output_frameno,
1389 fti,
1390 self.maybe_prev_log_base_q,
1391 log_isqrt_mean_scale,
1392 );
1393 frame_data.fi.set_quantizers(&qps);
1394 }
1395
1396 let data =
1397 encode_frame(&frame_data.fi, &mut frame_data.fs, &self.inter_cfg);
1398 #[cfg(feature = "dump_lookahead_data")]
1399 {
1400 let input_frameno = frame_data.fi.input_frameno;
1401 let data_location = Self::build_dump_properties();
1402 frame_data.fs.segmentation.dump_threshold(data_location, input_frameno);
1403 }
1404 let enc_stats = frame_data.fs.enc_stats.clone();
1405 self.maybe_prev_log_base_q = Some(qps.log_base_q);
1406 // TODO: Add support for dropping frames.
1407 self.rc_state.update_state(
1408 (data.len() * 8) as i64,
1409 fti,
1410 frame_data.fi.show_frame,
1411 qps.log_target_q,
1412 false,
1413 false,
1414 );
1415 self.packet_data.extend(data);
1416
1417 let planes =
1418 if frame_data.fi.sequence.chroma_sampling == Cs400 { 1 } else { 3 };
1419
1420 Arc::get_mut(&mut frame_data.fs.rec).unwrap().pad(
1421 frame_data.fi.width,
1422 frame_data.fi.height,
1423 planes,
1424 );
1425
1426 let (rec, source) = if frame_data.fi.show_frame {
1427 (Some(frame_data.fs.rec.clone()), Some(frame_data.fs.input.clone()))
1428 } else {
1429 (None, None)
1430 };
1431
1432 update_rec_buffer(cur_output_frameno, &mut frame_data.fi, &frame_data.fs);
1433
1434 // Copy persistent fields into subsequent FrameInvariants.
1435 let rec_buffer = frame_data.fi.rec_buffer.clone();
1436 for subsequent_fi in self
1437 .frame_data
1438 .iter_mut()
1439 .skip_while(|(&output_frameno, _)| output_frameno <= cur_output_frameno)
1440 // Here we want the next valid non-show-existing-frame inter frame.
1441 //
1442 // Copying to show-existing-frame frames isn't actually required
1443 // for correct encoding, but it's needed for the reconstruction to
1444 // work correctly.
1445 .filter_map(|(_, frame_data)| frame_data.as_mut().map(|fd| &mut fd.fi))
1446 .take_while(|fi| fi.frame_type != FrameType::KEY)
1447 {
1448 subsequent_fi.rec_buffer = rec_buffer.clone();
1449 subsequent_fi.set_ref_frame_sign_bias();
1450
1451 // Stop after the first non-show-existing-frame.
1452 if !subsequent_fi.is_show_existing_frame() {
1453 break;
1454 }
1455 }
1456
1457 self.frame_data.insert(cur_output_frameno, Some(frame_data));
1458 let frame_data =
1459 self.frame_data.get(&cur_output_frameno).unwrap().as_ref().unwrap();
1460 let fi = &frame_data.fi;
1461
1462 self.output_frameno += 1;
1463
1464 if fi.show_frame {
1465 let input_frameno = fi.input_frameno;
1466 let frame_type = fi.frame_type;
1467 let qp = fi.base_q_idx;
1468 self.finalize_packet(
1469 rec,
1470 source,
1471 input_frameno,
1472 frame_type,
1473 qp,
1474 enc_stats,
1475 )
1476 } else {
1477 Err(EncoderStatus::Encoded)
1478 }
1479 }
1480
1481 #[profiling::function]
1482 pub fn receive_packet(&mut self) -> Result<Packet<T>, EncoderStatus> {
1483 if self.done_processing() {
1484 return Err(EncoderStatus::LimitReached);
1485 }
1486
1487 if self.needs_more_fi_lookahead() {
1488 return Err(EncoderStatus::NeedMoreData);
1489 }
1490
1491 // Find the next output_frameno corresponding to a non-skipped frame.
1492 self.output_frameno = self
1493 .frame_data
1494 .iter()
1495 .skip_while(|(&output_frameno, _)| output_frameno < self.output_frameno)
1496 .find(|(_, data)| data.is_some())
1497 .map(|(&output_frameno, _)| output_frameno)
1498 .ok_or(EncoderStatus::NeedMoreData)?; // TODO: doesn't play well with the below check?
1499
1500 let input_frameno =
1501 self.frame_data[&self.output_frameno].as_ref().unwrap().fi.input_frameno;
1502 if !self.needs_more_frames(input_frameno) {
1503 return Err(EncoderStatus::LimitReached);
1504 }
1505
1506 if self.config.temporal_rdo() {
1507 // Compute the block importances for the current output frame.
1508 self.compute_block_importances();
1509 }
1510
1511 let cur_output_frameno = self.output_frameno;
1512
1513 let mut ret = self.encode_packet(cur_output_frameno);
1514
1515 if let Ok(ref mut pkt) = ret {
1516 self.garbage_collect(pkt.input_frameno);
1517 pkt.opaque = self.opaque_q.remove(&pkt.input_frameno);
1518 }
1519
1520 ret
1521 }
1522
1523 fn finalize_packet(
1524 &mut self, rec: Option<Arc<Frame<T>>>, source: Option<Arc<Frame<T>>>,
1525 input_frameno: u64, frame_type: FrameType, qp: u8,
1526 enc_stats: EncoderStats,
1527 ) -> Result<Packet<T>, EncoderStatus> {
1528 let data = self.packet_data.clone();
1529 self.packet_data.clear();
1530 if write_temporal_delimiter(&mut self.packet_data).is_err() {
1531 return Err(EncoderStatus::Failure);
1532 }
1533
1534 self.frames_processed += 1;
1535 Ok(Packet {
1536 data,
1537 rec,
1538 source,
1539 input_frameno,
1540 frame_type,
1541 qp,
1542 enc_stats,
1543 opaque: None,
1544 })
1545 }
1546
1547 #[profiling::function]
1548 fn garbage_collect(&mut self, cur_input_frameno: u64) {
1549 if cur_input_frameno == 0 {
1550 return;
1551 }
1552 let frame_q_start = self.frame_q.keys().next().cloned().unwrap_or(0);
1553 for i in frame_q_start..cur_input_frameno {
1554 self.frame_q.remove(&i);
1555 }
1556
1557 if self.output_frameno < 2 {
1558 return;
1559 }
1560 let fi_start = self.frame_data.keys().next().cloned().unwrap_or(0);
1561 for i in fi_start..(self.output_frameno - 1) {
1562 self.frame_data.remove(&i);
1563 self.gop_output_frameno_start.remove(&i);
1564 self.gop_input_frameno_start.remove(&i);
1565 }
1566 }
1567
1568 /// Counts the number of output frames of each subtype in the next
1569 /// `reservoir_frame_delay` temporal units (needed for rate control).
1570 /// Returns the number of output frames (excluding SEF frames) and output TUs
1571 /// until the last keyframe in the next `reservoir_frame_delay` temporal units,
1572 /// or the end of the interval, whichever comes first.
1573 /// The former is needed because it indicates the number of rate estimates we
1574 /// will make.
1575 /// The latter is needed because it indicates the number of times new bitrate
1576 /// is added to the buffer.
1577 pub(crate) fn guess_frame_subtypes(
1578 &self, nframes: &mut [i32; FRAME_NSUBTYPES + 1],
1579 reservoir_frame_delay: i32,
1580 ) -> (i32, i32) {
1581 for fti in 0..=FRAME_NSUBTYPES {
1582 nframes[fti] = 0;
1583 }
1584
1585 // Two-pass calls this function before receive_packet(), and in particular
1586 // before the very first send_frame(), when the following maps are empty.
1587 // In this case, return 0 as the default value.
1588 let mut prev_keyframe_input_frameno = *self
1589 .gop_input_frameno_start
1590 .get(&self.output_frameno)
1591 .unwrap_or_else(|| {
1592 assert!(self.output_frameno == 0);
1593 &0
1594 });
1595 let mut prev_keyframe_output_frameno = *self
1596 .gop_output_frameno_start
1597 .get(&self.output_frameno)
1598 .unwrap_or_else(|| {
1599 assert!(self.output_frameno == 0);
1600 &0
1601 });
1602
1603 let mut prev_keyframe_ntus = 0;
1604 // Does not include SEF frames.
1605 let mut prev_keyframe_nframes = 0;
1606 let mut acc: [i32; FRAME_NSUBTYPES + 1] = [0; FRAME_NSUBTYPES + 1];
1607 // Updates the frame counts with the accumulated values when we hit a
1608 // keyframe.
1609 fn collect_counts(
1610 nframes: &mut [i32; FRAME_NSUBTYPES + 1],
1611 acc: &mut [i32; FRAME_NSUBTYPES + 1],
1612 ) {
1613 for fti in 0..=FRAME_NSUBTYPES {
1614 nframes[fti] += acc[fti];
1615 acc[fti] = 0;
1616 }
1617 acc[FRAME_SUBTYPE_I] += 1;
1618 }
1619 let mut output_frameno = self.output_frameno;
1620 let mut ntus = 0;
1621 // Does not include SEF frames.
1622 let mut nframes_total = 0;
1623 while ntus < reservoir_frame_delay {
1624 let output_frameno_in_gop =
1625 output_frameno - prev_keyframe_output_frameno;
1626 let is_kf =
1627 if let Some(Some(frame_data)) = self.frame_data.get(&output_frameno) {
1628 if frame_data.fi.frame_type == FrameType::KEY {
1629 prev_keyframe_input_frameno = frame_data.fi.input_frameno;
1630 // We do not currently use forward keyframes, so they should always
1631 // end the current TU (thus we always increment ntus below).
1632 debug_assert!(frame_data.fi.show_frame);
1633 true
1634 } else {
1635 false
1636 }
1637 } else {
1638 // It is possible to be invoked for the first time from twopass_out()
1639 // before receive_packet() is called, in which case frame_invariants
1640 // will not be populated.
1641 // Force the first frame in each GOP to be a keyframe in that case.
1642 output_frameno_in_gop == 0
1643 };
1644 if is_kf {
1645 collect_counts(nframes, &mut acc);
1646 prev_keyframe_output_frameno = output_frameno;
1647 prev_keyframe_ntus = ntus;
1648 prev_keyframe_nframes = nframes_total;
1649 output_frameno += 1;
1650 ntus += 1;
1651 nframes_total += 1;
1652 continue;
1653 }
1654 let idx_in_group_output =
1655 self.inter_cfg.get_idx_in_group_output(output_frameno_in_gop);
1656 let input_frameno = prev_keyframe_input_frameno
1657 + self
1658 .inter_cfg
1659 .get_order_hint(output_frameno_in_gop, idx_in_group_output)
1660 as u64;
1661 // For rate control purposes, ignore any limit on frame count that has
1662 // been set.
1663 // We pretend that we will keep encoding frames forever to prevent the
1664 // control loop from driving us into the rails as we come up against a
1665 // hard stop (with no more chance to correct outstanding errors).
1666 let next_keyframe_input_frameno =
1667 self.next_keyframe_input_frameno(prev_keyframe_input_frameno, true);
1668 // If we are re-ordering, we may skip some output frames in the final
1669 // re-order group of the GOP.
1670 if input_frameno >= next_keyframe_input_frameno {
1671 // If we have encoded enough whole groups to reach the next keyframe,
1672 // then start the next keyframe gop.
1673 if 1
1674 + (output_frameno - prev_keyframe_output_frameno)
1675 / self.inter_cfg.group_output_len
1676 * self.inter_cfg.group_input_len
1677 >= next_keyframe_input_frameno - prev_keyframe_input_frameno
1678 {
1679 collect_counts(nframes, &mut acc);
1680 prev_keyframe_input_frameno = input_frameno;
1681 prev_keyframe_output_frameno = output_frameno;
1682 prev_keyframe_ntus = ntus;
1683 prev_keyframe_nframes = nframes_total;
1684 // We do not currently use forward keyframes, so they should always
1685 // end the current TU.
1686 output_frameno += 1;
1687 ntus += 1;
1688 }
1689 output_frameno += 1;
1690 continue;
1691 }
1692 if self.inter_cfg.get_show_existing_frame(idx_in_group_output) {
1693 acc[FRAME_SUBTYPE_SEF] += 1;
1694 } else {
1695 // TODO: Implement golden P-frames.
1696 let fti = FRAME_SUBTYPE_P
1697 + (self.inter_cfg.get_level(idx_in_group_output) as usize);
1698 acc[fti] += 1;
1699 nframes_total += 1;
1700 }
1701 if self.inter_cfg.get_show_frame(idx_in_group_output) {
1702 ntus += 1;
1703 }
1704 output_frameno += 1;
1705 }
1706 if prev_keyframe_output_frameno <= self.output_frameno {
1707 // If there were no keyframes at all, or only the first frame was a
1708 // keyframe, the accumulators never flushed and still contain counts for
1709 // the entire buffer.
1710 // In both cases, we return these counts.
1711 collect_counts(nframes, &mut acc);
1712 (nframes_total, ntus)
1713 } else {
1714 // Otherwise, we discard what remains in the accumulators as they contain
1715 // the counts from and past the last keyframe.
1716 (prev_keyframe_nframes, prev_keyframe_ntus)
1717 }
1718 }
1719}
1720