api.rs 17.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
27
28
29
30
31
32
33
34
#[derive(Debug, Clone, Copy)]
pub struct VideoDetails {
  pub width: usize,
  pub height: usize,
  pub bits: usize,
  pub bytes: usize,
  pub mono: bool,
  pub bit_depth: usize,
  pub chroma_sampling: ChromaSampling,
  pub chroma_sample_position: ChromaSamplePosition,
Luca Barbato's avatar
Luca Barbato committed
35
  pub time_base: Rational,
36
37
38
39
40
41
42
43
44
45
46
47
48
}

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

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

Luca Barbato's avatar
Luca Barbato committed
62
63
64
impl Rational {
  pub fn new(num: u64, den: u64) -> Self {
    Rational { num, den }
Luca Barbato's avatar
Luca Barbato committed
65
66
67
  }
}

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

impl Default for EncoderConfig {
  fn default() -> Self {
84
85
    const DEFAULT_SPEED: usize = 3;
    Self::with_speed_preset(DEFAULT_SPEED)
86
87
88
  }
}

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

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

192
193
arg_enum!{
  #[derive(Debug, Clone, Copy, PartialEq)]
194
  #[repr(C)]
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
  pub enum MatrixCoefficients {
      Identity = 0,
      BT709,
      Unspecified,
      BT470M = 4,
      BT470BG,
      ST170M,
      ST240M,
      YCgCo,
      BT2020NonConstantLuminance,
      BT2020ConstantLuminance,
      ST2085,
      ChromaticityDerivedNonConstantLuminance,
      ChromaticityDerivedConstantLuminance,
      ICtCp,
  }
211
212
213
214
215
216
217
218
}

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

219
220
arg_enum!{
  #[derive(Debug,Clone,Copy,PartialEq)]
221
  #[repr(C)]
222
223
224
225
226
227
228
229
230
231
232
233
234
235
  pub enum ColorPrimaries {
      BT709 = 1,
      Unspecified,
      BT470M = 4,
      BT470BG,
      ST170M,
      ST240M,
      Film,
      BT2020,
      ST428,
      P3DCI,
      P3Display,
      Tech3213 = 22,
  }
236
237
238
239
240
241
242
243
}

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

244
245
arg_enum!{
  #[derive(Debug,Clone,Copy,PartialEq)]
246
  #[repr(C)]
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
  pub enum TransferCharacteristics {
      BT1886 = 1,
      Unspecified,
      BT470M = 4,
      BT470BG,
      ST170M,
      ST240M,
      Linear,
      Logarithmic100,
      Logarithmic316,
      XVYCC,
      BT1361E,
      SRGB,
      BT2020Ten,
      BT2020Twelve,
      PerceptualQuantizer,
      ST428,
      HybridLogGamma,
  }
266
267
268
269
270
271
272
273
}

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

274
275
276
277
278
279
280
#[derive(Copy, Clone, Debug)]
pub struct ColorDescription {
    pub color_primaries: ColorPrimaries,
    pub transfer_characteristics: TransferCharacteristics,
    pub matrix_coefficients: MatrixCoefficients
}

281
282
283
/// Contain all the encoder configuration
#[derive(Clone, Copy, Debug)]
pub struct Config {
284
  pub video_info: VideoDetails,
Luca Barbato's avatar
Luca Barbato committed
285
286
287
288
  pub enc: EncoderConfig
}

impl Config {
Luca Barbato's avatar
Luca Barbato committed
289
290
  pub fn parse(&mut self, key: &str, value: &str) -> Result<(), EncoderStatus> {
    match key {
291
      "low_latency" => self.enc.low_latency = value.parse().map_err(|_e| ParseError)?,
Josh Holmer's avatar
Josh Holmer committed
292
293
      "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)?,
294
      "quantizer" => self.enc.quantizer = value.parse().map_err(|_e| ParseError)?,
295
      "speed" => self.enc.speed_settings = SpeedSettings::from_preset(value.parse().map_err(|_e| ParseError)?),
296
297
      "tune" => self.enc.tune = value.parse().map_err(|_e| ParseError)?,
      _ => return Err(InvalidKey)
Luca Barbato's avatar
Luca Barbato committed
298
299
300
301
302
    }

    Ok(())
  }

Luca Barbato's avatar
Luca Barbato committed
303
  pub fn new_context(&self) -> Context {
304
    #[cfg(feature = "aom")]
Luca Barbato's avatar
Luca Barbato committed
305
    unsafe {
Luca Barbato's avatar
Luca Barbato committed
306
307
      av1_rtcd();
      aom_dsp_rtcd();
Luca Barbato's avatar
Luca Barbato committed
308
309
    }

Josh Holmer's avatar
Josh Holmer committed
310
311
312
313
    Context {
      frame_count: 0,
      frames_to_be_coded: 0,
      idx: 0,
Josh Holmer's avatar
Josh Holmer committed
314
      frames_processed: 0,
Josh Holmer's avatar
Josh Holmer committed
315
      frame_q: BTreeMap::new(),
Josh Holmer's avatar
Josh Holmer committed
316
317
      frame_data: BTreeMap::new(),
      keyframes: BTreeSet::new(),
Josh Holmer's avatar
Josh Holmer committed
318
319
320
      packet_data: Vec::new(),
      segment_start_idx: 0,
      segment_start_frame: 0,
321
      keyframe_detector: SceneChangeDetector::new(&self.video_info),
Josh Holmer's avatar
Josh Holmer committed
322
      config: *self,
Josh Holmer's avatar
Josh Holmer committed
323
    }
Luca Barbato's avatar
Luca Barbato committed
324
325
326
327
  }
}

pub struct Context {
Luca Barbato's avatar
Luca Barbato committed
328
  //    timebase: Rational,
fbossen's avatar
fbossen committed
329
  frame_count: u64,
330
  frames_to_be_coded: u64,
fbossen's avatar
fbossen committed
331
  idx: u64,
Josh Holmer's avatar
Josh Holmer committed
332
333
  frames_processed: u64,
  /// Maps frame *number* to frames
334
  frame_q: BTreeMap<u64, Option<Arc<Frame>>>, //    packet_q: VecDeque<Packet>
Josh Holmer's avatar
Josh Holmer committed
335
336
337
338
339
  /// 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>,
Josh Holmer's avatar
Josh Holmer committed
340
341
342
343
  packet_data: Vec<u8>,
  segment_start_idx: u64,
  segment_start_frame: u64,
  keyframe_detector: SceneChangeDetector,
344
  pub config: Config,
Luca Barbato's avatar
Luca Barbato committed
345
346
347
348
349
350
351
352
353
}

#[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
354
355
356
  Failure,
  InvalidKey,
  ParseError
Luca Barbato's avatar
Luca Barbato committed
357
358
359
360
}

pub struct Packet {
  pub data: Vec<u8>,
fbossen's avatar
fbossen committed
361
362
  pub rec: Option<Frame>,
  pub number: u64,
363
364
365
  pub frame_type: FrameType,
  /// PSNR for Y, U, and V planes
  pub psnr: Option<(f64, f64, f64)>,
366
367
368
369
}

impl fmt::Display for Packet {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
Luca Barbato's avatar
Luca Barbato committed
370
371
372
373
374
375
376
    write!(
      f,
      "Frame {} - {} - {} bytes",
      self.number,
      self.frame_type,
      self.data.len()
    )
377
  }
Luca Barbato's avatar
Luca Barbato committed
378
379
380
381
}

impl Context {
  pub fn new_frame(&self) -> Arc<Frame> {
Josh Holmer's avatar
Josh Holmer committed
382
    Arc::new(Frame::new(
383
384
385
      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
386
    ))
Luca Barbato's avatar
Luca Barbato committed
387
388
389
390
391
392
  }

  pub fn send_frame<F>(&mut self, frame: F) -> Result<(), EncoderStatus>
  where
    F: Into<Option<Arc<Frame>>>
  {
Josh Holmer's avatar
Josh Holmer committed
393
394
    let idx = self.frame_count;
    self.frame_q.insert(idx, frame.into());
395
    self.frame_count += 1;
Luca Barbato's avatar
Luca Barbato committed
396
397
398
    Ok(())
  }

399
400
401
402
403
404
405
406
  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;
  }

Josh Holmer's avatar
Josh Holmer committed
407
408
409
410
  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)
  }

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

415
416
417
  pub fn container_sequence_header(&mut self) -> Vec<u8> {
    fn sequence_header_inner(seq: &Sequence) -> io::Result<Vec<u8>> {
      let mut buf = Vec::new();
418

419
      {
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
        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
437
      }
438

439
440
441
      Ok(buf)
    }

Josh Holmer's avatar
Josh Holmer committed
442
    sequence_header_inner(&self.frame_data[&0].sequence).unwrap()
443
444
  }

Josh Holmer's avatar
Josh Holmer committed
445
  fn next_keyframe(&self) -> u64 {
Josh Holmer's avatar
Josh Holmer committed
446
447
448
449
    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
450
451
452
453
454
    if next_detected.is_none() {
      return next_limit;
    }
    cmp::min(next_detected.unwrap(), next_limit)
  }
fbossen's avatar
fbossen committed
455

Josh Holmer's avatar
Josh Holmer committed
456
  fn set_frame_properties(&mut self, idx: u64) -> Result<(), ()> {
Josh Holmer's avatar
Josh Holmer committed
457
458
459
460
461
462
463
464
465
466
467
    let props = self.get_frame_properties(idx);
    let result = props.as_ref().map(|_| ()).map_err(|_| ());
    match props {
      Ok(props) | Err(props) => {
        self.frame_data.insert(idx, props);
      }
    }
    result
  }

  fn get_frame_properties(&mut self, idx: u64) -> Result<FrameInvariants, FrameInvariants> {
Josh Holmer's avatar
Josh Holmer committed
468
469
    if idx == 0 {
      // The first frame will always be a key frame
Josh Holmer's avatar
Josh Holmer committed
470
471
      let fi = FrameInvariants::new_key_frame(
        &FrameInvariants::new(
472
473
          self.config.video_info.width,
          self.config.video_info.height,
Josh Holmer's avatar
Josh Holmer committed
474
          self.config.enc,
475
          Sequence::new(&self.config.video_info)
Josh Holmer's avatar
Josh Holmer committed
476
477
478
479
        ),
        0
      );
      return Ok(fi);
Josh Holmer's avatar
Josh Holmer committed
480
481
    }

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

Josh Holmer's avatar
Josh Holmer committed
484
485
486
487
488
489
    // 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
490
      let next_keyframe = self.next_keyframe();
Josh Holmer's avatar
Josh Holmer committed
491
492
      let (fi_temp, success) = FrameInvariants::new_inter_frame(
        &fi,
493
494
495
496
        self.segment_start_frame,
        idx_in_segment,
        next_keyframe
      );
Josh Holmer's avatar
Josh Holmer committed
497
      fi = fi_temp;
Josh Holmer's avatar
Josh Holmer committed
498
      if !success {
Josh Holmer's avatar
Josh Holmer committed
499
500
501
        if !fi.inter_cfg.unwrap().reorder
          || ((idx_in_segment - 1) % fi.inter_cfg.unwrap().group_len == 0
          && fi.number == (next_keyframe - 1))
502
        {
Josh Holmer's avatar
Josh Holmer committed
503
          self.segment_start_idx = idx;
Josh Holmer's avatar
Josh Holmer committed
504
          self.segment_start_frame = next_keyframe;
Josh Holmer's avatar
Josh Holmer committed
505
          fi.number = next_keyframe;
Josh Holmer's avatar
Josh Holmer committed
506
        } else {
Josh Holmer's avatar
Josh Holmer committed
507
          return Err(fi);
Josh Holmer's avatar
Josh Holmer committed
508
509
        }
      }
fbossen's avatar
fbossen committed
510
    }
Luca Barbato's avatar
Luca Barbato committed
511

Josh Holmer's avatar
Josh Holmer committed
512
    // Now that we know the frame number, look up the correct frame type
Josh Holmer's avatar
Josh Holmer committed
513
514
515
516
517
518
519
    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
520

Josh Holmer's avatar
Josh Holmer committed
521
522
523
524
525
526
527
528
529
530
531
532
533
534
    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();
      let (fi_temp, success) = FrameInvariants::new_inter_frame(
        &fi,
        self.segment_start_frame,
        idx_in_segment,
        next_keyframe
      );
      fi = fi_temp;
      if !success {
        return Err(fi);
Josh Holmer's avatar
Josh Holmer committed
535
536
      }
    }
Josh Holmer's avatar
Josh Holmer committed
537
    Ok(fi)
fbossen's avatar
fbossen committed
538
  }
Luca Barbato's avatar
Luca Barbato committed
539

fbossen's avatar
fbossen committed
540
541
  pub fn receive_packet(&mut self) -> Result<Packet, EncoderStatus> {
    let mut idx = self.idx;
Josh Holmer's avatar
Josh Holmer committed
542
    while self.set_frame_properties(idx).is_err() {
Josh Holmer's avatar
Josh Holmer committed
543
      self.idx += 1;
fbossen's avatar
fbossen committed
544
545
      idx = self.idx;
    }
Luca Barbato's avatar
Luca Barbato committed
546

Josh Holmer's avatar
Josh Holmer committed
547
    if !self.needs_more_frames(self.frame_data.get(&idx).unwrap().number) {
Josh Holmer's avatar
Josh Holmer committed
548
      self.idx += 1;
549
550
551
      return Err(EncoderStatus::EnoughData)
    }

Josh Holmer's avatar
Josh Holmer committed
552
553
    let fi = self.frame_data.get_mut(&idx).unwrap();
    if fi.show_existing_frame {
Josh Holmer's avatar
Josh Holmer committed
554
      self.idx += 1;
Luca Barbato's avatar
Luca Barbato committed
555

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

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

Josh Holmer's avatar
Josh Holmer committed
560
      let rec = if fi.show_frame { Some(fs.rec) } else { None };
561
      let mut psnr = None;
Josh Holmer's avatar
Josh Holmer committed
562
      if self.config.enc.show_psnr {
563
        if let Some(ref rec) = rec {
Josh Holmer's avatar
Josh Holmer committed
564
          psnr = Some(calculate_frame_psnr(&*fs.input, rec, fi.sequence.bit_depth));
565
566
        }
      }
Luca Barbato's avatar
Luca Barbato committed
567

Josh Holmer's avatar
Josh Holmer committed
568
569
      self.frames_processed += 1;
      Ok(Packet { data, rec, number: fi.number, frame_type: fi.frame_type, psnr })
Luca Barbato's avatar
Luca Barbato committed
570
    } else {
Josh Holmer's avatar
Josh Holmer committed
571
      if let Some(f) = self.frame_q.get(&fi.number) {
Josh Holmer's avatar
Josh Holmer committed
572
        self.idx += 1;
fbossen's avatar
fbossen committed
573
574

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

Josh Holmer's avatar
Josh Holmer committed
577
          let data = encode_frame(fi, &mut fs);
578
          self.packet_data.extend(data);
fbossen's avatar
fbossen committed
579

Josh Holmer's avatar
Josh Holmer committed
580
          fs.rec.pad(fi.width, fi.height);
fbossen's avatar
fbossen committed
581
582

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

Josh Holmer's avatar
Josh Holmer committed
585
          update_rec_buffer(fi, fs);
fbossen's avatar
fbossen committed
586

Josh Holmer's avatar
Josh Holmer committed
587
          if fi.show_frame {
588
589
            let data = self.packet_data.clone();
            self.packet_data = Vec::new();
590
591

            let mut psnr = None;
Josh Holmer's avatar
Josh Holmer committed
592
            if self.config.enc.show_psnr {
593
              if let Some(ref rec) = rec {
Josh Holmer's avatar
Josh Holmer committed
594
                psnr = Some(calculate_frame_psnr(&*frame, rec, fi.sequence.bit_depth));
595
596
597
              }
            }

Josh Holmer's avatar
Josh Holmer committed
598
599
            self.frames_processed += 1;
            Ok(Packet { data, rec, number: fi.number, frame_type: fi.frame_type, psnr })
600
601
602
          } else {
            Err(EncoderStatus::NeedMoreData)
          }
fbossen's avatar
fbossen committed
603
604
605
606
607
608
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      } else {
        Err(EncoderStatus::NeedMoreData)
      }
Luca Barbato's avatar
Luca Barbato committed
609
610
611
    }
  }

Josh Holmer's avatar
Josh Holmer committed
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
  pub fn garbage_collect(&mut self, cur_frame: u64) {
    if cur_frame == 0 {
      return;
    }
    for i in 0..cur_frame {
      self.frame_q.remove(&i);
    }
    if self.idx < 2 {
      return;
    }
    for i in 0..(self.idx - 1) {
      self.frame_data.remove(&i);
    }
  }

Luca Barbato's avatar
Luca Barbato committed
627
  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
628
    self.frame_q.insert(self.frame_count, None);
629
    self.frame_count += 1;
Luca Barbato's avatar
Luca Barbato committed
630
  }
Josh Holmer's avatar
Josh Holmer committed
631

Josh Holmer's avatar
Josh Holmer committed
632
633
  fn determine_frame_type(&mut self, frame_number: u64) -> FrameType {
    if frame_number == 0 {
Josh Holmer's avatar
Josh Holmer committed
634
635
636
      return FrameType::KEY;
    }

Josh Holmer's avatar
Josh Holmer committed
637
638
639
640
641
642
643
644
    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
645
    if let Some(frame) = frame {
Josh Holmer's avatar
Josh Holmer committed
646
647
648
      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
649
650
          // 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
651
          self.keyframe_detector.detect_scene_change(frame, frame_number as usize);
Josh Holmer's avatar
Josh Holmer committed
652
653
654
        }
        return FrameType::INTER;
      }
Josh Holmer's avatar
Josh Holmer committed
655
      if distance >= self.config.enc.max_key_frame_interval {
Josh Holmer's avatar
Josh Holmer committed
656
657
        return FrameType::KEY;
      }
Josh Holmer's avatar
Josh Holmer committed
658
      if self.keyframe_detector.detect_scene_change(frame, frame_number as usize) {
Josh Holmer's avatar
Josh Holmer committed
659
660
661
662
663
        return FrameType::KEY;
      }
    }
    FrameType::INTER
  }
Luca Barbato's avatar
Luca Barbato committed
664
}