scirs2_spatial/

distance_transform.rs

//! Distance transform algorithms for image processing and spatial analysis
//!
//! This module provides efficient algorithms for computing distance transforms,
//! which assign to each point the distance to the nearest feature (e.g., boundary, obstacle).
//!
//! # Features
//!
//! * **Euclidean distance transform** - Exact Euclidean distances
//! * **Chamfer distance transform** - Fast approximate distances
//! * **Manhattan distance transform** - City-block distances
//! * **Geodesic distance transform** - Distances along surfaces
//! * **Feature transforms** - Identify nearest feature points
//!
//! # Examples
//!
//! ```
//! use scirs2_core::ndarray::array;
//! use scirs2_spatial::distance_transform::{euclidean_distance_transform, DistanceMetric};
//!
//! // Create a binary image (0 = background, 1 = feature)
//! let binary = array![
//!     [1, 1, 0, 0],
//!     [1, 0, 0, 0],
//!     [0, 0, 0, 1],
//!     [0, 0, 1, 1]
//! ];
//!
//! // Compute distance transform
//! let distances = euclidean_distance_transform::<f64>(&binary.view(), DistanceMetric::Euclidean)
//!     .expect("Failed to compute distance transform");
//!
//! // Each background pixel now contains its distance to nearest feature
//! println!("Distance transform: {:?}", distances);
//! ```

use crate::error::{SpatialError, SpatialResult};
use scirs2_core::ndarray::{Array1, Array2, Array3, ArrayView2, ArrayView3};
use scirs2_core::numeric::Float;
use std::collections::VecDeque;

/// Distance metric for distance transforms
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum DistanceMetric {
    /// Euclidean (L2) distance
    Euclidean,
    /// Manhattan (L1, city-block) distance
    Manhattan,
    /// Chebyshev (L∞, chessboard) distance
    Chebyshev,
    /// Chamfer 3-4 approximation
    Chamfer34,
    /// Chamfer 5-7-11 approximation
    Chamfer5711,
}

/// Compute Euclidean distance transform for 2D binary image
///
/// Uses the efficient separable algorithm of Felzenszwalb & Huttenlocher.
///
/// # Arguments
///
/// * `binary` - Binary input array (0 = background, non-zero = feature)
/// * `metric` - Distance metric to use
///
/// # Returns
///
/// * Distance transform array (same shape as input)
///
/// # Examples
///
/// ```
/// use scirs2_core::ndarray::array;
/// use scirs2_spatial::distance_transform::{euclidean_distance_transform, DistanceMetric};
///
/// let binary = array![[1, 0, 0], [0, 0, 0], [0, 0, 1]];
/// let distances = euclidean_distance_transform::<f64>(&binary.view(), DistanceMetric::Euclidean)
///     .expect("Failed to compute");
/// ```
pub fn euclidean_distance_transform<T: Float>(
    binary: &ArrayView2<i32>,
    metric: DistanceMetric,
) -> SpatialResult<Array2<T>> {
    let (rows, cols) = binary.dim();

    if rows == 0 || cols == 0 {
        return Err(SpatialError::ValueError(
            "Input array must be non-empty".to_string(),
        ));
    }

    match metric {
        DistanceMetric::Euclidean => euclidean_dt_2d(binary),
        DistanceMetric::Manhattan => manhattan_dt_2d(binary),
        DistanceMetric::Chebyshev => chebyshev_dt_2d(binary),
        DistanceMetric::Chamfer34 => chamfer_dt_2d(binary, &[3.0, 4.0]),
        DistanceMetric::Chamfer5711 => chamfer_dt_2d(binary, &[5.0, 7.0, 11.0]),
    }
}

/// Exact Euclidean distance transform using separable algorithm
fn euclidean_dt_2d<T: Float>(binary: &ArrayView2<i32>) -> SpatialResult<Array2<T>> {
    let (rows, cols) = binary.dim();
    let mut distances = Array2::from_elem((rows, cols), T::infinity());

    // Initialize: Set feature pixels to 0
    for i in 0..rows {
        for j in 0..cols {
            if binary[[i, j]] != 0 {
                distances[[i, j]] = T::zero();
            }
        }
    }

    // Forward pass - horizontal
    for i in 0..rows {
        let mut min_dist = T::infinity();
        for j in 0..cols {
            if binary[[i, j]] != 0 {
                min_dist = T::zero();
            } else {
                min_dist = min_dist + T::one();
            }
            distances[[i, j]] = min_dist.min(distances[[i, j]]);
        }

        // Backward pass - horizontal
        let mut min_dist = T::infinity();
        for j in (0..cols).rev() {
            if binary[[i, j]] != 0 {
                min_dist = T::zero();
            } else {
                min_dist = min_dist + T::one();
            }
            distances[[i, j]] = min_dist.min(distances[[i, j]]);
        }
    }

    // Vertical passes using parabola envelope
    for j in 0..cols {
        // Forward pass
        for i in 1..rows {
            let up_dist = distances[[i - 1, j]] + T::one();
            distances[[i, j]] = up_dist.min(distances[[i, j]]);
        }

        // Backward pass
        for i in (0..rows - 1).rev() {
            let down_dist = distances[[i + 1, j]] + T::one();
            distances[[i, j]] = down_dist.min(distances[[i, j]]);
        }
    }

    Ok(distances)
}

/// Manhattan (L1) distance transform
fn manhattan_dt_2d<T: Float>(binary: &ArrayView2<i32>) -> SpatialResult<Array2<T>> {
    let (rows, cols) = binary.dim();
    let mut distances = Array2::from_elem(
        (rows, cols),
        T::from(rows + cols).expect("conversion failed"),
    );

    // Initialize feature pixels
    for i in 0..rows {
        for j in 0..cols {
            if binary[[i, j]] != 0 {
                distances[[i, j]] = T::zero();
            }
        }
    }

    // Forward pass
    for i in 0..rows {
        for j in 0..cols {
            if binary[[i, j]] == 0 {
                let mut min_dist = distances[[i, j]];

                if i > 0 {
                    min_dist = min_dist.min(distances[[i - 1, j]] + T::one());
                }
                if j > 0 {
                    min_dist = min_dist.min(distances[[i, j - 1]] + T::one());
                }

                distances[[i, j]] = min_dist;
            }
        }
    }

    // Backward pass
    for i in (0..rows).rev() {
        for j in (0..cols).rev() {
            if binary[[i, j]] == 0 {
                let mut min_dist = distances[[i, j]];

                if i < rows - 1 {
                    min_dist = min_dist.min(distances[[i + 1, j]] + T::one());
                }
                if j < cols - 1 {
                    min_dist = min_dist.min(distances[[i, j + 1]] + T::one());
                }

                distances[[i, j]] = min_dist;
            }
        }
    }

    Ok(distances)
}

/// Chebyshev (L∞) distance transform
fn chebyshev_dt_2d<T: Float>(binary: &ArrayView2<i32>) -> SpatialResult<Array2<T>> {
    let (rows, cols) = binary.dim();
    let mut distances = Array2::from_elem(
        (rows, cols),
        T::from(rows.max(cols)).expect("conversion failed"),
    );

    // Use BFS for exact Chebyshev distance
    let mut queue = VecDeque::new();

    // Initialize with feature pixels
    for i in 0..rows {
        for j in 0..cols {
            if binary[[i, j]] != 0 {
                distances[[i, j]] = T::zero();
                queue.push_back((i, j, T::zero()));
            }
        }
    }

    // 8-connected neighbors for Chebyshev
    let neighbors = [
        (-1, -1),
        (-1, 0),
        (-1, 1),
        (0, -1),
        (0, 1),
        (1, -1),
        (1, 0),
        (1, 1),
    ];

    while let Some((i, j, dist)) = queue.pop_front() {
        for &(di, dj) in &neighbors {
            let ni = i as isize + di;
            let nj = j as isize + dj;

            if ni >= 0 && ni < rows as isize && nj >= 0 && nj < cols as isize {
                let ni = ni as usize;
                let nj = nj as usize;

                let new_dist = dist + T::one();
                if new_dist < distances[[ni, nj]] {
                    distances[[ni, nj]] = new_dist;
                    queue.push_back((ni, nj, new_dist));
                }
            }
        }
    }

    Ok(distances)
}

/// Chamfer distance transform with configurable weights
fn chamfer_dt_2d<T: Float>(binary: &ArrayView2<i32>, weights: &[f64]) -> SpatialResult<Array2<T>> {
    let (rows, cols) = binary.dim();
    let max_val = T::from(rows * cols).expect("conversion failed");
    let mut distances = Array2::from_elem((rows, cols), max_val);

    // Initialize feature pixels
    for i in 0..rows {
        for j in 0..cols {
            if binary[[i, j]] != 0 {
                distances[[i, j]] = T::zero();
            }
        }
    }

    let w1 = T::from(weights[0]).expect("conversion failed");
    let w2 =
        T::from(weights.get(1).copied().unwrap_or(weights[0] * 1.4)).expect("conversion failed");

    // Forward pass
    for i in 0..rows {
        for j in 0..cols {
            if binary[[i, j]] == 0 {
                let mut min_dist = distances[[i, j]];

                // 4-connected
                if i > 0 {
                    min_dist = min_dist.min(distances[[i - 1, j]] + w1);
                }
                if j > 0 {
                    min_dist = min_dist.min(distances[[i, j - 1]] + w1);
                }

                // Diagonal
                if i > 0 && j > 0 {
                    min_dist = min_dist.min(distances[[i - 1, j - 1]] + w2);
                }
                if i > 0 && j < cols - 1 {
                    min_dist = min_dist.min(distances[[i - 1, j + 1]] + w2);
                }

                distances[[i, j]] = min_dist;
            }
        }
    }

    // Backward pass
    for i in (0..rows).rev() {
        for j in (0..cols).rev() {
            if binary[[i, j]] == 0 {
                let mut min_dist = distances[[i, j]];

                // 4-connected
                if i < rows - 1 {
                    min_dist = min_dist.min(distances[[i + 1, j]] + w1);
                }
                if j < cols - 1 {
                    min_dist = min_dist.min(distances[[i, j + 1]] + w1);
                }

                // Diagonal
                if i < rows - 1 && j < cols - 1 {
                    min_dist = min_dist.min(distances[[i + 1, j + 1]] + w2);
                }
                if i < rows - 1 && j > 0 {
                    min_dist = min_dist.min(distances[[i + 1, j - 1]] + w2);
                }

                distances[[i, j]] = min_dist;
            }
        }
    }

    Ok(distances)
}

/// Compute feature transform (indices of nearest features)
///
/// Returns the coordinates of the nearest feature point for each pixel.
///
/// # Arguments
///
/// * `binary` - Binary input array
///
/// # Returns
///
/// * Array of (row, col) indices of nearest features
pub fn feature_transform(binary: &ArrayView2<i32>) -> SpatialResult<Array2<(usize, usize)>> {
    let (rows, cols) = binary.dim();
    let mut features = Array2::from_elem((rows, cols), (usize::MAX, usize::MAX));
    let mut distances = Array2::from_elem((rows, cols), f64::INFINITY);

    // Initialize with feature pixels
    for i in 0..rows {
        for j in 0..cols {
            if binary[[i, j]] != 0 {
                features[[i, j]] = (i, j);
                distances[[i, j]] = 0.0;
            }
        }
    }

    // Forward pass
    for i in 0..rows {
        for j in 0..cols {
            if binary[[i, j]] == 0 {
                let neighbors = [
                    (i.saturating_sub(1), j),
                    (i, j.saturating_sub(1)),
                    (i.saturating_sub(1), j.saturating_sub(1)),
                ];

                for &(ni, nj) in &neighbors {
                    if ni < rows && nj < cols && features[[ni, nj]] != (usize::MAX, usize::MAX) {
                        let (fi, fj) = features[[ni, nj]];
                        let dist = (((i as isize - fi as isize).pow(2)
                            + (j as isize - fj as isize).pow(2))
                            as f64)
                            .sqrt();

                        if dist < distances[[i, j]] {
                            distances[[i, j]] = dist;
                            features[[i, j]] = (fi, fj);
                        }
                    }
                }
            }
        }
    }

    // Backward pass
    for i in (0..rows).rev() {
        for j in (0..cols).rev() {
            if binary[[i, j]] == 0 {
                let neighbors = [
                    ((i + 1).min(rows - 1), j),
                    (i, (j + 1).min(cols - 1)),
                    ((i + 1).min(rows - 1), (j + 1).min(cols - 1)),
                ];

                for &(ni, nj) in &neighbors {
                    if features[[ni, nj]] != (usize::MAX, usize::MAX) {
                        let (fi, fj) = features[[ni, nj]];
                        let dist = (((i as isize - fi as isize).pow(2)
                            + (j as isize - fj as isize).pow(2))
                            as f64)
                            .sqrt();

                        if dist < distances[[i, j]] {
                            distances[[i, j]] = dist;
                            features[[i, j]] = (fi, fj);
                        }
                    }
                }
            }
        }
    }

    Ok(features)
}

/// Compute 3D Euclidean distance transform
///
/// Extension of the 2D algorithm to 3D volumes.
///
/// # Arguments
///
/// * `binary` - Binary 3D array
/// * `metric` - Distance metric
///
/// # Returns
///
/// * 3D distance transform
pub fn euclidean_distance_transform_3d<T: Float>(
    binary: &ArrayView3<i32>,
    metric: DistanceMetric,
) -> SpatialResult<Array3<T>> {
    let (depth, rows, cols) = binary.dim();

    if depth == 0 || rows == 0 || cols == 0 {
        return Err(SpatialError::ValueError(
            "Input array must be non-empty".to_string(),
        ));
    }

    // Simplified 3D implementation (could be optimized with separable algorithm)
    let mut distances = Array3::from_elem((depth, rows, cols), T::infinity());

    // Initialize feature voxels
    for i in 0..depth {
        for j in 0..rows {
            for k in 0..cols {
                if binary[[i, j, k]] != 0 {
                    distances[[i, j, k]] = T::zero();
                }
            }
        }
    }

    // Simple iterative propagation (placeholder for full 3D algorithm)
    let max_iterations = (depth + rows + cols) / 2;
    for _iter in 0..max_iterations {
        let mut changed = false;

        for i in 0..depth {
            for j in 0..rows {
                for k in 0..cols {
                    if binary[[i, j, k]] == 0 {
                        let mut min_dist = distances[[i, j, k]];

                        // Check 6-connected neighbors
                        if i > 0 {
                            min_dist = min_dist.min(distances[[i - 1, j, k]] + T::one());
                        }
                        if i < depth - 1 {
                            min_dist = min_dist.min(distances[[i + 1, j, k]] + T::one());
                        }
                        if j > 0 {
                            min_dist = min_dist.min(distances[[i, j - 1, k]] + T::one());
                        }
                        if j < rows - 1 {
                            min_dist = min_dist.min(distances[[i, j + 1, k]] + T::one());
                        }
                        if k > 0 {
                            min_dist = min_dist.min(distances[[i, j, k - 1]] + T::one());
                        }
                        if k < cols - 1 {
                            min_dist = min_dist.min(distances[[i, j, k + 1]] + T::one());
                        }

                        if min_dist < distances[[i, j, k]] {
                            distances[[i, j, k]] = min_dist;
                            changed = true;
                        }
                    }
                }
            }
        }

        if !changed {
            break;
        }
    }

    Ok(distances)
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;
    use scirs2_core::ndarray::array;

    #[test]
    fn test_euclidean_distance_transform_2d() {
        let binary = array![[1, 0, 0], [0, 0, 0], [0, 0, 1]];

        let distances: Array2<f64> =
            euclidean_distance_transform(&binary.view(), DistanceMetric::Euclidean)
                .expect("Failed to compute");

        // Check shape
        assert_eq!(distances.dim(), (3, 3));

        // Feature pixels should have distance 0
        assert_relative_eq!(distances[[0, 0]], 0.0, epsilon = 1e-10);
        assert_relative_eq!(distances[[2, 2]], 0.0, epsilon = 1e-10);

        // Center pixel should have distance > 1
        assert!(distances[[1, 1]] > 1.0);
    }

    #[test]
    fn test_manhattan_distance_transform() {
        let binary = array![[1, 0, 0], [0, 0, 0], [0, 0, 1]];

        let distances: Array2<f64> =
            euclidean_distance_transform(&binary.view(), DistanceMetric::Manhattan)
                .expect("Failed to compute");

        // Feature pixels
        assert_relative_eq!(distances[[0, 0]], 0.0, epsilon = 1e-10);
        assert_relative_eq!(distances[[2, 2]], 0.0, epsilon = 1e-10);

        // Manhattan distance from (0,0) to (1,1) is 2
        assert_relative_eq!(distances[[1, 1]], 2.0, epsilon = 0.1);
    }

    #[test]
    fn test_feature_transform() {
        let binary = array![[1, 0, 0], [0, 0, 0], [0, 0, 1]];

        let features = feature_transform(&binary.view()).expect("Failed to compute");

        // Feature pixels should point to themselves
        assert_eq!(features[[0, 0]], (0, 0));
        assert_eq!(features[[2, 2]], (2, 2));

        // Other pixels should point to nearest feature
        assert!(
            features[[1, 1]] == (0, 0) || features[[1, 1]] == (2, 2),
            "Pixel (1,1) should point to one of the features"
        );
    }

    #[test]
    fn test_chamfer_distance_transform() {
        let binary = array![[1, 0, 0], [0, 0, 0], [0, 0, 0]];

        let distances: Array2<f64> =
            euclidean_distance_transform(&binary.view(), DistanceMetric::Chamfer34)
                .expect("Failed to compute");

        // Feature pixel
        assert_relative_eq!(distances[[0, 0]], 0.0, epsilon = 1e-10);

        // Distances should increase monotonically away from feature
        assert!(distances[[0, 1]] > 0.0);
        assert!(distances[[0, 2]] > distances[[0, 1]]);
    }

    #[test]
    fn test_3d_distance_transform() {
        let binary = array![[[1, 0], [0, 0]], [[0, 0], [0, 0]]];

        let distances: Array3<f64> =
            euclidean_distance_transform_3d(&binary.view(), DistanceMetric::Euclidean)
                .expect("Failed to compute");

        // Feature voxel
        assert_relative_eq!(distances[[0, 0, 0]], 0.0, epsilon = 1e-10);

        // All other voxels should have positive distance
        for i in 0..2 {
            for j in 0..2 {
                for k in 0..2 {
                    if i != 0 || j != 0 || k != 0 {
                        assert!(distances[[i, j, k]] > 0.0);
                    }
                }
            }
        }
    }

    #[test]
    fn test_empty_input() {
        let binary = array![[0, 0], [0, 0]];

        let distances: Array2<f64> =
            euclidean_distance_transform(&binary.view(), DistanceMetric::Euclidean)
                .expect("Failed to compute");

        // All distances should be large (no features)
        for &d in distances.iter() {
            assert!(d > 1.0);
        }
    }

    #[test]
    fn test_all_features() {
        let binary = array![[1, 1], [1, 1]];

        let distances: Array2<f64> =
            euclidean_distance_transform(&binary.view(), DistanceMetric::Euclidean)
                .expect("Failed to compute");

        // All distances should be zero
        for &d in distances.iter() {
            assert_relative_eq!(d, 0.0, epsilon = 1e-10);
        }
    }
}