Commit 8057ee7c authored by Romain Vimont's avatar Romain Vimont Committed by Luca Barbato

Make plane generic over pixel component type

In order to support both u8 and u16 for plane components, make the Plane
structure generic over the component type. As a consequence, many other
structures and functions also become generic.

Some functions are not u8-compatible yet, although they have been make
generic over the component type to make the compilation work. They
assert that the size of the generic parameter is 16 bits wide.

For this reason, the root context structure is unconditionally created
as Context<u16> for now.
parent bc358bf7
......@@ -49,7 +49,7 @@ fn write_b_bench(b: &mut Bencher, tx_size: TxSize, qindex: usize) {
..Default::default()
};
let sequence = Sequence::new(&Default::default());
let mut fi = FrameInvariants::new(config, sequence);
let mut fi = FrameInvariants::<u16>::new(config, sequence);
let mut w = ec::WriterEncoder::new();
let fc = CDFContext::new(fi.base_q_idx);
let bc = BlockContext::new(fi.sb_width * 16, fi.sb_height * 16);
......@@ -120,7 +120,7 @@ fn cdef_frame_bench(b: &mut Bencher, width: usize, height: usize) {
..Default::default()
};
let sequence = Sequence::new(&Default::default());
let fi = FrameInvariants::new(config, sequence);
let fi = FrameInvariants::<u16>::new(config, sequence);
let mut bc = BlockContext::new(fi.sb_width * 16, fi.sb_height * 16);
let mut fs = FrameState::new(&fi);
......@@ -148,7 +148,7 @@ fn cfl_rdo_bench(b: &mut Bencher, bsize: BlockSize) {
..Default::default()
};
let sequence = Sequence::new(&Default::default());
let fi = FrameInvariants::new(config, sequence );
let fi = FrameInvariants::<u16>::new(config, sequence);
let mut fs = FrameState::new(&fi);
let offset = BlockOffset { x: 1, y: 1 };
b.iter(|| rdo_cfl_alpha(&mut fs, &offset, bsize, fi.sequence.bit_depth, fi.sequence.chroma_sampling))
......
......@@ -13,18 +13,19 @@ use crate::partition::BlockSize::*;
use crate::plane::*;
use rand::{ChaChaRng, Rng, SeedableRng};
use rav1e::me;
use rav1e::Pixel;
fn fill_plane(ra: &mut ChaChaRng, plane: &mut Plane) {
fn fill_plane<T: Pixel>(ra: &mut ChaChaRng, plane: &mut Plane<T>) {
let stride = plane.cfg.stride;
for row in plane.data.chunks_mut(stride) {
for pixel in row {
let v: u8 = ra.gen();
*pixel = v as u16;
*pixel = T::cast_from(v);
}
}
}
fn new_plane(ra: &mut ChaChaRng, width: usize, height: usize) -> Plane {
fn new_plane<T: Pixel>(ra: &mut ChaChaRng, width: usize, height: usize) -> Plane<T> {
let mut p = Plane::new(width, height, 0, 0, 128 + 8, 128 + 8);
fill_plane(ra, &mut p);
......@@ -39,8 +40,8 @@ fn bench_get_sad(b: &mut Bencher, bs: &BlockSize) {
let w = 640;
let h = 480;
let bit_depth = 10;
let input_plane = new_plane(&mut ra, w, h);
let rec_plane = new_plane(&mut ra, w, h);
let input_plane = new_plane::<u16>(&mut ra, w, h);
let rec_plane = new_plane::<u16>(&mut ra, w, h);
let po = PlaneOffset { x: 0, y: 0 };
let plane_org = input_plane.slice(&po);
......
......@@ -16,6 +16,7 @@ use crate::rate::FRAME_NSUBTYPES;
use crate::rate::FRAME_SUBTYPE_I;
use crate::rate::FRAME_SUBTYPE_P;
use crate::scenechange::SceneChangeDetector;
use crate::util::Pixel;
use self::EncoderStatus::*;
use std::{cmp, fmt, io};
......@@ -437,7 +438,7 @@ impl Config {
Ok(())
}
pub fn new_context(&self) -> Context {
pub fn new_context<T: Pixel>(&self) -> Context<T> {
#[cfg(feature = "aom")]
unsafe {
av1_rtcd();
......@@ -483,16 +484,16 @@ impl Config {
}
}
pub struct Context {
pub struct Context<T: Pixel> {
// timebase: Rational,
frame_count: u64,
limit: u64,
pub(crate) idx: u64,
frames_processed: u64,
/// Maps frame *number* to frames
frame_q: BTreeMap<u64, Option<Arc<Frame>>>, // packet_q: VecDeque<Packet>
frame_q: BTreeMap<u64, Option<Arc<Frame<T>>>>, // packet_q: VecDeque<Packet>
/// Maps frame *idx* to frame data
frame_data: BTreeMap<u64, FrameInvariants>,
frame_data: BTreeMap<u64, FrameInvariants<T>>,
/// A list of keyframe *numbers* in this encode. Needed so that we don't
/// need to keep all of the frame_data in memory for the whole life of the encode.
keyframes: BTreeSet<u64>,
......@@ -500,7 +501,7 @@ pub struct Context {
packet_data: Vec<u8>,
segment_start_idx: u64,
segment_start_frame: u64,
keyframe_detector: SceneChangeDetector,
keyframe_detector: SceneChangeDetector<T>,
pub config: Config,
rc_state: RCState,
maybe_prev_log_base_q: Option<i64>,
......@@ -519,16 +520,16 @@ pub enum EncoderStatus {
ParseError
}
pub struct Packet {
pub struct Packet<T: Pixel> {
pub data: Vec<u8>,
pub rec: Option<Frame>,
pub rec: Option<Frame<T>>,
pub number: u64,
pub frame_type: FrameType,
/// PSNR for Y, U, and V planes
pub psnr: Option<(f64, f64, f64)>,
}
impl fmt::Display for Packet {
impl<T: Pixel> fmt::Display for Packet<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
......@@ -540,8 +541,8 @@ impl fmt::Display for Packet {
}
}
impl Context {
pub fn new_frame(&self) -> Arc<Frame> {
impl<T: Pixel> Context<T> {
pub fn new_frame(&self) -> Arc<Frame<T>> {
Arc::new(Frame::new(
self.config.enc.width,
self.config.enc.height,
......@@ -551,7 +552,8 @@ impl Context {
pub fn send_frame<F>(&mut self, frame: F) -> Result<(), EncoderStatus>
where
F: Into<Option<Arc<Frame>>>
F: Into<Option<Arc<Frame<T>>>>,
T: Pixel,
{
let idx = self.frame_count;
self.frame_q.insert(idx, frame.into());
......@@ -559,7 +561,7 @@ impl Context {
Ok(())
}
pub fn get_frame(&self, frame_number: u64) -> Arc<Frame> {
pub fn get_frame(&self, frame_number: u64) -> Arc<Frame<T>> {
// Clones only the arc, so low cost overhead
self.frame_q.get(&frame_number).as_ref().unwrap().as_ref().unwrap().clone()
}
......@@ -628,7 +630,7 @@ impl Context {
end_of_subgop
}
fn build_frame_properties(&mut self, idx: u64) -> (FrameInvariants, bool) {
fn build_frame_properties(&mut self, idx: u64) -> (FrameInvariants<T>, bool) {
if idx == 0 {
let seq = Sequence::new(&self.config.enc);
......@@ -711,7 +713,7 @@ impl Context {
(fi, true)
}
pub fn receive_packet(&mut self) -> Result<Packet, EncoderStatus> {
pub fn receive_packet(&mut self) -> Result<Packet<T>, EncoderStatus> {
let idx = {
let mut idx = self.idx;
while !self.set_frame_properties(idx) {
......@@ -794,7 +796,7 @@ impl Context {
ret
}
fn finalize_packet(&mut self, rec: Option<Frame>, fi: &FrameInvariants) -> Result<Packet, EncoderStatus> {
fn finalize_packet(&mut self, rec: Option<Frame<T>>, fi: &FrameInvariants<T>) -> Result<Packet<T>, EncoderStatus> {
let data = self.packet_data.clone();
self.packet_data.clear();
if write_temporal_delimiter(&mut self.packet_data).is_err() {
......@@ -972,8 +974,8 @@ pub struct FirstPassFrame {
frame_type: FrameType,
}
impl From<&FrameInvariants> for FirstPassFrame {
fn from(fi: &FrameInvariants) -> FirstPassFrame {
impl<T: Pixel> From<&FrameInvariants<T>> for FirstPassFrame {
fn from(fi: &FrameInvariants<T>) -> FirstPassFrame {
FirstPassFrame {
number: fi.number,
frame_type: fi.frame_type,
......
......@@ -463,8 +463,8 @@ pub struct FrameSummary {
pub psnr: Option<(f64, f64, f64)>,
}
impl From<Packet> for FrameSummary {
fn from(packet: Packet) -> Self {
impl<T: Pixel> From<Packet<T>> for FrameSummary {
fn from(packet: Packet<T>) -> Self {
Self {
size: packet.data.len(),
number: packet.number,
......
......@@ -6,7 +6,7 @@ pub mod y4m;
pub trait Decoder {
fn get_video_details(&self) -> VideoDetails;
fn read_frame(&mut self, cfg: &VideoDetails) -> Result<Frame, DecodeError>;
fn read_frame<T: Pixel>(&mut self, cfg: &VideoDetails) -> Result<Frame<T>, DecodeError>;
}
#[derive(Debug)]
......
......@@ -7,6 +7,7 @@ use crate::decoder::VideoDetails;
use crate::ChromaSamplePosition;
use crate::ChromaSampling;
use crate::encoder::Frame;
use rav1e::*;
impl Decoder for y4m::Decoder<'_, Box<dyn Read>> {
fn get_video_details(&self) -> VideoDetails {
......@@ -28,11 +29,11 @@ impl Decoder for y4m::Decoder<'_, Box<dyn Read>> {
}
}
fn read_frame(&mut self, cfg: &VideoDetails) -> Result<Frame, DecodeError> {
fn read_frame<T: Pixel>(&mut self, cfg: &VideoDetails) -> Result<Frame<T>, DecodeError> {
let bytes = self.get_bytes_per_sample();
self.read_frame()
.map(|frame| {
let mut f = Frame::new(cfg.width, cfg.height, cfg.chroma_sampling);
let mut f: Frame<T> = Frame::new(cfg.width, cfg.height, cfg.chroma_sampling);
let (chroma_period, _) = cfg.chroma_sampling.sampling_period();
......
......@@ -10,10 +10,11 @@
use crate::decoder::VideoDetails;
use std::io::Write;
use std::slice;
use rav1e::*;
pub use ivf::*;
pub fn write_y4m_frame(y4m_enc: &mut y4m::Encoder<'_, Box<dyn Write>>, rec: &rav1e::Frame, y4m_details: VideoDetails) {
pub fn write_y4m_frame<T: Pixel>(y4m_enc: &mut y4m::Encoder<'_, Box<dyn Write>>, rec: &rav1e::Frame<T>, y4m_details: VideoDetails) {
let pitch_y = if y4m_details.bit_depth > 8 { y4m_details.width * 2 } else { y4m_details.width };
let chroma_sampling_period = y4m_details.chroma_sampling.sampling_period();
let (pitch_uv, height_uv) = (
......@@ -47,7 +48,7 @@ pub fn write_y4m_frame(y4m_enc: &mut y4m::Encoder<'_, Box<dyn Write>>, rec: &rav
}
} else {
line_out.copy_from_slice(
&line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_y]
&line.iter().map(|&v| u8::cast_from(v)).collect::<Vec<u8>>()[..pitch_y]
);
}
}
......@@ -65,7 +66,7 @@ pub fn write_y4m_frame(y4m_enc: &mut y4m::Encoder<'_, Box<dyn Write>>, rec: &rav
}
} else {
line_out.copy_from_slice(
&line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
&line.iter().map(|&v| u8::cast_from(v)).collect::<Vec<u8>>()[..pitch_uv]
);
}
}
......@@ -83,7 +84,7 @@ pub fn write_y4m_frame(y4m_enc: &mut y4m::Encoder<'_, Box<dyn Write>>, rec: &rav
}
} else {
line_out.copy_from_slice(
&line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
&line.iter().map(|&v| u8::cast_from(v)).collect::<Vec<u8>>()[..pitch_uv]
);
}
}
......
......@@ -29,7 +29,7 @@ use crate::decoder::VideoDetails;
use std::fs::File;
use std::io::BufWriter;
fn read_frame_batch<D: Decoder>(ctx: &mut Context, decoder: &mut D, video_info: VideoDetails) {
fn read_frame_batch<T: Pixel, D: Decoder>(ctx: &mut Context<T>, decoder: &mut D, video_info: VideoDetails) {
loop {
if ctx.needs_more_lookahead() {
match decoder.read_frame(&video_info) {
......@@ -58,8 +58,8 @@ fn read_frame_batch<D: Decoder>(ctx: &mut Context, decoder: &mut D, video_info:
// Encode and write a frame.
// Returns frame information in a `Result`.
fn process_frame(
ctx: &mut Context,
fn process_frame<T: Pixel>(
ctx: &mut Context<T>,
output_file: &mut dyn Write,
y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>,
......@@ -78,7 +78,7 @@ fn process_frame(
Ok(frame_summaries)
}
fn write_stats_file(ctx: &Context, filename: &Path) -> Result<(), io::Error> {
fn write_stats_file<T: Pixel>(ctx: &Context<T>, filename: &Path) -> Result<(), io::Error> {
let file = File::create(filename)?;
let writer = BufWriter::new(file);
serde_json::to_writer(writer, &ctx.first_pass_data).expect("Serialization should not fail");
......@@ -112,7 +112,8 @@ fn main() {
enc: cli.enc
};
let mut ctx = cfg.new_context();
// FIXME for now, unconditionally create Context<u16>
let mut ctx: Context<u16> = cfg.new_context();
let stderr = io::stderr();
let mut err = stderr.lock();
......
This diff is collapsed.
......@@ -25,7 +25,7 @@ use crate::lrf::*;
use crate::plane::*;
use crate::scan_order::*;
use crate::token_cdfs::*;
use crate::util::{clamp, msb};
use crate::util::{clamp, msb, Pixel};
use std::*;
......@@ -2137,12 +2137,12 @@ impl ContextWriter {
}
}
fn add_extra_mv_candidate(
fn add_extra_mv_candidate<T: Pixel>(
&self,
blk: &Block,
ref_frames: [usize; 2],
mv_stack: &mut Vec<CandidateMV>,
fi: &FrameInvariants,
fi: &FrameInvariants<T>,
is_compound: bool,
ref_id_count: &mut [usize; 2],
ref_id_mvs: &mut [[MotionVector; 2]; 2],
......@@ -2320,8 +2320,10 @@ impl ContextWriter {
}
}
fn setup_mvref_list(&mut self, bo: &BlockOffset, ref_frames: [usize; 2], mv_stack: &mut Vec<CandidateMV>,
bsize: BlockSize, fi: &FrameInvariants, is_compound: bool) -> usize {
fn setup_mvref_list<T: Pixel>(
&mut self, bo: &BlockOffset, ref_frames: [usize; 2], mv_stack: &mut Vec<CandidateMV>,
bsize: BlockSize, fi: &FrameInvariants<T>, is_compound: bool
) -> usize {
let (_rf, _rf_num) = self.get_mvref_ref_frames(INTRA_FRAME);
let target_n4_h = bsize.height_mi();
......@@ -2536,9 +2538,11 @@ impl ContextWriter {
mode_context
}
pub fn find_mvrefs(&mut self, bo: &BlockOffset, ref_frames: [usize; 2],
mv_stack: &mut Vec<CandidateMV>, bsize: BlockSize,
fi: &FrameInvariants, is_compound: bool) -> usize {
pub fn find_mvrefs<T: Pixel>(
&mut self, bo: &BlockOffset, ref_frames: [usize; 2],
mv_stack: &mut Vec<CandidateMV>, bsize: BlockSize,
fi: &FrameInvariants<T>, is_compound: bool
) -> usize {
assert!(ref_frames[0] != NONE_FRAME);
if ref_frames[0] < REF_FRAMES {
if ref_frames[0] != INTRA_FRAME {
......@@ -2753,7 +2757,7 @@ impl ContextWriter {
}
}
pub fn write_ref_frames(&mut self, w: &mut dyn Writer, fi: &FrameInvariants, bo: &BlockOffset) {
pub fn write_ref_frames<T: Pixel>(&mut self, w: &mut dyn Writer, fi: &FrameInvariants<T>, bo: &BlockOffset) {
let rf = self.bc.at(bo).ref_frames;
let sz = self.bc.at(bo).n4_w.min(self.bc.at(bo).n4_h);
......@@ -3008,8 +3012,9 @@ impl ContextWriter {
symbol_with_update!(self, w, coded_id as u32, &mut self.fc.spatial_segmentation_cdfs[cdf_index as usize]);
}
pub fn write_lrf(&mut self, w: &mut dyn Writer, fi: &FrameInvariants, rs: &mut RestorationState,
sbo: &SuperBlockOffset) {
pub fn write_lrf<T: Pixel>(
&mut self, w: &mut dyn Writer, fi: &FrameInvariants<T>, rs: &mut RestorationState, sbo: &SuperBlockOffset
) {
if !fi.allow_intrabc { // TODO: also disallow if lossless
for pli in 0..PLANES {
let code;
......
This diff is collapsed.
This diff is collapsed.
......@@ -12,6 +12,7 @@ use crate::context::*;
use crate::ec::*;
use crate::lrf::*;
use crate::partition::*;
use crate::util::Pixel;
use crate::SegmentationState;
use crate::DeblockState;
......@@ -135,37 +136,37 @@ pub trait UncompressedHeader {
fn write_metadata_obu(
&mut self, obu_meta_type: ObuMetaType, seq: Sequence
) -> io::Result<()>;
fn write_sequence_header_obu(
&mut self, fi: &mut FrameInvariants
fn write_sequence_header_obu<T: Pixel>(
&mut self, fi: &mut FrameInvariants<T>
) -> io::Result<()>;
fn write_frame_header_obu(
&mut self, fi: &FrameInvariants, fs: &FrameState
fn write_frame_header_obu<T: Pixel>(
&mut self, fi: &FrameInvariants<T>, fs: &FrameState<T>
) -> io::Result<()>;
fn write_sequence_header(
&mut self, fi: &mut FrameInvariants
fn write_sequence_header<T: Pixel>(
&mut self, fi: &mut FrameInvariants<T>
) -> io::Result<()>;
fn write_color_config(
&mut self, seq: &mut Sequence
) -> io::Result<()>;
// End of OBU Headers
fn write_frame_size(
&mut self, fi: &FrameInvariants
fn write_frame_size<T: Pixel>(
&mut self, fi: &FrameInvariants<T>
) -> io::Result<()>;
fn write_deblock_filter_a(
&mut self, fi: &FrameInvariants, deblock: &DeblockState
fn write_deblock_filter_a<T: Pixel>(
&mut self, fi: &FrameInvariants<T>, deblock: &DeblockState
) -> io::Result<()>;
fn write_deblock_filter_b(
&mut self, fi: &FrameInvariants, deblock: &DeblockState
fn write_deblock_filter_b<T: Pixel>(
&mut self, fi: &FrameInvariants<T>, deblock: &DeblockState
) -> io::Result<()>;
fn write_frame_cdef(
&mut self, fi: &FrameInvariants
fn write_frame_cdef<T: Pixel>(
&mut self, fi: &FrameInvariants<T>
) -> io::Result<()>;
fn write_frame_lrf(
&mut self, fi: &FrameInvariants, rs: &RestorationState
fn write_frame_lrf<T: Pixel>(
&mut self, fi: &FrameInvariants<T>, rs: &RestorationState
) -> io::Result<()>;
fn write_segment_data(
&mut self, fi: &FrameInvariants, segmentation: &SegmentationState
fn write_segment_data<T: Pixel>(
&mut self, fi: &FrameInvariants<T>, segmentation: &SegmentationState
) -> io::Result<()>;
fn write_delta_q(
&mut self, delta_q: i8
......@@ -237,8 +238,8 @@ impl<W: io::Write> UncompressedHeader for BitWriter<W, BigEndian> {
Ok(())
}
fn write_sequence_header_obu(
&mut self, fi: &mut FrameInvariants
fn write_sequence_header_obu<T: Pixel>(
&mut self, fi: &mut FrameInvariants<T>
) -> io::Result<()> {
self.write(3, fi.sequence.profile)?; // profile
self.write_bit(false)?; // still_picture
......@@ -262,8 +263,8 @@ impl<W: io::Write> UncompressedHeader for BitWriter<W, BigEndian> {
Ok(())
}
fn write_sequence_header(
&mut self, fi: &mut FrameInvariants
fn write_sequence_header<T: Pixel>(
&mut self, fi: &mut FrameInvariants<T>
) -> io::Result<()> {
self.write_frame_size(fi)?;
......@@ -405,8 +406,8 @@ impl<W: io::Write> UncompressedHeader for BitWriter<W, BigEndian> {
}
#[allow(unused)]
fn write_frame_header_obu(
&mut self, fi: &FrameInvariants, fs: &FrameState
fn write_frame_header_obu<T: Pixel>(
&mut self, fi: &FrameInvariants<T>, fs: &FrameState<T>
) -> io::Result<()> {
if fi.sequence.reduced_still_picture_hdr {
assert!(fi.show_existing_frame);
......@@ -740,7 +741,7 @@ impl<W: io::Write> UncompressedHeader for BitWriter<W, BigEndian> {
}
// End of OBU Headers
fn write_frame_size(&mut self, fi: &FrameInvariants) -> io::Result<()> {
fn write_frame_size<T: Pixel>(&mut self, fi: &FrameInvariants<T>) -> io::Result<()> {
// width_bits and height_bits will have to be moved to the sequence header OBU
// when we add support for it.
let width_bits = 32 - (fi.width as u32).leading_zeros();
......@@ -754,8 +755,8 @@ impl<W: io::Write> UncompressedHeader for BitWriter<W, BigEndian> {
Ok(())
}
fn write_deblock_filter_a(
&mut self, fi: &FrameInvariants, deblock: &DeblockState
fn write_deblock_filter_a<T: Pixel>(
&mut self, fi: &FrameInvariants<T>, deblock: &DeblockState
) -> io::Result<()> {
if fi.delta_q_present {
if !fi.allow_intrabc {
......@@ -769,8 +770,8 @@ impl<W: io::Write> UncompressedHeader for BitWriter<W, BigEndian> {
Ok(())
}
fn write_deblock_filter_b(
&mut self, fi: &FrameInvariants, deblock: &DeblockState
fn write_deblock_filter_b<T: Pixel>(
&mut self, fi: &FrameInvariants<T>, deblock: &DeblockState
) -> io::Result<()> {
assert!(deblock.levels[0] < 64);
self.write(6, deblock.levels[0])?; // loop deblocking filter level 0
......@@ -822,7 +823,7 @@ impl<W: io::Write> UncompressedHeader for BitWriter<W, BigEndian> {
Ok(())
}
fn write_frame_cdef(&mut self, fi: &FrameInvariants) -> io::Result<()> {
fn write_frame_cdef<T: Pixel>(&mut self, fi: &FrameInvariants<T>) -> io::Result<()> {
if fi.sequence.enable_cdef {
assert!(fi.cdef_damping >= 3);
assert!(fi.cdef_damping <= 6);
......@@ -840,8 +841,8 @@ impl<W: io::Write> UncompressedHeader for BitWriter<W, BigEndian> {
Ok(())
}
fn write_frame_lrf(
&mut self, fi: &FrameInvariants, rs: &RestorationState
fn write_frame_lrf<T: Pixel>(
&mut self, fi: &FrameInvariants<T>, rs: &RestorationState
) -> io::Result<()> {
if fi.sequence.enable_restoration && !fi.allow_intrabc {
// && !self.lossless
......@@ -882,8 +883,8 @@ impl<W: io::Write> UncompressedHeader for BitWriter<W, BigEndian> {
Ok(())
}
fn write_segment_data(
&mut self, fi: &FrameInvariants, segmentation: &SegmentationState
fn write_segment_data<T: Pixel>(
&mut self, fi: &FrameInvariants<T>, segmentation: &SegmentationState
) -> io::Result<()> {
self.write_bit(segmentation.enabled)?;
if segmentation.enabled {
......
......@@ -49,6 +49,7 @@ mod header;
pub use crate::api::*;
pub use crate::encoder::*;
pub use crate::header::*;
pub use crate::util::{CastFromPrimitive, Pixel};
#[cfg(all(test, feature="decode_test"))]
mod test_encode_decode_aom;
......
......@@ -20,6 +20,8 @@ use crate::plane::PlaneOffset;
use crate::plane::PlaneConfig;
use std::cmp;
use crate::util::clamp;
use crate::util::CastFromPrimitive;
use crate::util::Pixel;
pub const RESTORATION_TILESIZE_MAX_LOG2: usize = 8;
......@@ -71,12 +73,20 @@ impl RestorationFilter {
}
}
fn sgrproj_box_ab(af: &mut[i32; 64+2],
bf: &mut[i32; 64+2],
r: isize, eps: isize,
crop_w: usize, crop_h: usize, stripe_h: isize,
stripe_x: isize, stripe_y: isize,
cdeffed: &Plane, deblocked: &Plane, bit_depth: usize) {
fn sgrproj_box_ab<T: Pixel>(
af: &mut[i32; 64+2],
bf: &mut[i32; 64+2],
r: isize,
eps: isize,
crop_w: usize,
crop_h: usize,
stripe_h: isize,
stripe_x: isize,
stripe_y: isize,
cdeffed: &Plane<T>,
deblocked: &Plane<T>,
bit_depth: usize,
) {
let n = ((2*r + 1) * (2*r + 1)) as i32;
let one_over_n = ((1 << SGRPROJ_RECIP_BITS) + n/2 ) / n;
let n2e = n*n*eps as i32;
......@@ -87,7 +97,7 @@ fn sgrproj_box_ab(af: &mut[i32; 64+2],
let mut b:i32 = 0;
for yi in stripe_y+row-r..=stripe_y+row+r {
let src_plane: &Plane;
let src_plane: &Plane<T>;
let ly;
if yi < stripe_y {
ly = cmp::max(cmp::max(yi, 0), stripe_y - 2) as usize;
......@@ -101,17 +111,17 @@ fn sgrproj_box_ab(af: &mut[i32; 64+2],
}
for _xi in -r..-stripe_x {
let c = src_plane.p(0, ly) as i32;
let c: i32 = src_plane.p(0, ly).as_();
a += c*c;
b += c;
}
for xi in cmp::max(-r, -stripe_x)..xn {
let c = src_plane.p((xi + stripe_x) as usize, ly) as i32;
let c: i32 = src_plane.p((xi + stripe_x) as usize, ly).as_();
a += c*c;
b += c;
}
for _xi in xn..=r {
let c = src_plane.p(crop_w - 1, ly) as i32;
let c: i32 = src_plane.p(crop_w - 1, ly).as_();
a += c*c;
b += c;
}
......@@ -133,9 +143,16 @@ fn sgrproj_box_ab(af: &mut[i32; 64+2],
}
}
fn sgrproj_box_f(af: &[&[i32; 64+2]; 3], bf: &[&[i32; 64+2]; 3], f: &mut[i32; 64],
x: usize, y: isize, h: isize, cdeffed: &Plane, pass: usize) {
fn sgrproj_box_f<T: Pixel>(
af: &[&[i32; 64+2]; 3],
bf: &[&[i32; 64+2]; 3],
f: &mut[i32; 64],
x: usize,
y: isize,
h: isize,
cdeffed: &Plane<T>,
pass: usize,
) {
for i in 0..h as usize {
let shift = if pass == 0 && (i&1) == 1 {4} else {5} + SGRPROJ_SGR_BITS - SGRPROJ_RST_BITS;
let mut a = 0;
......@@ -163,17 +180,25 @@ fn sgrproj_box_f(af: &[&[i32; 64+2]; 3], bf: &[&[i32; 64+2]; 3], f: &mut[i32; 64
b += weight * bf[dx][i+dy];
}
}
let v = a * cdeffed.p(x, (y+i as isize) as usize) as i32 + b;
let v = a * i32::cast_from(cdeffed.p(x, (y+i as isize) as usize)) + b;
f[i as usize] = v + (1 << shift >> 1) >> shift;
}
}
fn sgrproj_stripe_rdu(set: u8, xqd: [i8; 2], fi: &FrameInvariants,
crop_w: usize, crop_h: usize,
stripe_w: usize, stripe_h: isize,
stripe_x: usize, stripe_y: isize,
cdeffed: &Plane, deblocked: &Plane, out: &mut Plane) {
fn sgrproj_stripe_rdu<T: Pixel>(
set: u8,
xqd: [i8; 2],
fi: &FrameInvariants<T>,
crop_w: usize,
crop_h: usize,
stripe_w: usize,
stripe_h: isize,
stripe_x: usize,
stripe_y: isize,
cdeffed: &Plane<T>,
deblocked: &Plane<T>,
out: &mut Plane<T>,
) {
assert!(stripe_h <= 64);
let mut a0: [[i32; 64+2]; 3] = [[0; 64+2]; 3];
let mut a1: [[i32; 64+2]; 3] = [[0; 64+2]; 3];
......@@ -243,19 +268,19 @@ fn sgrproj_stripe_rdu(set: u8, xqd: [i8; 2], fi: &FrameInvariants,
let w1 = xqd[1] as i32;
let w2 = (1 << SGRPROJ_PRJ_BITS) - w0 - w1;
for yi in start..stripe_h as usize {
let u = (cdeffed_slice.p(xi, yi) as i32) << SGRPROJ_RST_BITS;
let u = i32::cast_from(cdeffed_slice.p(xi, yi)) << SGRPROJ_RST_BITS;
let v = w0*f0[yi] + w1*u + w2*f1[yi];
let s = v + (1 << SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS >> 1) >> SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS;
out_data[xi + yi*outstride] = clamp(s as u16, 0, (1 << bit_depth) - 1);
out_data[xi + yi*outstride] = T::cast_from(clamp(u16::cast_from(s), 0, (1 << bit_depth) - 1));
}
} else {
let w0 = xqd[0] as i32;
let w = (1 << SGRPROJ_PRJ_BITS) - w0;
for yi in start..stripe_h as usize {
let u = (cdeffed_slice.p(xi, yi) as i32) << SGRPROJ_RST_BITS;
let u = i32::cast_from(cdeffed_slice.p(xi, yi)) << SGRPROJ_RST_BITS;