api.rs 6.28 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 14

use std::collections::VecDeque;
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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105
    }

    Context { fi, seq, frame_q: VecDeque::new() }
  }
}

pub struct Context {
  fi: FrameInvariants,
  seq: Sequence,
  //    timebase: Ratio,
  frame_q: VecDeque<Option<Arc<Frame>>> //    packet_q: VecDeque<Packet>
}

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

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

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>>>
  {
    self.frame_q.push_back(frame.into());
    Ok(())
  }

  pub fn receive_packet(&mut self) -> Result<Packet, EncoderStatus> {
    let f = self.frame_q.pop_front().ok_or(EncoderStatus::NeedMoreData)?;
    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),
Luca Barbato's avatar
Luca Barbato committed
151
        deblock: Default::default()
Luca Barbato's avatar
Luca Barbato committed
152 153
      };

fbossen's avatar
fbossen committed
154 155 156 157 158
      let frame_number_in_segment = self.fi.number % 30;

      self.fi.order_hint = frame_number_in_segment as u32;

      self.fi.frame_type = if frame_number_in_segment == 0 {
Luca Barbato's avatar
Luca Barbato committed
159 160 161 162 163
        FrameType::KEY
      } else {
        FrameType::INTER
      };

164 165
      let slot_idx = frame_number_in_segment % REF_FRAMES as u64;

Luca Barbato's avatar
Luca Barbato committed
166 167 168
      self.fi.refresh_frame_flags = if self.fi.frame_type == FrameType::KEY {
        ALL_REF_FRAMES_MASK
      } else {
169
        1 << slot_idx
Luca Barbato's avatar
Luca Barbato committed
170 171 172 173 174 175 176
      };

      self.fi.intra_only = self.fi.frame_type == FrameType::KEY
        || self.fi.frame_type == FrameType::INTRA_ONLY;
      // self.fi.use_prev_frame_mvs =
      //  !(self.fi.intra_only || self.fi.error_resilient);

177 178 179 180 181 182 183 184 185 186
      let use_multiple_ref_frames = self.fi.config.speed <= 2;

      let log_boost_frequency = if use_multiple_ref_frames {
        2 // Higher quality frame every 4 frames
      } else {
        0 // No boosting with single reference frame
      };

      assert!(log_boost_frequency >= 0 && log_boost_frequency <= 2);
      let boost_frequency = 1 << log_boost_frequency;
Luca Barbato's avatar
Luca Barbato committed
187 188 189
      self.fi.base_q_idx = if self.fi.frame_type == FrameType::KEY {
        let q_boost = 15;
        self.fi.config.quantizer.max(1 + q_boost).min(255 + q_boost) - q_boost
190
      } else if slot_idx & (boost_frequency - 1) == 0 {
Luca Barbato's avatar
Luca Barbato committed
191
        self.fi.config.quantizer.max(1).min(255)
192 193 194
      } else {
        let q_drop = 15;
        self.fi.config.quantizer.min(255 - q_drop) + q_drop
Luca Barbato's avatar
Luca Barbato committed
195 196
      } as u8;

197
      let first_ref_frame = LAST_FRAME;
Luca Barbato's avatar
Luca Barbato committed
198 199
      let second_ref_frame =
        if use_multiple_ref_frames { ALTREF_FRAME } else { NONE_FRAME };
200

Luca Barbato's avatar
Luca Barbato committed
201 202 203 204 205 206
      self.fi.primary_ref_frame =
        if self.fi.intra_only || self.fi.error_resilient {
          PRIMARY_REF_NONE
        } else {
          (first_ref_frame - LAST_FRAME) as u32
        };
Luca Barbato's avatar
Luca Barbato committed
207

208 209 210 211 212 213 214 215
      for i in 0..INTER_REFS_PER_FRAME {
        self.fi.ref_frames[i] = if i == second_ref_frame - LAST_FRAME {
          (REF_FRAMES + slot_idx as usize - 2) & boost_frequency as usize
        } else {
          (REF_FRAMES + slot_idx as usize - 1) & (REF_FRAMES - 1)
        };
      }

Luca Barbato's avatar
Luca Barbato committed
216 217 218 219 220 221 222 223 224 225 226 227 228
      let data = encode_frame(&mut self.seq, &mut self.fi, &mut fs);

      let number = self.fi.number as usize;

      self.fi.number += 1;

      fs.rec.pad();

      // TODO avoid the clone by having rec Arc.
      let rec = fs.rec.clone();

      update_rec_buffer(&mut self.fi, fs);

229
      Ok(Packet { data, rec, number, frame_type: self.fi.frame_type })
Luca Barbato's avatar
Luca Barbato committed
230 231 232 233 234 235 236 237 238
    } else {
      unimplemented!("Flushing not implemented")
    }
  }

  pub fn flush(&mut self) {
    self.frame_q.push_back(None);
  }
}
Luca Barbato's avatar
Luca Barbato committed
239 240 241 242 243 244

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