Add filtering functionalityHEADmaster

2 files changed, 102 insertions, 6 deletions

diff --git a/Cargo.toml b/Cargo.toml
index 3e06d25..7d340d9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "label_toolkit"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2024"
 
 [dependencies]
diff --git a/src/analysis.rs b/src/analysis.rs
index aabe4e5..41a4f30 100644
--- a/src/analysis.rs
+++ b/src/analysis.rs
@@ -2,7 +2,7 @@ use crate::{LabelFormat,LabelU16};
 
 use avl_tree::*;
 
-#[derive(Debug)]
+#[derive(Debug,PartialEq)]
 pub struct Data {
     pub label: u16,
     pub area: usize,
@@ -99,6 +99,32 @@ impl LabelFormat<u16> {
         }
         avt
     }
+
+    pub fn filter<T: Fn(&Data) -> bool>(&self, filter: T) -> LabelFormat<u16> {
+        let mut output = LabelFormat::<u16> {
+            buffer: vec![0u16; self.width*self.height],
+            width: self.width,
+            height: self.height,
+        };
+        let tree = self.get_tree();
+        for y in 0..self.height {
+            for x in 0..self.width {
+                let index = x + y * self.width;
+                let key = self.buffer[index];
+                if key.is_zero() {
+                    continue;
+                }
+                if let Some(data) = tree.get(key) {
+                    if filter(data) {
+                        output.buffer[index] = self.buffer[index];
+                    }
+                } else {
+                    panic!("Tree should contain all non-zero values!");
+                }
+            }
+        }
+        output
+    }
 }
 
 #[cfg(test)]
@@ -118,9 +144,79 @@ mod tests {
         test_data.buffer[4+3*DIM] = 2;
         test_data.idilate();
 
-        test_data.display();
-        for (key,val) in test_data.get_tree().iter() {
-            println!("{} -> {:?}",key,val)
-        }
+        let expected: [Vec<u16>;1] = [
+            [
+                0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0,
+                0, 0, 1, 1, 2, 0,
+                0, 1, 1, 1, 2, 2,
+                0, 0, 1, 1, 2, 0,
+                0, 0, 0, 0, 0, 0,
+            ].to_vec(),
+        ];
+
+        assert_eq!(test_data.buffer, expected[0]);
+        let tree = test_data.get_tree();
+        let data = tree.get(0);
+        assert_eq!(data, None);
+        let data = tree.get(1).expect("Key should exist");
+        assert_eq!(*data, Data::new(1, 7, 6, 3, 5));
+        let data = tree.get(2).expect("Key should exist");
+        assert_eq!(*data, Data::new(2, 4, 4, 3, 3));
+        let data = tree.get(3);
+        assert_eq!(data, None);
+    }
+
+    #[test]
+    fn filter_test() {
+        const DIM: usize = 6;
+        let mut test_data = LabelFormat::<u16> {
+            buffer: vec![0u16; DIM*DIM],
+            width: DIM,
+            height: DIM,
+        };
+        test_data.buffer[2+3*DIM] = 1;
+        test_data.buffer[3+3*DIM] = 1;
+        test_data.buffer[4+3*DIM] = 2;
+        test_data.idilate();
+
+        let expected: [Vec<u16>;2] = [
+            [
+                0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0,
+                0, 0, 1, 1, 2, 0,
+                0, 1, 1, 1, 2, 2,
+                0, 0, 1, 1, 2, 0,
+                0, 0, 0, 0, 0, 0,
+            ].to_vec(),
+            [
+                0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0,
+                0, 0, 1, 1, 0, 0,
+                0, 1, 1, 1, 0, 0,
+                0, 0, 1, 1, 0, 0,
+                0, 0, 0, 0, 0, 0,
+            ].to_vec(),
+        ];
+
+        let tree = test_data.get_tree();
+        let data = tree.get(0);
+        assert_eq!(data, None);
+        let data = tree.get(1).expect("Key should exist");
+        assert_eq!(*data, Data::new(1, 7, 6, 3, 5));
+        let data = tree.get(2).expect("Key should exist");
+        assert_eq!(*data, Data::new(2, 4, 4, 3, 3));
+        let data = tree.get(3);
+        assert_eq!(data, None);
+
+        let filtered_data = test_data.filter(|x| x.area > 5);
+        assert_eq!(filtered_data.buffer, expected[1]);
+        let tree = filtered_data.get_tree();
+        let data = tree.get(0);
+        assert_eq!(data, None);
+        let data = tree.get(1).expect("Key should exist");
+        assert_eq!(*data, Data::new(1, 7, 6, 0, 6));
+        let data = tree.get(2);
+        assert_eq!(data, None);
     }
 }