Implement edge hysteresis for Canny

src/filter/canny/hysteresis.rs

@@ -0,0 +1,189 @@
+use super::double_threshold::Candidate;
+use crate::image::*;
+
+#[derive(Debug)]
+pub struct DisjointSet {
+    parent: Vec<usize>,
+    rank: Vec<usize>,
+}
+
+impl DisjointSet {
+    pub fn new() -> Self {
+        Self {
+            parent: Vec::new(),
+            rank: Vec::new(),
+        }
+    }
+
+    pub fn make_set(&mut self) -> usize {
+        let x = self.parent.len();
+        self.parent.push(x);
+        self.rank.push(0);
+        x
+    }
+
+    pub fn find(&mut self, x: usize) -> usize {
+        if self.parent[x] != x {
+            self.parent[x] = self.find(self.parent[x]);
+        }
+        self.parent[x]
+    }
+
+    pub fn union(&mut self, x: usize, y: usize) {
+        let root_x = self.find(x);
+        let root_y = self.find(y);
+        if root_x != root_y {
+            if self.rank[root_x] < self.rank[root_y] {
+                self.parent[root_x] = root_y;
+            } else {
+                self.parent[root_y] = root_x;
+                if self.rank[root_y] == self.rank[root_x] {
+                    self.rank[root_x] += 1;
+                }
+            }
+        }
+    }
+}
+
+pub fn hysteresis(src: Image<Candidate>) -> Image<bool> {
+    let mut links = DisjointSet::new();
+    let mut map = Image::<Option<usize>>::with_extent(src.extent);
+    let mut neighbor_labels = std::collections::BTreeSet::<usize>::new();
+
+    for y in 0..src.extent.y {
+        for x in 0..src.extent.x {
+            if src.at(x, y) == Candidate::None {
+                continue;
+            }
+
+            get_neighbor_labels(x, y, &src, &map, &mut neighbor_labels);
+            *map.at_mut(x, y) = Some(match neighbor_labels.first() {
+                None => links.make_set(),
+                Some(&label) => {
+                    for neighbor in &neighbor_labels {
+                        links.union(*neighbor, label);
+                    }
+                    label
+                }
+            });
+        }
+    }
+
+    let mut alive = std::collections::BTreeSet::<usize>::new();
+    for y in 0..src.extent.y {
+        for x in 0..src.extent.x {
+            if let Some(label) = map.at(x, y) {
+                let final_label = links.find(label);
+                *map.at_mut(x, y) = Some(final_label);
+                if src.at(x, y) == Candidate::Strong {
+                    alive.insert(final_label);
+                }
+            }
+        }
+    }
+
+    map.map(|p| match p {
+        None => false,
+        Some(label) => alive.contains(&label),
+    })
+}
+
+fn get_neighbor_labels(
+    x: usize,
+    y: usize,
+    src: &Image<Candidate>,
+    map: &Image<Option<usize>>,
+    out: &mut std::collections::BTreeSet<usize>,
+) {
+    out.clear();
+    for (nx, ny) in neighbors(x, y, src.extent) {
+        if src.at(nx, ny) == Candidate::None {
+            continue;
+        }
+        if let Some(label) = map.at(nx, ny) {
+            out.insert(label);
+        }
+    }
+}
+
+fn neighbors(x: usize, y: usize, extent: Extent) -> Vec<(usize, usize)> {
+    let start_x = if x <= 0 { 0 } else { x - 1 };
+    let start_y = if y <= 0 { 0 } else { y - 1 };
+    let end_x = if x + 1 >= extent.x { x + 1 } else { x + 2 };
+    let end_y = if y + 1 >= extent.y { y + 1 } else { y + 2 };
+
+    let mut res = Vec::new();
+    for i in start_y..end_y {
+        for j in start_x..end_x {
+            if i != y || j != x {
+                res.push((j, i));
+            }
+        }
+    }
+
+    res
+}
+
+#[rustfmt::skip]
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[test]
+	fn simple_horiz_lines() {
+		let img = Image {
+			extent: Extent { x: 3, y: 3 },
+			pixels: vec![
+				Candidate::Weak, Candidate::Weak, Candidate::Weak,
+				Candidate::None, Candidate::None, Candidate::None,
+				Candidate::Weak, Candidate::Weak, Candidate::Strong,
+			],
+		};
+		let exp = vec![
+			false, false, false,
+			false, false, false,
+			true, true, true,
+		];
+		assert_eq!(hysteresis(img).pixels, exp);
+	}
+
+	#[test]
+	fn simple_diagonal() {
+		let img = Image {
+			extent: Extent { x: 3, y: 3 },
+			pixels: vec![
+				Candidate::Weak, Candidate::None, Candidate::None,
+				Candidate::None, Candidate::Weak, Candidate::None,
+				Candidate::None, Candidate::None, Candidate::Strong,
+			],
+		};
+		let exp = vec![
+			true, false, false,
+			false, true, false,
+			false, false, true,
+		];
+		assert_eq!(hysteresis(img).pixels, exp);
+	}
+
+	#[test]
+	fn complex_single_group() {
+		let img = Image {
+			extent: Extent { x: 5, y: 5 },
+			pixels: vec![
+				Candidate::None, Candidate::Weak, Candidate::None, Candidate::None, Candidate::Weak,
+				Candidate::Weak, Candidate::None, Candidate::None, Candidate::Weak, Candidate::None,
+				Candidate::None, Candidate::Weak, Candidate::Weak, Candidate::None, Candidate::None,
+				Candidate::Weak, Candidate::None, Candidate::Weak, Candidate::None, Candidate::Strong,
+				Candidate::None, Candidate::None, Candidate::None, Candidate::Weak, Candidate::None,
+			],
+		};
+		let exp = vec![
+			false, true, false, false, true,
+			true, false, false, true, false,
+			false, true, true, false, false,
+			true, false, true, false, true,
+			false, false, false, true, false,
+		];
+		assert_eq!(hysteresis(img).pixels, exp);
+	}
+}