api.rs 10.3 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
    }

88
    Context { fi, seq, frame_count: 0, idx: 0, frame_q: BTreeMap::new(), packet_data: Vec::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
  frame_count: u64,
  idx: u64,
98 99
  frame_q: BTreeMap<u64, Option<Arc<Frame>>>, //    packet_q: VecDeque<Packet>
  packet_data: Vec<u8>
Luca Barbato's avatar
Luca Barbato committed
100 101 102 103 104 105 106 107 108
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

fbossen's avatar
fbossen committed
290 291 292 293 294 295 296 297 298
      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
299

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

fbossen's avatar
fbossen committed
303
      Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type })
Luca Barbato's avatar
Luca Barbato committed
304
    } else {
fbossen's avatar
fbossen committed
305 306 307 308 309 310 311 312 313 314 315 316 317
      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);
318
          self.packet_data.extend(data);
fbossen's avatar
fbossen committed
319 320 321 322 323 324 325 326

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

327 328 329 330 331 332 333
          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
334 335 336 337 338 339
        } else {
          Err(EncoderStatus::NeedMoreData)
        }
      } else {
        Err(EncoderStatus::NeedMoreData)
      }
Luca Barbato's avatar
Luca Barbato committed
340 341 342 343
    }
  }

  pub fn flush(&mut self) {
fbossen's avatar
fbossen committed
344 345
    self.frame_q.insert(self.frame_count, None);
    self.frame_count = self.frame_count + 1;
Luca Barbato's avatar
Luca Barbato committed
346 347
  }
}
Luca Barbato's avatar
Luca Barbato committed
348 349 350 351 352 353

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