diff options
Diffstat (limited to 'bv1/codec/src/encode.rs')
| -rw-r--r-- | bv1/codec/src/encode.rs | 190 |
1 files changed, 190 insertions, 0 deletions
diff --git a/bv1/codec/src/encode.rs b/bv1/codec/src/encode.rs new file mode 100644 index 0000000..564b0b3 --- /dev/null +++ b/bv1/codec/src/encode.rs @@ -0,0 +1,190 @@ +use crate::diff::{diff, pixel_diff}; +use crate::huff::write_huff; +use crate::impls::join; +use crate::split::split; +use crate::{decode::decode_block, Block, Frame, Pixel, Ref, View, P2}; +use std::io::{BufReader, BufWriter, Read, Write}; +use std::time::Instant; + +#[derive(Debug, Clone)] +pub struct EncodeConfig { + pub threshold: f32, + pub max_block_size: usize, + pub attention_split: u32, + pub keyframe_interval: usize, +} + +pub fn encode( + config: EncodeConfig, + size: P2, + input: impl Read, + output: impl Write, +) -> std::io::Result<()> { + let mut input = BufReader::new(input); + let mut output = BufWriter::new(output); + + let mut last_frame = Frame::new(size); + for frame_number in 0.. { + let mut frame = Frame::read(&mut input, size)?; + + let mut config = config.clone(); + if frame_number % config.keyframe_interval != 0 { + config.threshold = std::f32::INFINITY; + } + + let t = Instant::now(); + let b: Block = encode_block(&last_frame, &frame, View::all(size), &config); + let time_encode = t.elapsed(); + + let t = Instant::now(); + decode_block(&last_frame, &mut frame, View::all(size), &b); + last_frame = frame; + let time_decode = t.elapsed(); + + if true { + let mut buf = vec![]; + let mut bufw = std::io::Cursor::new(&mut buf); + b.write(&mut bufw)?; + drop(bufw); + let t = Instant::now(); + let bits_raw = buf.len() * 8; + let bits_huff = write_huff(&buf, &mut output)?; + let time_huff = t.elapsed(); + drop(buf); + + eprintln!( + "frame {frame_number}: {:?}", + time_decode + time_huff + time_encode + ); + eprintln!( + "\tencode {time_encode:?} ({:.2}%)", + (bits_raw as f32 / (size.area() * 24) as f32) * 100.0 + ); + eprintln!( + "\thuff {time_huff:?} ({:.2}%)", + (bits_huff as f32 / bits_raw as f32) * 100.0 + ); + eprintln!("\tdecode {time_decode:?}"); + } else { + b.write(&mut output)?; + } + } + Ok(()) +} + +pub fn encode_block(last_frame: &Frame, frame: &Frame, view: View, config: &EncodeConfig) -> Block { + let view_area = view.size().area(); + if view_area > config.max_block_size + || (view_area > 64 && attention(frame, view) > config.attention_split) + { + let [av, bv] = split(view); + let (ab, bb) = join( + || Box::new(encode_block(last_frame, frame, av, config)), + || Box::new(encode_block(last_frame, frame, bv, config)), + ); + return Block::Split(ab, bb); + } + + let mut r = Ref::default(); + let mut d = diff([last_frame, frame], view, r); + + // let att = 1. - attention(frame, view) as f32 * 0.000001; + // let thres = (config.threshold as f32 * att.clamp(0.2, 1.0)) as u32; + let thres = (config.threshold * view_area as f32) as u32; + + let target_average = average_color(frame, view); + + for granularity in [2, 1, 2, 1, 2, 1, 2, 1] { + let (nd, nrp) = optimize_ref(last_frame, frame, view, r, granularity, target_average); + if nd < d { + r = nrp; + d = nd; + } else { + break; + } + } + + if d < thres { + return Block::Ref(r); + } else { + Block::Lit(frame.export(view)) + } +} + +pub fn optimize_ref( + last_frame: &Frame, + frame: &Frame, + view: View, + r: Ref, + g: i32, + target_average: Pixel, +) -> (u32, Ref) { + let g2 = g * 2; + [ + Some(r.apply(|r| r.pos_off += P2 { x: g, y: 0 })), + Some(r.apply(|r| r.pos_off += P2 { x: g, y: g })), + Some(r.apply(|r| r.pos_off += P2 { x: 0, y: g })), + Some(r.apply(|r| r.pos_off += P2 { x: -g, y: g })), + Some(r.apply(|r| r.pos_off += P2 { x: -g, y: 0 })), + Some(r.apply(|r| r.pos_off += P2 { x: -g, y: -g })), + Some(r.apply(|r| r.pos_off += P2 { x: 0, y: -g })), + Some(r.apply(|r| r.pos_off += P2 { x: g, y: -g })), + Some(r.apply(|r| r.pos_off += P2 { x: g2, y: 0 })), + Some(r.apply(|r| r.pos_off += P2 { x: g2, y: g2 })), + Some(r.apply(|r| r.pos_off += P2 { x: 0, y: g2 })), + Some(r.apply(|r| r.pos_off += P2 { x: -g2, y: g2 })), + Some(r.apply(|r| r.pos_off += P2 { x: -g2, y: 0 })), + Some(r.apply(|r| r.pos_off += P2 { x: -g2, y: -g2 })), + Some(r.apply(|r| r.pos_off += P2 { x: 0, y: -g2 })), + Some(r.apply(|r| r.pos_off += P2 { x: g2, y: -g2 })), + { + let mut r = r; + let last_avr = average_color(last_frame, view); + let diff = target_average - last_avr; + r.color_off = diff; + if diff != Pixel::BLACK { + Some(r) + } else { + None + } + }, + ] + .into_iter() + .flatten() + .map(|r| (diff([last_frame, frame], view, r), r)) + .min_by_key(|e| e.0) + .unwrap() +} + +pub fn attention(frame: &Frame, view: View) -> u32 { + let mut k = 0; + for y in view.a.y..view.b.y - 1 { + for x in view.a.x..view.b.x - 1 { + let p = P2 { x, y }; + k += pixel_diff(frame[p], frame[p + P2::X]).pow(2); + k += pixel_diff(frame[p], frame[p + P2::Y]).pow(2); + } + } + k +} + +pub fn average_color(frame: &Frame, view: View) -> Pixel { + let mut r = 0u32; + let mut g = 0u32; + let mut b = 0u32; + + for y in view.a.y..view.b.y { + for x in view.a.x..view.b.x { + let p = frame[P2 { x, y }]; + r += p.r as u32; + g += p.g as u32; + b += p.b as u32; + } + } + let area = view.size().area() as u32; + Pixel { + r: (r / area) as i16, + g: (g / area) as i16, + b: (b / area) as i16, + } +} |