glint_mask_tools/core/

postprocessor.rs

/// Post-processing abstraction for mask refinement.
///
/// This module defines the [`PostProcessor`] trait which provides a uniform
/// interface for post-processing generated masks. The primary post-processor
/// is pixel buffering, but the trait allows for composable operations.
use crate::error::Result;
use ndarray::Array2;

/// Trait for post-processing glint masks.
///
/// Post-processors take binary masks and apply various refinement operations
/// such as morphological operations, buffering, or filtering.
pub trait PostProcessor: Send + Sync {
    /// Apply post-processing to a binary mask
    ///
    /// # Arguments
    ///
    /// * `mask` - Binary mask with shape (height, width)
    ///
    /// # Returns
    ///
    /// A processed binary mask with the same shape
    fn process_mask(&self, mask: &Array2<u8>) -> Result<Array2<u8>>;

    /// Get the name of this post-processor
    fn name(&self) -> &'static str;

    /// Get a description of this post-processor
    fn description(&self) -> &'static str;

    /// Validate that the post-processor parameters are valid
    fn validate_parameters(&self) -> Result<()> {
        Ok(())
    }
}

/// Composite post-processor that applies multiple processors in sequence
pub struct CompositePostProcessor {
    processors: Vec<Box<dyn PostProcessor>>,
}

impl CompositePostProcessor {
    /// Create a new composite post-processor
    pub fn new() -> Self {
        Self {
            processors: Vec::new(),
        }
    }

    /// Add a post-processor to the pipeline
    pub fn add_processor(mut self, processor: Box<dyn PostProcessor>) -> Self {
        self.processors.push(processor);
        self
    }

    /// Get the number of processors in the pipeline
    pub fn len(&self) -> usize {
        self.processors.len()
    }

    /// Check if the pipeline is empty
    pub fn is_empty(&self) -> bool {
        self.processors.is_empty()
    }
}

impl Default for CompositePostProcessor {
    fn default() -> Self {
        Self::new()
    }
}

impl PostProcessor for CompositePostProcessor {
    fn process_mask(&self, mask: &Array2<u8>) -> Result<Array2<u8>> {
        let mut result = mask.clone();

        for processor in &self.processors {
            result = processor.process_mask(&result)?;
        }

        Ok(result)
    }

    fn name(&self) -> &'static str {
        "Composite"
    }

    fn description(&self) -> &'static str {
        "Applies multiple post-processors in sequence"
    }

    fn validate_parameters(&self) -> Result<()> {
        for processor in &self.processors {
            processor.validate_parameters()?;
        }
        Ok(())
    }
}

/// Pixel buffer post-processor that dilates the mask by a specified radius
#[derive(Debug, Clone)]
pub struct PixelBufferProcessor {
    radius: usize,
    kernel: Array2<bool>,
}

impl PixelBufferProcessor {
    /// Create a new pixel buffer processor
    pub fn new(radius: usize) -> Self {
        let kernel = create_circular_kernel(radius);
        Self { radius, kernel }
    }

    /// Get the buffer radius
    pub fn radius(&self) -> usize {
        self.radius
    }

    /// Create a kernel for morphological operations
    fn create_kernel(&self) -> &Array2<bool> {
        &self.kernel
    }
}

impl PostProcessor for PixelBufferProcessor {
    fn process_mask(&self, mask: &Array2<u8>) -> Result<Array2<u8>> {
        if self.radius == 0 {
            return Ok(mask.clone());
        }

        let (height, width) = mask.dim();
        let mut result = Array2::zeros((height, width));
        let kernel = self.create_kernel();
        let (kernel_height, kernel_width) = kernel.dim();
        let kernel_center_y = kernel_height / 2;
        let kernel_center_x = kernel_width / 2;

        for y in 0..height {
            for x in 0..width {
                let mut should_mask = false;

                // Check the kernel area around this pixel
                for ky in 0..kernel_height {
                    for kx in 0..kernel_width {
                        if !kernel[[ky, kx]] {
                            continue;
                        }

                        let img_y = y as i32 + ky as i32 - kernel_center_y as i32;
                        let img_x = x as i32 + kx as i32 - kernel_center_x as i32;

                        if img_y >= 0
                            && img_y < height as i32
                            && img_x >= 0
                            && img_x < width as i32
                            && mask[[img_y as usize, img_x as usize]] > 0
                        {
                            should_mask = true;
                            break;
                        }
                    }
                    if should_mask {
                        break;
                    }
                }

                result[[y, x]] = if should_mask { 1 } else { 0 };
            }
        }

        Ok(result)
    }

    fn name(&self) -> &'static str {
        "PixelBuffer"
    }

    fn description(&self) -> &'static str {
        "Dilates the mask by a specified pixel radius"
    }

    fn validate_parameters(&self) -> Result<()> {
        // Radius validation could be added here if needed
        Ok(())
    }
}

/// Metashape format converter that inverts the mask and converts to the expected format
#[derive(Debug, Clone)]
pub struct MetashapeConverter;

impl MetashapeConverter {
    /// Create a new Metashape converter
    pub fn new() -> Self {
        Self
    }
}

impl Default for MetashapeConverter {
    fn default() -> Self {
        Self::new()
    }
}

impl PostProcessor for MetashapeConverter {
    fn process_mask(&self, mask: &Array2<u8>) -> Result<Array2<u8>> {
        // Invert the mask and scale to 0-255 range
        // In our convention: 1 = mask (glint), 0 = keep
        // In Metashape: 0 = mask (ignore), 255 = keep
        let result = mask.map(|&pixel| if pixel > 0 { 0 } else { 255 });
        Ok(result)
    }

    fn name(&self) -> &'static str {
        "Metashape"
    }

    fn description(&self) -> &'static str {
        "Converts mask to Metashape format (inverted, 0-255 range)"
    }
}

/// Create a circular kernel for morphological operations
fn create_circular_kernel(radius: usize) -> Array2<bool> {
    let size = 2 * radius + 1;
    let mut kernel = Array2::from_elem((size, size), false);
    let center = radius as i32;

    for y in 0..size {
        for x in 0..size {
            let dy = y as i32 - center;
            let dx = x as i32 - center;
            let distance_sq = dx * dx + dy * dy;

            if distance_sq <= (radius as i32) * (radius as i32) {
                kernel[[y, x]] = true;
            }
        }
    }

    kernel
}

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

    #[test]
    fn test_circular_kernel() {
        let kernel = create_circular_kernel(1);
        assert_eq!(kernel.dim(), (3, 3));
        // Center should be true
        assert!(kernel[[1, 1]]);
        // Adjacent pixels should be true
        assert!(kernel[[0, 1]]);
        assert!(kernel[[1, 0]]);
        assert!(kernel[[2, 1]]);
        assert!(kernel[[1, 2]]);
        // Corners should be false for radius 1
        assert!(!kernel[[0, 0]]);
        assert!(!kernel[[2, 2]]);
    }

    #[test]
    fn test_pixel_buffer_processor() {
        let mut mask = ndarray::Array2::zeros((5, 5));
        mask[[2, 2]] = 1; // Single pixel in center

        let processor = PixelBufferProcessor::new(1);
        let result = processor.process_mask(&mask).unwrap();

        // Should have expanded to adjacent pixels
        assert_eq!(result[[2, 2]], 1);
        assert_eq!(result[[1, 2]], 1);
        assert_eq!(result[[3, 2]], 1);
        assert_eq!(result[[2, 1]], 1);
        assert_eq!(result[[2, 3]], 1);

        // Corners should still be 0
        assert_eq!(result[[0, 0]], 0);
        assert_eq!(result[[4, 4]], 0);
    }

    #[test]
    fn test_metashape_converter() {
        let mut mask = ndarray::Array2::zeros((3, 3));
        mask[[1, 1]] = 1; // Single masked pixel

        let converter = MetashapeConverter::new();
        let result = converter.process_mask(&mask).unwrap();

        // Masked pixel should become 0
        assert_eq!(result[[1, 1]], 0);
        // Unmasked pixels should become 255
        assert_eq!(result[[0, 0]], 255);
        assert_eq!(result[[2, 2]], 255);
    }

    #[test]
    fn test_composite_processor() {
        let mut mask = ndarray::Array2::zeros((5, 5));
        mask[[2, 2]] = 1;

        let processor = CompositePostProcessor::new()
            .add_processor(Box::new(PixelBufferProcessor::new(1)))
            .add_processor(Box::new(MetashapeConverter::new()));

        let result = processor.process_mask(&mask).unwrap();

        // Should have buffered and then converted to Metashape format
        assert_eq!(result[[2, 2]], 0); // Center (was masked, buffered, then inverted)
        assert_eq!(result[[1, 2]], 0); // Adjacent (was buffered, then inverted)
        assert_eq!(result[[0, 0]], 255); // Corner (was not buffered, inverted to 255)
    }
}