more magic

10 files changed, 94 insertions, 24 deletions

diff --git a/lvc/src/debug.rs b/lvc/src/debug.rs
index 0e6fb7a..d6c37c2 100644
--- a/lvc/src/debug.rs
+++ b/lvc/src/debug.rs
@@ -8,7 +8,18 @@ pub fn draw_debug(frame: &mut Frame, view: View, block: &Block) {
             draw_debug(frame, av, &a);
             draw_debug(frame, bv, &b);
         }
-        Block::Ref(_r) => {}
+        Block::Ref(r) => {
+            let v = View {
+                a: view.a,
+                b: view.a + P2 { x: 2, y: 2 },
+            };
+            if r.pos_off != P2::ZERO {
+                fill_rect(frame, v + P2 { x: 0, y: 0 }, Pixel::BLUE)
+            }
+            if r.color_off != Pixel::BLACK {
+                fill_rect(frame, v + P2 { x: 2, y: 0 }, Pixel::RED)
+            }
+        }
     }
 }
 
@@ -22,9 +33,10 @@ fn rect(frame: &mut Frame, view: View, color: Pixel) {
         frame[P2 { y, x: view.b.x - 1 }] = color;
     }
 }
-
-impl Pixel {
-    pub const RED: Pixel = Pixel { r: 255, g: 0, b: 0 };
-    pub const GREEN: Pixel = Pixel { r: 0, g: 255, b: 0 };
-    pub const BLUE: Pixel = Pixel { r: 0, g: 0, b: 255 };
+fn fill_rect(frame: &mut Frame, view: View, color: Pixel) {
+    for y in view.a.y..view.b.y {
+        for x in view.a.x..view.b.x {
+            frame[P2 { x, y }] = color;
+        }
+    }
 }
diff --git a/lvc/src/diff.rs b/lvc/src/diff.rs
index 354f146..fbb6b0a 100644
--- a/lvc/src/diff.rs
+++ b/lvc/src/diff.rs
@@ -1,4 +1,4 @@
-use crate::{Frame, Ref, View, P2};
+use crate::{Frame, Pixel, Ref, View, P2};
 
 // 4ms
 pub fn diff([frame1, frame2]: [&Frame; 2], view: View, rp: Ref) -> u32 {
@@ -8,14 +8,17 @@ pub fn diff([frame1, frame2]: [&Frame; 2], view: View, rp: Ref) -> u32 {
             let pos = P2 { x, y };
             let p1 = frame1[pos + rp.pos_off] + rp.color_off;
             let p2 = frame2[pos];
-            k += p1.r.abs_diff(p2.r) as u32
-                + p1.g.abs_diff(p2.g) as u32
-                + p1.b.abs_diff(p2.b) as u32;
+            k += pixel_diff(p1, p2)
         }
     }
     k
 }
 
+#[inline(always)]
+pub fn pixel_diff(p1: Pixel, p2: Pixel) -> u32 {
+    p1.r.abs_diff(p2.r) as u32 + p1.g.abs_diff(p2.g) as u32 + p1.b.abs_diff(p2.b) as u32
+}
+
 // pub fn fast_diff([frame1, frame2]: [&Frame; 2], view: View, rp: Ref) -> u32 {
 //     assert!(view.size().x % 5 == 0);
 
diff --git a/lvc/src/encode.rs b/lvc/src/encode.rs
index 935aaba..10f8c50 100644
--- a/lvc/src/encode.rs
+++ b/lvc/src/encode.rs
@@ -1,14 +1,16 @@
-use crate::diff::diff;
+use crate::diff::{diff, pixel_diff};
 use crate::split::split;
-use crate::{Block, Frame, Ref, View};
+use crate::{Block, Frame, Ref, View, P2};
 
 pub struct EncodeConfig {
     pub threshold: u32,
     pub max_block_size: usize,
+    pub iters: usize,
 }
 
 pub fn encode(last_frame: &Frame, frame: &Frame, view: View, config: &EncodeConfig) -> Block {
-    if view.size().area() > config.max_block_size {
+    let view_area = view.size().area();
+    if view_area > config.max_block_size {
         let [av, bv] = split(view);
         let (ab, bb) = rayon::join(
             || Box::new(encode(last_frame, frame, av, config)),
@@ -20,8 +22,11 @@ pub fn encode(last_frame: &Frame, frame: &Frame, view: View, config: &EncodeConf
     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;
+
     for granularity in [4, 2, 1] {
-        for _ in 0..10 {
+        for _ in 0..config.iters {
             let (nd, nrp) = optimize_ref(last_frame, frame, view, d, r, granularity);
             if nd < d {
                 r = nrp;
@@ -32,7 +37,7 @@ pub fn encode(last_frame: &Frame, frame: &Frame, view: View, config: &EncodeConf
         }
     }
 
-    if d < config.threshold {
+    if d < thres {
         return Block::Ref(r);
     } else {
         Block::Lit(frame.export(view))
@@ -71,3 +76,15 @@ pub fn optimize_ref(
 
     (d, r)
 }
+
+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
+}
diff --git a/lvc/src/impls.rs b/lvc/src/impls.rs
index 00614c9..2a5a497 100644
--- a/lvc/src/impls.rs
+++ b/lvc/src/impls.rs
@@ -27,11 +27,18 @@ impl Frame {
     }
 }
 impl Ref {
+    #[inline]
     pub fn apply<F: Fn(&mut Self)>(mut self, f: F) -> Self {
         f(&mut self);
         self
     }
 }
+impl Pixel {
+    pub const BLACK: Pixel = Pixel { r: 0, g: 0, b: 0 };
+    pub const RED: Pixel = Pixel { r: 255, g: 0, b: 0 };
+    pub const GREEN: Pixel = Pixel { r: 0, g: 255, b: 0 };
+    pub const BLUE: Pixel = Pixel { r: 0, g: 0, b: 255 };
+}
 impl Add for Pixel {
     type Output = Pixel;
     #[inline]
@@ -45,6 +52,9 @@ impl Add for Pixel {
 }
 impl P2 {
     pub const ZERO: P2 = P2 { x: 0, y: 0 };
+    pub const X: P2 = P2 { x: 1, y: 0 };
+    pub const Y: P2 = P2 { x: 0, y: 1 };
+
     #[inline]
     pub fn area(&self) -> usize {
         (self.x * self.y) as usize
@@ -63,6 +73,16 @@ impl View {
         self.b - self.a
     }
 }
+impl Add<P2> for View {
+    type Output = View;
+    #[inline]
+    fn add(self, rhs: P2) -> Self::Output {
+        View {
+            a: self.a + rhs,
+            b: self.b + rhs,
+        }
+    }
+}
 impl Add for P2 {
     type Output = P2;
     #[inline]
@@ -107,6 +127,7 @@ pub trait ToArray {
 }
 impl<A> ToArray for (A, A) {
     type Output = A;
+    #[inline]
     fn to_array(self) -> [A; 2] {
         [self.0, self.1]
     }
diff --git a/lvc/src/lib.rs b/lvc/src/lib.rs
index 33a8cfb..94a432d 100644
--- a/lvc/src/lib.rs
+++ b/lvc/src/lib.rs
@@ -13,14 +13,14 @@ pub mod split;
 
 pub type PixelValue = i16;
 
-#[derive(Debug, Clone, Copy, Default, Encode, Decode)]
+#[derive(Debug, Clone, Copy, Default, Encode, Decode, PartialEq, Eq)]
 pub struct Pixel {
     pub r: PixelValue,
     pub g: PixelValue,
     pub b: PixelValue,
 }
 
-#[derive(Debug, Clone, Copy, Default, Encode, Decode)]
+#[derive(Debug, Clone, Copy, Default, Encode, Decode, PartialEq, Eq)]
 pub struct P2 {
     pub x: i32,
     pub y: i32,
@@ -40,7 +40,7 @@ pub struct View {
 #[derive(Debug, Clone, Encode, Decode)]
 pub enum Block {
     Lit(Vec<Pixel>),
-    Split(Box<Block>,Box<Block>),
+    Split(Box<Block>, Box<Block>),
     Ref(Ref),
 }
 
diff --git a/lvc/src/main.rs b/lvc/src/main.rs
index 48df911..7ecd4b8 100644
--- a/lvc/src/main.rs
+++ b/lvc/src/main.rs
@@ -6,7 +6,10 @@ use lvc::{
     encode::{encode, EncodeConfig},
     Block, Frame, Pixel, PixelValue, View, P2,
 };
-use std::io::{stdin, stdout, BufReader, BufWriter, Read, Write};
+use std::{
+    io::{stdin, stdout, BufReader, BufWriter, Read, Write},
+    time::Instant,
+};
 
 #[derive(Parser)]
 #[clap(about, version)]
@@ -25,8 +28,10 @@ enum Action {
     // Compress video
     Encode {
         #[arg(short, long, default_value_t = 1024)]
-        block_size: usize,
-        #[arg(short, long, default_value_t = 10_000)]
+        max_block_size: usize,
+        #[arg(short, long, default_value_t = 10)]
+        iters: usize,
+        #[arg(short, long, default_value_t = 20_000)]
         threshold: u32,
     },
     // Decompress video
@@ -68,12 +73,14 @@ fn main() {
             }
         }
         Action::Encode {
-            block_size,
+            max_block_size,
             threshold,
+            iters,
         } => {
             let config = EncodeConfig {
                 threshold,
-                max_block_size: block_size,
+                max_block_size,
+                iters,
             };
 
             let mut last_frame = Frame::new(size);
@@ -82,10 +89,12 @@ fn main() {
             let mut stdout = BufWriter::new(stdout());
 
             for frame_number in 0.. {
-                eprintln!("encode frame {frame_number}");
                 let mut frame = read_frame(&mut stdin, size);
 
+                let t = Instant::now();
                 let b: Block = encode(&last_frame, &frame, View::all(size), &config);
+                eprintln!("frame {frame_number} took {:?}", t.elapsed());
+
                 bincode::encode_into_std_write(&b, &mut stdout, standard()).unwrap();
 
                 decode(&last_frame, &mut frame, View::all(size), &b);
diff --git a/lvc/tools/d-disp b/lvc/tools/d-disp
new file mode 100755
index 0000000..5879bf3
--- /dev/null
+++ b/lvc/tools/d-disp
@@ -0,0 +1,2 @@
+#!/bin/fish
+cargo run --release -- -w 1920 -h 1080 decode | ffplay -video_size 1920x1080 -pixel_format rgb24 -f rawvideo pipe:0
\ No newline at end of file
diff --git a/lvc/tools/d-dispd b/lvc/tools/d-dispd
new file mode 100755
index 0000000..c4b435d
--- /dev/null
+++ b/lvc/tools/d-dispd
@@ -0,0 +1,2 @@
+#!/bin/fish
+cargo run --release -- -w 1920 -h 1080 decode --debug | ffplay -video_size 1920x1080 -pixel_format rgb24 -f rawvideo pipe:0
\ No newline at end of file
diff --git a/lvc/tools/d-save b/lvc/tools/d-save
new file mode 100755
index 0000000..72e89c2
--- /dev/null
+++ b/lvc/tools/d-save
@@ -0,0 +1,2 @@
+#!/bin/fish
+ffmpeg -pixel_format rgb24 -video_size 1920x1080 -f rawvideo -i pipe:0 -i data/input.webm -map '0:v' -map '1:a'  -c:v libsvtav1 -y data/output.webm
diff --git a/lvc/tools/e b/lvc/tools/e
new file mode 100755
index 0000000..89d822a
--- /dev/null
+++ b/lvc/tools/e
@@ -0,0 +1,2 @@
+#!/bin/fish
+ffmpeg -loglevel quiet -i $argv[1] -vf format=rgb24 -f rawvideo pipe:1 | cargo run --release -- -w 1920 -h 1080 encode $argv[2..]