api.rs 11.5 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 36 37 38 39 40 41
  pub quantizer: usize,
  pub speed: usize,
  pub tune: Tune
}

impl Default for EncoderConfig {
  fn default() -> Self {
Frank Bossen's avatar
Frank Bossen committed
42
    EncoderConfig { key_frame_interval: 30, low_latency: true, quantizer: 100, speed: 0, tune: Tune::Psnr  }
43 44 45 46
  }
}

/// Frame-specific information
Luca Barbato's avatar
Luca Barbato committed
47
#[derive(Clone, Copy, Debug)]
48
pub struct FrameInfo {
Luca Barbato's avatar
Luca Barbato committed
49 50 51
  pub width: usize,
  pub height: usize,
  pub bit_depth: usize,
52 53 54 55 56 57 58
  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
59
  pub timebase: Rational,
Luca Barbato's avatar
Luca Barbato committed
60 61 62 63
  pub enc: EncoderConfig
}

impl Config {
Luca Barbato's avatar
Luca Barbato committed
64 65 66
  pub fn parse(&mut self, key: &str, value: &str) -> Result<(), EncoderStatus> {
    use self::EncoderStatus::*;
    match key {
67 68 69 70 71 72
      "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)?,
      "speed" => self.enc.speed = value.parse().map_err(|_e| ParseError)?,
      "tune" => self.enc.tune = value.parse().map_err(|_e| ParseError)?,
      _ => return Err(InvalidKey)
Luca Barbato's avatar
Luca Barbato committed
73 74 75 76 77
    }

    Ok(())
  }

Luca Barbato's avatar
Luca Barbato committed
78
  pub fn new_context(&self) -> Context {
79 80 81
    let fi = FrameInvariants::new(
      self.frame_info.width,
      self.frame_info.height,
82
      self.enc
Luca Barbato's avatar
Luca Barbato committed
83
    );
84
    let seq = Sequence::new(&self.frame_info);
Luca Barbato's avatar
Luca Barbato committed
85

86
    #[cfg(feature = "aom")]
Luca Barbato's avatar
Luca Barbato committed
87
    unsafe {
Luca Barbato's avatar
Luca Barbato committed
88 89
      av1_rtcd();
      aom_dsp_rtcd();
Luca Barbato's avatar
Luca Barbato committed
90 91
    }

92
    Context { fi, seq, frame_count: 0, idx: 0, frame_q: BTreeMap::new(), packet_data: Vec::new() }
Luca Barbato's avatar
Luca Barbato committed
93 94 95 96 97 98
  }
}

pub struct Context {
  fi: FrameInvariants,
  seq: Sequence,
Luca Barbato's avatar
Luca Barbato committed
99
  //    timebase: Rational,
fbossen's avatar
fbossen committed
100 101
  frame_count: u64,
  idx: u64,
102 103
  frame_q: BTreeMap<u64, Option<Arc<Frame>>>, //    packet_q: VecDeque<Packet>
  packet_data: Vec<u8>
Luca Barbato's avatar
Luca Barbato committed
104 105 106 107 108 109 110 111 112
}

#[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
113 114 115
  Failure,
  InvalidKey,
  ParseError
Luca Barbato's avatar
Luca Barbato committed
116 117 118 119
}

pub struct Packet {
  pub data: Vec<u8>,
fbossen's avatar
fbossen committed
120 121
  pub rec: Option<Frame>,
  pub number: u64,
122 123 124 125 126
  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
127 128 129 130 131 132 133
    write!(
      f,
      "Frame {} - {} - {} bytes",
      self.number,
      self.frame_type,
      self.data.len()
    )
134
  }
Luca Barbato's avatar
Luca Barbato committed
135 136 137 138 139 140 141 142 143 144 145
}

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
146 147
    self.frame_q.insert(self.frame_count, frame.into());
    self.frame_count = self.frame_count + 1;
Luca Barbato's avatar
Luca Barbato committed
148 149 150
    Ok(())
  }

151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180
  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
181
  pub fn frame_properties(&mut self, idx: u64) -> bool {
182
    let key_frame_interval: u64 = self.fi.config.key_frame_interval;
fbossen's avatar
fbossen committed
183

Frank Bossen's avatar
Frank Bossen committed
184
    let reorder = !self.fi.config.low_latency;
fbossen's avatar
fbossen committed
185 186
    let multiref = reorder || self.fi.config.speed <= 2;

fbossen's avatar
fbossen committed
187
    let pyramid_depth = if reorder { 2 } else { 0 };
fbossen's avatar
fbossen committed
188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212
    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
213

fbossen's avatar
fbossen committed
214 215
      self.fi.frame_type = FrameType::INTER;
      self.fi.intra_only = false;
fbossen's avatar
fbossen committed
216

fbossen's avatar
fbossen committed
217 218 219 220 221 222 223 224 225
      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;
      }
226

227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252
      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
253 254 255 256 257 258 259
      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
260
      } else {
261
        1 << slot_idx
Luca Barbato's avatar
Luca Barbato committed
262 263
      };

fbossen's avatar
fbossen committed
264 265
      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
266

fbossen's avatar
fbossen committed
267 268 269 270 271 272 273 274
      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;
275

276 277
      // 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
278

279
      for i in 0..INTER_REFS_PER_FRAME {
280 281 282 283 284 285
        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
          }
286
        } else {
287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310
          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
            }
          }
        }
311 312
      }

fbossen's avatar
fbossen committed
313 314 315 316 317
      self.fi.reference_mode = if multiref && reorder && idx_in_group != 0 {
        ReferenceMode::SELECT
      } else {
        ReferenceMode::SINGLE
      };
fbossen's avatar
fbossen committed
318
      self.fi.number = segment_idx * key_frame_interval + self.fi.order_hint as u64;
fbossen's avatar
fbossen committed
319
      self.fi.me_range_scale = (group_src_len >> lvl) as u8;
fbossen's avatar
fbossen committed
320
    }
Luca Barbato's avatar
Luca Barbato committed
321

fbossen's avatar
fbossen committed
322 323
    true
  }
Luca Barbato's avatar
Luca Barbato committed
324

fbossen's avatar
fbossen committed
325 326 327 328 329 330
  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
331

fbossen's avatar
fbossen committed
332 333
    if self.fi.show_existing_frame {
      self.idx = self.idx + 1;
Luca Barbato's avatar
Luca Barbato committed
334

335
      let mut fs = FrameState::new(&self.fi);
fbossen's avatar
fbossen committed
336 337

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

fbossen's avatar
fbossen committed
339 340
      // 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
341

fbossen's avatar
fbossen committed
342
      Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type })
Luca Barbato's avatar
Luca Barbato committed
343
    } else {
fbossen's avatar
fbossen committed
344 345 346 347
      if let Some(f) = self.frame_q.remove(&self.fi.number) {
        self.idx = self.idx + 1;

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

          let data = encode_frame(&mut self.seq, &mut self.fi, &mut fs);
351
          self.packet_data.extend(data);
fbossen's avatar
fbossen committed
352 353 354 355 356 357 358 359

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

360 361 362 363 364 365 366
          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
367 368 369 370 371 372
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      } else {
        Err(EncoderStatus::NeedMoreData)
      }
Luca Barbato's avatar
Luca Barbato committed
373 374 375 376
    }
  }

  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
377 378
    self.frame_q.insert(self.frame_count, None);
    self.frame_count = self.frame_count + 1;
Luca Barbato's avatar
Luca Barbato committed
379 380
  }
}
Luca Barbato's avatar
Luca Barbato committed
381 382 383 384 385 386

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