Commit 3202bf70 authored by Josh Holmer's avatar Josh Holmer Committed by Thomas Daede

Implement lookahead queue

Allows rav1e to read frames into the frame queue ahead of the current
one, to allow for future features to be added that depend on a lookahead
pass. Currently the lookahead distance is hardcoded to 10. In the
future, this would probably be configurable as a CLI option.
parent be245fc0
This diff is collapsed.
......@@ -17,6 +17,7 @@ use std::io::prelude::*;
use std::sync::Arc;
use std::time::Instant;
use y4m;
use y4m::Colorspace;
pub struct EncoderIO {
pub input: Box<dyn Read>,
......@@ -257,21 +258,9 @@ impl fmt::Display for FrameSummary {
}
}
// Encode and write a frame.
// Returns frame information in a `Result`.
pub fn process_frame(
ctx: &mut Context, output_file: &mut dyn Write,
y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
) -> Result<Vec<FrameSummary>, ()> {
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 y4m_color_space = y4m_dec.get_colorspace();
let bit_depth = y4m_color_space.get_bit_depth();
if ctx.needs_more_frames(ctx.get_frame_count()) {
fn read_frame_batch(ctx: &mut Context, y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>, y4m_details: Y4MDetails) {
loop {
if ctx.needs_more_lookahead() {
match y4m_dec.read_frame() {
Ok(y4m_frame) => {
let y4m_y = y4m_frame.get_y_plane();
......@@ -280,25 +269,26 @@ pub fn process_frame(
let mut input = ctx.new_frame();
{
let input = Arc::get_mut(&mut input).unwrap();
input.planes[0].copy_from_raw_u8(&y4m_y, width * y4m_bytes, y4m_bytes);
input.planes[0].copy_from_raw_u8(&y4m_y, y4m_details.width * y4m_details.bytes, y4m_details.bytes);
input.planes[1].copy_from_raw_u8(
&y4m_u,
width * y4m_bytes / 2,
y4m_bytes
y4m_details.width * y4m_details.bytes / 2,
y4m_details.bytes
);
input.planes[2].copy_from_raw_u8(
&y4m_v,
width * y4m_bytes / 2,
y4m_bytes
y4m_details.width * y4m_details.bytes / 2,
y4m_details.bytes
);
}
match y4m_bits {
match y4m_details.bits {
8 | 10 | 12 => {}
_ => panic!("unknown input bit depth!")
}
let _ = ctx.send_frame(input);
let _ = ctx.send_frame(Some(input));
continue;
}
_ => {
let frames_to_be_coded = ctx.get_frame_count();
......@@ -306,27 +296,66 @@ pub fn process_frame(
ctx.flush();
}
}
} else {
} else if !ctx.needs_more_frames(ctx.get_frame_count()) {
ctx.flush();
};
}
break;
}
}
#[derive(Debug, Clone, Copy)]
struct Y4MDetails {
width: usize,
height: usize,
bits: usize,
bytes: usize,
color_space: Colorspace,
bit_depth: usize,
}
impl Y4MDetails {
fn new(y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>) -> Self {
let width = y4m_dec.get_width();
let height = y4m_dec.get_height();
let bits = y4m_dec.get_bit_depth();
let bytes = y4m_dec.get_bytes_per_sample();
let color_space = y4m_dec.get_colorspace();
let bit_depth = color_space.get_bit_depth();
Y4MDetails {
width,
height,
bits,
bytes,
color_space,
bit_depth,
}
}
}
// Encode and write a frame.
// Returns frame information in a `Result`.
pub fn process_frame(
ctx: &mut Context, output_file: &mut dyn Write,
y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
) -> Result<Vec<FrameSummary>, ()> {
let y4m_details = Y4MDetails::new(y4m_dec);
let mut frame_summaries = Vec::new();
loop {
read_frame_batch(ctx, y4m_dec, y4m_details);
let pkt_wrapped = ctx.receive_packet();
match pkt_wrapped {
Ok(pkt) => {
if let Ok(pkt) = pkt_wrapped {
write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
if let Some(y4m_enc_uw) = y4m_enc.as_mut() {
if let Some(ref rec) = pkt.rec {
let pitch_y = if bit_depth > 8 { width * 2 } else { width };
let chroma_sampling_period = map_y4m_color_space(y4m_color_space).0.sampling_period();
let pitch_y = if y4m_details.bit_depth > 8 { y4m_details.width * 2 } else { y4m_details.width };
let chroma_sampling_period = map_y4m_color_space(y4m_details.color_space).0.sampling_period();
let (pitch_uv, height_uv) = (
pitch_y / chroma_sampling_period.0,
height / chroma_sampling_period.1
y4m_details.height / chroma_sampling_period.1
);
let (mut rec_y, mut rec_u, mut rec_v) = (
vec![128u8; pitch_y * height],
vec![128u8; pitch_y * y4m_details.height],
vec![128u8; pitch_uv * height_uv],
vec![128u8; pitch_uv * height_uv]
);
......@@ -342,7 +371,7 @@ pub fn process_frame(
.chunks(stride_y)
.zip(rec_y.chunks_mut(pitch_y))
{
if bit_depth > 8 {
if y4m_details.bit_depth > 8 {
unsafe {
line_out.copy_from_slice(slice::from_raw_parts::<u8>(
line.as_ptr() as (*const u8),
......@@ -360,7 +389,7 @@ pub fn process_frame(
.chunks(stride_u)
.zip(rec_u.chunks_mut(pitch_uv))
{
if bit_depth > 8 {
if y4m_details.bit_depth > 8 {
unsafe {
line_out.copy_from_slice(slice::from_raw_parts::<u8>(
line.as_ptr() as (*const u8),
......@@ -378,7 +407,7 @@ pub fn process_frame(
.chunks(stride_v)
.zip(rec_v.chunks_mut(pitch_uv))
{
if bit_depth > 8 {
if y4m_details.bit_depth > 8 {
unsafe {
line_out.copy_from_slice(slice::from_raw_parts::<u8>(
line.as_ptr() as (*const u8),
......@@ -396,10 +425,8 @@ pub fn process_frame(
y4m_enc_uw.write_frame(&rec_frame).unwrap();
}
}
ctx.garbage_collect(pkt.number);
frame_summaries.push(pkt.into());
},
_ => { break; }
}
}
Ok(frame_summaries)
}
......
......@@ -306,7 +306,7 @@ fn encode_decode(
let mut input = ctx.new_frame();
fill_frame(&mut ra, Arc::get_mut(&mut input).unwrap());
let _ = ctx.send_frame(input);
let _ = ctx.send_frame(Some(input));
let mut done = false;
let mut corrupted_count = 0;
......
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