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.rs104
1 files changed, 48 insertions, 56 deletions
diff --git a/mtree-test/src/encode.rs b/mtree-test/src/encode.rs
index d746e7f..0aae8ee 100644
--- a/mtree-test/src/encode.rs
+++ b/mtree-test/src/encode.rs
@@ -1,6 +1,6 @@
use crate::{AbsRef, BLOCK_SIZE, Frame, LastFrames, frame_to_frame_rect_copy};
use framework::BitstreamFilter;
-use glam::{I16Vec2, i16vec2};
+use glam::{IVec2, ivec2};
use rayon::{
iter::{IntoParallelRefIterator, ParallelIterator},
join,
@@ -8,7 +8,7 @@ use rayon::{
use std::{collections::VecDeque, time::Instant};
pub struct Enc {
- res: I16Vec2,
+ res: IVec2,
last: LastFrames,
tree: MTree,
frame_num: u64,
@@ -21,13 +21,13 @@ impl BitstreamFilter for Enc {
fn new(width: u32, height: u32) -> Self {
Self {
- res: i16vec2(width as i16, height as i16),
+ res: ivec2(width as i32, height as i32),
last: LastFrames {
frame_offset: 0,
frames: VecDeque::new(),
},
tree: MTree::Leaf(AbsRef {
- off: i16vec2(0, 0),
+ off: ivec2(0, 0),
frame: u64::MAX,
}),
frame_num: 0,
@@ -35,21 +35,23 @@ impl BitstreamFilter for Enc {
}
}
fn process_block(&mut self, frame: Vec<u8>) -> Vec<u8> {
- let mut frame = Frame(frame);
+ let mut frame = Frame {
+ data: frame,
+ res: self.res,
+ };
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 boff = ivec2(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,
@@ -59,20 +61,19 @@ impl BitstreamFilter for Enc {
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());
+ out.extend((r.off.x as i16).to_le_bytes());
+ out.extend((r.off.y as i16).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,
+ &mut frame,
+ IVec2::splat(BLOCK_SIZE),
r.off,
+ boff,
);
num_refs += 1;
} else {
out.push(0);
- frame.export_rect(self.res, boff..boff + I16Vec2::splat(BLOCK_SIZE), &mut out);
+ frame.export_rect(boff, IVec2::splat(BLOCK_SIZE), &mut out);
}
}
}
@@ -89,7 +90,7 @@ impl BitstreamFilter for Enc {
for yo in 0..self.res.y - BLOCK_SIZE {
rs.push(AbsRef {
frame: frame_index,
- off: i16vec2(xo, yo),
+ off: ivec2(xo, yo),
});
}
}
@@ -102,7 +103,7 @@ impl BitstreamFilter for Enc {
&self.last,
AbsRef {
frame: frame_index,
- off: i16vec2(xo, yo),
+ off: ivec2(xo, yo),
},
)
}
@@ -135,14 +136,14 @@ impl MTree {
MTree::Leaf(r) => *r,
}
}
- pub fn bulk_insert(&mut self, res: I16Vec2, last: &LastFrames, mut rs: Vec<AbsRef>) {
+ pub fn bulk_insert(&mut self, res: IVec2, 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)
+ distance_absref(last, b[0].center(), *r)
+ < distance_absref(last, b[1].center(), *r)
})
.collect::<Vec<bool>>();
@@ -178,11 +179,10 @@ impl MTree {
}
}
}
- pub fn insert(&mut self, res: I16Vec2, last: &LastFrames, r: AbsRef) {
+ pub fn insert(&mut self, res: IVec2, 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)
+ if distance_absref(last, b[0].center(), r) < distance_absref(last, b[1].center(), r)
{
b[0].insert(res, last, r)
} else {
@@ -198,26 +198,18 @@ impl MTree {
}
}
}
- pub fn search(
- &mut self,
- res: I16Vec2,
- last: &LastFrames,
- frame: &Frame,
- off: I16Vec2,
- ) -> AbsRef {
+ pub fn search(&mut self, res: IVec2, last: &LastFrames, frame: &Frame, off: IVec2) -> 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,
@@ -234,39 +226,39 @@ impl MTree {
}
}
-fn distance_absref(res: I16Vec2, last: &LastFrames, a: AbsRef, b: AbsRef) -> u32 {
+fn distance_absref(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)
+ distance(aframe, bframe, a.off, b.off)
}
-fn distance(res: I16Vec2, aframe: &Frame, bframe: &Frame, aoff: I16Vec2, boff: I16Vec2) -> u32 {
+fn distance(aframe: &Frame, bframe: &Frame, aoff: IVec2, boff: IVec2) -> u32 {
+ assert_eq!(aframe.res, bframe.res);
+ let res = aframe.res;
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);
+ // Luma
+ for y in 0..BLOCK_SIZE {
+ let ay_off = aoff.x + (y + aoff.y) * res.x;
+ let by_off = boff.x + (y + boff.y) * res.x;
- 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 x in 0..BLOCK_SIZE {
+ diff += aframe.data[(ay_off + x) as usize].abs_diff(bframe.data[(by_off + x) as usize])
+ 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;
+ }
+
+ // Chroma
+ let uvplane_off = res.x * res.y;
+ for y in 0..BLOCK_SIZE / 2 {
+ let ay_off = uvplane_off + (aoff.x & !1) + (y + aoff.y / 2) * res.x;
+ let by_off = uvplane_off + (boff.x & !1) + (y + boff.y / 2) * res.x;
+
+ for x in 0..BLOCK_SIZE {
+ diff += aframe.data[(ay_off + x) as usize].abs_diff(bframe.data[(by_off + x) as usize])
+ as u32
+ * 2;
}
}
+
diff
}
generated by cgit v1.2.3-70-g09d2 (git 2.51.0) at 2025-10-07 21:26:18 +0000