api.rs 10.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

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

fbossen's avatar
fbossen committed
227 228 229 230 231 232 233 234
      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
235
      } else {
236
        1 << slot_idx
Luca Barbato's avatar
Luca Barbato committed
237 238
      };

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

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

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

260
      assert!(group_src_len <= REF_FRAMES as u64);
261
      for i in 0..INTER_REFS_PER_FRAME {
262 263 264
        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
265
        } else if i == ref_in_previous_group - LAST_FRAME {
266
          REF_FRAMES as u64 - group_src_len
267
        } else {
268 269
          REF_FRAMES as u64 - (group_src_len >> lvl)
        } as u8) & 7;
270 271
      }

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

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

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

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

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

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

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

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

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

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

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

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

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

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