OSUOSL/Nero are experiencing Internet connectivity problems. This affects us as we're hosted with OSUOSL. We apologize for the inconvenience.

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

Luca Barbato's avatar
Luca Barbato committed
10
use encoder::*;
11
use partition::*;
Luca Barbato's avatar
Luca Barbato committed
12

Josh Holmer's avatar
Josh Holmer committed
13
use std::cmp;
fbossen's avatar
fbossen committed
14
use std::collections::BTreeMap;
Luca Barbato's avatar
Luca Barbato committed
15
use std::fmt;
Luca Barbato's avatar
Luca Barbato committed
16
use std::sync::Arc;
Josh Holmer's avatar
Josh Holmer committed
17
use scenechange::SceneChangeDetector;
18
use metrics::calculate_frame_psnr;
Luca Barbato's avatar
Luca Barbato committed
19 20 21

// TODO: use the num crate?
#[derive(Clone, Copy, Debug)]
22
#[repr(C)]
Luca Barbato's avatar
Luca Barbato committed
23 24 25
pub struct Rational {
  pub num: u64,
  pub den: u64
Luca Barbato's avatar
Luca Barbato committed
26 27
}

Luca Barbato's avatar
Luca Barbato committed
28 29 30
impl Rational {
  pub fn new(num: u64, den: u64) -> Self {
    Rational { num, den }
Luca Barbato's avatar
Luca Barbato committed
31 32 33
  }
}

34 35
#[derive(Copy, Clone, Debug)]
pub struct EncoderConfig {
Josh Holmer's avatar
Josh Holmer committed
36 37 38 39
  /// The *minimum* interval between two keyframes
  pub min_key_frame_interval: u64,
  /// The *maximum* interval between two keyframes
  pub max_key_frame_interval: u64,
40
  pub low_latency: bool,
41
  pub quantizer: usize,
42 43
  pub tune: Tune,
  pub speed_settings: SpeedSettings,
44
  pub show_psnr: bool,
45 46 47 48
}

impl Default for EncoderConfig {
  fn default() -> Self {
49 50
    const DEFAULT_SPEED: usize = 3;
    Self::with_speed_preset(DEFAULT_SPEED)
51 52 53
  }
}

54 55 56
impl EncoderConfig {
  pub fn with_speed_preset(speed: usize) -> Self {
    EncoderConfig {
Josh Holmer's avatar
Josh Holmer committed
57 58
      min_key_frame_interval: 12,
      max_key_frame_interval: 240,
59 60 61
      low_latency: true,
      quantizer: 100,
      tune: Tune::Psnr,
62 63
      speed_settings: SpeedSettings::from_preset(speed),
      show_psnr: false,
64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 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
    }
  }
}

#[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 {
    if speed <= 1 {
      PredictionModesSetting::ComplexAll
    } else if speed <= 3 {
      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,
}

156
/// Frame-specific information
Luca Barbato's avatar
Luca Barbato committed
157
#[derive(Clone, Copy, Debug)]
158
pub struct FrameInfo {
Luca Barbato's avatar
Luca Barbato committed
159 160 161
  pub width: usize,
  pub height: usize,
  pub bit_depth: usize,
162 163 164 165 166 167 168
  pub chroma_sampling: ChromaSampling
}

/// Contain all the encoder configuration
#[derive(Clone, Copy, Debug)]
pub struct Config {
  pub frame_info: FrameInfo,
Luca Barbato's avatar
Luca Barbato committed
169
  pub timebase: Rational,
Luca Barbato's avatar
Luca Barbato committed
170 171 172 173
  pub enc: EncoderConfig
}

impl Config {
Luca Barbato's avatar
Luca Barbato committed
174 175 176
  pub fn parse(&mut self, key: &str, value: &str) -> Result<(), EncoderStatus> {
    use self::EncoderStatus::*;
    match key {
177
      "low_latency" => self.enc.low_latency = value.parse().map_err(|_e| ParseError)?,
Josh Holmer's avatar
Josh Holmer committed
178 179
      "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)?,
180
      "quantizer" => self.enc.quantizer = value.parse().map_err(|_e| ParseError)?,
181
      "speed" => self.enc.speed_settings = SpeedSettings::from_preset(value.parse().map_err(|_e| ParseError)?),
182 183
      "tune" => self.enc.tune = value.parse().map_err(|_e| ParseError)?,
      _ => return Err(InvalidKey)
Luca Barbato's avatar
Luca Barbato committed
184 185 186 187 188
    }

    Ok(())
  }

Luca Barbato's avatar
Luca Barbato committed
189
  pub fn new_context(&self) -> Context {
190 191 192
    let fi = FrameInvariants::new(
      self.frame_info.width,
      self.frame_info.height,
193
      self.enc
Luca Barbato's avatar
Luca Barbato committed
194
    );
195
    let seq = Sequence::new(&self.frame_info);
Luca Barbato's avatar
Luca Barbato committed
196

197
    #[cfg(feature = "aom")]
Luca Barbato's avatar
Luca Barbato committed
198
    unsafe {
Luca Barbato's avatar
Luca Barbato committed
199 200
      av1_rtcd();
      aom_dsp_rtcd();
Luca Barbato's avatar
Luca Barbato committed
201 202
    }

Josh Holmer's avatar
Josh Holmer committed
203 204 205 206 207 208 209 210 211 212 213 214 215
    Context {
      fi,
      seq,
      frame_count: 0,
      frames_to_be_coded: 0,
      idx: 0,
      frame_q: BTreeMap::new(),
      packet_data: Vec::new(),
      segment_start_idx: 0,
      segment_start_frame: 0,
      frame_types: BTreeMap::new(),
      keyframe_detector: SceneChangeDetector::new(&self.frame_info),
    }
Luca Barbato's avatar
Luca Barbato committed
216 217 218 219 220 221
  }
}

pub struct Context {
  fi: FrameInvariants,
  seq: Sequence,
Luca Barbato's avatar
Luca Barbato committed
222
  //    timebase: Rational,
fbossen's avatar
fbossen committed
223
  frame_count: u64,
224
  frames_to_be_coded: u64,
fbossen's avatar
fbossen committed
225
  idx: u64,
226
  frame_q: BTreeMap<u64, Option<Arc<Frame>>>, //    packet_q: VecDeque<Packet>
Josh Holmer's avatar
Josh Holmer committed
227 228 229 230 231
  packet_data: Vec<u8>,
  segment_start_idx: u64,
  segment_start_frame: u64,
  frame_types: BTreeMap<u64, FrameType>,
  keyframe_detector: SceneChangeDetector,
Luca Barbato's avatar
Luca Barbato committed
232 233 234 235 236 237 238 239 240
}

#[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
241 242 243
  Failure,
  InvalidKey,
  ParseError
Luca Barbato's avatar
Luca Barbato committed
244 245 246 247
}

pub struct Packet {
  pub data: Vec<u8>,
fbossen's avatar
fbossen committed
248 249
  pub rec: Option<Frame>,
  pub number: u64,
250 251 252
  pub frame_type: FrameType,
  /// PSNR for Y, U, and V planes
  pub psnr: Option<(f64, f64, f64)>,
253 254 255 256
}

impl fmt::Display for Packet {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
Luca Barbato's avatar
Luca Barbato committed
257 258 259 260 261 262 263
    write!(
      f,
      "Frame {} - {} - {} bytes",
      self.number,
      self.frame_type,
      self.data.len()
    )
264
  }
Luca Barbato's avatar
Luca Barbato committed
265 266
}

Josh Holmer's avatar
Josh Holmer committed
267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282
#[derive(Debug, Clone, Copy)]
enum FramePropsStatus {
  FramePastSegmentEnd,
}

#[derive(Debug, Clone, Copy)]
struct InterPropsConfig {
  reorder: bool,
  multiref: bool,
  pyramid_depth: u64,
  group_src_len: u64,
  group_len: u64,
  idx_in_group: u64,
  group_idx: u64,
}

Luca Barbato's avatar
Luca Barbato committed
283 284 285 286 287 288 289 290 291
impl Context {
  pub fn new_frame(&self) -> Arc<Frame> {
    Arc::new(Frame::new(self.fi.padded_w, self.fi.padded_h))
  }

  pub fn send_frame<F>(&mut self, frame: F) -> Result<(), EncoderStatus>
  where
    F: Into<Option<Arc<Frame>>>
  {
Josh Holmer's avatar
Josh Holmer committed
292 293 294
    let idx = self.frame_count;
    self.frame_q.insert(idx, frame.into());
    self.save_frame_type(idx);
fbossen's avatar
fbossen committed
295
    self.frame_count = self.frame_count + 1;
Luca Barbato's avatar
Luca Barbato committed
296 297 298
    Ok(())
  }

299 300 301 302 303 304 305 306 307 308 309 310
  pub fn get_frame_count(&self) -> u64 {
    self.frame_count
  }

  pub fn set_frames_to_be_coded(&mut self, frames_to_be_coded: u64) {
    self.frames_to_be_coded = frames_to_be_coded;
  }

  pub fn needs_more_frames(&self, frame_count: u64) -> bool {
    self.frames_to_be_coded == 0 || frame_count < self.frames_to_be_coded
  }

311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340
  pub fn container_sequence_header(&mut self) -> Vec<u8> {
    use bitstream_io::*;
    use std::io;
    fn sequence_header_inner(seq: &Sequence) -> io::Result<Vec<u8>> {
      let mut buf = Vec::new();
      {
      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
      }
      Ok(buf)
    }

    sequence_header_inner(&self.seq).unwrap()
  }

Josh Holmer's avatar
Josh Holmer committed
341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356
  fn setup_key_frame_properties(&mut self) {
    self.fi.frame_type = FrameType::KEY;
    self.fi.intra_only = true;
    self.fi.order_hint = 0;
    self.fi.refresh_frame_flags = ALL_REF_FRAMES_MASK;
    self.fi.show_frame = true;
    self.fi.show_existing_frame = false;
    self.fi.frame_to_show_map_idx = 0;
    let q_boost = 15;
    self.fi.base_q_idx = (self.fi.config.quantizer.max(1 + q_boost).min(255 + q_boost) - q_boost) as u8;
    self.fi.primary_ref_frame = PRIMARY_REF_NONE;
    self.fi.number = self.segment_start_frame;
    for i in 0..INTER_REFS_PER_FRAME {
      self.fi.ref_frames[i] = 0;
    }
  }
fbossen's avatar
fbossen committed
357

Josh Holmer's avatar
Josh Holmer committed
358
  fn get_inter_props_cfg(&mut self, idx_in_segment: u64) -> InterPropsConfig {
Frank Bossen's avatar
Frank Bossen committed
359
    let reorder = !self.fi.config.low_latency;
360
    let multiref = reorder || self.fi.config.speed_settings.multiref;
fbossen's avatar
fbossen committed
361

362
    let pyramid_depth = if reorder { 2 } else { 0 };
fbossen's avatar
fbossen committed
363 364 365
    let group_src_len = 1 << pyramid_depth;
    let group_len = group_src_len + if reorder { pyramid_depth } else { 0 };

Josh Holmer's avatar
Josh Holmer committed
366 367 368 369 370 371 372 373 374 375 376 377 378
    let idx_in_group = (idx_in_segment - 1) % group_len;
    let group_idx = (idx_in_segment - 1) / group_len;

    InterPropsConfig {
      reorder,
      multiref,
      pyramid_depth,
      group_src_len,
      group_len,
      idx_in_group,
      group_idx,
    }
  }
fbossen's avatar
fbossen committed
379

Josh Holmer's avatar
Josh Holmer committed
380 381 382 383 384 385 386 387 388 389
  fn next_keyframe(&self) -> u64 {
    let next_detected = self.frame_types.iter()
      .find(|(&i, &ty)| ty == FrameType::KEY && i > self.segment_start_frame)
      .map(|(&i, _)| i);
    let next_limit = self.segment_start_frame + self.fi.config.max_key_frame_interval;
    if next_detected.is_none() {
      return next_limit;
    }
    cmp::min(next_detected.unwrap(), next_limit)
  }
fbossen's avatar
fbossen committed
390

Josh Holmer's avatar
Josh Holmer committed
391 392 393
  fn setup_inter_frame_properties(&mut self, cfg: &InterPropsConfig) -> Result<(), FramePropsStatus> {
    self.fi.frame_type = FrameType::INTER;
    self.fi.intra_only = false;
fbossen's avatar
fbossen committed
394

Josh Holmer's avatar
Josh Holmer committed
395 396 397 398 399 400 401 402 403 404 405 406
    self.fi.order_hint = (cfg.group_src_len * cfg.group_idx +
      if cfg.reorder && cfg.idx_in_group < cfg.pyramid_depth {
        cfg.group_src_len >> cfg.idx_in_group
      } else {
        cfg.idx_in_group - cfg.pyramid_depth + 1
      }) as u32;
    let number = self.segment_start_frame + self.fi.order_hint as u64;
    if number >= self.next_keyframe() {
      self.fi.show_existing_frame = false;
      self.fi.show_frame = false;
      return Err(FramePropsStatus::FramePastSegmentEnd);
    }
407

Josh Holmer's avatar
Josh Holmer committed
408 409 410 411 412 413 414 415 416 417
    fn pos_to_lvl(pos: u64, pyramid_depth: u64) -> u64 {
      // Derive level within pyramid for a frame with a given coding order position
      // For example, with a pyramid of depth 2, the 2 least significant bits of the
      // position determine the level:
      // 00 -> 0
      // 01 -> 2
      // 10 -> 1
      // 11 -> 2
      pyramid_depth - (pos | (1 << pyramid_depth)).trailing_zeros() as u64
    }
418

Josh Holmer's avatar
Josh Holmer committed
419 420 421 422 423 424 425
    let lvl = if !cfg.reorder {
      0
    } else if cfg.idx_in_group < cfg.pyramid_depth {
      cfg.idx_in_group
    } else {
      pos_to_lvl(cfg.idx_in_group - cfg.pyramid_depth + 1, cfg.pyramid_depth)
    };
426

Josh Holmer's avatar
Josh Holmer committed
427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443
    // Frames with lvl == 0 are stored in slots 0..4 and frames with higher values
    // of lvl in slots 4..8
    let slot_idx = if lvl == 0 {
      (self.fi.order_hint >> cfg.pyramid_depth) % 4 as u32
    } else {
      3 + lvl as u32
    };
    self.fi.show_frame = !cfg.reorder || cfg.idx_in_group >= cfg.pyramid_depth;
    self.fi.show_existing_frame = self.fi.show_frame && cfg.reorder &&
      (cfg.idx_in_group - cfg.pyramid_depth + 1).count_ones() == 1 &&
      cfg.idx_in_group != cfg.pyramid_depth;
    self.fi.frame_to_show_map_idx = slot_idx;
    self.fi.refresh_frame_flags = if self.fi.show_existing_frame {
      0
    } else {
      1 << slot_idx
    };
Luca Barbato's avatar
Luca Barbato committed
444

Josh Holmer's avatar
Josh Holmer committed
445 446
    let q_drop = 15 * lvl as usize;
    self.fi.base_q_idx = (self.fi.config.quantizer.min(255 - q_drop) + q_drop) as u8;
Luca Barbato's avatar
Luca Barbato committed
447

Josh Holmer's avatar
Josh Holmer committed
448 449 450 451 452 453 454 455
    let second_ref_frame = if !cfg.multiref {
      NONE_FRAME
    } else if !cfg.reorder || cfg.idx_in_group == 0 {
      LAST2_FRAME
    } else {
      ALTREF_FRAME
    };
    let ref_in_previous_group = LAST3_FRAME;
456

Josh Holmer's avatar
Josh Holmer committed
457 458
    // reuse probability estimates from previous frames only in top level frames
    self.fi.primary_ref_frame = if lvl > 0 { PRIMARY_REF_NONE } else { (ref_in_previous_group - LAST_FRAME) as u32 };
Luca Barbato's avatar
Luca Barbato committed
459

Josh Holmer's avatar
Josh Holmer committed
460 461 462 463 464 465 466 467 468 469 470 471 472
    for i in 0..INTER_REFS_PER_FRAME {
      self.fi.ref_frames[i] = if lvl == 0 {
        if i == second_ref_frame - LAST_FRAME {
          (slot_idx + 4 - 2) as u8 % 4
        } else {
          (slot_idx + 4 - 1) as u8 % 4
        }
      } else {
        if i == second_ref_frame - LAST_FRAME {
          let oh = self.fi.order_hint + (cfg.group_src_len as u32 >> lvl);
          let lvl2 = pos_to_lvl(oh as u64, cfg.pyramid_depth);
          if lvl2 == 0 {
            ((oh >> cfg.pyramid_depth) % 4) as u8
473
          } else {
Josh Holmer's avatar
Josh Holmer committed
474 475 476 477
            3 + lvl2 as u8
          }
        } else if i == ref_in_previous_group - LAST_FRAME {
          if lvl == 0 {
478
            (slot_idx + 4 - 1) as u8 % 4
Josh Holmer's avatar
Josh Holmer committed
479 480
          } else {
            slot_idx as u8
481
          }
482
        } else {
Josh Holmer's avatar
Josh Holmer committed
483 484 485 486
          let oh = self.fi.order_hint - (cfg.group_src_len as u32 >> lvl);
          let lvl1 = pos_to_lvl(oh as u64, cfg.pyramid_depth);
          if lvl1 == 0 {
            ((oh >> cfg.pyramid_depth) % 4) as u8
487
          } else {
Josh Holmer's avatar
Josh Holmer committed
488
            3 + lvl1 as u8
489 490
          }
        }
491
      }
Josh Holmer's avatar
Josh Holmer committed
492
    }
493

Josh Holmer's avatar
Josh Holmer committed
494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527
    self.fi.reference_mode = if cfg.multiref && cfg.reorder && cfg.idx_in_group != 0 {
      ReferenceMode::SELECT
    } else {
      ReferenceMode::SINGLE
    };
    self.fi.number = number;
    self.fi.me_range_scale = (cfg.group_src_len >> lvl) as u8;
    Ok(())
  }

  fn frame_properties(&mut self, idx: u64) -> Result<(), FramePropsStatus> {
    if idx == 0 {
      // The first frame will always be a key frame
      self.setup_key_frame_properties();
      return Ok(());
    }

    // 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 {
      let inter_props = self.get_inter_props_cfg(idx_in_segment);
      if let Err(FramePropsStatus::FramePastSegmentEnd) = self.setup_inter_frame_properties(&inter_props) {
        let start_frame = self.next_keyframe();
        if !inter_props.reorder || ((idx_in_segment - 1) % inter_props.group_len == 0 && self.fi.number == (start_frame - 1)) {
          self.segment_start_idx = idx;
          self.segment_start_frame = start_frame;
          self.fi.number = start_frame;
        } else {
          return Err(FramePropsStatus::FramePastSegmentEnd);
        }
      }
fbossen's avatar
fbossen committed
528
    }
Luca Barbato's avatar
Luca Barbato committed
529

Josh Holmer's avatar
Josh Holmer committed
530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547
    // Now that we know the frame number, look up the correct frame type
    let frame_type = self.frame_types.get(&self.fi.number).cloned();
    if let Some(frame_type) = frame_type {
      if frame_type == FrameType::KEY {
        self.segment_start_idx = idx;
        self.segment_start_frame = self.fi.number;
      }
      self.fi.frame_type = frame_type;

      let idx_in_segment = idx - self.segment_start_idx;
      if idx_in_segment == 0 {
        self.setup_key_frame_properties();
      } else {
        let inter_props = self.get_inter_props_cfg(idx_in_segment);
        self.setup_inter_frame_properties(&inter_props)?;
      }
    }
    Ok(())
fbossen's avatar
fbossen committed
548
  }
Luca Barbato's avatar
Luca Barbato committed
549

fbossen's avatar
fbossen committed
550 551
  pub fn receive_packet(&mut self) -> Result<Packet, EncoderStatus> {
    let mut idx = self.idx;
Josh Holmer's avatar
Josh Holmer committed
552 553
    while self.frame_properties(idx).is_err() {
      self.idx += 1;
fbossen's avatar
fbossen committed
554 555
      idx = self.idx;
    }
Luca Barbato's avatar
Luca Barbato committed
556

557
    if !self.needs_more_frames(self.fi.number) {
Josh Holmer's avatar
Josh Holmer committed
558
      self.idx += 1;
559 560 561
      return Err(EncoderStatus::EnoughData)
    }

fbossen's avatar
fbossen committed
562
    if self.fi.show_existing_frame {
Josh Holmer's avatar
Josh Holmer committed
563
      self.idx += 1;
Luca Barbato's avatar
Luca Barbato committed
564

565
      let mut fs = FrameState::new(&self.fi);
fbossen's avatar
fbossen committed
566 567

      let data = encode_frame(&mut self.seq, &mut self.fi, &mut fs);
Luca Barbato's avatar
Luca Barbato committed
568

fbossen's avatar
fbossen committed
569 570
      // TODO avoid the clone by having rec Arc.
      let rec = if self.fi.show_frame { Some(fs.rec.clone()) } else { None };
571 572 573 574 575 576
      let mut psnr = None;
      if self.fi.config.show_psnr {
        if let Some(ref rec) = rec {
          psnr = Some(calculate_frame_psnr(&*fs.input, rec, self.seq.bit_depth));
        }
      }
Luca Barbato's avatar
Luca Barbato committed
577

578
      Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type, psnr })
Luca Barbato's avatar
Luca Barbato committed
579
    } else {
fbossen's avatar
fbossen committed
580
      if let Some(f) = self.frame_q.remove(&self.fi.number) {
Josh Holmer's avatar
Josh Holmer committed
581
        self.idx += 1;
fbossen's avatar
fbossen committed
582 583

        if let Some(frame) = f {
584
          let mut fs = FrameState::new_with_frame(&self.fi, frame.clone());
fbossen's avatar
fbossen committed
585 586

          let data = encode_frame(&mut self.seq, &mut self.fi, &mut fs);
587
          self.packet_data.extend(data);
fbossen's avatar
fbossen committed
588

589
          fs.rec.pad(self.fi.width, self.fi.height);
fbossen's avatar
fbossen committed
590 591 592 593 594 595

          // TODO avoid the clone by having rec Arc.
          let rec = if self.fi.show_frame { Some(fs.rec.clone()) } else { None };

          update_rec_buffer(&mut self.fi, fs);

596 597 598
          if self.fi.show_frame {
            let data = self.packet_data.clone();
            self.packet_data = Vec::new();
599 600 601 602 603 604 605 606 607

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

            Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type, psnr })
608 609 610
          } else {
            Err(EncoderStatus::NeedMoreData)
          }
fbossen's avatar
fbossen committed
611 612 613 614 615 616
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      } else {
        Err(EncoderStatus::NeedMoreData)
      }
Luca Barbato's avatar
Luca Barbato committed
617 618 619 620
    }
  }

  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
621 622
    self.frame_q.insert(self.frame_count, None);
    self.frame_count = self.frame_count + 1;
Luca Barbato's avatar
Luca Barbato committed
623
  }
Josh Holmer's avatar
Josh Holmer committed
624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655

  fn save_frame_type(&mut self, idx: u64) {
    let frame_type = self.determine_frame_type(idx);
    self.frame_types.insert(idx, frame_type);
  }

  fn determine_frame_type(&mut self, idx: u64) -> FrameType {
    if idx == 0 {
      return FrameType::KEY;
    }

    let prev_keyframe = *self.frame_types.iter().rfind(|(_, &ty)| ty == FrameType::KEY).unwrap().0;
    let frame = self.frame_q.get(&idx).cloned().unwrap();
    if let Some(frame) = frame {
      let distance = idx - prev_keyframe;
      if distance < self.fi.config.min_key_frame_interval {
        if distance + 1 == self.fi.config.min_key_frame_interval {
          // 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.
          self.keyframe_detector.detect_scene_change(frame, idx as usize);
        }
        return FrameType::INTER;
      }
      if distance >= self.fi.config.max_key_frame_interval {
        return FrameType::KEY;
      }
      if self.keyframe_detector.detect_scene_change(frame, idx as usize) {
        return FrameType::KEY;
      }
    }
    FrameType::INTER
  }
Luca Barbato's avatar
Luca Barbato committed
656
}
Luca Barbato's avatar
Luca Barbato committed
657 658 659 660 661 662

impl fmt::Display for Context {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "Frame {} - {}", self.fi.number, self.fi.frame_type)
  }
}