image_debug_utils/
region_labelling.rs

1//! Tools for visualizing and analyzing connected components results from `imageproc::region_labelling`.
2
3use crate::colors::generate_contrasting_colors;
4use image::{ImageBuffer, Luma, Rgba, RgbaImage};
5use std::collections::HashMap;
6
7/// Draws the n largest connected components with contrasting colors.
8///
9/// # Arguments
10/// * `labelled_image` - The labelled image generated by a function like `imageproc::region_labelling::connected_components`.
11/// * `n` - The number of largest components to keep and color.
12/// * `background_color` - The color for the background and smaller, unselected components.
13///
14/// # Examples
15///
16/// ```
17/// use image::{ImageBuffer, Luma, Rgba};
18/// use image_debug_utils::region_labelling::draw_principal_connected_components;
19///
20/// let mut labelled_image = ImageBuffer::<Luma<u32>, Vec<u32>>::new(10, 10);
21/// // Simulate a large component (label 1) and a small one (label 2)
22/// labelled_image.put_pixel(0, 0, Luma([1]));
23/// labelled_image.put_pixel(0, 1, Luma([1]));
24/// labelled_image.put_pixel(5, 5, Luma([2]));
25///
26/// // Keep top 1 component, use transparent black for background
27/// let colored = draw_principal_connected_components(&labelled_image, 1, Rgba([0, 0, 0, 0]));
28/// ```
29///
30/// # Returns
31/// An `RgbaImage` where the `n` largest components are colored and the rest is background.
32pub fn draw_principal_connected_components(
33    labelled_image: &ImageBuffer<Luma<u32>, Vec<u32>>,
34    n: usize,
35    background_color: Rgba<u8>,
36) -> RgbaImage {
37    let mut counts = HashMap::new();
38    for label in labelled_image.pixels() {
39        if label[0] != 0 {
40            *counts.entry(label[0]).or_insert(0) += 1;
41        }
42    }
43
44    let mut sorted_counts: Vec<_> = counts.into_iter().collect();
45    sorted_counts.sort_by(|a, b| b.1.cmp(&a.1));
46
47    let principal_labels: Vec<u32> = sorted_counts
48        .iter()
49        .take(n)
50        .map(|(label, _)| *label)
51        .collect();
52
53    let colors = generate_contrasting_colors(n, 255);
54    let color_map: HashMap<u32, Rgba<u8>> = principal_labels.into_iter().zip(colors).collect();
55
56    let (width, height) = labelled_image.dimensions();
57    let mut output_image = RgbaImage::from_pixel(width, height, background_color);
58
59    for (x, y, pixel) in labelled_image.enumerate_pixels() {
60        let label = pixel[0];
61        if let Some(color) = color_map.get(&label) {
62            output_image.put_pixel(x, y, *color);
63        }
64    }
65
66    output_image
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72
73    #[test]
74    fn test_draw_principal_connected_components() {
75        // Arrange:
76        // Define a labelled image with multiple components of different sizes.
77        // This simulates the output of `connected_components`.
78        // 1. A large 4x4 square (16 pixels), label 1.
79        // 2. A medium 2x3 rectangle (6 pixels), label 2.
80        // 3. A small 1x2 line (2 pixels), label 3.
81        #[rustfmt::skip]
82        let labels_vec = vec![
83            0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
84            0, 1, 1, 1, 1, 0, 0, 0, 0, 0,
85            0, 1, 1, 1, 1, 0, 2, 2, 2, 0,
86            0, 1, 1, 1, 1, 0, 2, 2, 2, 0,
87            0, 1, 1, 1, 1, 0, 0, 0, 0, 0,
88            0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
89            0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
90            0, 3, 0, 0, 0, 0, 0, 0, 0, 0,
91            0, 3, 0, 0, 0, 0, 0, 0, 0, 0,
92            0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
93        ];
94
95        let labelled_image = ImageBuffer::<Luma<u32>, _>::from_raw(10, 10, labels_vec).unwrap();
96
97        // We want to draw the two largest components (n=2).
98        let n = 2;
99        let background = Rgba([0, 0, 0, 255]);
100
101        // Act:
102        let result_image = draw_principal_connected_components(&labelled_image, n, background);
103
104        // Assert:
105        // -- Create the expected output image --
106
107        // The function will identify the largest components.
108        // Counts: {1: 16, 2: 6, 3: 2}.
109        // The two largest labels are 1 and 2.
110
111        // Get the colors that will be generated for these two labels.
112        let colors = generate_contrasting_colors(n, 255);
113        let color_for_label1 = colors[0]; // The largest component gets the first color.
114        let color_for_label2 = colors[1]; // The second largest component gets the second color.
115
116        let mut expected_image = RgbaImage::from_pixel(10, 10, background);
117
118        // Draw the largest component (label 1).
119        for y in 1..=4 {
120            for x in 1..=4 {
121                expected_image.put_pixel(x, y, color_for_label1);
122            }
123        }
124
125        // Draw the second largest component (label 2).
126        for y in 2..=3 {
127            for x in 6..=8 {
128                expected_image.put_pixel(x, y, color_for_label2);
129            }
130        }
131
132        // The smallest component (label 3) should remain as the background color,
133        // so we don't need to draw it on the expected_image.
134
135        assert_eq!(result_image, expected_image);
136    }
137
138    #[test]
139    fn test_draw_principal_connected_components_empty() {
140        // All background image.
141        let labelled_image = ImageBuffer::<Luma<u32>, Vec<u32>>::new(5, 5);
142        let background = Rgba([255, 255, 255, 255]);
143
144        // n = 0
145        let result = draw_principal_connected_components(&labelled_image, 0, background);
146        assert_eq!(result.dimensions(), (5, 5));
147        assert!(result.pixels().all(|p| *p == background));
148    }
149}
image_debug_utils/region_labelling.rs

image_debug_utils/
region_labelling.rs
