api.rs 17.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::*;
Luca Barbato's avatar
Luca Barbato committed
11
use encoder::*;
Raphaël Zumer's avatar
Raphaël Zumer committed
12
use metrics::calculate_frame_psnr;
13
use partition::*;
Raphaël Zumer's avatar
Raphaël Zumer committed
14 15
use scenechange::SceneChangeDetector;
use self::EncoderStatus::*;
Luca Barbato's avatar
Luca Barbato committed
16

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

const LOOKAHEAD_FRAMES: u64 = 10;
Luca Barbato's avatar
Luca Barbato committed
24

25 26 27 28 29 30 31 32
#[derive(Debug, Clone, Copy)]
pub struct VideoDetails {
  pub width: usize,
  pub height: usize,
  pub mono: bool,
  pub bit_depth: usize,
  pub chroma_sampling: ChromaSampling,
  pub chroma_sample_position: ChromaSamplePosition,
Luca Barbato's avatar
Luca Barbato committed
33
  pub time_base: Rational,
34 35 36 37 38 39 40 41 42 43 44
}

impl Default for VideoDetails {
  fn default() -> Self {
    VideoDetails {
      width: 640,
      height: 480,
      mono: false,
      bit_depth: 8,
      chroma_sampling: ChromaSampling::Cs420,
      chroma_sample_position: ChromaSamplePosition::Unknown,
Luca Barbato's avatar
Luca Barbato committed
45
      time_base: Rational { num: 30, den: 1 }
46 47 48 49
    }
  }
}

Luca Barbato's avatar
Luca Barbato committed
50 51
// TODO: use the num crate?
#[derive(Clone, Copy, Debug)]
52
#[repr(C)]
Luca Barbato's avatar
Luca Barbato committed
53 54 55
pub struct Rational {
  pub num: u64,
  pub den: u64
Luca Barbato's avatar
Luca Barbato committed
56 57
}

Luca Barbato's avatar
Luca Barbato committed
58 59 60
impl Rational {
  pub fn new(num: u64, den: u64) -> Self {
    Rational { num, den }
Luca Barbato's avatar
Luca Barbato committed
61 62 63
  }
}

64 65
#[derive(Copy, Clone, Debug)]
pub struct EncoderConfig {
Josh Holmer's avatar
Josh Holmer committed
66 67 68 69
  /// The *minimum* interval between two keyframes
  pub min_key_frame_interval: u64,
  /// The *maximum* interval between two keyframes
  pub max_key_frame_interval: u64,
70
  pub low_latency: bool,
71
  pub quantizer: usize,
72
  pub tune: Tune,
73
  pub pixel_range: PixelRange,
74
  pub color_description: Option<ColorDescription>,
75
  pub speed_settings: SpeedSettings,
76
  pub show_psnr: bool,
77 78 79 80
}

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

86 87 88
impl EncoderConfig {
  pub fn with_speed_preset(speed: usize) -> Self {
    EncoderConfig {
Josh Holmer's avatar
Josh Holmer committed
89 90
      min_key_frame_interval: 12,
      max_key_frame_interval: 240,
91 92 93
      low_latency: true,
      quantizer: 100,
      tune: Tune::Psnr,
94
      pixel_range: PixelRange::Unspecified,
95
      color_description: None,
96 97
      speed_settings: SpeedSettings::from_preset(speed),
      show_psnr: false,
98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168
    }
  }
}

#[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,
}

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),
    }
  }

  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 {
    speed >= 1
  }

  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 {
169
    if speed <= 1 {
170
      PredictionModesSetting::ComplexAll
171
    } else if speed <= 3 {
172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189
      PredictionModesSetting::ComplexKeyframes
    } else {
      PredictionModesSetting::Simple
    }
  }

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

#[derive(Clone, Copy, Debug, PartialOrd, PartialEq)]
pub enum PredictionModesSetting {
  Simple,
  ComplexKeyframes,
  ComplexAll,
}

190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205
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
    }
}

206 207
arg_enum!{
  #[derive(Debug, Clone, Copy, PartialEq)]
208
  #[repr(C)]
209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224
  pub enum MatrixCoefficients {
      Identity = 0,
      BT709,
      Unspecified,
      BT470M = 4,
      BT470BG,
      ST170M,
      ST240M,
      YCgCo,
      BT2020NonConstantLuminance,
      BT2020ConstantLuminance,
      ST2085,
      ChromaticityDerivedNonConstantLuminance,
      ChromaticityDerivedConstantLuminance,
      ICtCp,
  }
225 226 227 228 229 230 231 232
}

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

233 234
arg_enum!{
  #[derive(Debug,Clone,Copy,PartialEq)]
235
  #[repr(C)]
236 237 238 239 240 241 242 243 244 245 246 247 248 249
  pub enum ColorPrimaries {
      BT709 = 1,
      Unspecified,
      BT470M = 4,
      BT470BG,
      ST170M,
      ST240M,
      Film,
      BT2020,
      ST428,
      P3DCI,
      P3Display,
      Tech3213 = 22,
  }
250 251 252 253 254 255 256 257
}

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

258 259
arg_enum!{
  #[derive(Debug,Clone,Copy,PartialEq)]
260
  #[repr(C)]
261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279
  pub enum TransferCharacteristics {
      BT1886 = 1,
      Unspecified,
      BT470M = 4,
      BT470BG,
      ST170M,
      ST240M,
      Linear,
      Logarithmic100,
      Logarithmic316,
      XVYCC,
      BT1361E,
      SRGB,
      BT2020Ten,
      BT2020Twelve,
      PerceptualQuantizer,
      ST428,
      HybridLogGamma,
  }
280 281 282 283 284 285 286 287
}

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

288 289 290 291 292 293 294
#[derive(Copy, Clone, Debug)]
pub struct ColorDescription {
    pub color_primaries: ColorPrimaries,
    pub transfer_characteristics: TransferCharacteristics,
    pub matrix_coefficients: MatrixCoefficients
}

295 296 297
/// Contain all the encoder configuration
#[derive(Clone, Copy, Debug)]
pub struct Config {
298
  pub video_info: VideoDetails,
Luca Barbato's avatar
Luca Barbato committed
299 300 301 302
  pub enc: EncoderConfig
}

impl Config {
Luca Barbato's avatar
Luca Barbato committed
303 304
  pub fn parse(&mut self, key: &str, value: &str) -> Result<(), EncoderStatus> {
    match key {
305
      "low_latency" => self.enc.low_latency = value.parse().map_err(|_e| ParseError)?,
Josh Holmer's avatar
Josh Holmer committed
306 307
      "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)?,
308
      "quantizer" => self.enc.quantizer = value.parse().map_err(|_e| ParseError)?,
309
      "speed" => self.enc.speed_settings = SpeedSettings::from_preset(value.parse().map_err(|_e| ParseError)?),
310 311
      "tune" => self.enc.tune = value.parse().map_err(|_e| ParseError)?,
      _ => return Err(InvalidKey)
Luca Barbato's avatar
Luca Barbato committed
312 313 314 315 316
    }

    Ok(())
  }

Luca Barbato's avatar
Luca Barbato committed
317
  pub fn new_context(&self) -> Context {
318
    #[cfg(feature = "aom")]
Luca Barbato's avatar
Luca Barbato committed
319
    unsafe {
Luca Barbato's avatar
Luca Barbato committed
320 321
      av1_rtcd();
      aom_dsp_rtcd();
Luca Barbato's avatar
Luca Barbato committed
322 323
    }

Josh Holmer's avatar
Josh Holmer committed
324 325
    Context {
      frame_count: 0,
326
      limit: 0,
Josh Holmer's avatar
Josh Holmer committed
327
      idx: 0,
Josh Holmer's avatar
Josh Holmer committed
328
      frames_processed: 0,
Josh Holmer's avatar
Josh Holmer committed
329
      frame_q: BTreeMap::new(),
Josh Holmer's avatar
Josh Holmer committed
330 331
      frame_data: BTreeMap::new(),
      keyframes: BTreeSet::new(),
Josh Holmer's avatar
Josh Holmer committed
332 333 334
      packet_data: Vec::new(),
      segment_start_idx: 0,
      segment_start_frame: 0,
335
      keyframe_detector: SceneChangeDetector::new(&self.video_info),
Josh Holmer's avatar
Josh Holmer committed
336
      config: *self,
Josh Holmer's avatar
Josh Holmer committed
337
    }
Luca Barbato's avatar
Luca Barbato committed
338 339 340 341
  }
}

pub struct Context {
Luca Barbato's avatar
Luca Barbato committed
342
  //    timebase: Rational,
fbossen's avatar
fbossen committed
343
  frame_count: u64,
344
  limit: u64,
fbossen's avatar
fbossen committed
345
  idx: u64,
Josh Holmer's avatar
Josh Holmer committed
346 347
  frames_processed: u64,
  /// Maps frame *number* to frames
348
  frame_q: BTreeMap<u64, Option<Arc<Frame>>>, //    packet_q: VecDeque<Packet>
Josh Holmer's avatar
Josh Holmer committed
349 350 351 352 353
  /// 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>,
354
  /// A storage space for reordered frames.
Josh Holmer's avatar
Josh Holmer committed
355 356 357 358
  packet_data: Vec<u8>,
  segment_start_idx: u64,
  segment_start_frame: u64,
  keyframe_detector: SceneChangeDetector,
359
  pub config: Config,
Luca Barbato's avatar
Luca Barbato committed
360 361 362 363 364 365 366 367 368
}

#[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
369 370 371
  Failure,
  InvalidKey,
  ParseError
Luca Barbato's avatar
Luca Barbato committed
372 373 374 375
}

pub struct Packet {
  pub data: Vec<u8>,
fbossen's avatar
fbossen committed
376 377
  pub rec: Option<Frame>,
  pub number: u64,
378 379 380
  pub frame_type: FrameType,
  /// PSNR for Y, U, and V planes
  pub psnr: Option<(f64, f64, f64)>,
381 382 383 384
}

impl fmt::Display for Packet {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
Luca Barbato's avatar
Luca Barbato committed
385 386 387 388 389 390 391
    write!(
      f,
      "Frame {} - {} - {} bytes",
      self.number,
      self.frame_type,
      self.data.len()
    )
392
  }
Luca Barbato's avatar
Luca Barbato committed
393 394 395 396
}

impl Context {
  pub fn new_frame(&self) -> Arc<Frame> {
Josh Holmer's avatar
Josh Holmer committed
397
    Arc::new(Frame::new(
398 399 400
      self.config.video_info.width.align_power_of_two(3),
      self.config.video_info.height.align_power_of_two(3),
      self.config.video_info.chroma_sampling
Josh Holmer's avatar
Josh Holmer committed
401
    ))
Luca Barbato's avatar
Luca Barbato committed
402 403 404 405 406 407
  }

  pub fn send_frame<F>(&mut self, frame: F) -> Result<(), EncoderStatus>
  where
    F: Into<Option<Arc<Frame>>>
  {
Josh Holmer's avatar
Josh Holmer committed
408 409
    let idx = self.frame_count;
    self.frame_q.insert(idx, frame.into());
410
    self.frame_count += 1;
Luca Barbato's avatar
Luca Barbato committed
411 412 413
    Ok(())
  }

414 415 416 417
  pub fn get_frame_count(&self) -> u64 {
    self.frame_count
  }

418 419
  pub fn set_limit(&mut self, limit: u64) {
    self.limit = limit;
420 421
  }

Josh Holmer's avatar
Josh Holmer committed
422 423 424 425
  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)
  }

426
  pub fn needs_more_frames(&self, frame_count: u64) -> bool {
427
    self.limit == 0 || frame_count < self.limit
428 429
  }

430 431 432
  pub fn container_sequence_header(&mut self) -> Vec<u8> {
    fn sequence_header_inner(seq: &Sequence) -> io::Result<Vec<u8>> {
      let mut buf = Vec::new();
433

434
      {
435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451
        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
452
      }
453

454 455 456
      Ok(buf)
    }

Josh Holmer's avatar
Josh Holmer committed
457
    sequence_header_inner(&self.frame_data[&0].sequence).unwrap()
458 459
  }

Josh Holmer's avatar
Josh Holmer committed
460
  fn next_keyframe(&self) -> u64 {
Josh Holmer's avatar
Josh Holmer committed
461 462 463 464
    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
465 466 467 468 469
    if next_detected.is_none() {
      return next_limit;
    }
    cmp::min(next_detected.unwrap(), next_limit)
  }
fbossen's avatar
fbossen committed
470

471 472 473 474 475
  fn set_frame_properties(&mut self, idx: u64) -> bool {
    let (fi, end_of_subgop) = self.get_frame_properties(idx);
    self.frame_data.insert(idx, fi);

    end_of_subgop
Josh Holmer's avatar
Josh Holmer committed
476 477
  }

478
  fn get_frame_properties(&mut self, idx: u64) -> (FrameInvariants, bool) {
Josh Holmer's avatar
Josh Holmer committed
479
    if idx == 0 {
480
      let mut seq = Sequence::new(&self.config.video_info);
481
      seq.pixel_range = self.config.enc.pixel_range;
482 483
      seq.color_description = self.config.enc.color_description;

Josh Holmer's avatar
Josh Holmer committed
484
      // The first frame will always be a key frame
Josh Holmer's avatar
Josh Holmer committed
485 486
      let fi = FrameInvariants::new_key_frame(
        &FrameInvariants::new(
487 488
          self.config.video_info.width,
          self.config.video_info.height,
Josh Holmer's avatar
Josh Holmer committed
489
          self.config.enc,
490
          seq
Josh Holmer's avatar
Josh Holmer committed
491 492 493
        ),
        0
      );
494
      return (fi, true);
Josh Holmer's avatar
Josh Holmer committed
495 496
    }

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

Josh Holmer's avatar
Josh Holmer committed
499 500 501 502 503 504
    // 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 {
Josh Holmer's avatar
Josh Holmer committed
505
      let next_keyframe = self.next_keyframe();
506
      let (fi_temp, end_of_subgop) = FrameInvariants::new_inter_frame(
Josh Holmer's avatar
Josh Holmer committed
507
        &fi,
508 509 510 511
        self.segment_start_frame,
        idx_in_segment,
        next_keyframe
      );
Josh Holmer's avatar
Josh Holmer committed
512
      fi = fi_temp;
513
      if !end_of_subgop {
Josh Holmer's avatar
Josh Holmer committed
514 515 516
        if !fi.inter_cfg.unwrap().reorder
          || ((idx_in_segment - 1) % fi.inter_cfg.unwrap().group_len == 0
          && fi.number == (next_keyframe - 1))
517
        {
Josh Holmer's avatar
Josh Holmer committed
518
          self.segment_start_idx = idx;
Josh Holmer's avatar
Josh Holmer committed
519
          self.segment_start_frame = next_keyframe;
Josh Holmer's avatar
Josh Holmer committed
520
          fi.number = next_keyframe;
Josh Holmer's avatar
Josh Holmer committed
521
        } else {
522
          return (fi, false);
Josh Holmer's avatar
Josh Holmer committed
523 524
        }
      }
fbossen's avatar
fbossen committed
525
    }
Luca Barbato's avatar
Luca Barbato committed
526

527 528
    match self.frame_q.get(&fi.number) {
      Some(Some(_)) => {},
Josh Holmer's avatar
Josh Holmer committed
529
      _ => { return (fi, false); }
530 531
    }

Josh Holmer's avatar
Josh Holmer committed
532
    // Now that we know the frame number, look up the correct frame type
Josh Holmer's avatar
Josh Holmer committed
533 534 535 536 537 538 539
    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
540

Josh Holmer's avatar
Josh Holmer committed
541 542 543 544 545
    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();
546
      let (fi_temp, end_of_subgop) = FrameInvariants::new_inter_frame(
Josh Holmer's avatar
Josh Holmer committed
547 548 549 550 551 552
        &fi,
        self.segment_start_frame,
        idx_in_segment,
        next_keyframe
      );
      fi = fi_temp;
553 554
      if !end_of_subgop {
        return (fi, false);
Josh Holmer's avatar
Josh Holmer committed
555 556
      }
    }
557
    (fi, true)
fbossen's avatar
fbossen committed
558
  }
Luca Barbato's avatar
Luca Barbato committed
559

fbossen's avatar
fbossen committed
560 561
  pub fn receive_packet(&mut self) -> Result<Packet, EncoderStatus> {
    let mut idx = self.idx;
562
    while !self.set_frame_properties(idx) {
Josh Holmer's avatar
Josh Holmer committed
563
      self.idx += 1;
fbossen's avatar
fbossen committed
564 565
      idx = self.idx;
    }
Luca Barbato's avatar
Luca Barbato committed
566

Josh Holmer's avatar
Josh Holmer committed
567
    if !self.needs_more_frames(self.frame_data.get(&idx).unwrap().number) {
Josh Holmer's avatar
Josh Holmer committed
568
      self.idx += 1;
569 570 571
      return Err(EncoderStatus::EnoughData)
    }

Josh Holmer's avatar
Josh Holmer committed
572 573
    let fi = self.frame_data.get_mut(&idx).unwrap();
    if fi.show_existing_frame {
Josh Holmer's avatar
Josh Holmer committed
574
      self.idx += 1;
Luca Barbato's avatar
Luca Barbato committed
575

Josh Holmer's avatar
Josh Holmer committed
576
      let mut fs = FrameState::new(fi);
fbossen's avatar
fbossen committed
577

Josh Holmer's avatar
Josh Holmer committed
578
      let data = encode_frame(fi, &mut fs);
Luca Barbato's avatar
Luca Barbato committed
579

Josh Holmer's avatar
Josh Holmer committed
580
      let rec = if fi.show_frame { Some(fs.rec) } else { None };
581
      let mut psnr = None;
Josh Holmer's avatar
Josh Holmer committed
582
      if self.config.enc.show_psnr {
583
        if let Some(ref rec) = rec {
Josh Holmer's avatar
Josh Holmer committed
584
          psnr = Some(calculate_frame_psnr(&*fs.input, rec, fi.sequence.bit_depth));
585 586
        }
      }
Luca Barbato's avatar
Luca Barbato committed
587

Josh Holmer's avatar
Josh Holmer committed
588 589
      self.frames_processed += 1;
      Ok(Packet { data, rec, number: fi.number, frame_type: fi.frame_type, psnr })
Luca Barbato's avatar
Luca Barbato committed
590
    } else {
Josh Holmer's avatar
Josh Holmer committed
591
      if let Some(f) = self.frame_q.get(&fi.number) {
Josh Holmer's avatar
Josh Holmer committed
592
        self.idx += 1;
fbossen's avatar
fbossen committed
593 594

        if let Some(frame) = f {
Josh Holmer's avatar
Josh Holmer committed
595
          let mut fs = FrameState::new_with_frame(fi, frame.clone());
fbossen's avatar
fbossen committed
596

Josh Holmer's avatar
Josh Holmer committed
597
          let data = encode_frame(fi, &mut fs);
598
          self.packet_data.extend(data);
fbossen's avatar
fbossen committed
599

Josh Holmer's avatar
Josh Holmer committed
600
          fs.rec.pad(fi.width, fi.height);
fbossen's avatar
fbossen committed
601 602

          // TODO avoid the clone by having rec Arc.
Josh Holmer's avatar
Josh Holmer committed
603
          let rec = if fi.show_frame { Some(fs.rec.clone()) } else { None };
fbossen's avatar
fbossen committed
604

Josh Holmer's avatar
Josh Holmer committed
605
          update_rec_buffer(fi, fs);
fbossen's avatar
fbossen committed
606

Josh Holmer's avatar
Josh Holmer committed
607
          if fi.show_frame {
608
            let data = self.packet_data.clone();
609
            self.packet_data.clear();
610 611

            let mut psnr = None;
Josh Holmer's avatar
Josh Holmer committed
612
            if self.config.enc.show_psnr {
613
              if let Some(ref rec) = rec {
Josh Holmer's avatar
Josh Holmer committed
614
                psnr = Some(calculate_frame_psnr(&*frame, rec, fi.sequence.bit_depth));
615 616 617
              }
            }

Josh Holmer's avatar
Josh Holmer committed
618 619
            self.frames_processed += 1;
            Ok(Packet { data, rec, number: fi.number, frame_type: fi.frame_type, psnr })
620 621 622
          } else {
            Err(EncoderStatus::NeedMoreData)
          }
fbossen's avatar
fbossen committed
623 624 625 626 627 628
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      } else {
        Err(EncoderStatus::NeedMoreData)
      }
Luca Barbato's avatar
Luca Barbato committed
629 630 631
    }
  }

Josh Holmer's avatar
Josh Holmer committed
632 633 634 635 636 637 638 639 640 641 642 643 644 645 646
  pub fn garbage_collect(&mut self, cur_frame: u64) {
    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
647
  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
648
    self.frame_q.insert(self.frame_count, None);
649
    self.frame_count += 1;
Luca Barbato's avatar
Luca Barbato committed
650
  }
Josh Holmer's avatar
Josh Holmer committed
651

Josh Holmer's avatar
Josh Holmer committed
652 653
  fn determine_frame_type(&mut self, frame_number: u64) -> FrameType {
    if frame_number == 0 {
Josh Holmer's avatar
Josh Holmer committed
654 655 656
      return FrameType::KEY;
    }

Josh Holmer's avatar
Josh Holmer committed
657 658 659 660 661 662 663 664
    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
665
    if let Some(frame) = frame {
Josh Holmer's avatar
Josh Holmer committed
666 667 668
      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
669 670
          // 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
671
          self.keyframe_detector.detect_scene_change(frame, frame_number as usize);
Josh Holmer's avatar
Josh Holmer committed
672 673 674
        }
        return FrameType::INTER;
      }
Josh Holmer's avatar
Josh Holmer committed
675
      if distance >= self.config.enc.max_key_frame_interval {
Josh Holmer's avatar
Josh Holmer committed
676 677
        return FrameType::KEY;
      }
Josh Holmer's avatar
Josh Holmer committed
678
      if self.keyframe_detector.detect_scene_change(frame, frame_number as usize) {
Josh Holmer's avatar
Josh Holmer committed
679 680 681 682 683
        return FrameType::KEY;
      }
    }
    FrameType::INTER
  }
Luca Barbato's avatar
Luca Barbato committed
684
}