api.rs 19.1 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
#[derive(Debug, Clone, Copy)]
pub struct VideoDetails {
  pub width: usize,
  pub height: usize,
  pub bit_depth: usize,
  pub chroma_sampling: ChromaSampling,
  pub chroma_sample_position: ChromaSamplePosition,
Luca Barbato's avatar
Luca Barbato committed
32
  pub time_base: Rational,
33 34 35 36 37 38 39 40 41 42
}

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

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

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

62 63 64 65 66 67 68
#[derive(Clone, Copy, Debug)]
#[repr(C)]
pub struct Point {
  pub x: u16,
  pub y: u16
}

69 70
#[derive(Copy, Clone, Debug)]
pub struct EncoderConfig {
Josh Holmer's avatar
Josh Holmer committed
71 72 73 74
  /// The *minimum* interval between two keyframes
  pub min_key_frame_interval: u64,
  /// The *maximum* interval between two keyframes
  pub max_key_frame_interval: u64,
75
  pub low_latency: bool,
76
  pub quantizer: usize,
77
  pub tune: Tune,
78
  pub pixel_range: PixelRange,
79
  pub color_description: Option<ColorDescription>,
80 81
  pub mastering_display: Option<MasteringDisplay>,
  pub content_light: Option<ContentLight>,
82
  pub speed_settings: SpeedSettings,
83
  pub show_psnr: bool,
84 85 86 87
}

impl Default for EncoderConfig {
  fn default() -> Self {
88 89
    const DEFAULT_SPEED: usize = 3;
    Self::with_speed_preset(DEFAULT_SPEED)
90 91 92
  }
}

93 94 95
impl EncoderConfig {
  pub fn with_speed_preset(speed: usize) -> Self {
    EncoderConfig {
Josh Holmer's avatar
Josh Holmer committed
96 97
      min_key_frame_interval: 12,
      max_key_frame_interval: 240,
98
      low_latency: false,
99 100
      quantizer: 100,
      tune: Tune::Psnr,
101
      pixel_range: PixelRange::Unspecified,
102
      color_description: None,
103 104
      mastering_display: None,
      content_light: None,
105 106
      speed_settings: SpeedSettings::from_preset(speed),
      show_psnr: false,
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 169 170 171 172 173 174 175 176 177
    }
  }
}

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

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

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

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

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

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

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

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

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

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

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

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

    Ok(())
  }

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

Josh Holmer's avatar
Josh Holmer committed
347 348
    Context {
      frame_count: 0,
349
      limit: 0,
Josh Holmer's avatar
Josh Holmer committed
350
      idx: 0,
Josh Holmer's avatar
Josh Holmer committed
351
      frames_processed: 0,
Josh Holmer's avatar
Josh Holmer committed
352
      frame_q: BTreeMap::new(),
Josh Holmer's avatar
Josh Holmer committed
353 354
      frame_data: BTreeMap::new(),
      keyframes: BTreeSet::new(),
Josh Holmer's avatar
Josh Holmer committed
355 356 357
      packet_data: Vec::new(),
      segment_start_idx: 0,
      segment_start_frame: 0,
358
      keyframe_detector: SceneChangeDetector::new(&self.video_info),
Josh Holmer's avatar
Josh Holmer committed
359
      config: *self,
Josh Holmer's avatar
Josh Holmer committed
360
    }
Luca Barbato's avatar
Luca Barbato committed
361 362 363 364
  }
}

pub struct Context {
Luca Barbato's avatar
Luca Barbato committed
365
  //    timebase: Rational,
fbossen's avatar
fbossen committed
366
  frame_count: u64,
367
  limit: u64,
fbossen's avatar
fbossen committed
368
  idx: u64,
Josh Holmer's avatar
Josh Holmer committed
369 370
  frames_processed: u64,
  /// Maps frame *number* to frames
371
  frame_q: BTreeMap<u64, Option<Arc<Frame>>>, //    packet_q: VecDeque<Packet>
Josh Holmer's avatar
Josh Holmer committed
372 373 374 375 376
  /// 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>,
377
  /// A storage space for reordered frames.
Josh Holmer's avatar
Josh Holmer committed
378 379 380 381
  packet_data: Vec<u8>,
  segment_start_idx: u64,
  segment_start_frame: u64,
  keyframe_detector: SceneChangeDetector,
382
  pub config: Config,
Luca Barbato's avatar
Luca Barbato committed
383 384 385 386 387 388 389 390 391
}

#[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
392 393 394
  Failure,
  InvalidKey,
  ParseError
Luca Barbato's avatar
Luca Barbato committed
395 396 397 398
}

pub struct Packet {
  pub data: Vec<u8>,
fbossen's avatar
fbossen committed
399 400
  pub rec: Option<Frame>,
  pub number: u64,
401 402 403
  pub frame_type: FrameType,
  /// PSNR for Y, U, and V planes
  pub psnr: Option<(f64, f64, f64)>,
404 405 406 407
}

impl fmt::Display for Packet {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
Luca Barbato's avatar
Luca Barbato committed
408 409 410 411 412 413 414
    write!(
      f,
      "Frame {} - {} - {} bytes",
      self.number,
      self.frame_type,
      self.data.len()
    )
415
  }
Luca Barbato's avatar
Luca Barbato committed
416 417 418 419
}

impl Context {
  pub fn new_frame(&self) -> Arc<Frame> {
Josh Holmer's avatar
Josh Holmer committed
420
    Arc::new(Frame::new(
421 422 423
      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
424
    ))
Luca Barbato's avatar
Luca Barbato committed
425 426 427 428 429 430
  }

  pub fn send_frame<F>(&mut self, frame: F) -> Result<(), EncoderStatus>
  where
    F: Into<Option<Arc<Frame>>>
  {
Josh Holmer's avatar
Josh Holmer committed
431 432
    let idx = self.frame_count;
    self.frame_q.insert(idx, frame.into());
433
    self.frame_count += 1;
Luca Barbato's avatar
Luca Barbato committed
434 435 436
    Ok(())
  }

437 438 439 440
  pub fn get_frame_count(&self) -> u64 {
    self.frame_count
  }

441 442
  pub fn set_limit(&mut self, limit: u64) {
    self.limit = limit;
443 444
  }

Josh Holmer's avatar
Josh Holmer committed
445 446 447 448
  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)
  }

449
  pub fn needs_more_frames(&self, frame_count: u64) -> bool {
450
    self.limit == 0 || frame_count < self.limit
451 452
  }

453 454 455
  pub fn container_sequence_header(&mut self) -> Vec<u8> {
    fn sequence_header_inner(seq: &Sequence) -> io::Result<Vec<u8>> {
      let mut buf = Vec::new();
456

457
      {
458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474
        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
475
      }
476

477 478 479
      Ok(buf)
    }

Josh Holmer's avatar
Josh Holmer committed
480
    sequence_header_inner(&self.frame_data[&0].sequence).unwrap()
481 482
  }

Josh Holmer's avatar
Josh Holmer committed
483
  fn next_keyframe(&self) -> u64 {
Josh Holmer's avatar
Josh Holmer committed
484 485 486 487
    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
488 489 490 491 492
    if next_detected.is_none() {
      return next_limit;
    }
    cmp::min(next_detected.unwrap(), next_limit)
  }
fbossen's avatar
fbossen committed
493

494
  fn set_frame_properties(&mut self, idx: u64) -> bool {
495
    let (fi, end_of_subgop) = self.build_frame_properties(idx);
496 497 498
    self.frame_data.insert(idx, fi);

    end_of_subgop
Josh Holmer's avatar
Josh Holmer committed
499 500
  }

501
  fn build_frame_properties(&mut self, idx: u64) -> (FrameInvariants, bool) {
Josh Holmer's avatar
Josh Holmer committed
502
    if idx == 0 {
503
      let mut seq = Sequence::new(&self.config.video_info);
504
      seq.pixel_range = self.config.enc.pixel_range;
505
      seq.color_description = self.config.enc.color_description;
506 507
      seq.mastering_display = self.config.enc.mastering_display;
      seq.content_light = self.config.enc.content_light;
508

Josh Holmer's avatar
Josh Holmer committed
509
      // The first frame will always be a key frame
Josh Holmer's avatar
Josh Holmer committed
510 511
      let fi = FrameInvariants::new_key_frame(
        &FrameInvariants::new(
512 513
          self.config.video_info.width,
          self.config.video_info.height,
Josh Holmer's avatar
Josh Holmer committed
514
          self.config.enc,
515
          seq
Josh Holmer's avatar
Josh Holmer committed
516 517 518
        ),
        0
      );
519
      return (fi, true);
Josh Holmer's avatar
Josh Holmer committed
520 521
    }

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

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

552 553
    match self.frame_q.get(&fi.number) {
      Some(Some(_)) => {},
Josh Holmer's avatar
Josh Holmer committed
554
      _ => { return (fi, false); }
555 556
    }

Josh Holmer's avatar
Josh Holmer committed
557
    // Now that we know the frame number, look up the correct frame type
Josh Holmer's avatar
Josh Holmer committed
558 559 560 561 562 563 564
    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
565

Josh Holmer's avatar
Josh Holmer committed
566 567 568 569 570
    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();
571
      let (fi_temp, end_of_subgop) = FrameInvariants::new_inter_frame(
Josh Holmer's avatar
Josh Holmer committed
572 573 574 575 576 577
        &fi,
        self.segment_start_frame,
        idx_in_segment,
        next_keyframe
      );
      fi = fi_temp;
578 579
      if !end_of_subgop {
        return (fi, false);
Josh Holmer's avatar
Josh Holmer committed
580 581
      }
    }
582
    (fi, true)
fbossen's avatar
fbossen committed
583
  }
Luca Barbato's avatar
Luca Barbato committed
584

fbossen's avatar
fbossen committed
585
  pub fn receive_packet(&mut self) -> Result<Packet, EncoderStatus> {
586 587 588 589 590 591
    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
592

593 594 595 596 597 598
      if !self.needs_more_frames(self.frame_data.get(&idx).unwrap().number) {
        self.idx += 1;
        return Err(EncoderStatus::EnoughData);
      }
      idx
    };
599

600
    let ret = {
Luca Barbato's avatar
Luca Barbato committed
601 602 603
      let fi = self.frame_data.get_mut(&idx).unwrap();
      if fi.show_existing_frame {
        self.idx += 1;
Luca Barbato's avatar
Luca Barbato committed
604

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

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

Luca Barbato's avatar
Luca Barbato committed
609 610 611 612 613 614 615 616 617 618
        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
            ));
          }
619
        }
Luca Barbato's avatar
Luca Barbato committed
620

Luca Barbato's avatar
Luca Barbato committed
621 622 623 624 625 626 627 628 629 630 631
        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
632

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

Luca Barbato's avatar
Luca Barbato committed
636 637
            let data = encode_frame(fi, &mut fs);
            self.packet_data.extend(data);
fbossen's avatar
fbossen committed
638

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

Luca Barbato's avatar
Luca Barbato committed
641 642
            // 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
643

Luca Barbato's avatar
Luca Barbato committed
644
            update_rec_buffer(fi, fs);
fbossen's avatar
fbossen committed
645

Luca Barbato's avatar
Luca Barbato committed
646 647 648
            if fi.show_frame {
              let data = self.packet_data.clone();
              self.packet_data.clear();
649

Luca Barbato's avatar
Luca Barbato committed
650 651 652 653 654 655 656 657 658
              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
                  ));
                }
659 660
              }

Luca Barbato's avatar
Luca Barbato committed
661 662 663 664 665 666 667 668 669 670 671
              self.frames_processed += 1;
              Ok(Packet {
                data,
                rec,
                number: fi.number,
                frame_type: fi.frame_type,
                psnr
              })
            } else {
              Err(EncoderStatus::NeedMoreData)
            }
672 673 674
          } else {
            Err(EncoderStatus::NeedMoreData)
          }
fbossen's avatar
fbossen committed
675 676 677 678
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      }
679 680 681
    };

    if let Ok(ref pkt) = ret {
Luca Barbato's avatar
Luca Barbato committed
682
      self.garbage_collect(pkt.number);
683 684 685
    }

    ret
Luca Barbato's avatar
Luca Barbato committed
686 687
  }

688
  fn garbage_collect(&mut self, cur_frame: u64) {
Josh Holmer's avatar
Josh Holmer committed
689 690 691 692 693 694 695 696 697 698 699 700 701 702
    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
703
  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
704
    self.frame_q.insert(self.frame_count, None);
705
    self.frame_count += 1;
Luca Barbato's avatar
Luca Barbato committed
706
  }
Josh Holmer's avatar
Josh Holmer committed
707

Josh Holmer's avatar
Josh Holmer committed
708 709
  fn determine_frame_type(&mut self, frame_number: u64) -> FrameType {
    if frame_number == 0 {
Josh Holmer's avatar
Josh Holmer committed
710 711 712
      return FrameType::KEY;
    }

Josh Holmer's avatar
Josh Holmer committed
713 714 715 716 717 718 719 720
    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
721
    if let Some(frame) = frame {
Josh Holmer's avatar
Josh Holmer committed
722 723 724
      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
725 726
          // 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
727
          self.keyframe_detector.detect_scene_change(frame, frame_number as usize);
Josh Holmer's avatar
Josh Holmer committed
728 729 730
        }
        return FrameType::INTER;
      }
Josh Holmer's avatar
Josh Holmer committed
731
      if distance >= self.config.enc.max_key_frame_interval {
Josh Holmer's avatar
Josh Holmer committed
732 733
        return FrameType::KEY;
      }
Josh Holmer's avatar
Josh Holmer committed
734
      if self.keyframe_detector.detect_scene_change(frame, frame_number as usize) {
Josh Holmer's avatar
Josh Holmer committed
735 736 737 738 739
        return FrameType::KEY;
      }
    }
    FrameType::INTER
  }
Luca Barbato's avatar
Luca Barbato committed
740
}