api.rs 14.2 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 267 268 269 270 271 272 273 274 275
}

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
276 277 278
    let idx = self.frame_count;
    self.frame_q.insert(idx, frame.into());
    self.save_frame_type(idx);
fbossen's avatar
fbossen committed
279
    self.frame_count = self.frame_count + 1;
Luca Barbato's avatar
Luca Barbato committed
280 281 282
    Ok(())
  }

283 284 285 286 287 288 289 290 291 292 293 294
  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
  }

295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324
  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
325 326 327 328 329 330 331 332 333 334
  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
335

Josh Holmer's avatar
Josh Holmer committed
336
  fn set_frame_properties(&mut self, idx: u64) -> Result<(), ()> {
Josh Holmer's avatar
Josh Holmer committed
337 338
    if idx == 0 {
      // The first frame will always be a key frame
Josh Holmer's avatar
Josh Holmer committed
339
      self.fi = FrameInvariants::new_key_frame(&self.fi,0);
Josh Holmer's avatar
Josh Holmer committed
340 341 342 343 344 345 346 347 348
      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 {
Josh Holmer's avatar
Josh Holmer committed
349 350 351 352 353
      let next_keyframe = self.next_keyframe();
      let (fi, success) = FrameInvariants::new_inter_frame(&self.fi, self.segment_start_frame, idx_in_segment, next_keyframe);
      self.fi = fi;
      if !success {
        if !self.fi.inter_cfg.unwrap().reorder || ((idx_in_segment - 1) % self.fi.inter_cfg.unwrap().group_len == 0 && self.fi.number == (next_keyframe - 1)) {
Josh Holmer's avatar
Josh Holmer committed
354
          self.segment_start_idx = idx;
Josh Holmer's avatar
Josh Holmer committed
355 356
          self.segment_start_frame = next_keyframe;
          self.fi.number = next_keyframe;
Josh Holmer's avatar
Josh Holmer committed
357
        } else {
Josh Holmer's avatar
Josh Holmer committed
358
          return Err(());
Josh Holmer's avatar
Josh Holmer committed
359 360
        }
      }
fbossen's avatar
fbossen committed
361
    }
Luca Barbato's avatar
Luca Barbato committed
362

Josh Holmer's avatar
Josh Holmer committed
363 364 365 366 367 368 369 370 371 372 373
    // 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 {
Josh Holmer's avatar
Josh Holmer committed
374
        self.fi = FrameInvariants::new_key_frame(&self.fi, self.segment_start_frame);
Josh Holmer's avatar
Josh Holmer committed
375
      } else {
Josh Holmer's avatar
Josh Holmer committed
376 377 378 379 380 381
        let next_keyframe = self.next_keyframe();
        let (fi, success) = FrameInvariants::new_inter_frame(&self.fi, self.segment_start_frame, idx_in_segment, next_keyframe);
        self.fi = fi;
        if !success {
          return Err(());
        }
Josh Holmer's avatar
Josh Holmer committed
382 383 384
      }
    }
    Ok(())
fbossen's avatar
fbossen committed
385
  }
Luca Barbato's avatar
Luca Barbato committed
386

fbossen's avatar
fbossen committed
387 388
  pub fn receive_packet(&mut self) -> Result<Packet, EncoderStatus> {
    let mut idx = self.idx;
Josh Holmer's avatar
Josh Holmer committed
389
    while self.set_frame_properties(idx).is_err() {
Josh Holmer's avatar
Josh Holmer committed
390
      self.idx += 1;
fbossen's avatar
fbossen committed
391 392
      idx = self.idx;
    }
Luca Barbato's avatar
Luca Barbato committed
393

394
    if !self.needs_more_frames(self.fi.number) {
Josh Holmer's avatar
Josh Holmer committed
395
      self.idx += 1;
396 397 398
      return Err(EncoderStatus::EnoughData)
    }

fbossen's avatar
fbossen committed
399
    if self.fi.show_existing_frame {
Josh Holmer's avatar
Josh Holmer committed
400
      self.idx += 1;
Luca Barbato's avatar
Luca Barbato committed
401

402
      let mut fs = FrameState::new(&self.fi);
fbossen's avatar
fbossen committed
403 404

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

fbossen's avatar
fbossen committed
406 407
      // TODO avoid the clone by having rec Arc.
      let rec = if self.fi.show_frame { Some(fs.rec.clone()) } else { None };
408 409 410 411 412 413
      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
414

415
      Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type, psnr })
Luca Barbato's avatar
Luca Barbato committed
416
    } else {
fbossen's avatar
fbossen committed
417
      if let Some(f) = self.frame_q.remove(&self.fi.number) {
Josh Holmer's avatar
Josh Holmer committed
418
        self.idx += 1;
fbossen's avatar
fbossen committed
419 420

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

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

426
          fs.rec.pad(self.fi.width, self.fi.height);
fbossen's avatar
fbossen committed
427 428 429 430 431 432

          // 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);

433 434 435
          if self.fi.show_frame {
            let data = self.packet_data.clone();
            self.packet_data = Vec::new();
436 437 438 439 440 441 442 443 444

            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 })
445 446 447
          } else {
            Err(EncoderStatus::NeedMoreData)
          }
fbossen's avatar
fbossen committed
448 449 450 451 452 453
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      } else {
        Err(EncoderStatus::NeedMoreData)
      }
Luca Barbato's avatar
Luca Barbato committed
454 455 456 457
    }
  }

  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
458 459
    self.frame_q.insert(self.frame_count, None);
    self.frame_count = self.frame_count + 1;
Luca Barbato's avatar
Luca Barbato committed
460
  }
Josh Holmer's avatar
Josh Holmer committed
461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492

  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
493
}
Luca Barbato's avatar
Luca Barbato committed
494 495 496 497 498 499

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