api.rs 10.1 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 82
    let fi = FrameInvariants::new(
      self.frame_info.width,
      self.frame_info.height,
      self.enc.clone()
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

    unsafe {
Luca Barbato's avatar
Luca Barbato committed
87 88
      av1_rtcd();
      aom_dsp_rtcd();
Luca Barbato's avatar
Luca Barbato committed
89 90
    }

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

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

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

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

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

150 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
  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
180
  pub fn frame_properties(&mut self, idx: u64) -> bool {
181
    let key_frame_interval: u64 = self.fi.config.key_frame_interval;
fbossen's avatar
fbossen committed
182

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

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

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

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

fbossen's avatar
fbossen committed
226 227 228 229 230 231 232 233
      let slot_idx = self.fi.order_hint % REF_FRAMES as u32;
      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
234
      } else {
235
        1 << slot_idx
Luca Barbato's avatar
Luca Barbato committed
236 237
      };

fbossen's avatar
fbossen committed
238 239 240 241
      let lvl = if !reorder {
        0
      } else if idx_in_group < pyramid_depth {
        idx_in_group
242
      } else {
fbossen's avatar
fbossen committed
243
        pyramid_depth - (idx_in_group - pyramid_depth + 1).trailing_zeros() as u64
244
      };
fbossen's avatar
fbossen committed
245 246
      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
247

fbossen's avatar
fbossen committed
248 249 250 251 252 253 254 255
      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;
256

fbossen's avatar
fbossen committed
257
      self.fi.primary_ref_frame = (ref_in_previous_group - LAST_FRAME) as u32;
Luca Barbato's avatar
Luca Barbato committed
258

259
      assert!(group_src_len <= REF_FRAMES as u64);
260
      for i in 0..INTER_REFS_PER_FRAME {
261 262 263
        self.fi.ref_frames[i] = (slot_idx as u8 +
        if i == second_ref_frame - LAST_FRAME {
          if lvl == 0 { (REF_FRAMES as u64 - 2) * group_src_len } else { group_src_len >> lvl }
fbossen's avatar
fbossen committed
264
        } else if i == ref_in_previous_group - LAST_FRAME {
265
          REF_FRAMES as u64 - group_src_len
266
        } else {
267 268
          REF_FRAMES as u64 - (group_src_len >> lvl)
        } as u8) & 7;
269 270
      }

fbossen's avatar
fbossen committed
271 272 273 274 275
      self.fi.reference_mode = if multiref && reorder && idx_in_group != 0 {
        ReferenceMode::SELECT
      } else {
        ReferenceMode::SINGLE
      };
fbossen's avatar
fbossen committed
276
      self.fi.number = segment_idx * key_frame_interval + self.fi.order_hint as u64;
fbossen's avatar
fbossen committed
277
      self.fi.me_range_scale = (group_src_len >> lvl) as u8;
fbossen's avatar
fbossen committed
278
    }
Luca Barbato's avatar
Luca Barbato committed
279

fbossen's avatar
fbossen committed
280 281
    true
  }
Luca Barbato's avatar
Luca Barbato committed
282

fbossen's avatar
fbossen committed
283 284 285 286 287 288
  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
289

fbossen's avatar
fbossen committed
290 291
    if self.fi.show_existing_frame {
      self.idx = self.idx + 1;
Luca Barbato's avatar
Luca Barbato committed
292

293
      let mut fs = FrameState::new(&self.fi);
fbossen's avatar
fbossen committed
294 295

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

fbossen's avatar
fbossen committed
297 298
      // 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
299

fbossen's avatar
fbossen committed
300
      Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type })
Luca Barbato's avatar
Luca Barbato committed
301
    } else {
fbossen's avatar
fbossen committed
302 303 304 305
      if let Some(f) = self.frame_q.remove(&self.fi.number) {
        self.idx = self.idx + 1;

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

          let data = encode_frame(&mut self.seq, &mut self.fi, &mut fs);
309
          self.packet_data.extend(data);
fbossen's avatar
fbossen committed
310 311 312 313 314 315 316 317

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

318 319 320 321 322 323 324
          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
325 326 327 328 329 330
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      } else {
        Err(EncoderStatus::NeedMoreData)
      }
Luca Barbato's avatar
Luca Barbato committed
331 332 333 334
    }
  }

  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
335 336
    self.frame_q.insert(self.frame_count, None);
    self.frame_count = self.frame_count + 1;
Luca Barbato's avatar
Luca Barbato committed
337 338
  }
}
Luca Barbato's avatar
Luca Barbato committed
339 340 341 342 343 344

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