api.rs 14.4 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, frames_to_be_coded: 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
  frame_count: u64,
202
  frames_to_be_coded: u64,
fbossen's avatar
fbossen committed
203
  idx: u64,
204 205
  frame_q: BTreeMap<u64, Option<Arc<Frame>>>, //    packet_q: VecDeque<Packet>
  packet_data: Vec<u8>
Luca Barbato's avatar
Luca Barbato committed
206 207 208 209 210 211 212 213 214
}

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

pub struct Packet {
  pub data: Vec<u8>,
fbossen's avatar
fbossen committed
222 223
  pub rec: Option<Frame>,
  pub number: u64,
224 225 226 227 228
  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
229 230 231 232 233 234 235
    write!(
      f,
      "Frame {} - {} - {} bytes",
      self.number,
      self.frame_type,
      self.data.len()
    )
236
  }
Luca Barbato's avatar
Luca Barbato committed
237 238 239 240 241 242 243 244 245 246 247
}

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

253 254 255 256 257 258 259 260 261 262 263 264
  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
  }

265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294
  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
295
  pub fn frame_properties(&mut self, idx: u64) -> bool {
296
    let key_frame_interval: u64 = self.fi.config.key_frame_interval;
fbossen's avatar
fbossen committed
297

Frank Bossen's avatar
Frank Bossen committed
298
    let reorder = !self.fi.config.low_latency;
299
    let multiref = reorder || self.fi.config.speed_settings.multiref;
fbossen's avatar
fbossen committed
300

301
    let pyramid_depth = if reorder { 2 } else { 0 };
fbossen's avatar
fbossen committed
302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326
    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
327

fbossen's avatar
fbossen committed
328 329
      self.fi.frame_type = FrameType::INTER;
      self.fi.intra_only = false;
fbossen's avatar
fbossen committed
330

fbossen's avatar
fbossen committed
331 332 333 334 335 336 337 338 339
      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;
      }
340

341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366
      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
367 368 369 370 371 372 373
      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
374
      } else {
375
        1 << slot_idx
Luca Barbato's avatar
Luca Barbato committed
376 377
      };

fbossen's avatar
fbossen committed
378 379
      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
380

fbossen's avatar
fbossen committed
381 382 383 384 385 386 387 388
      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;
389

390 391
      // 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
392

393
      for i in 0..INTER_REFS_PER_FRAME {
394 395 396 397 398 399
        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
          }
400
        } else {
401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424
          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
            }
          }
        }
425 426
      }

427 428 429 430 431
      self.fi.reference_mode = if multiref && reorder && idx_in_group != 0 {
        ReferenceMode::SELECT
      } else {
        ReferenceMode::SINGLE
      };
fbossen's avatar
fbossen committed
432
      self.fi.number = segment_idx * key_frame_interval + self.fi.order_hint as u64;
433
      self.fi.me_range_scale = (group_src_len >> lvl) as u8;
fbossen's avatar
fbossen committed
434
    }
Luca Barbato's avatar
Luca Barbato committed
435

fbossen's avatar
fbossen committed
436 437
    true
  }
Luca Barbato's avatar
Luca Barbato committed
438

fbossen's avatar
fbossen committed
439 440 441 442 443 444
  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
445

446 447 448 449 450
    if !self.needs_more_frames(self.fi.number) {
      self.idx = self.idx + 1;
      return Err(EncoderStatus::EnoughData)
    }

fbossen's avatar
fbossen committed
451 452
    if self.fi.show_existing_frame {
      self.idx = self.idx + 1;
Luca Barbato's avatar
Luca Barbato committed
453

454
      let mut fs = FrameState::new(&self.fi);
fbossen's avatar
fbossen committed
455 456

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

fbossen's avatar
fbossen committed
458 459
      // 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
460

fbossen's avatar
fbossen committed
461
      Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type })
Luca Barbato's avatar
Luca Barbato committed
462
    } else {
fbossen's avatar
fbossen committed
463 464 465 466
      if let Some(f) = self.frame_q.remove(&self.fi.number) {
        self.idx = self.idx + 1;

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

          let data = encode_frame(&mut self.seq, &mut self.fi, &mut fs);
470
          self.packet_data.extend(data);
fbossen's avatar
fbossen committed
471 472 473 474 475 476 477 478

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

479 480 481 482 483 484 485
          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
486 487 488 489 490 491
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      } else {
        Err(EncoderStatus::NeedMoreData)
      }
Luca Barbato's avatar
Luca Barbato committed
492 493 494 495
    }
  }

  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
496 497
    self.frame_q.insert(self.frame_count, None);
    self.frame_count = self.frame_count + 1;
Luca Barbato's avatar
Luca Barbato committed
498 499
  }
}
Luca Barbato's avatar
Luca Barbato committed
500 501 502 503 504 505

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