api.rs 21.9 KB
Newer Older
rzumer's avatar
rzumer committed
1 2 3 4 5 6 7 8 9
// Copyright (c) 2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

Raphaël Zumer's avatar
Raphaël Zumer committed
10
use bitstream_io::*;
11 12 13
use crate::encoder::*;
use crate::metrics::calculate_frame_psnr;
use crate::partition::*;
14
use crate::rate::RCState;
15
use crate::scenechange::SceneChangeDetector;
Raphaël Zumer's avatar
Raphaël Zumer committed
16
use self::EncoderStatus::*;
Luca Barbato's avatar
Luca Barbato committed
17

Raphaël Zumer's avatar
Raphaël Zumer committed
18
use std::{cmp, fmt, io};
fbossen's avatar
fbossen committed
19
use std::collections::BTreeMap;
Luca Barbato's avatar
Luca Barbato committed
20
use std::sync::Arc;
Josh Holmer's avatar
Josh Holmer committed
21
use std::collections::BTreeSet;
22
use std::path::PathBuf;
Josh Holmer's avatar
Josh Holmer committed
23 24

const LOOKAHEAD_FRAMES: u64 = 10;
Luca Barbato's avatar
Luca Barbato committed
25 26 27

// TODO: use the num crate?
#[derive(Clone, Copy, Debug)]
28
#[repr(C)]
Luca Barbato's avatar
Luca Barbato committed
29 30 31
pub struct Rational {
  pub num: u64,
  pub den: u64
Luca Barbato's avatar
Luca Barbato committed
32 33
}

Luca Barbato's avatar
Luca Barbato committed
34 35 36
impl Rational {
  pub fn new(num: u64, den: u64) -> Self {
    Rational { num, den }
Luca Barbato's avatar
Luca Barbato committed
37 38 39
  }
}

40 41 42 43 44 45 46
#[derive(Clone, Copy, Debug)]
#[repr(C)]
pub struct Point {
  pub x: u16,
  pub y: u16
}

47
#[derive(Clone, Debug)]
48
pub struct EncoderConfig {
49
  // output size
50 51
  pub width: usize,
  pub height: usize,
52 53

  // data format and ancillary color information
54 55 56
  pub bit_depth: usize,
  pub chroma_sampling: ChromaSampling,
  pub chroma_sample_position: ChromaSamplePosition,
57 58 59 60 61 62
  pub pixel_range: PixelRange,
  pub color_description: Option<ColorDescription>,
  pub mastering_display: Option<MasteringDisplay>,
  pub content_light: Option<ContentLight>,

  // encoder configuration
63
  pub time_base: Rational,
Josh Holmer's avatar
Josh Holmer committed
64 65 66 67
  /// The *minimum* interval between two keyframes
  pub min_key_frame_interval: u64,
  /// The *maximum* interval between two keyframes
  pub max_key_frame_interval: u64,
68
  pub low_latency: bool,
69
  pub quantizer: usize,
70 71
  pub tune: Tune,
  pub speed_settings: SpeedSettings,
72 73
  /// `None` for one-pass encode. `Some(1)` or `Some(2)` for two-pass encoding.
  pub pass: Option<u8>,
74
  pub show_psnr: bool,
75
  pub stats_file: Option<PathBuf>,
76 77 78 79
}

impl Default for EncoderConfig {
  fn default() -> Self {
80 81
    const DEFAULT_SPEED: usize = 3;
    Self::with_speed_preset(DEFAULT_SPEED)
82 83 84
  }
}

85 86 87
impl EncoderConfig {
  pub fn with_speed_preset(speed: usize) -> Self {
    EncoderConfig {
88 89
      width: 640,
      height: 480,
90

91 92 93
      bit_depth: 8,
      chroma_sampling: ChromaSampling::Cs420,
      chroma_sample_position: ChromaSamplePosition::Unknown,
94 95 96 97 98
      pixel_range: PixelRange::Unspecified,
      color_description: None,
      mastering_display: None,
      content_light: None,

99
      time_base: Rational { num: 30, den: 1 },
Josh Holmer's avatar
Josh Holmer committed
100 101
      min_key_frame_interval: 12,
      max_key_frame_interval: 240,
102
      low_latency: false,
103 104
      quantizer: 100,
      tune: Tune::Psnr,
105
      speed_settings: SpeedSettings::from_preset(speed),
106
      pass: None,
107
      show_psnr: false,
108
      stats_file: None,
109 110 111 112 113 114 115 116 117 118 119 120 121 122 123
    }
  }
}

#[derive(Clone, Copy, Debug)]
pub struct SpeedSettings {
  pub min_block_size: BlockSize,
  pub multiref: bool,
  pub fast_deblock: bool,
  pub reduced_tx_set: bool,
  pub tx_domain_distortion: bool,
  pub encode_bottomup: bool,
  pub rdo_tx_decision: bool,
  pub prediction_modes: PredictionModesSetting,
  pub include_near_mvs: bool,
124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141
  pub no_scene_detection: bool,
}

impl Default for SpeedSettings {
  fn default() -> Self {
    SpeedSettings {
      min_block_size: BlockSize::BLOCK_32X32,
      multiref: false,
      fast_deblock: false,
      reduced_tx_set: false,
      tx_domain_distortion: false,
      encode_bottomup: false,
      rdo_tx_decision: false,
      prediction_modes: PredictionModesSetting::Simple,
      include_near_mvs: false,
      no_scene_detection: false,
    }
  }
142 143 144 145 146 147 148 149 150 151 152 153 154 155
}

impl SpeedSettings {
  pub fn from_preset(speed: usize) -> Self {
    SpeedSettings {
      min_block_size: Self::min_block_size_preset(speed),
      multiref: Self::multiref_preset(speed),
      fast_deblock: Self::fast_deblock_preset(speed),
      reduced_tx_set: Self::reduced_tx_set_preset(speed),
      tx_domain_distortion: Self::tx_domain_distortion_preset(speed),
      encode_bottomup: Self::encode_bottomup_preset(speed),
      rdo_tx_decision: Self::rdo_tx_decision_preset(speed),
      prediction_modes: Self::prediction_modes_preset(speed),
      include_near_mvs: Self::include_near_mvs_preset(speed),
156
      no_scene_detection: Self::no_scene_detection_preset(speed),
157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186
    }
  }

  fn min_block_size_preset(speed: usize) -> BlockSize {
    if speed <= 1 {
      BlockSize::BLOCK_4X4
    } else if speed <= 2 {
      BlockSize::BLOCK_8X8
    } else if speed <= 3 {
      BlockSize::BLOCK_16X16
    } else if speed <= 4 {
      BlockSize::BLOCK_32X32
    } else {
      BlockSize::BLOCK_64X64
    }
  }

  fn multiref_preset(speed: usize) -> bool {
    speed <= 2
  }

  fn fast_deblock_preset(speed: usize) -> bool {
    speed >= 4
  }

  fn reduced_tx_set_preset(speed: usize) -> bool {
    speed >= 2
  }

  fn tx_domain_distortion_preset(speed: usize) -> bool {
187
    speed >= 4
188 189 190 191 192 193 194 195 196 197 198
  }

  fn encode_bottomup_preset(speed: usize) -> bool {
    speed == 0
  }

  fn rdo_tx_decision_preset(speed: usize) -> bool {
    speed <= 3
  }

  fn prediction_modes_preset(speed: usize) -> PredictionModesSetting {
199
    if speed <= 1 {
200
      PredictionModesSetting::ComplexAll
201
    } else if speed <= 3 {
202 203 204 205 206 207 208 209 210
      PredictionModesSetting::ComplexKeyframes
    } else {
      PredictionModesSetting::Simple
    }
  }

  fn include_near_mvs_preset(speed: usize) -> bool {
    speed <= 2
  }
211 212 213 214

  fn no_scene_detection_preset(speed: usize) -> bool {
    speed == 10
  }
215 216
}

217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238
#[allow(dead_code, non_camel_case_types)]
#[derive(Debug, PartialEq, Clone, Copy, Serialize, Deserialize)]
#[repr(C)]
pub enum FrameType {
  KEY,
  INTER,
  INTRA_ONLY,
  SWITCH
}

impl fmt::Display for FrameType {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    use self::FrameType::*;
    match self {
      KEY => write!(f, "Key frame"),
      INTER => write!(f, "Inter frame"),
      INTRA_ONLY => write!(f, "Intra only frame"),
      SWITCH => write!(f, "Switching frame"),
    }
  }
}

239 240 241 242 243 244 245
#[derive(Clone, Copy, Debug, PartialOrd, PartialEq)]
pub enum PredictionModesSetting {
  Simple,
  ComplexKeyframes,
  ComplexAll,
}

246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287
#[derive(Copy, Clone, Debug, PartialEq)]
#[repr(C)]
pub enum ChromaSampling {
  Cs420,
  Cs422,
  Cs444,
  Cs400,
}

impl Default for ChromaSampling {
  fn default() -> Self {
    ChromaSampling::Cs420
  }
}

impl ChromaSampling {
  // Provides the sampling period in the horizontal and vertical axes.
  pub fn sampling_period(self) -> (usize, usize) {
    use self::ChromaSampling::*;
    match self {
      Cs420 => (2, 2),
      Cs422 => (2, 1),
      Cs444 => (1, 1),
      Cs400 => (2, 2),
    }
  }
}

#[derive(Copy, Clone, Debug, PartialEq)]
#[repr(C)]
pub enum ChromaSamplePosition {
  Unknown,
  Vertical,
  Colocated
}

impl Default for ChromaSamplePosition {
  fn default() -> Self {
    ChromaSamplePosition::Unknown
  }
}

288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303
arg_enum!{
  #[derive(Debug, Clone, Copy, PartialEq)]
  #[repr(C)]
  pub enum PixelRange {
      Unspecified = 0,
      Limited,
      Full,
  }
}

impl Default for PixelRange {
    fn default() -> Self {
        PixelRange::Unspecified
    }
}

304 305
arg_enum!{
  #[derive(Debug, Clone, Copy, PartialEq)]
306
  #[repr(C)]
307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322
  pub enum MatrixCoefficients {
      Identity = 0,
      BT709,
      Unspecified,
      BT470M = 4,
      BT470BG,
      ST170M,
      ST240M,
      YCgCo,
      BT2020NonConstantLuminance,
      BT2020ConstantLuminance,
      ST2085,
      ChromaticityDerivedNonConstantLuminance,
      ChromaticityDerivedConstantLuminance,
      ICtCp,
  }
323 324 325 326 327 328 329 330
}

impl Default for MatrixCoefficients {
    fn default() -> Self {
        MatrixCoefficients::Unspecified
    }
}

331
arg_enum!{
332
  #[derive(Debug, Clone, Copy, PartialEq)]
333
  #[repr(C)]
334 335 336 337 338 339 340 341 342 343 344 345 346 347
  pub enum ColorPrimaries {
      BT709 = 1,
      Unspecified,
      BT470M = 4,
      BT470BG,
      ST170M,
      ST240M,
      Film,
      BT2020,
      ST428,
      P3DCI,
      P3Display,
      Tech3213 = 22,
  }
348 349 350 351 352 353 354 355
}

impl Default for ColorPrimaries {
    fn default() -> Self {
        ColorPrimaries::Unspecified
    }
}

356
arg_enum!{
357
  #[derive(Debug, Clone, Copy, PartialEq)]
358
  #[repr(C)]
359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377
  pub enum TransferCharacteristics {
      BT1886 = 1,
      Unspecified,
      BT470M = 4,
      BT470BG,
      ST170M,
      ST240M,
      Linear,
      Logarithmic100,
      Logarithmic316,
      XVYCC,
      BT1361E,
      SRGB,
      BT2020Ten,
      BT2020Twelve,
      PerceptualQuantizer,
      ST428,
      HybridLogGamma,
  }
378 379 380 381 382 383 384 385
}

impl Default for TransferCharacteristics {
    fn default() -> Self {
        TransferCharacteristics::Unspecified
    }
}

386 387 388 389 390 391 392
#[derive(Copy, Clone, Debug)]
pub struct ColorDescription {
    pub color_primaries: ColorPrimaries,
    pub transfer_characteristics: TransferCharacteristics,
    pub matrix_coefficients: MatrixCoefficients
}

393 394 395 396 397 398 399 400 401 402 403 404 405 406
#[derive(Copy, Clone, Debug)]
pub struct MasteringDisplay {
    pub primaries: [Point; 3],
    pub white_point: Point,
    pub max_luminance: u32,
    pub min_luminance: u32,
}

#[derive(Copy, Clone, Debug)]
pub struct ContentLight {
    pub max_content_light_level: u16,
    pub max_frame_average_light_level: u16,
}

407
/// Contain all the encoder configuration
408
#[derive(Clone, Debug)]
409
pub struct Config {
Luca Barbato's avatar
Luca Barbato committed
410 411 412 413
  pub enc: EncoderConfig
}

impl Config {
Luca Barbato's avatar
Luca Barbato committed
414 415
  pub fn parse(&mut self, key: &str, value: &str) -> Result<(), EncoderStatus> {
    match key {
416
      "low_latency" => self.enc.low_latency = value.parse().map_err(|_e| ParseError)?,
Josh Holmer's avatar
Josh Holmer committed
417 418
      "min_key_frame_interval" => self.enc.min_key_frame_interval = value.parse().map_err(|_e| ParseError)?,
      "key_frame_interval" => self.enc.max_key_frame_interval = value.parse().map_err(|_e| ParseError)?,
419
      "quantizer" => self.enc.quantizer = value.parse().map_err(|_e| ParseError)?,
420
      "speed" => self.enc.speed_settings = SpeedSettings::from_preset(value.parse().map_err(|_e| ParseError)?),
421 422
      "tune" => self.enc.tune = value.parse().map_err(|_e| ParseError)?,
      _ => return Err(InvalidKey)
Luca Barbato's avatar
Luca Barbato committed
423 424 425 426 427
    }

    Ok(())
  }

Luca Barbato's avatar
Luca Barbato committed
428
  pub fn new_context(&self) -> Context {
429
    #[cfg(feature = "aom")]
Luca Barbato's avatar
Luca Barbato committed
430
    unsafe {
Luca Barbato's avatar
Luca Barbato committed
431 432
      av1_rtcd();
      aom_dsp_rtcd();
Luca Barbato's avatar
Luca Barbato committed
433 434
    }

435 436 437
    // initialize with temporal delimiter
    let packet_data = TEMPORAL_DELIMITER.to_vec();

Josh Holmer's avatar
Josh Holmer committed
438 439
    Context {
      frame_count: 0,
440
      limit: 0,
Josh Holmer's avatar
Josh Holmer committed
441
      idx: 0,
Josh Holmer's avatar
Josh Holmer committed
442
      frames_processed: 0,
Josh Holmer's avatar
Josh Holmer committed
443
      frame_q: BTreeMap::new(),
Josh Holmer's avatar
Josh Holmer committed
444 445
      frame_data: BTreeMap::new(),
      keyframes: BTreeSet::new(),
446
      packet_data,
Josh Holmer's avatar
Josh Holmer committed
447 448
      segment_start_idx: 0,
      segment_start_frame: 0,
449
      keyframe_detector: SceneChangeDetector::new(self.enc.bit_depth),
450 451 452 453 454
      config: self.clone(),
      rc_state: RCState::new(),
      first_pass_data: FirstPassData {
        frames: Vec::new(),
      },
Josh Holmer's avatar
Josh Holmer committed
455
    }
Luca Barbato's avatar
Luca Barbato committed
456 457 458 459
  }
}

pub struct Context {
Luca Barbato's avatar
Luca Barbato committed
460
  //    timebase: Rational,
fbossen's avatar
fbossen committed
461
  frame_count: u64,
462
  limit: u64,
fbossen's avatar
fbossen committed
463
  idx: u64,
Josh Holmer's avatar
Josh Holmer committed
464 465
  frames_processed: u64,
  /// Maps frame *number* to frames
466
  frame_q: BTreeMap<u64, Option<Arc<Frame>>>, //    packet_q: VecDeque<Packet>
Josh Holmer's avatar
Josh Holmer committed
467 468 469 470 471
  /// Maps frame *idx* to frame data
  frame_data: BTreeMap<u64, FrameInvariants>,
  /// A list of keyframe *numbers* in this encode. Needed so that we don't
  /// need to keep all of the frame_data in memory for the whole life of the encode.
  keyframes: BTreeSet<u64>,
472
  /// A storage space for reordered frames.
Josh Holmer's avatar
Josh Holmer committed
473 474 475 476
  packet_data: Vec<u8>,
  segment_start_idx: u64,
  segment_start_frame: u64,
  keyframe_detector: SceneChangeDetector,
477
  pub config: Config,
478 479
  rc_state: RCState,
  pub first_pass_data: FirstPassData,
Luca Barbato's avatar
Luca Barbato committed
480 481 482 483 484 485 486 487 488
}

#[derive(Clone, Copy, Debug)]
pub enum EncoderStatus {
  /// The encoder needs more Frames to produce an output Packet
  NeedMoreData,
  /// There are enough Frames queue
  EnoughData,
  ///
Luca Barbato's avatar
Luca Barbato committed
489 490 491
  Failure,
  InvalidKey,
  ParseError
Luca Barbato's avatar
Luca Barbato committed
492 493 494 495
}

pub struct Packet {
  pub data: Vec<u8>,
fbossen's avatar
fbossen committed
496 497
  pub rec: Option<Frame>,
  pub number: u64,
498 499 500
  pub frame_type: FrameType,
  /// PSNR for Y, U, and V planes
  pub psnr: Option<(f64, f64, f64)>,
501 502 503 504
}

impl fmt::Display for Packet {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
Luca Barbato's avatar
Luca Barbato committed
505 506 507 508 509 510 511
    write!(
      f,
      "Frame {} - {} - {} bytes",
      self.number,
      self.frame_type,
      self.data.len()
    )
512
  }
Luca Barbato's avatar
Luca Barbato committed
513 514 515 516
}

impl Context {
  pub fn new_frame(&self) -> Arc<Frame> {
Josh Holmer's avatar
Josh Holmer committed
517
    Arc::new(Frame::new(
518 519
      self.config.enc.width,
      self.config.enc.height,
520
      self.config.enc.chroma_sampling
Josh Holmer's avatar
Josh Holmer committed
521
    ))
Luca Barbato's avatar
Luca Barbato committed
522 523 524 525 526 527
  }

  pub fn send_frame<F>(&mut self, frame: F) -> Result<(), EncoderStatus>
  where
    F: Into<Option<Arc<Frame>>>
  {
Josh Holmer's avatar
Josh Holmer committed
528 529
    let idx = self.frame_count;
    self.frame_q.insert(idx, frame.into());
530
    self.frame_count += 1;
Luca Barbato's avatar
Luca Barbato committed
531 532 533
    Ok(())
  }

534 535 536 537
  pub fn get_frame_count(&self) -> u64 {
    self.frame_count
  }

538 539
  pub fn set_limit(&mut self, limit: u64) {
    self.limit = limit;
540 541
  }

Josh Holmer's avatar
Josh Holmer committed
542 543 544 545
  pub fn needs_more_lookahead(&self) -> bool {
    self.needs_more_frames(self.frame_count) && self.frames_processed + LOOKAHEAD_FRAMES > self.frame_q.keys().last().cloned().unwrap_or(0)
  }

546
  pub fn needs_more_frames(&self, frame_count: u64) -> bool {
547
    self.limit == 0 || frame_count < self.limit
548 549
  }

550 551 552
  pub fn container_sequence_header(&mut self) -> Vec<u8> {
    fn sequence_header_inner(seq: &Sequence) -> io::Result<Vec<u8>> {
      let mut buf = Vec::new();
553

554
      {
555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571
        let mut bw = BitWriter::endian(&mut buf, BigEndian);
        bw.write_bit(true)?; // marker
        bw.write(7, 1)?; // version
        bw.write(3, seq.profile)?;
        bw.write(5, 32)?; // level
        bw.write_bit(false)?; // tier
        bw.write_bit(seq.bit_depth > 8)?; // high_bitdepth
        bw.write_bit(seq.bit_depth == 12)?; // twelve_bit
        bw.write_bit(seq.bit_depth == 1)?; // monochrome
        bw.write_bit(seq.bit_depth == 12)?; // twelve_bit
        bw.write_bit(seq.chroma_sampling != ChromaSampling::Cs444)?; // chroma_subsampling_x
        bw.write_bit(seq.chroma_sampling == ChromaSampling::Cs420)?; // chroma_subsampling_y
        bw.write(2, 0)?; // sample_position
        bw.write(3, 0)?; // reserved
        bw.write_bit(false)?; // initial_presentation_delay_present

        bw.write(4, 0)?; // reserved
572
      }
573

574 575 576
      Ok(buf)
    }

Josh Holmer's avatar
Josh Holmer committed
577
    sequence_header_inner(&self.frame_data[&0].sequence).unwrap()
578 579
  }

Josh Holmer's avatar
Josh Holmer committed
580
  fn next_keyframe(&self) -> u64 {
Josh Holmer's avatar
Josh Holmer committed
581 582 583 584
    let next_detected = self.frame_data.values()
      .find(|fi| fi.frame_type == FrameType::KEY && fi.number > self.segment_start_frame)
      .map(|fi| fi.number);
    let next_limit = self.segment_start_frame + self.config.enc.max_key_frame_interval;
Josh Holmer's avatar
Josh Holmer committed
585 586 587 588 589
    if next_detected.is_none() {
      return next_limit;
    }
    cmp::min(next_detected.unwrap(), next_limit)
  }
fbossen's avatar
fbossen committed
590

591
  fn set_frame_properties(&mut self, idx: u64) -> bool {
592
    let (fi, end_of_subgop) = self.build_frame_properties(idx);
593 594 595
    self.frame_data.insert(idx, fi);

    end_of_subgop
Josh Holmer's avatar
Josh Holmer committed
596 597
  }

598
  fn build_frame_properties(&mut self, idx: u64) -> (FrameInvariants, bool) {
Josh Holmer's avatar
Josh Holmer committed
599
    if idx == 0 {
600
      let seq = Sequence::new(&self.config.enc);
601

Josh Holmer's avatar
Josh Holmer committed
602
      // The first frame will always be a key frame
Josh Holmer's avatar
Josh Holmer committed
603 604
      let fi = FrameInvariants::new_key_frame(
        &FrameInvariants::new(
605
          self.config.enc.clone(),
606
          seq
Josh Holmer's avatar
Josh Holmer committed
607 608 609
        ),
        0
      );
610
      return (fi, true);
Josh Holmer's avatar
Josh Holmer committed
611 612
    }

Josh Holmer's avatar
Josh Holmer committed
613 614
    let mut fi = self.frame_data[&(idx - 1)].clone();

615 616 617 618 619
    // FIXME: inter unsupported with 4:2:2 and 4:4:4 chroma sampling
    let chroma_sampling = self.config.enc.chroma_sampling;
    let keyframe_only = chroma_sampling == ChromaSampling::Cs444 ||
      chroma_sampling == ChromaSampling::Cs422;

Josh Holmer's avatar
Josh Holmer committed
620 621 622 623 624 625
    // Initially set up the frame as an inter frame.
    // We need to determine what the frame number is before we can
    // look up the frame type. If reordering is enabled, the idx
    // may not match the frame number.
    let idx_in_segment = idx - self.segment_start_idx;
    if idx_in_segment > 0 {
626
      let next_keyframe = if keyframe_only { self.segment_start_frame + 1 } else { self.next_keyframe() };
627
      let (fi_temp, end_of_subgop) = FrameInvariants::new_inter_frame(
Josh Holmer's avatar
Josh Holmer committed
628
        &fi,
629 630 631 632
        self.segment_start_frame,
        idx_in_segment,
        next_keyframe
      );
Josh Holmer's avatar
Josh Holmer committed
633
      fi = fi_temp;
634
      if !end_of_subgop {
Josh Holmer's avatar
Josh Holmer committed
635 636 637
        if !fi.inter_cfg.unwrap().reorder
          || ((idx_in_segment - 1) % fi.inter_cfg.unwrap().group_len == 0
          && fi.number == (next_keyframe - 1))
638
        {
Josh Holmer's avatar
Josh Holmer committed
639
          self.segment_start_idx = idx;
Josh Holmer's avatar
Josh Holmer committed
640
          self.segment_start_frame = next_keyframe;
Josh Holmer's avatar
Josh Holmer committed
641
          fi.number = next_keyframe;
Josh Holmer's avatar
Josh Holmer committed
642
        } else {
643
          return (fi, false);
Josh Holmer's avatar
Josh Holmer committed
644 645
        }
      }
fbossen's avatar
fbossen committed
646
    }
Luca Barbato's avatar
Luca Barbato committed
647

648 649
    match self.frame_q.get(&fi.number) {
      Some(Some(_)) => {},
Josh Holmer's avatar
Josh Holmer committed
650
      _ => { return (fi, false); }
651 652
    }

Josh Holmer's avatar
Josh Holmer committed
653
    // Now that we know the frame number, look up the correct frame type
Josh Holmer's avatar
Josh Holmer committed
654 655 656 657 658 659 660
    let frame_type = self.determine_frame_type(fi.number);
    if frame_type == FrameType::KEY {
      self.segment_start_idx = idx;
      self.segment_start_frame = fi.number;
      self.keyframes.insert(fi.number);
    }
    fi.frame_type = frame_type;
Josh Holmer's avatar
Josh Holmer committed
661

Josh Holmer's avatar
Josh Holmer committed
662 663 664 665 666
    let idx_in_segment = idx - self.segment_start_idx;
    if idx_in_segment == 0 {
      fi = FrameInvariants::new_key_frame(&fi, self.segment_start_frame);
    } else {
      let next_keyframe = self.next_keyframe();
667
      let (fi_temp, end_of_subgop) = FrameInvariants::new_inter_frame(
Josh Holmer's avatar
Josh Holmer committed
668 669 670 671 672 673
        &fi,
        self.segment_start_frame,
        idx_in_segment,
        next_keyframe
      );
      fi = fi_temp;
674 675
      if !end_of_subgop {
        return (fi, false);
Josh Holmer's avatar
Josh Holmer committed
676 677
      }
    }
678
    (fi, true)
fbossen's avatar
fbossen committed
679
  }
Luca Barbato's avatar
Luca Barbato committed
680

fbossen's avatar
fbossen committed
681
  pub fn receive_packet(&mut self) -> Result<Packet, EncoderStatus> {
682 683 684 685 686 687
    let idx = {
      let mut idx = self.idx;
      while !self.set_frame_properties(idx) {
        self.idx += 1;
        idx = self.idx;
      }
Luca Barbato's avatar
Luca Barbato committed
688

689 690 691 692 693 694
      if !self.needs_more_frames(self.frame_data.get(&idx).unwrap().number) {
        self.idx += 1;
        return Err(EncoderStatus::EnoughData);
      }
      idx
    };
695

696
    let ret = {
Luca Barbato's avatar
Luca Barbato committed
697 698 699
      let fi = self.frame_data.get_mut(&idx).unwrap();
      if fi.show_existing_frame {
        self.idx += 1;
Luca Barbato's avatar
Luca Barbato committed
700

Luca Barbato's avatar
Luca Barbato committed
701
        let mut fs = FrameState::new(fi);
fbossen's avatar
fbossen committed
702

703 704
        let sef_data = encode_frame(fi, &mut fs);
        self.packet_data.extend(sef_data);
Luca Barbato's avatar
Luca Barbato committed
705

Luca Barbato's avatar
Luca Barbato committed
706
        let rec = if fi.show_frame { Some(fs.rec) } else { None };
707 708
        let fi = fi.clone();
        self.finalize_packet(&*fs.input, rec, &fi)
Luca Barbato's avatar
Luca Barbato committed
709 710 711
      } else {
        if let Some(f) = self.frame_q.get(&fi.number) {
          self.idx += 1;
fbossen's avatar
fbossen committed
712

713
          if let Some(frame) = f.clone() {
714 715 716
            let fti = fi.get_frame_subtype();
            let qps = self.rc_state.select_qi(fi, fti);
            fi.set_quantizers(&qps);
Luca Barbato's avatar
Luca Barbato committed
717
            let mut fs = FrameState::new_with_frame(fi, frame.clone());
fbossen's avatar
fbossen committed
718

Luca Barbato's avatar
Luca Barbato committed
719 720
            let data = encode_frame(fi, &mut fs);
            self.packet_data.extend(data);
fbossen's avatar
fbossen committed
721

Luca Barbato's avatar
Luca Barbato committed
722
            fs.rec.pad(fi.width, fi.height);
fbossen's avatar
fbossen committed
723

Luca Barbato's avatar
Luca Barbato committed
724 725
            // TODO avoid the clone by having rec Arc.
            let rec = if fi.show_frame { Some(fs.rec.clone()) } else { None };
fbossen's avatar
fbossen committed
726

Luca Barbato's avatar
Luca Barbato committed
727
            update_rec_buffer(fi, fs);
fbossen's avatar
fbossen committed
728

Luca Barbato's avatar
Luca Barbato committed
729
            if fi.show_frame {
730 731
              let fi = fi.clone();
              self.finalize_packet(&*frame, rec, &fi)
Luca Barbato's avatar
Luca Barbato committed
732 733 734
            } else {
              Err(EncoderStatus::NeedMoreData)
            }
735 736 737
          } else {
            Err(EncoderStatus::NeedMoreData)
          }
fbossen's avatar
fbossen committed
738 739 740 741
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      }
742 743 744
    };

    if let Ok(ref pkt) = ret {
Luca Barbato's avatar
Luca Barbato committed
745
      self.garbage_collect(pkt.number);
746 747 748
    }

    ret
Luca Barbato's avatar
Luca Barbato committed
749 750
  }

751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768
  fn finalize_packet(&mut self, original_frame: &Frame, rec: Option<Frame>, fi: &FrameInvariants) -> Result<Packet, EncoderStatus> {
    let data = self.packet_data.clone();
    self.packet_data.clear();
    if write_temporal_delimiter(&mut self.packet_data).is_err() {
      return Err(EncoderStatus::Failure);
    }

    let mut psnr = None;
    if self.config.enc.show_psnr {
      if let Some(ref rec) = rec {
        psnr = Some(calculate_frame_psnr(
          &*original_frame,
          rec,
          fi.sequence.bit_depth
        ));
      }
    }

769 770 771 772
    if self.config.enc.pass == Some(1) {
      self.first_pass_data.frames.push(FirstPassFrame::from(fi));
    }

773 774 775 776 777 778 779 780 781 782
    self.frames_processed += 1;
    Ok(Packet {
      data,
      rec,
      number: fi.number,
      frame_type: fi.frame_type,
      psnr
    })
  }

783
  fn garbage_collect(&mut self, cur_frame: u64) {
Josh Holmer's avatar
Josh Holmer committed
784 785 786 787 788 789 790 791 792 793 794 795 796 797
    if cur_frame == 0 {
      return;
    }
    for i in 0..cur_frame {
      self.frame_q.remove(&i);
    }
    if self.idx < 2 {
      return;
    }
    for i in 0..(self.idx - 1) {
      self.frame_data.remove(&i);
    }
  }

Luca Barbato's avatar
Luca Barbato committed
798
  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
799
    self.frame_q.insert(self.frame_count, None);
800
    self.frame_count += 1;
Luca Barbato's avatar
Luca Barbato committed
801
  }
Josh Holmer's avatar
Josh Holmer committed
802

Josh Holmer's avatar
Josh Holmer committed
803 804
  fn determine_frame_type(&mut self, frame_number: u64) -> FrameType {
    if frame_number == 0 {
Josh Holmer's avatar
Josh Holmer committed
805 806
      return FrameType::KEY;
    }
807 808 809 810 811 812 813
    if self.config.enc.speed_settings.no_scene_detection {
      if frame_number % self.config.enc.max_key_frame_interval == 0 {
        return FrameType::KEY;
      } else {
        return FrameType::INTER;
      }
    }
Josh Holmer's avatar
Josh Holmer committed
814

Josh Holmer's avatar
Josh Holmer committed
815 816 817 818 819 820 821 822
    let prev_keyframe = self.keyframes.iter()
      .rfind(|&&keyframe| keyframe < frame_number)
      .cloned()
      .unwrap_or(0);
    let frame = match self.frame_q.get(&frame_number).cloned() {
      Some(frame) => frame,
      None => { return FrameType::KEY; }
    };
Josh Holmer's avatar
Josh Holmer committed
823
    if let Some(frame) = frame {
Josh Holmer's avatar
Josh Holmer committed
824 825 826
      let distance = frame_number - prev_keyframe;
      if distance < self.config.enc.min_key_frame_interval {
        if distance + 1 == self.config.enc.min_key_frame_interval {
Josh Holmer's avatar
Josh Holmer committed
827 828
          // Run the detector for the current frame, so that it will contain this frame's information
          // to compare against the next frame. We can ignore the results for this frame.
Josh Holmer's avatar
Josh Holmer committed
829
          self.keyframe_detector.detect_scene_change(frame, frame_number as usize);
Josh Holmer's avatar
Josh Holmer committed
830 831 832
        }
        return FrameType::INTER;
      }
Josh Holmer's avatar
Josh Holmer committed
833
      if distance >= self.config.enc.max_key_frame_interval {
Josh Holmer's avatar
Josh Holmer committed
834 835
        return FrameType::KEY;
      }
Josh Holmer's avatar
Josh Holmer committed
836
      if self.keyframe_detector.detect_scene_change(frame, frame_number as usize) {
Josh Holmer's avatar
Josh Holmer committed
837 838 839 840 841
        return FrameType::KEY;
      }
    }
    FrameType::INTER
  }
Luca Barbato's avatar
Luca Barbato committed
842
}
843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FirstPassData {
  frames: Vec<FirstPassFrame>,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct FirstPassFrame {
  number: u64,
  frame_type: FrameType,
}

impl From<&FrameInvariants> for FirstPassFrame {
  fn from(fi: &FrameInvariants) -> FirstPassFrame {
    FirstPassFrame {
      number: fi.number,
      frame_type: fi.frame_type,
    }
  }
862
}