image_debug_utils/

region_labelling.rs

//! Tools for visualizing and analyzing connected components results from `imageproc::region_labelling`.

use crate::colors::generate_contrasting_colors;
use image::{ImageBuffer, Luma, Rgba, RgbaImage};
use std::collections::HashMap;

/// Draws the n largest connected components with contrasting colors.
///
/// # Arguments
/// * `labelled_image` - The labelled image generated by a function like `imageproc::region_labelling::connected_components`.
/// * `n` - The number of largest components to keep and color.
/// * `background_color` - The color for the background and smaller, unselected components.
///
/// # Examples
///
/// ```
/// use image::{ImageBuffer, Luma, Rgba};
/// use image_debug_utils::region_labelling::draw_principal_connected_components;
///
/// let mut labelled_image = ImageBuffer::<Luma<u32>, Vec<u32>>::new(10, 10);
/// // Simulate a large component (label 1) and a small one (label 2)
/// labelled_image.put_pixel(0, 0, Luma([1]));
/// labelled_image.put_pixel(0, 1, Luma([1]));
/// labelled_image.put_pixel(5, 5, Luma([2]));
///
/// // Keep top 1 component, use transparent black for background
/// let colored = draw_principal_connected_components(&labelled_image, 1, Rgba([0, 0, 0, 0]));
/// ```
///
/// # Returns
/// An `RgbaImage` where the `n` largest components are colored and the rest is background.
pub fn draw_principal_connected_components(
    labelled_image: &ImageBuffer<Luma<u32>, Vec<u32>>,
    n: usize,
    background_color: Rgba<u8>,
) -> RgbaImage {
    let mut counts = HashMap::new();
    for label in labelled_image.pixels() {
        if label[0] != 0 {
            *counts.entry(label[0]).or_insert(0) += 1;
        }
    }

    let mut sorted_counts: Vec<_> = counts.into_iter().collect();
    sorted_counts.sort_by(|a, b| b.1.cmp(&a.1));

    let principal_labels: Vec<u32> = sorted_counts
        .iter()
        .take(n)
        .map(|(label, _)| *label)
        .collect();

    let colors = generate_contrasting_colors(n, 255);
    let color_map: HashMap<u32, Rgba<u8>> = principal_labels.into_iter().zip(colors).collect();

    let (width, height) = labelled_image.dimensions();
    let mut output_image = RgbaImage::from_pixel(width, height, background_color);

    for (x, y, pixel) in labelled_image.enumerate_pixels() {
        let label = pixel[0];
        if let Some(color) = color_map.get(&label) {
            output_image.put_pixel(x, y, *color);
        }
    }

    output_image
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_draw_principal_connected_components() {
        // Arrange:
        // Define a labelled image with multiple components of different sizes.
        // This simulates the output of `connected_components`.
        // 1. A large 4x4 square (16 pixels), label 1.
        // 2. A medium 2x3 rectangle (6 pixels), label 2.
        // 3. A small 1x2 line (2 pixels), label 3.
        #[rustfmt::skip]
        let labels_vec = vec![
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 1, 1, 1, 1, 0, 0, 0, 0, 0,
            0, 1, 1, 1, 1, 0, 2, 2, 2, 0,
            0, 1, 1, 1, 1, 0, 2, 2, 2, 0,
            0, 1, 1, 1, 1, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 3, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 3, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        ];

        let labelled_image = ImageBuffer::<Luma<u32>, _>::from_raw(10, 10, labels_vec).unwrap();

        // We want to draw the two largest components (n=2).
        let n = 2;
        let background = Rgba([0, 0, 0, 255]);

        // Act:
        let result_image = draw_principal_connected_components(&labelled_image, n, background);

        // Assert:
        // -- Create the expected output image --

        // The function will identify the largest components.
        // Counts: {1: 16, 2: 6, 3: 2}.
        // The two largest labels are 1 and 2.

        // Get the colors that will be generated for these two labels.
        let colors = generate_contrasting_colors(n, 255);
        let color_for_label1 = colors[0]; // The largest component gets the first color.
        let color_for_label2 = colors[1]; // The second largest component gets the second color.

        let mut expected_image = RgbaImage::from_pixel(10, 10, background);

        // Draw the largest component (label 1).
        for y in 1..=4 {
            for x in 1..=4 {
                expected_image.put_pixel(x, y, color_for_label1);
            }
        }

        // Draw the second largest component (label 2).
        for y in 2..=3 {
            for x in 6..=8 {
                expected_image.put_pixel(x, y, color_for_label2);
            }
        }

        // The smallest component (label 3) should remain as the background color,
        // so we don't need to draw it on the expected_image.

        assert_eq!(result_image, expected_image);
    }

    #[test]
    fn test_draw_principal_connected_components_empty() {
        // All background image.
        let labelled_image = ImageBuffer::<Luma<u32>, Vec<u32>>::new(5, 5);
        let background = Rgba([255, 255, 255, 255]);

        // n = 0
        let result = draw_principal_connected_components(&labelled_image, 0, background);
        assert_eq!(result.dimensions(), (5, 5));
        assert!(result.pixels().all(|p| *p == background));
    }
}