Commit f593320a authored by Luca Barbato's avatar Luca Barbato Committed by Luca Barbato

Initial user api

parent cf90a229
use encoder::*;
use context::CDFContext;
use partition::LAST_FRAME;
use std::collections::VecDeque;
use std::sync::Arc;
// TODO: use the num crate?
#[derive(Clone, Copy, Debug)]
pub struct Ratio {
pub num: usize,
pub den: usize
}
impl Ratio {
pub fn new(num: usize, den: usize) -> Self {
Ratio { num, den }
}
}
/// Here we store all the information we might receive from the cli
#[derive(Clone, Copy, Debug)]
pub struct Config {
pub width: usize,
pub height: usize,
pub bit_depth: usize,
pub chroma_sampling: ChromaSampling,
pub timebase: Ratio,
pub enc: EncoderConfig
}
impl Config {
pub fn new_context(&self) -> Context {
let fi = FrameInvariants::new(self.width, self.height, self.enc.clone());
let seq = Sequence::new(
self.width,
self.height,
self.bit_depth,
self.chroma_sampling
);
unsafe {
av1_rtcd();
aom_dsp_rtcd();
}
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,
///
Failure
}
pub struct Packet {
pub data: Vec<u8>,
pub rec: Frame,
pub number: usize
}
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),
};
self.fi.frame_type = if self.fi.number % 30 == 0 {
FrameType::KEY
} else {
FrameType::INTER
};
self.fi.refresh_frame_flags = if self.fi.frame_type == FrameType::KEY {
ALL_REF_FRAMES_MASK
} else {
1
};
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);
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
} else {
self.fi.config.quantizer.max(1).min(255)
} as u8;
self.fi.primary_ref_frame = if self.fi.intra_only || self.fi.error_resilient {
PRIMARY_REF_NONE
} else {
(LAST_FRAME - LAST_FRAME) as u32
};
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);
Ok(Packet { data, rec, number })
} else {
unimplemented!("Flushing not implemented")
}
}
pub fn flush(&mut self) {
self.frame_q.push_back(None);
}
}
......@@ -94,7 +94,7 @@ impl FromCli for EncoderConfig {
}
/// Encode and write a frame.
pub fn process_frame(sequence: &mut Sequence, fi: &mut FrameInvariants,
pub fn process_frame(ctx: &mut Context,
output_file: &mut dyn Write,
y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>) -> bool {
......@@ -102,32 +102,20 @@ pub fn process_frame(sequence: &mut Sequence, fi: &mut FrameInvariants,
av1_rtcd();
aom_dsp_rtcd();
}
let width = fi.width;
let height = fi.height;
let width = y4m_dec.get_width();
let height = y4m_dec.get_height();
let y4m_bits = y4m_dec.get_bit_depth();
let y4m_bytes = y4m_dec.get_bytes_per_sample();
let csp = y4m_dec.get_colorspace();
match csp {
y4m::Colorspace::C420 |
y4m::Colorspace::C420jpeg |
y4m::Colorspace::C420paldv |
y4m::Colorspace::C420mpeg2 |
y4m::Colorspace::C420p10 |
y4m::Colorspace::C420p12 => {},
_ => {
panic!("Colorspace {:?} is not supported yet.", csp);
},
}
let bit_depth = y4m_dec.get_colorspace().get_bit_depth();
match y4m_dec.read_frame() {
Ok(y4m_frame) => {
let y4m_y = y4m_frame.get_y_plane();
let y4m_u = y4m_frame.get_u_plane();
let y4m_v = y4m_frame.get_v_plane();
eprintln!("{}", fi);
let mut fs = FrameState::new(&fi);
let mut input = ctx.new_frame();
{
let input = Arc::get_mut(&mut fs.input).unwrap();
let input = Arc::get_mut(&mut input).unwrap();
input.planes[0].copy_from_raw_u8(&y4m_y, width * y4m_bytes, y4m_bytes);
input.planes[1].copy_from_raw_u8(&y4m_u, width * y4m_bytes / 2, y4m_bytes);
input.planes[2].copy_from_raw_u8(&y4m_v, width * y4m_bytes / 2, y4m_bytes);
......@@ -138,10 +126,12 @@ pub fn process_frame(sequence: &mut Sequence, fi: &mut FrameInvariants,
_ => panic! ("unknown input bit depth!"),
}
let packet = encode_frame(sequence, fi, &mut fs);
write_ivf_frame(output_file, fi.number, packet.as_ref());
let _ = ctx.send_frame(input);
let pkt = ctx.receive_packet().unwrap();
write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
if let Some(mut y4m_enc) = y4m_enc {
let pitch_y = if sequence.bit_depth > 8 {
let pitch_y = if bit_depth > 8 {
width * 2
} else {
width
......@@ -154,12 +144,12 @@ pub fn process_frame(sequence: &mut Sequence, fi: &mut FrameInvariants,
vec![128u8; pitch_uv * (height / 2)]);
let (stride_y, stride_u, stride_v) = (
fs.rec.planes[0].cfg.stride,
fs.rec.planes[1].cfg.stride,
fs.rec.planes[2].cfg.stride);
pkt.rec.planes[0].cfg.stride,
pkt.rec.planes[1].cfg.stride,
pkt.rec.planes[2].cfg.stride);
for (line, line_out) in fs.rec.planes[0].data_origin().chunks(stride_y).zip(rec_y.chunks_mut(pitch_y)) {
if sequence.bit_depth > 8 {
for (line, line_out) in pkt.rec.planes[0].data_origin().chunks(stride_y).zip(rec_y.chunks_mut(pitch_y)) {
if bit_depth > 8 {
unsafe {
line_out.copy_from_slice(
slice::from_raw_parts::<u8>(line.as_ptr() as (*const u8), pitch_y));
......@@ -169,8 +159,8 @@ pub fn process_frame(sequence: &mut Sequence, fi: &mut FrameInvariants,
&line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_y]);
}
}
for (line, line_out) in fs.rec.planes[1].data_origin().chunks(stride_u).zip(rec_u.chunks_mut(pitch_uv)) {
if sequence.bit_depth > 8 {
for (line, line_out) in pkt.rec.planes[1].data_origin().chunks(stride_u).zip(rec_u.chunks_mut(pitch_uv)) {
if bit_depth > 8 {
unsafe {
line_out.copy_from_slice(
slice::from_raw_parts::<u8>(line.as_ptr() as (*const u8), pitch_uv));
......@@ -180,8 +170,8 @@ pub fn process_frame(sequence: &mut Sequence, fi: &mut FrameInvariants,
&line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]);
}
}
for (line, line_out) in fs.rec.planes[2].data_origin().chunks(stride_v).zip(rec_v.chunks_mut(pitch_uv)) {
if sequence.bit_depth > 8 {
for (line, line_out) in pkt.rec.planes[2].data_origin().chunks(stride_v).zip(rec_v.chunks_mut(pitch_uv)) {
if bit_depth > 8 {
unsafe {
line_out.copy_from_slice(
slice::from_raw_parts::<u8>(line.as_ptr() as (*const u8), pitch_uv));
......@@ -196,8 +186,6 @@ pub fn process_frame(sequence: &mut Sequence, fi: &mut FrameInvariants,
y4m_enc.write_frame(&rec_frame).unwrap();
}
fs.rec.pad();
update_rec_buffer(fi, fs);
true
},
_ => false
......
......@@ -14,16 +14,17 @@ extern crate y4m;
mod common;
use common::*;
use rav1e::partition::LAST_FRAME;
use rav1e::*;
fn main() {
let (mut io, config) = EncoderConfig::from_cli();
let (mut io, enc) = EncoderConfig::from_cli();
let mut y4m_dec = y4m::decode(&mut io.input).unwrap();
let width = y4m_dec.get_width();
let height = y4m_dec.get_height();
let framerate = y4m_dec.get_framerate();
let color_space = y4m_dec.get_colorspace();
let mut limit = enc.limit;
let mut y4m_enc = match io.rec.as_mut() {
Some(rec) => Some(
y4m::encode(width, height, framerate)
......@@ -34,8 +35,6 @@ fn main() {
None => None
};
let mut fi = FrameInvariants::new(width, height, config);
let chroma_sampling = match color_space {
y4m::Colorspace::C420
| y4m::Colorspace::C420jpeg
......@@ -52,8 +51,12 @@ fn main() {
_ => panic!("Chroma sampling unknown for the specified color space.")
};
let mut sequence =
Sequence::new(width, height, color_space.get_bit_depth(), chroma_sampling);
let bit_depth = color_space.get_bit_depth();
let cfg = Config {width, height, bit_depth, chroma_sampling, timebase: Ratio::new(framerate.den, framerate.num), enc };
let mut ctx = cfg.new_context();
write_ivf_header(
&mut io.output,
width,
......@@ -63,40 +66,21 @@ fn main() {
);
loop {
fi.frame_type =
if fi.number % 30 == 0 { FrameType::KEY } else { FrameType::INTER };
fi.base_q_idx = if fi.frame_type == FrameType::KEY {
let q_boost = 15;
fi.config.quantizer.max(1 + q_boost).min(255 + q_boost) - q_boost
} else {
fi.config.quantizer.max(1).min(255)
} as u8;
fi.refresh_frame_flags =
if fi.frame_type == FrameType::KEY { ALL_REF_FRAMES_MASK } else { 1 };
fi.intra_only = fi.frame_type == FrameType::KEY
|| fi.frame_type == FrameType::INTRA_ONLY;
fi.primary_ref_frame = if fi.intra_only || fi.error_resilient {
PRIMARY_REF_NONE
} else {
(LAST_FRAME - LAST_FRAME) as u32
};
if !process_frame(
&mut sequence,
&mut fi,
&mut ctx,
&mut io.output,
&mut y4m_dec,
y4m_enc.as_mut()
y4m_enc.as_mut(),
) {
break;
}
fi.number += 1;
//fi.show_existing_frame = fi.number % 2 == 1;
if fi.number == config.limit {
limit -= 1;
if limit == 0 {
break;
}
io.output.flush().unwrap();
}
}
......@@ -22,26 +22,48 @@ use rustyline::error::ReadlineError;
use rustyline::Editor;
fn main() {
let (mut io, config) = EncoderConfig::from_cli();
let (mut io, enc) = EncoderConfig::from_cli();
let mut y4m_dec = y4m::decode(&mut io.input).unwrap();
let width = y4m_dec.get_width();
let height = y4m_dec.get_height();
let framerate = y4m_dec.get_framerate();
let color_space = y4m_dec.get_colorspace();
let mut y4m_enc = match io.rec.as_mut() {
Some(rec) =>
Some(y4m::encode(width, height, framerate).write_header(rec).unwrap()),
None => None
};
let mut fi = FrameInvariants::new(width, height, &config);
let mut sequence = Sequence::new(width, height);
write_ivf_header(
&mut io.output,
fi.padded_w,
fi.padded_h,
width,
height,
framerate.num,
framerate.den
);
let chroma_sampling = match color_space {
y4m::Colorspace::C420 |
y4m::Colorspace::C420jpeg |
y4m::Colorspace::C420paldv |
y4m::Colorspace::C420mpeg2 |
y4m::Colorspace::C420p10 |
y4m::Colorspace::C420p12 => ChromaSampling::Cs420,
y4m::Colorspace::C422 |
y4m::Colorspace::C422p10 |
y4m::Colorspace::C422p12 => ChromaSampling::Cs422,
y4m::Colorspace::C444 |
y4m::Colorspace::C444p10 |
y4m::Colorspace::C444p12 => ChromaSampling::Cs444,
_ => {
panic!("Chroma sampling unknown for the specified color space.")
}
};
let bit_depth = color_space.get_bit_depth();
let cfg = Config {width, height, bit_depth, chroma_sampling, timebase: Ratio::new(framerate.den, framerate.num), enc };
let mut ctx = cfg.new_context();
let mut rl = Editor::<()>::new();
let _ = rl.load_history(".rav1e-history");
let mut last_rec: Option<Frame> = None;
......@@ -52,17 +74,16 @@ fn main() {
rl.add_history_entry(&line);
match line.split_whitespace().next() {
Some("process_frame") => {
process_frame(
&mut sequence,
&mut fi,
if !process_frame(
&mut ctx,
&mut io.output,
&mut y4m_dec,
y4m_enc.as_mut()
);
fi.number += 1;
if fi.number == config.limit {
y4m_enc.as_mut(),
) {
break;
}
io.output.flush().unwrap();
}
Some("quit") => break,
Some("exit") => break,
......
......@@ -32,6 +32,9 @@ pub mod cdef;
pub mod encoder;
pub mod me;
mod api;
pub use api::*;
pub use encoder::*;
// #[cfg(test)]
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment