api.rs 13.9 KB
Newer Older
rzumer's avatar
rzumer committed
1 2 3 4 5 6 7 8 9
// Copyright (c) 2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

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

fbossen's avatar
fbossen committed
13
use std::collections::BTreeMap;
Luca Barbato's avatar
Luca Barbato committed
14
use std::fmt;
Luca Barbato's avatar
Luca Barbato committed
15 16 17 18
use std::sync::Arc;

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

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

31 32
#[derive(Copy, Clone, Debug)]
pub struct EncoderConfig {
33 34
  pub key_frame_interval: u64,
  pub low_latency: bool,
35
  pub quantizer: usize,
36 37
  pub tune: Tune,
  pub speed_settings: SpeedSettings,
38 39 40 41
}

impl Default for EncoderConfig {
  fn default() -> Self {
42 43
    const DEFAULT_SPEED: usize = 3;
    Self::with_speed_preset(DEFAULT_SPEED)
44 45 46
  }
}

47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 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
impl EncoderConfig {
  pub fn with_speed_preset(speed: usize) -> Self {
    EncoderConfig {
      key_frame_interval: 30,
      low_latency: true,
      quantizer: 100,
      tune: Tune::Psnr,
      speed_settings: SpeedSettings::from_preset(speed)
    }
  }
}

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

147
/// Frame-specific information
Luca Barbato's avatar
Luca Barbato committed
148
#[derive(Clone, Copy, Debug)]
149
pub struct FrameInfo {
Luca Barbato's avatar
Luca Barbato committed
150 151 152
  pub width: usize,
  pub height: usize,
  pub bit_depth: usize,
153 154 155 156 157 158 159
  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
160
  pub timebase: Rational,
Luca Barbato's avatar
Luca Barbato committed
161 162 163 164
  pub enc: EncoderConfig
}

impl Config {
Luca Barbato's avatar
Luca Barbato committed
165 166 167
  pub fn parse(&mut self, key: &str, value: &str) -> Result<(), EncoderStatus> {
    use self::EncoderStatus::*;
    match key {
168 169 170
      "low_latency" => self.enc.low_latency = value.parse().map_err(|_e| ParseError)?,
      "key_frame_interval" => self.enc.key_frame_interval = value.parse().map_err(|_e| ParseError)?,
      "quantizer" => self.enc.quantizer = value.parse().map_err(|_e| ParseError)?,
171
      "speed" => self.enc.speed_settings = SpeedSettings::from_preset(value.parse().map_err(|_e| ParseError)?),
172 173
      "tune" => self.enc.tune = value.parse().map_err(|_e| ParseError)?,
      _ => return Err(InvalidKey)
Luca Barbato's avatar
Luca Barbato committed
174 175 176 177 178
    }

    Ok(())
  }

Luca Barbato's avatar
Luca Barbato committed
179
  pub fn new_context(&self) -> Context {
180 181 182
    let fi = FrameInvariants::new(
      self.frame_info.width,
      self.frame_info.height,
183
      self.enc
Luca Barbato's avatar
Luca Barbato committed
184
    );
185
    let seq = Sequence::new(&self.frame_info);
Luca Barbato's avatar
Luca Barbato committed
186

187
    #[cfg(feature = "aom")]
Luca Barbato's avatar
Luca Barbato committed
188
    unsafe {
Luca Barbato's avatar
Luca Barbato committed
189 190
      av1_rtcd();
      aom_dsp_rtcd();
Luca Barbato's avatar
Luca Barbato committed
191 192
    }

193
    Context { fi, seq, frame_count: 0, idx: 0, frame_q: BTreeMap::new(), packet_data: Vec::new() }
Luca Barbato's avatar
Luca Barbato committed
194 195 196 197 198 199
  }
}

pub struct Context {
  fi: FrameInvariants,
  seq: Sequence,
Luca Barbato's avatar
Luca Barbato committed
200
  //    timebase: Rational,
fbossen's avatar
fbossen committed
201 202
  frame_count: u64,
  idx: u64,
203 204
  frame_q: BTreeMap<u64, Option<Arc<Frame>>>, //    packet_q: VecDeque<Packet>
  packet_data: Vec<u8>
Luca Barbato's avatar
Luca Barbato committed
205 206 207 208 209 210 211 212 213
}

#[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
214 215 216
  Failure,
  InvalidKey,
  ParseError
Luca Barbato's avatar
Luca Barbato committed
217 218 219 220
}

pub struct Packet {
  pub data: Vec<u8>,
fbossen's avatar
fbossen committed
221 222
  pub rec: Option<Frame>,
  pub number: u64,
223 224 225 226 227
  pub frame_type: FrameType
}

impl fmt::Display for Packet {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
Luca Barbato's avatar
Luca Barbato committed
228 229 230 231 232 233 234
    write!(
      f,
      "Frame {} - {} - {} bytes",
      self.number,
      self.frame_type,
      self.data.len()
    )
235
  }
Luca Barbato's avatar
Luca Barbato committed
236 237 238 239 240 241 242 243 244 245 246
}

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>>>
  {
fbossen's avatar
fbossen committed
247 248
    self.frame_q.insert(self.frame_count, frame.into());
    self.frame_count = self.frame_count + 1;
Luca Barbato's avatar
Luca Barbato committed
249 250 251
    Ok(())
  }

252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281
  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()
  }

fbossen's avatar
fbossen committed
282
  pub fn frame_properties(&mut self, idx: u64) -> bool {
283
    let key_frame_interval: u64 = self.fi.config.key_frame_interval;
fbossen's avatar
fbossen committed
284

Frank Bossen's avatar
Frank Bossen committed
285
    let reorder = !self.fi.config.low_latency;
286
    let multiref = reorder || self.fi.config.speed_settings.multiref;
fbossen's avatar
fbossen committed
287

288
    let pyramid_depth = if reorder { 2 } else { 0 };
fbossen's avatar
fbossen committed
289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313
    let group_src_len = 1 << pyramid_depth;
    let group_len = group_src_len + if reorder { pyramid_depth } else { 0 };
    let segment_len = 1 + (key_frame_interval - 1 + group_src_len - 1) / group_src_len * group_len;

    let idx_in_segment = idx % segment_len;
    let segment_idx = idx / segment_len;

    if idx_in_segment == 0 {
      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 = segment_idx * key_frame_interval;
      for i in 0..INTER_REFS_PER_FRAME {
        self.fi.ref_frames[i] = 0;
      }
    } else {
      let idx_in_group = (idx_in_segment - 1) % group_len;
      let group_idx = (idx_in_segment - 1) / group_len;
fbossen's avatar
fbossen committed
314

fbossen's avatar
fbossen committed
315 316
      self.fi.frame_type = FrameType::INTER;
      self.fi.intra_only = false;
fbossen's avatar
fbossen committed
317

fbossen's avatar
fbossen committed
318 319 320 321 322 323 324 325 326
      self.fi.order_hint = (group_src_len * group_idx +
        if reorder && idx_in_group < pyramid_depth {
          group_src_len >> idx_in_group
        } else {
          idx_in_group - pyramid_depth + 1
        }) as u32;
      if self.fi.order_hint >= key_frame_interval as u32 {
        return false;
      }
327

328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353
      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
      }

      let lvl = if !reorder {
        0
      } else if idx_in_group < pyramid_depth {
        idx_in_group
      } else {
        pos_to_lvl(idx_in_group - pyramid_depth + 1, pyramid_depth)
      };

      // 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 >> pyramid_depth) % 4 as u32
      } else {
        3 + lvl as u32
      };
fbossen's avatar
fbossen committed
354 355 356 357 358 359 360
      self.fi.show_frame = !reorder || idx_in_group >= pyramid_depth;
      self.fi.show_existing_frame = self.fi.show_frame && reorder &&
        (idx_in_group - pyramid_depth + 1).count_ones() == 1 &&
        idx_in_group != pyramid_depth;
      self.fi.frame_to_show_map_idx = slot_idx;
      self.fi.refresh_frame_flags = if self.fi.show_existing_frame {
        0
Luca Barbato's avatar
Luca Barbato committed
361
      } else {
362
        1 << slot_idx
Luca Barbato's avatar
Luca Barbato committed
363 364
      };

fbossen's avatar
fbossen committed
365 366
      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
367

fbossen's avatar
fbossen committed
368 369 370 371 372 373 374 375
      let second_ref_frame = if !multiref {
        NONE_FRAME
      } else if !reorder || idx_in_group == 0 {
        LAST2_FRAME
      } else {
        ALTREF_FRAME
      };
      let ref_in_previous_group = LAST3_FRAME;
376

377 378
      // 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
379

380
      for i in 0..INTER_REFS_PER_FRAME {
381 382 383 384 385 386
        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
          }
387
        } else {
388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411
          if i == second_ref_frame - LAST_FRAME {
            let oh = self.fi.order_hint + (group_src_len as u32 >> lvl);
            let lvl2 = pos_to_lvl(oh as u64, pyramid_depth);
            if lvl2 == 0 {
              ((oh >> pyramid_depth) % 4) as u8
            } else {
              3 + lvl2 as u8
            }
          } else if i == ref_in_previous_group - LAST_FRAME {
            if lvl == 0 {
              (slot_idx + 4 - 1) as u8 % 4
            } else {
              slot_idx as u8
            }
          } else {
            let oh = self.fi.order_hint - (group_src_len as u32 >> lvl);
            let lvl1 = pos_to_lvl(oh as u64, pyramid_depth);
            if lvl1 == 0 {
              ((oh >> pyramid_depth) % 4) as u8
            } else {
              3 + lvl1 as u8
            }
          }
        }
412 413
      }

414 415 416 417 418
      self.fi.reference_mode = if multiref && reorder && idx_in_group != 0 {
        ReferenceMode::SELECT
      } else {
        ReferenceMode::SINGLE
      };
fbossen's avatar
fbossen committed
419
      self.fi.number = segment_idx * key_frame_interval + self.fi.order_hint as u64;
420
      self.fi.me_range_scale = (group_src_len >> lvl) as u8;
fbossen's avatar
fbossen committed
421
    }
Luca Barbato's avatar
Luca Barbato committed
422

fbossen's avatar
fbossen committed
423 424
    true
  }
Luca Barbato's avatar
Luca Barbato committed
425

fbossen's avatar
fbossen committed
426 427 428 429 430 431
  pub fn receive_packet(&mut self) -> Result<Packet, EncoderStatus> {
    let mut idx = self.idx;
    while !self.frame_properties(idx) {
      self.idx = self.idx + 1;
      idx = self.idx;
    }
Luca Barbato's avatar
Luca Barbato committed
432

fbossen's avatar
fbossen committed
433 434
    if self.fi.show_existing_frame {
      self.idx = self.idx + 1;
Luca Barbato's avatar
Luca Barbato committed
435

436
      let mut fs = FrameState::new(&self.fi);
fbossen's avatar
fbossen committed
437 438

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

fbossen's avatar
fbossen committed
440 441
      // TODO avoid the clone by having rec Arc.
      let rec = if self.fi.show_frame { Some(fs.rec.clone()) } else { None };
Luca Barbato's avatar
Luca Barbato committed
442

fbossen's avatar
fbossen committed
443
      Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type })
Luca Barbato's avatar
Luca Barbato committed
444
    } else {
fbossen's avatar
fbossen committed
445 446 447 448
      if let Some(f) = self.frame_q.remove(&self.fi.number) {
        self.idx = self.idx + 1;

        if let Some(frame) = f {
449
          let mut fs = FrameState::new_with_frame(&self.fi, frame);
fbossen's avatar
fbossen committed
450 451

          let data = encode_frame(&mut self.seq, &mut self.fi, &mut fs);
452
          self.packet_data.extend(data);
fbossen's avatar
fbossen committed
453 454 455 456 457 458 459 460

          fs.rec.pad();

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

461 462 463 464 465 466 467
          if self.fi.show_frame {
            let data = self.packet_data.clone();
            self.packet_data = Vec::new();
            Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type })
          } else {
            Err(EncoderStatus::NeedMoreData)
          }
fbossen's avatar
fbossen committed
468 469 470 471 472 473
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      } else {
        Err(EncoderStatus::NeedMoreData)
      }
Luca Barbato's avatar
Luca Barbato committed
474 475 476 477
    }
  }

  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
478 479
    self.frame_q.insert(self.frame_count, None);
    self.frame_count = self.frame_count + 1;
Luca Barbato's avatar
Luca Barbato committed
480 481
  }
}
Luca Barbato's avatar
Luca Barbato committed
482 483 484 485 486 487

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