aboutsummaryrefslogtreecommitdiff
path: root/mtree-test/src/encode.rs
diff options
context:
space:
mode:
Diffstat (limited to 'mtree-test/src/encode.rs')
-rw-r--r--mtree-test/src/encode.rs272
1 files changed, 272 insertions, 0 deletions
diff --git a/mtree-test/src/encode.rs b/mtree-test/src/encode.rs
new file mode 100644
index 0000000..d746e7f
--- /dev/null
+++ b/mtree-test/src/encode.rs
@@ -0,0 +1,272 @@
+use crate::{AbsRef, BLOCK_SIZE, Frame, LastFrames, frame_to_frame_rect_copy};
+use framework::BitstreamFilter;
+use glam::{I16Vec2, i16vec2};
+use rayon::{
+ iter::{IntoParallelRefIterator, ParallelIterator},
+ join,
+};
+use std::{collections::VecDeque, time::Instant};
+
+pub struct Enc {
+ res: I16Vec2,
+ last: LastFrames,
+ tree: MTree,
+ frame_num: u64,
+
+ use_bulk_insert: bool,
+}
+impl BitstreamFilter for Enc {
+ const INPUT_CODEC_ID: &str = "V_UNCOMPRESSED";
+ const OUTPUT_CODEC_ID: &str = "V_VCEMTREE";
+
+ fn new(width: u32, height: u32) -> Self {
+ Self {
+ res: i16vec2(width as i16, height as i16),
+ last: LastFrames {
+ frame_offset: 0,
+ frames: VecDeque::new(),
+ },
+ tree: MTree::Leaf(AbsRef {
+ off: i16vec2(0, 0),
+ frame: u64::MAX,
+ }),
+ frame_num: 0,
+ use_bulk_insert: true,
+ }
+ }
+ fn process_block(&mut self, frame: Vec<u8>) -> Vec<u8> {
+ let mut frame = Frame(frame);
+ let mut out = Vec::new();
+ let mut num_refs = 0;
+
+ let t = Instant::now();
+ for bx in 0..self.res.x / BLOCK_SIZE {
+ for by in 0..self.res.y / BLOCK_SIZE {
+ let boff = i16vec2(bx * BLOCK_SIZE, by * BLOCK_SIZE);
+ let r = self.tree.search(self.res, &self.last, &frame, boff);
+
+ let rdist = if r.frame == u64::MAX {
+ u32::MAX
+ } else {
+ distance(
+ self.res,
+ &self.last.frames[(r.frame - self.last.frame_offset) as usize],
+ &frame,
+ r.off,
+ boff,
+ )
+ };
+ if rdist < BLOCK_SIZE as u32 * BLOCK_SIZE as u32 * 50 {
+ out.push(1);
+ out.extend(r.frame.to_le_bytes());
+ out.extend(r.off.x.to_le_bytes());
+ out.extend(r.off.y.to_le_bytes());
+ frame_to_frame_rect_copy(
+ self.res,
+ &mut frame,
+ &self.last.frames[(r.frame - self.last.frame_offset) as usize],
+ I16Vec2::splat(BLOCK_SIZE),
+ boff,
+ r.off,
+ );
+ num_refs += 1;
+ } else {
+ out.push(0);
+ frame.export_rect(self.res, boff..boff + I16Vec2::splat(BLOCK_SIZE), &mut out);
+ }
+ }
+ }
+ eprintln!("search\t{:?}", t.elapsed());
+
+ self.last.frames.push_back(frame);
+ let frame_index = self.last.frame_offset + self.last.frames.len() as u64 - 1;
+
+ if self.frame_num % 30 == 0 {
+ let t = Instant::now();
+ if self.use_bulk_insert {
+ let mut rs = Vec::new();
+ for xo in 0..self.res.x - BLOCK_SIZE {
+ for yo in 0..self.res.y - BLOCK_SIZE {
+ rs.push(AbsRef {
+ frame: frame_index,
+ off: i16vec2(xo, yo),
+ });
+ }
+ }
+ self.tree.bulk_insert(self.res, &self.last, rs);
+ } else {
+ for xo in 0..self.res.x - BLOCK_SIZE {
+ for yo in 0..self.res.y - BLOCK_SIZE {
+ self.tree.insert(
+ self.res,
+ &self.last,
+ AbsRef {
+ frame: frame_index,
+ off: i16vec2(xo, yo),
+ },
+ )
+ }
+ }
+ }
+ eprintln!("insert\t{:?}", t.elapsed());
+ }
+
+ eprintln!("depth={} refs={num_refs}", self.tree.depth());
+
+ self.frame_num += 1;
+ out
+ }
+}
+
+enum MTree {
+ Branch(Box<[MTree; 2]>),
+ Leaf(AbsRef),
+}
+impl MTree {
+ pub fn depth(&self) -> usize {
+ (match self {
+ MTree::Branch(b) => b[0].depth().max(b[1].depth()),
+ MTree::Leaf(_) => 0,
+ }) + 1
+ }
+ pub fn center(&self) -> AbsRef {
+ match self {
+ MTree::Branch(b) => b[0].center(),
+ MTree::Leaf(r) => *r,
+ }
+ }
+ pub fn bulk_insert(&mut self, res: I16Vec2, last: &LastFrames, mut rs: Vec<AbsRef>) {
+ match self {
+ MTree::Branch(b) => {
+ let dirs = rs
+ .par_iter()
+ .map(|r| {
+ distance_absref(res, last, b[0].center(), *r)
+ < distance_absref(res, last, b[1].center(), *r)
+ })
+ .collect::<Vec<bool>>();
+
+ let mut rs0 = Vec::new();
+ let mut rs1 = Vec::new();
+ for (r, d) in rs.into_iter().zip(dirs) {
+ if d {
+ rs0.push(r);
+ } else {
+ rs1.push(r);
+ }
+ }
+ let [b0, b1] = b.as_mut();
+ join(
+ || b0.bulk_insert(res, last, rs0),
+ || b1.bulk_insert(res, last, rs1),
+ );
+ }
+ MTree::Leaf(l) => {
+ if rs.is_empty() {
+ return;
+ }
+ if l.frame == u64::MAX {
+ *l = rs.pop().unwrap();
+ if rs.is_empty() {
+ return;
+ }
+ }
+ let mut b =
+ MTree::Branch(Box::new([MTree::Leaf(*l), MTree::Leaf(rs.pop().unwrap())]));
+ b.bulk_insert(res, last, rs);
+ *self = b;
+ }
+ }
+ }
+ pub fn insert(&mut self, res: I16Vec2, last: &LastFrames, r: AbsRef) {
+ match self {
+ MTree::Branch(b) => {
+ if distance_absref(res, last, b[0].center(), r)
+ < distance_absref(res, last, b[1].center(), r)
+ {
+ b[0].insert(res, last, r)
+ } else {
+ b[1].insert(res, last, r)
+ }
+ }
+ MTree::Leaf(l) => {
+ if l.frame == u64::MAX {
+ *l = r
+ } else {
+ *self = MTree::Branch(Box::new([MTree::Leaf(*l), MTree::Leaf(r)]))
+ }
+ }
+ }
+ }
+ pub fn search(
+ &mut self,
+ res: I16Vec2,
+ last: &LastFrames,
+ frame: &Frame,
+ off: I16Vec2,
+ ) -> AbsRef {
+ match self {
+ MTree::Branch(b) => {
+ let acenter = b[0].center();
+ let bcenter = b[1].center();
+ let adist = distance(
+ res,
+ &last.frames[(acenter.frame - last.frame_offset) as usize],
+ frame,
+ acenter.off,
+ off,
+ );
+ let bdist = distance(
+ res,
+ &last.frames[(bcenter.frame - last.frame_offset) as usize],
+ frame,
+ bcenter.off,
+ off,
+ );
+ if adist < bdist {
+ b[0].search(res, last, frame, off)
+ } else {
+ b[1].search(res, last, frame, off)
+ }
+ }
+ MTree::Leaf(l) => *l,
+ }
+ }
+}
+
+fn distance_absref(res: I16Vec2, last: &LastFrames, a: AbsRef, b: AbsRef) -> u32 {
+ let aframe = &last.frames[(a.frame - last.frame_offset) as usize];
+ let bframe = &last.frames[(b.frame - last.frame_offset) as usize];
+ distance(res, aframe, bframe, a.off, b.off)
+}
+fn distance(res: I16Vec2, aframe: &Frame, bframe: &Frame, aoff: I16Vec2, boff: I16Vec2) -> u32 {
+ let mut diff = 0;
+ for yo in 0..BLOCK_SIZE {
+ // for xo in 0..BLOCK_SIZE {
+ // let off = i16vec2(xo, yo);
+ // let asample = aframe.get(res, aoff + off);
+ // let bsample = bframe.get(res, boff + off);
+
+ // diff += asample[0].abs_diff(bsample[0]) as u32;
+ // diff += asample[1].abs_diff(bsample[1]) as u32;
+ // diff += asample[2].abs_diff(bsample[2]) as u32;
+ // }
+
+ let aslice = aframe.slice(res, yo + aoff.y, aoff.x..aoff.x + BLOCK_SIZE);
+ let bslice = bframe.slice(res, yo + boff.y, boff.x..boff.x + BLOCK_SIZE);
+
+ for i in 0..BLOCK_SIZE as usize {
+ diff += aslice[0][i].abs_diff(bslice[0][i]) as u32;
+ diff += aslice[0][i].abs_diff(bslice[0][i]) as u32;
+ }
+ for i in 0..BLOCK_SIZE as usize / 2 {
+ diff += aslice[1][i].abs_diff(bslice[1][i]) as u32;
+ diff += aslice[1][i].abs_diff(bslice[1][i]) as u32;
+ }
+ for i in 0..BLOCK_SIZE as usize / 2 {
+ diff += aslice[2][i].abs_diff(bslice[2][i]) as u32;
+ diff += aslice[2][i].abs_diff(bslice[2][i]) as u32;
+ }
+ }
+ diff
+}
generated by cgit v1.2.3-70-g09d2 (git 2.51.0) at 2025-10-07 19:50:44 +0000