From 0c44ce639609caa114fd5f61118c445308e47f86 Mon Sep 17 00:00:00 2001 From: Camden Dixie O'Brien Date: Tue, 27 Oct 2026 19:16:23 +0000 Subject: [PATCH] Implement edge hysteresis for Canny --- src/filter/canny/hysteresis.rs | 189 +++++++++++++++++++++++++++++++++ 1 file changed, 189 insertions(+) create mode 100644 src/filter/canny/hysteresis.rs diff --git a/src/filter/canny/hysteresis.rs b/src/filter/canny/hysteresis.rs new file mode 100644 index 0000000..ec50f99 --- /dev/null +++ b/src/filter/canny/hysteresis.rs @@ -0,0 +1,189 @@ +use super::double_threshold::Candidate; +use crate::image::*; + +#[derive(Debug)] +pub struct DisjointSet { + parent: Vec, + rank: Vec, +} + +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) -> Image { + let mut links = DisjointSet::new(); + let mut map = Image::>::with_extent(src.extent); + let mut neighbor_labels = std::collections::BTreeSet::::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::::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, + map: &Image>, + out: &mut std::collections::BTreeSet, +) { + 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); + } +}