api.rs 9.96 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 context::CDFContext;
Luca Barbato's avatar
Luca Barbato committed
11
use encoder::*;
12
use partition::*;
Luca Barbato's avatar
Luca Barbato committed
13

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

// TODO: use the num crate?
#[derive(Clone, Copy, Debug)]
20
#[repr(C)]
Luca Barbato's avatar
Luca Barbato committed
21 22 23 24 25 26 27 28 29 30 31
pub struct Ratio {
  pub num: usize,
  pub den: usize
}

impl Ratio {
  pub fn new(num: usize, den: usize) -> Self {
    Ratio { num, den }
  }
}

32 33 34 35 36 37 38 39 40 41 42 43 44 45
#[derive(Copy, Clone, Debug)]
pub struct EncoderConfig {
  pub quantizer: usize,
  pub speed: usize,
  pub tune: Tune
}

impl Default for EncoderConfig {
  fn default() -> Self {
    EncoderConfig { quantizer: 100, speed: 0, tune: Tune::Psnr }
  }
}

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

impl Config {
Luca Barbato's avatar
Luca Barbato committed
63 64 65 66 67 68 69 70 71 72 73 74
  pub fn parse(&mut self, key: &str, value: &str) -> Result<(), EncoderStatus> {
    use self::EncoderStatus::*;
    match key {
        "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)
    }

    Ok(())
  }

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

    unsafe {
Luca Barbato's avatar
Luca Barbato committed
84 85
      av1_rtcd();
      aom_dsp_rtcd();
Luca Barbato's avatar
Luca Barbato committed
86 87
    }

fbossen's avatar
fbossen committed
88
    Context { fi, seq, frame_count: 0, idx: 0, frame_q: BTreeMap::new() }
Luca Barbato's avatar
Luca Barbato committed
89 90 91 92 93 94 95
  }
}

pub struct Context {
  fi: FrameInvariants,
  seq: Sequence,
  //    timebase: Ratio,
fbossen's avatar
fbossen committed
96 97 98
  frame_count: u64,
  idx: u64,
  frame_q: BTreeMap<u64, Option<Arc<Frame>>> //    packet_q: VecDeque<Packet>
Luca Barbato's avatar
Luca Barbato committed
99 100 101 102 103 104 105 106 107
}

#[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
108 109 110
  Failure,
  InvalidKey,
  ParseError
Luca Barbato's avatar
Luca Barbato committed
111 112 113 114
}

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

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
141 142
    self.frame_q.insert(self.frame_count, frame.into());
    self.frame_count = self.frame_count + 1;
Luca Barbato's avatar
Luca Barbato committed
143 144 145
    Ok(())
  }

146 147 148 149 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
  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
176 177 178 179 180 181
  pub fn frame_properties(&mut self, idx: u64) -> bool {
    let key_frame_interval: u64 = 30;

    let reorder = false;
    let multiref = reorder || self.fi.config.speed <= 2;

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

fbossen's avatar
fbossen committed
209 210
      self.fi.frame_type = FrameType::INTER;
      self.fi.intra_only = false;
fbossen's avatar
fbossen committed
211

fbossen's avatar
fbossen committed
212 213 214 215 216 217 218 219 220
      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;
      }
221

fbossen's avatar
fbossen committed
222 223 224 225 226 227 228 229
      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
230
      } else {
231
        1 << slot_idx
Luca Barbato's avatar
Luca Barbato committed
232 233
      };

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

fbossen's avatar
fbossen committed
244 245 246 247 248 249 250 251
      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;
252

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

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

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

fbossen's avatar
fbossen committed
276 277
    true
  }
Luca Barbato's avatar
Luca Barbato committed
278

fbossen's avatar
fbossen committed
279 280 281 282 283 284
  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
285

fbossen's avatar
fbossen committed
286 287
    if self.fi.show_existing_frame {
      self.idx = self.idx + 1;
Luca Barbato's avatar
Luca Barbato committed
288

fbossen's avatar
fbossen committed
289 290 291 292 293 294 295 296 297
      let mut fs = FrameState {
        input: Arc::new(Frame::new(self.fi.padded_w, self.fi.padded_h)), // dummy
        rec: Frame::new(self.fi.padded_w, self.fi.padded_h),
        qc: Default::default(),
        cdfs: CDFContext::new(0),
        deblock: Default::default(),
      };

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

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

fbossen's avatar
fbossen committed
302
      Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type })
Luca Barbato's avatar
Luca Barbato committed
303
    } else {
fbossen's avatar
fbossen committed
304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331
      if let Some(f) = self.frame_q.remove(&self.fi.number) {
        self.idx = self.idx + 1;

        if let Some(frame) = f {
          let mut fs = FrameState {
            input: frame,
            rec: Frame::new(self.fi.padded_w, self.fi.padded_h),
            qc: Default::default(),
            cdfs: CDFContext::new(0),
            deblock: Default::default(),
          };

          let data = encode_frame(&mut self.seq, &mut self.fi, &mut fs);

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

          Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type })
        } 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)
  }
}