api.rs 18.6 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

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

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

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

46 47
#[derive(Copy, Clone, Debug)]
pub struct EncoderConfig {
48 49 50 51 52 53
  pub width: usize,
  pub height: usize,
  pub bit_depth: usize,
  pub chroma_sampling: ChromaSampling,
  pub chroma_sample_position: ChromaSamplePosition,
  pub time_base: Rational,
Josh Holmer's avatar
Josh Holmer committed
54 55 56 57
  /// The *minimum* interval between two keyframes
  pub min_key_frame_interval: u64,
  /// The *maximum* interval between two keyframes
  pub max_key_frame_interval: u64,
58
  pub low_latency: bool,
59
  pub quantizer: usize,
60
  pub tune: Tune,
61
  pub pixel_range: PixelRange,
62
  pub color_description: Option<ColorDescription>,
63 64
  pub mastering_display: Option<MasteringDisplay>,
  pub content_light: Option<ContentLight>,
65
  pub speed_settings: SpeedSettings,
66
  pub show_psnr: bool,
67 68 69 70
}

impl Default for EncoderConfig {
  fn default() -> Self {
71 72
    const DEFAULT_SPEED: usize = 3;
    Self::with_speed_preset(DEFAULT_SPEED)
73 74 75
  }
}

76 77 78
impl EncoderConfig {
  pub fn with_speed_preset(speed: usize) -> Self {
    EncoderConfig {
79 80 81 82 83 84
      width: 640,
      height: 480,
      bit_depth: 8,
      chroma_sampling: ChromaSampling::Cs420,
      chroma_sample_position: ChromaSamplePosition::Unknown,
      time_base: Rational { num: 30, den: 1 },
Josh Holmer's avatar
Josh Holmer committed
85 86
      min_key_frame_interval: 12,
      max_key_frame_interval: 240,
87
      low_latency: false,
88 89
      quantizer: 100,
      tune: Tune::Psnr,
90
      pixel_range: PixelRange::Unspecified,
91
      color_description: None,
92 93
      mastering_display: None,
      content_light: None,
94 95
      speed_settings: SpeedSettings::from_preset(speed),
      show_psnr: false,
96 97 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
    }
  }
}

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

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

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

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

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

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

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

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

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

293 294 295 296 297 298 299 300 301 302 303 304 305 306
#[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,
}

307 308 309
/// Contain all the encoder configuration
#[derive(Clone, Copy, Debug)]
pub struct Config {
Luca Barbato's avatar
Luca Barbato committed
310 311 312 313
  pub enc: EncoderConfig
}

impl Config {
Luca Barbato's avatar
Luca Barbato committed
314 315
  pub fn parse(&mut self, key: &str, value: &str) -> Result<(), EncoderStatus> {
    match key {
316
      "low_latency" => self.enc.low_latency = value.parse().map_err(|_e| ParseError)?,
Josh Holmer's avatar
Josh Holmer committed
317 318
      "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)?,
319
      "quantizer" => self.enc.quantizer = value.parse().map_err(|_e| ParseError)?,
320
      "speed" => self.enc.speed_settings = SpeedSettings::from_preset(value.parse().map_err(|_e| ParseError)?),
321 322
      "tune" => self.enc.tune = value.parse().map_err(|_e| ParseError)?,
      _ => return Err(InvalidKey)
Luca Barbato's avatar
Luca Barbato committed
323 324 325 326 327
    }

    Ok(())
  }

Luca Barbato's avatar
Luca Barbato committed
328
  pub fn new_context(&self) -> Context {
329
    #[cfg(feature = "aom")]
Luca Barbato's avatar
Luca Barbato committed
330
    unsafe {
Luca Barbato's avatar
Luca Barbato committed
331 332
      av1_rtcd();
      aom_dsp_rtcd();
Luca Barbato's avatar
Luca Barbato committed
333 334
    }

Josh Holmer's avatar
Josh Holmer committed
335 336
    Context {
      frame_count: 0,
337
      limit: 0,
Josh Holmer's avatar
Josh Holmer committed
338
      idx: 0,
Josh Holmer's avatar
Josh Holmer committed
339
      frames_processed: 0,
Josh Holmer's avatar
Josh Holmer committed
340
      frame_q: BTreeMap::new(),
Josh Holmer's avatar
Josh Holmer committed
341 342
      frame_data: BTreeMap::new(),
      keyframes: BTreeSet::new(),
Josh Holmer's avatar
Josh Holmer committed
343 344 345
      packet_data: Vec::new(),
      segment_start_idx: 0,
      segment_start_frame: 0,
346
      keyframe_detector: SceneChangeDetector::new(self.enc.bit_depth),
Josh Holmer's avatar
Josh Holmer committed
347
      config: *self,
Josh Holmer's avatar
Josh Holmer committed
348
    }
Luca Barbato's avatar
Luca Barbato committed
349 350 351 352
  }
}

pub struct Context {
Luca Barbato's avatar
Luca Barbato committed
353
  //    timebase: Rational,
fbossen's avatar
fbossen committed
354
  frame_count: u64,
355
  limit: u64,
fbossen's avatar
fbossen committed
356
  idx: u64,
Josh Holmer's avatar
Josh Holmer committed
357 358
  frames_processed: u64,
  /// Maps frame *number* to frames
359
  frame_q: BTreeMap<u64, Option<Arc<Frame>>>, //    packet_q: VecDeque<Packet>
Josh Holmer's avatar
Josh Holmer committed
360 361 362 363 364
  /// 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>,
365
  /// A storage space for reordered frames.
Josh Holmer's avatar
Josh Holmer committed
366 367 368 369
  packet_data: Vec<u8>,
  segment_start_idx: u64,
  segment_start_frame: u64,
  keyframe_detector: SceneChangeDetector,
370
  pub config: Config,
Luca Barbato's avatar
Luca Barbato committed
371 372 373 374 375 376 377 378 379
}

#[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
380 381 382
  Failure,
  InvalidKey,
  ParseError
Luca Barbato's avatar
Luca Barbato committed
383 384 385 386
}

pub struct Packet {
  pub data: Vec<u8>,
fbossen's avatar
fbossen committed
387 388
  pub rec: Option<Frame>,
  pub number: u64,
389 390 391
  pub frame_type: FrameType,
  /// PSNR for Y, U, and V planes
  pub psnr: Option<(f64, f64, f64)>,
392 393 394 395
}

impl fmt::Display for Packet {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
Luca Barbato's avatar
Luca Barbato committed
396 397 398 399 400 401 402
    write!(
      f,
      "Frame {} - {} - {} bytes",
      self.number,
      self.frame_type,
      self.data.len()
    )
403
  }
Luca Barbato's avatar
Luca Barbato committed
404 405 406 407
}

impl Context {
  pub fn new_frame(&self) -> Arc<Frame> {
Josh Holmer's avatar
Josh Holmer committed
408
    Arc::new(Frame::new(
409 410 411
      self.config.enc.width.align_power_of_two(3),
      self.config.enc.height.align_power_of_two(3),
      self.config.enc.chroma_sampling
Josh Holmer's avatar
Josh Holmer committed
412
    ))
Luca Barbato's avatar
Luca Barbato committed
413 414 415 416 417 418
  }

  pub fn send_frame<F>(&mut self, frame: F) -> Result<(), EncoderStatus>
  where
    F: Into<Option<Arc<Frame>>>
  {
Josh Holmer's avatar
Josh Holmer committed
419 420
    let idx = self.frame_count;
    self.frame_q.insert(idx, frame.into());
421
    self.frame_count += 1;
Luca Barbato's avatar
Luca Barbato committed
422 423 424
    Ok(())
  }

425 426 427 428
  pub fn get_frame_count(&self) -> u64 {
    self.frame_count
  }

429 430
  pub fn set_limit(&mut self, limit: u64) {
    self.limit = limit;
431 432
  }

Josh Holmer's avatar
Josh Holmer committed
433 434 435 436
  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)
  }

437
  pub fn needs_more_frames(&self, frame_count: u64) -> bool {
438
    self.limit == 0 || frame_count < self.limit
439 440
  }

441 442 443
  pub fn container_sequence_header(&mut self) -> Vec<u8> {
    fn sequence_header_inner(seq: &Sequence) -> io::Result<Vec<u8>> {
      let mut buf = Vec::new();
444

445
      {
446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462
        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
463
      }
464

465 466 467
      Ok(buf)
    }

Josh Holmer's avatar
Josh Holmer committed
468
    sequence_header_inner(&self.frame_data[&0].sequence).unwrap()
469 470
  }

Josh Holmer's avatar
Josh Holmer committed
471
  fn next_keyframe(&self) -> u64 {
Josh Holmer's avatar
Josh Holmer committed
472 473 474 475
    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
476 477 478 479 480
    if next_detected.is_none() {
      return next_limit;
    }
    cmp::min(next_detected.unwrap(), next_limit)
  }
fbossen's avatar
fbossen committed
481

482
  fn set_frame_properties(&mut self, idx: u64) -> bool {
483
    let (fi, end_of_subgop) = self.build_frame_properties(idx);
484 485 486
    self.frame_data.insert(idx, fi);

    end_of_subgop
Josh Holmer's avatar
Josh Holmer committed
487 488
  }

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

Josh Holmer's avatar
Josh Holmer committed
493
      // The first frame will always be a key frame
Josh Holmer's avatar
Josh Holmer committed
494 495 496
      let fi = FrameInvariants::new_key_frame(
        &FrameInvariants::new(
          self.config.enc,
497
          seq
Josh Holmer's avatar
Josh Holmer committed
498 499 500
        ),
        0
      );
501
      return (fi, true);
Josh Holmer's avatar
Josh Holmer committed
502 503
    }

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

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

534 535
    match self.frame_q.get(&fi.number) {
      Some(Some(_)) => {},
Josh Holmer's avatar
Josh Holmer committed
536
      _ => { return (fi, false); }
537 538
    }

Josh Holmer's avatar
Josh Holmer committed
539
    // Now that we know the frame number, look up the correct frame type
Josh Holmer's avatar
Josh Holmer committed
540 541 542 543 544 545 546
    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
547

Josh Holmer's avatar
Josh Holmer committed
548 549 550 551 552
    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();
553
      let (fi_temp, end_of_subgop) = FrameInvariants::new_inter_frame(
Josh Holmer's avatar
Josh Holmer committed
554 555 556 557 558 559
        &fi,
        self.segment_start_frame,
        idx_in_segment,
        next_keyframe
      );
      fi = fi_temp;
560 561
      if !end_of_subgop {
        return (fi, false);
Josh Holmer's avatar
Josh Holmer committed
562 563
      }
    }
564
    (fi, true)
fbossen's avatar
fbossen committed
565
  }
Luca Barbato's avatar
Luca Barbato committed
566

fbossen's avatar
fbossen committed
567
  pub fn receive_packet(&mut self) -> Result<Packet, EncoderStatus> {
568 569 570 571 572 573
    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
574

575 576 577 578 579 580
      if !self.needs_more_frames(self.frame_data.get(&idx).unwrap().number) {
        self.idx += 1;
        return Err(EncoderStatus::EnoughData);
      }
      idx
    };
581

582
    let ret = {
Luca Barbato's avatar
Luca Barbato committed
583 584 585
      let fi = self.frame_data.get_mut(&idx).unwrap();
      if fi.show_existing_frame {
        self.idx += 1;
Luca Barbato's avatar
Luca Barbato committed
586

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

Luca Barbato's avatar
Luca Barbato committed
589
        let data = encode_frame(fi, &mut fs);
Luca Barbato's avatar
Luca Barbato committed
590

Luca Barbato's avatar
Luca Barbato committed
591 592 593 594 595 596 597 598 599 600
        let rec = if fi.show_frame { Some(fs.rec) } else { None };
        let mut psnr = None;
        if self.config.enc.show_psnr {
          if let Some(ref rec) = rec {
            psnr = Some(calculate_frame_psnr(
              &*fs.input,
              rec,
              fi.sequence.bit_depth
            ));
          }
601
        }
Luca Barbato's avatar
Luca Barbato committed
602

Luca Barbato's avatar
Luca Barbato committed
603 604 605 606 607 608 609 610 611 612 613
        self.frames_processed += 1;
        Ok(Packet {
          data,
          rec,
          number: fi.number,
          frame_type: fi.frame_type,
          psnr
        })
      } else {
        if let Some(f) = self.frame_q.get(&fi.number) {
          self.idx += 1;
fbossen's avatar
fbossen committed
614

Luca Barbato's avatar
Luca Barbato committed
615 616
          if let Some(frame) = f {
            let mut fs = FrameState::new_with_frame(fi, frame.clone());
fbossen's avatar
fbossen committed
617

Luca Barbato's avatar
Luca Barbato committed
618 619
            let data = encode_frame(fi, &mut fs);
            self.packet_data.extend(data);
fbossen's avatar
fbossen committed
620

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

Luca Barbato's avatar
Luca Barbato committed
623 624
            // 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
625

Luca Barbato's avatar
Luca Barbato committed
626
            update_rec_buffer(fi, fs);
fbossen's avatar
fbossen committed
627

Luca Barbato's avatar
Luca Barbato committed
628 629 630
            if fi.show_frame {
              let data = self.packet_data.clone();
              self.packet_data.clear();
631

Luca Barbato's avatar
Luca Barbato committed
632 633 634 635 636 637 638 639 640
              let mut psnr = None;
              if self.config.enc.show_psnr {
                if let Some(ref rec) = rec {
                  psnr = Some(calculate_frame_psnr(
                    &*frame,
                    rec,
                    fi.sequence.bit_depth
                  ));
                }
641 642
              }

Luca Barbato's avatar
Luca Barbato committed
643 644 645 646 647 648 649 650 651 652 653
              self.frames_processed += 1;
              Ok(Packet {
                data,
                rec,
                number: fi.number,
                frame_type: fi.frame_type,
                psnr
              })
            } else {
              Err(EncoderStatus::NeedMoreData)
            }
654 655 656
          } else {
            Err(EncoderStatus::NeedMoreData)
          }
fbossen's avatar
fbossen committed
657 658 659 660
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      }
661 662 663
    };

    if let Ok(ref pkt) = ret {
Luca Barbato's avatar
Luca Barbato committed
664
      self.garbage_collect(pkt.number);
665 666 667
    }

    ret
Luca Barbato's avatar
Luca Barbato committed
668 669
  }

670
  fn garbage_collect(&mut self, cur_frame: u64) {
Josh Holmer's avatar
Josh Holmer committed
671 672 673 674 675 676 677 678 679 680 681 682 683 684
    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
685
  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
686
    self.frame_q.insert(self.frame_count, None);
687
    self.frame_count += 1;
Luca Barbato's avatar
Luca Barbato committed
688
  }
Josh Holmer's avatar
Josh Holmer committed
689

Josh Holmer's avatar
Josh Holmer committed
690 691
  fn determine_frame_type(&mut self, frame_number: u64) -> FrameType {
    if frame_number == 0 {
Josh Holmer's avatar
Josh Holmer committed
692 693 694
      return FrameType::KEY;
    }

Josh Holmer's avatar
Josh Holmer committed
695 696 697 698 699 700 701 702
    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
703
    if let Some(frame) = frame {
Josh Holmer's avatar
Josh Holmer committed
704 705 706
      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
707 708
          // 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
709
          self.keyframe_detector.detect_scene_change(frame, frame_number as usize);
Josh Holmer's avatar
Josh Holmer committed
710 711 712
        }
        return FrameType::INTER;
      }
Josh Holmer's avatar
Josh Holmer committed
713
      if distance >= self.config.enc.max_key_frame_interval {
Josh Holmer's avatar
Josh Holmer committed
714 715
        return FrameType::KEY;
      }
Josh Holmer's avatar
Josh Holmer committed
716
      if self.keyframe_detector.detect_scene_change(frame, frame_number as usize) {
Josh Holmer's avatar
Josh Holmer committed
717 718 719 720 721
        return FrameType::KEY;
      }
    }
    FrameType::INTER
  }
Luca Barbato's avatar
Luca Barbato committed
722
}