api.rs 19.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::*;
Raphaël Zumer's avatar
Raphaël Zumer committed
11 12 13
use crate::encoder::*;
use crate::metrics::calculate_frame_psnr;
use crate::partition::*;
14
use crate::rate::RCState;
Raphaël Zumer's avatar
Raphaël Zumer committed
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 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
    }
  }
}

#[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 {
168
    speed >= 4
169 170 171 172 173 174 175 176 177 178 179
  }

  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 {
180
    if speed <= 1 {
181
      PredictionModesSetting::ComplexAll
182
    } else if speed <= 3 {
183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200
      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,
}

201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216
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
    }
}

217 218
arg_enum!{
  #[derive(Debug, Clone, Copy, PartialEq)]
219
  #[repr(C)]
220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235
  pub enum MatrixCoefficients {
      Identity = 0,
      BT709,
      Unspecified,
      BT470M = 4,
      BT470BG,
      ST170M,
      ST240M,
      YCgCo,
      BT2020NonConstantLuminance,
      BT2020ConstantLuminance,
      ST2085,
      ChromaticityDerivedNonConstantLuminance,
      ChromaticityDerivedConstantLuminance,
      ICtCp,
  }
236 237 238 239 240 241 242 243
}

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

244
arg_enum!{
245
  #[derive(Debug, Clone, Copy, PartialEq)]
246
  #[repr(C)]
247 248 249 250 251 252 253 254 255 256 257 258 259 260
  pub enum ColorPrimaries {
      BT709 = 1,
      Unspecified,
      BT470M = 4,
      BT470BG,
      ST170M,
      ST240M,
      Film,
      BT2020,
      ST428,
      P3DCI,
      P3Display,
      Tech3213 = 22,
  }
261 262 263 264 265 266 267 268
}

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

269
arg_enum!{
270
  #[derive(Debug, Clone, Copy, PartialEq)]
271
  #[repr(C)]
272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290
  pub enum TransferCharacteristics {
      BT1886 = 1,
      Unspecified,
      BT470M = 4,
      BT470BG,
      ST170M,
      ST240M,
      Linear,
      Logarithmic100,
      Logarithmic316,
      XVYCC,
      BT1361E,
      SRGB,
      BT2020Ten,
      BT2020Twelve,
      PerceptualQuantizer,
      ST428,
      HybridLogGamma,
  }
291 292 293 294 295 296 297 298
}

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

299 300 301 302 303 304 305
#[derive(Copy, Clone, Debug)]
pub struct ColorDescription {
    pub color_primaries: ColorPrimaries,
    pub transfer_characteristics: TransferCharacteristics,
    pub matrix_coefficients: MatrixCoefficients
}

306 307 308 309 310 311 312 313 314 315 316 317 318 319
#[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,
}

320
/// Contain all the encoder configuration
321
#[derive(Clone, Debug)]
322
pub struct Config {
Luca Barbato's avatar
Luca Barbato committed
323 324 325 326
  pub enc: EncoderConfig
}

impl Config {
Luca Barbato's avatar
Luca Barbato committed
327 328
  pub fn parse(&mut self, key: &str, value: &str) -> Result<(), EncoderStatus> {
    match key {
329
      "low_latency" => self.enc.low_latency = value.parse().map_err(|_e| ParseError)?,
Josh Holmer's avatar
Josh Holmer committed
330 331
      "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)?,
332
      "quantizer" => self.enc.quantizer = value.parse().map_err(|_e| ParseError)?,
333
      "speed" => self.enc.speed_settings = SpeedSettings::from_preset(value.parse().map_err(|_e| ParseError)?),
334 335
      "tune" => self.enc.tune = value.parse().map_err(|_e| ParseError)?,
      _ => return Err(InvalidKey)
Luca Barbato's avatar
Luca Barbato committed
336 337 338 339 340
    }

    Ok(())
  }

Luca Barbato's avatar
Luca Barbato committed
341
  pub fn new_context(&self) -> Context {
342
    #[cfg(feature = "aom")]
Luca Barbato's avatar
Luca Barbato committed
343
    unsafe {
Luca Barbato's avatar
Luca Barbato committed
344 345
      av1_rtcd();
      aom_dsp_rtcd();
Luca Barbato's avatar
Luca Barbato committed
346 347
    }

348 349 350
    // initialize with temporal delimiter
    let packet_data = TEMPORAL_DELIMITER.to_vec();

Josh Holmer's avatar
Josh Holmer committed
351 352
    Context {
      frame_count: 0,
353
      limit: 0,
Josh Holmer's avatar
Josh Holmer committed
354
      idx: 0,
Josh Holmer's avatar
Josh Holmer committed
355
      frames_processed: 0,
Josh Holmer's avatar
Josh Holmer committed
356
      frame_q: BTreeMap::new(),
Josh Holmer's avatar
Josh Holmer committed
357 358
      frame_data: BTreeMap::new(),
      keyframes: BTreeSet::new(),
359
      packet_data,
Josh Holmer's avatar
Josh Holmer committed
360 361
      segment_start_idx: 0,
      segment_start_frame: 0,
362
      keyframe_detector: SceneChangeDetector::new(self.enc.bit_depth),
363 364 365 366 367
      config: self.clone(),
      rc_state: RCState::new(),
      first_pass_data: FirstPassData {
        frames: Vec::new(),
      },
Josh Holmer's avatar
Josh Holmer committed
368
    }
Luca Barbato's avatar
Luca Barbato committed
369 370 371 372
  }
}

pub struct Context {
Luca Barbato's avatar
Luca Barbato committed
373
  //    timebase: Rational,
fbossen's avatar
fbossen committed
374
  frame_count: u64,
375
  limit: u64,
fbossen's avatar
fbossen committed
376
  idx: u64,
Josh Holmer's avatar
Josh Holmer committed
377 378
  frames_processed: u64,
  /// Maps frame *number* to frames
379
  frame_q: BTreeMap<u64, Option<Arc<Frame>>>, //    packet_q: VecDeque<Packet>
Josh Holmer's avatar
Josh Holmer committed
380 381 382 383 384
  /// 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>,
385
  /// A storage space for reordered frames.
Josh Holmer's avatar
Josh Holmer committed
386 387 388 389
  packet_data: Vec<u8>,
  segment_start_idx: u64,
  segment_start_frame: u64,
  keyframe_detector: SceneChangeDetector,
390
  pub config: Config,
391 392
  rc_state: RCState,
  pub first_pass_data: FirstPassData,
Luca Barbato's avatar
Luca Barbato committed
393 394 395 396 397 398 399 400 401
}

#[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
402 403 404
  Failure,
  InvalidKey,
  ParseError
Luca Barbato's avatar
Luca Barbato committed
405 406 407 408
}

pub struct Packet {
  pub data: Vec<u8>,
fbossen's avatar
fbossen committed
409 410
  pub rec: Option<Frame>,
  pub number: u64,
411 412 413
  pub frame_type: FrameType,
  /// PSNR for Y, U, and V planes
  pub psnr: Option<(f64, f64, f64)>,
414 415 416 417
}

impl fmt::Display for Packet {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
Luca Barbato's avatar
Luca Barbato committed
418 419 420 421 422 423 424
    write!(
      f,
      "Frame {} - {} - {} bytes",
      self.number,
      self.frame_type,
      self.data.len()
    )
425
  }
Luca Barbato's avatar
Luca Barbato committed
426 427 428 429
}

impl Context {
  pub fn new_frame(&self) -> Arc<Frame> {
Josh Holmer's avatar
Josh Holmer committed
430
    Arc::new(Frame::new(
431 432
      self.config.enc.width,
      self.config.enc.height,
433
      self.config.enc.chroma_sampling
Josh Holmer's avatar
Josh Holmer committed
434
    ))
Luca Barbato's avatar
Luca Barbato committed
435 436 437 438 439 440
  }

  pub fn send_frame<F>(&mut self, frame: F) -> Result<(), EncoderStatus>
  where
    F: Into<Option<Arc<Frame>>>
  {
Josh Holmer's avatar
Josh Holmer committed
441 442
    let idx = self.frame_count;
    self.frame_q.insert(idx, frame.into());
443
    self.frame_count += 1;
Luca Barbato's avatar
Luca Barbato committed
444 445 446
    Ok(())
  }

447 448 449 450
  pub fn get_frame_count(&self) -> u64 {
    self.frame_count
  }

451 452
  pub fn set_limit(&mut self, limit: u64) {
    self.limit = limit;
453 454
  }

Josh Holmer's avatar
Josh Holmer committed
455 456 457 458
  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)
  }

459
  pub fn needs_more_frames(&self, frame_count: u64) -> bool {
460
    self.limit == 0 || frame_count < self.limit
461 462
  }

463 464 465
  pub fn container_sequence_header(&mut self) -> Vec<u8> {
    fn sequence_header_inner(seq: &Sequence) -> io::Result<Vec<u8>> {
      let mut buf = Vec::new();
466

467
      {
468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484
        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
485
      }
486

487 488 489
      Ok(buf)
    }

Josh Holmer's avatar
Josh Holmer committed
490
    sequence_header_inner(&self.frame_data[&0].sequence).unwrap()
491 492
  }

Josh Holmer's avatar
Josh Holmer committed
493
  fn next_keyframe(&self) -> u64 {
Josh Holmer's avatar
Josh Holmer committed
494 495 496 497
    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
498 499 500 501 502
    if next_detected.is_none() {
      return next_limit;
    }
    cmp::min(next_detected.unwrap(), next_limit)
  }
fbossen's avatar
fbossen committed
503

504
  fn set_frame_properties(&mut self, idx: u64) -> bool {
505
    let (fi, end_of_subgop) = self.build_frame_properties(idx);
506 507 508
    self.frame_data.insert(idx, fi);

    end_of_subgop
Josh Holmer's avatar
Josh Holmer committed
509 510
  }

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

Josh Holmer's avatar
Josh Holmer committed
515
      // The first frame will always be a key frame
Josh Holmer's avatar
Josh Holmer committed
516 517
      let fi = FrameInvariants::new_key_frame(
        &FrameInvariants::new(
518
          self.config.enc.clone(),
519
          seq
Josh Holmer's avatar
Josh Holmer committed
520 521 522
        ),
        0
      );
523
      return (fi, true);
Josh Holmer's avatar
Josh Holmer committed
524 525
    }

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

528 529 530 531 532
    // 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
533 534 535 536 537 538
    // 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 {
539
      let next_keyframe = if keyframe_only { self.segment_start_frame + 1 } else { self.next_keyframe() };
540
      let (fi_temp, end_of_subgop) = FrameInvariants::new_inter_frame(
Josh Holmer's avatar
Josh Holmer committed
541
        &fi,
542 543 544 545
        self.segment_start_frame,
        idx_in_segment,
        next_keyframe
      );
Josh Holmer's avatar
Josh Holmer committed
546
      fi = fi_temp;
547
      if !end_of_subgop {
Josh Holmer's avatar
Josh Holmer committed
548 549 550
        if !fi.inter_cfg.unwrap().reorder
          || ((idx_in_segment - 1) % fi.inter_cfg.unwrap().group_len == 0
          && fi.number == (next_keyframe - 1))
551
        {
Josh Holmer's avatar
Josh Holmer committed
552
          self.segment_start_idx = idx;
Josh Holmer's avatar
Josh Holmer committed
553
          self.segment_start_frame = next_keyframe;
Josh Holmer's avatar
Josh Holmer committed
554
          fi.number = next_keyframe;
Josh Holmer's avatar
Josh Holmer committed
555
        } else {
556
          return (fi, false);
Josh Holmer's avatar
Josh Holmer committed
557 558
        }
      }
fbossen's avatar
fbossen committed
559
    }
Luca Barbato's avatar
Luca Barbato committed
560

561 562
    match self.frame_q.get(&fi.number) {
      Some(Some(_)) => {},
Josh Holmer's avatar
Josh Holmer committed
563
      _ => { return (fi, false); }
564 565
    }

Josh Holmer's avatar
Josh Holmer committed
566
    // Now that we know the frame number, look up the correct frame type
Josh Holmer's avatar
Josh Holmer committed
567 568 569 570 571 572 573
    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
574

Josh Holmer's avatar
Josh Holmer committed
575 576 577 578 579
    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();
580
      let (fi_temp, end_of_subgop) = FrameInvariants::new_inter_frame(
Josh Holmer's avatar
Josh Holmer committed
581 582 583 584 585 586
        &fi,
        self.segment_start_frame,
        idx_in_segment,
        next_keyframe
      );
      fi = fi_temp;
587 588
      if !end_of_subgop {
        return (fi, false);
Josh Holmer's avatar
Josh Holmer committed
589 590
      }
    }
591
    (fi, true)
fbossen's avatar
fbossen committed
592
  }
Luca Barbato's avatar
Luca Barbato committed
593

fbossen's avatar
fbossen committed
594
  pub fn receive_packet(&mut self) -> Result<Packet, EncoderStatus> {
595 596 597 598 599 600
    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
601

602 603 604 605 606 607
      if !self.needs_more_frames(self.frame_data.get(&idx).unwrap().number) {
        self.idx += 1;
        return Err(EncoderStatus::EnoughData);
      }
      idx
    };
608

609
    let ret = {
Luca Barbato's avatar
Luca Barbato committed
610 611 612
      let fi = self.frame_data.get_mut(&idx).unwrap();
      if fi.show_existing_frame {
        self.idx += 1;
Luca Barbato's avatar
Luca Barbato committed
613

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

616 617
        let sef_data = encode_frame(fi, &mut fs);
        self.packet_data.extend(sef_data);
Luca Barbato's avatar
Luca Barbato committed
618

Luca Barbato's avatar
Luca Barbato committed
619
        let rec = if fi.show_frame { Some(fs.rec) } else { None };
620 621
        let fi = fi.clone();
        self.finalize_packet(&*fs.input, rec, &fi)
Luca Barbato's avatar
Luca Barbato committed
622 623 624
      } else {
        if let Some(f) = self.frame_q.get(&fi.number) {
          self.idx += 1;
fbossen's avatar
fbossen committed
625

626
          if let Some(frame) = f.clone() {
627 628 629
            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
630
            let mut fs = FrameState::new_with_frame(fi, frame.clone());
fbossen's avatar
fbossen committed
631

Luca Barbato's avatar
Luca Barbato committed
632 633
            let data = encode_frame(fi, &mut fs);
            self.packet_data.extend(data);
fbossen's avatar
fbossen committed
634

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

Luca Barbato's avatar
Luca Barbato committed
637 638
            // 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
639

Luca Barbato's avatar
Luca Barbato committed
640
            update_rec_buffer(fi, fs);
fbossen's avatar
fbossen committed
641

Luca Barbato's avatar
Luca Barbato committed
642
            if fi.show_frame {
643 644
              let fi = fi.clone();
              self.finalize_packet(&*frame, rec, &fi)
Luca Barbato's avatar
Luca Barbato committed
645 646 647
            } else {
              Err(EncoderStatus::NeedMoreData)
            }
648 649 650
          } else {
            Err(EncoderStatus::NeedMoreData)
          }
fbossen's avatar
fbossen committed
651 652 653 654
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      }
655 656 657
    };

    if let Ok(ref pkt) = ret {
Luca Barbato's avatar
Luca Barbato committed
658
      self.garbage_collect(pkt.number);
659 660 661
    }

    ret
Luca Barbato's avatar
Luca Barbato committed
662 663
  }

664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681
  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
        ));
      }
    }

682 683 684 685
    if self.config.enc.pass == Some(1) {
      self.first_pass_data.frames.push(FirstPassFrame::from(fi));
    }

686 687 688 689 690 691 692 693 694 695
    self.frames_processed += 1;
    Ok(Packet {
      data,
      rec,
      number: fi.number,
      frame_type: fi.frame_type,
      psnr
    })
  }

696
  fn garbage_collect(&mut self, cur_frame: u64) {
Josh Holmer's avatar
Josh Holmer committed
697 698 699 700 701 702 703 704 705 706 707 708 709 710
    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
711
  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
712
    self.frame_q.insert(self.frame_count, None);
713
    self.frame_count += 1;
Luca Barbato's avatar
Luca Barbato committed
714
  }
Josh Holmer's avatar
Josh Holmer committed
715

Josh Holmer's avatar
Josh Holmer committed
716 717
  fn determine_frame_type(&mut self, frame_number: u64) -> FrameType {
    if frame_number == 0 {
Josh Holmer's avatar
Josh Holmer committed
718 719 720
      return FrameType::KEY;
    }

Josh Holmer's avatar
Josh Holmer committed
721 722 723 724 725 726 727 728
    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
729
    if let Some(frame) = frame {
Josh Holmer's avatar
Josh Holmer committed
730 731 732
      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
733 734
          // 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
735
          self.keyframe_detector.detect_scene_change(frame, frame_number as usize);
Josh Holmer's avatar
Josh Holmer committed
736 737 738
        }
        return FrameType::INTER;
      }
Josh Holmer's avatar
Josh Holmer committed
739
      if distance >= self.config.enc.max_key_frame_interval {
Josh Holmer's avatar
Josh Holmer committed
740 741
        return FrameType::KEY;
      }
Josh Holmer's avatar
Josh Holmer committed
742
      if self.keyframe_detector.detect_scene_change(frame, frame_number as usize) {
Josh Holmer's avatar
Josh Holmer committed
743 744 745 746 747
        return FrameType::KEY;
      }
    }
    FrameType::INTER
  }
Luca Barbato's avatar
Luca Barbato committed
748
}
749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768

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