基于ndarray实现边缘检测的更惯用写法及抽象优化问询

问题概述 #

我有一个Array2<u32>类型的图像：

const SIZE: usize = 8;
let image = Array2::from_elem((SIZE, SIZE), 0_u32);

需要生成一个Array2<bool>，其中每个位置为true的条件是：该像素的4连通邻域（上下左右）中至少有一个像素与自身值不同，结果允许各边缩小2个像素。

当前手动拼接邻域切片的实现方式重复代码多，不够优雅，希望构建更简洁的抽象层。

优化方案1：自定义邻域迭代器（推荐） #

封装一个扩展trait，让Zip可以直接返回包含中心像素和邻域迭代器的结构，避免重复编写邻域参数：

use ndarray::{Array2, ArrayView2, Zip, s};
use std::iter::IntoIterator;

trait ZipNeighborsExt {
    fn with_4_neighbors(self, image: &Array2<u32>) -> Zip<(Self::Item, impl Iterator<Item = &u32>)>
    where
        Self: Sized;
}

impl<'a> ZipNeighborsExt for Zip<ArrayView2<'a, u32>> {
    fn with_4_neighbors(self, image: &Array2<u32>) -> Zip<(Self::Item, impl Iterator<Item = &u32>)> {
        let rows = image.nrows();
        let cols = image.ncols();
        let top = image.slice(s![0..rows-2, 1..cols-1]);
        let bottom = image.slice(s![2..rows, 1..cols-1]);
        let left = image.slice(s![1..rows-1, 0..cols-2]);
        let right = image.slice(s![1..rows-1, 2..cols]);

        Zip::from((self.into_inner(), top, bottom, left, right))
            .map(|(c, n, s, w, e)| (c, [n, s, w, e].into_iter()))
    }
}

使用时只需关注中心像素和邻域迭代器，写法大幅简化：

let center = image.slice(s![1..image.nrows()-1, 1..image.ncols()-1]);
let edges = Zip::from(center)
    .with_4_neighbors(&image)
    .map_collect(|(c, neighbors)| neighbors.any(|&v| v != c));

优化方案2：封装邻域视图结构体 #

自定义结构体包装四个邻域视图，实现IntoIterator让它能被Zip直接处理，接近你想要的“带额外邻域轴”的效果：

use ndarray::{Array2, ArrayView2, Zip, s};

struct FourNeighbors<'a> {
    top: ArrayView2<'a, u32>,
    bottom: ArrayView2<'a, u32>,
    left: ArrayView2<'a, u32>,
    right: ArrayView2<'a, u32>,
}

impl<'a> FourNeighbors<'a> {
    fn new(image: &'a Array2<u32>) -> Self {
        let rows = image.nrows();
        let cols = image.ncols();
        Self {
            top: image.slice(s![0..rows-2, 1..cols-1]),
            bottom: image.slice(s![2..rows, 1..cols-1]),
            left: image.slice(s![1..rows-1, 0..cols-2]),
            right: image.slice(s![1..rows-1, 2..cols]),
        }
    }
}

impl<'a> IntoIterator for FourNeighbors<'a> {
    type Item = (&'a u32, &'a u32, &'a u32, &'a u32);
    type IntoIter = Zip<(ArrayView2<'a, u32>, ArrayView2<'a, u32>, ArrayView2<'a, u32>, ArrayView2<'a, u32>)>;

    fn into_iter(self) -> Self::IntoIter {
        Zip::from((self.top, self.bottom, self.left, self.right))
    }
}

结合Zip使用的示例：

let center = image.slice(s![1..image.nrows()-1, 1..image.ncols()-1]);
let neighbors = FourNeighbors::new(&image);

let edges = Zip::from((center, neighbors))
    .map_collect(|(&c, (n, s, w, e))| [n, s, w, e].any(|&v| v != c));

方案说明 #

• 两种方案都将邻域切片逻辑封装为可复用结构，避免重复代码
• 方案1更简洁，闭包直接处理迭代器，适合多数场景
• 方案2的结构体可扩展支持8连通等更多邻域类型，灵活性更强

内容的提问来源于stack exchange，提问作者Thomas