Python转JavaScript：Numpy相关非极大值抑制函数转换难题

刚好我之前也处理过类似的需求——把依赖Numpy的非极大值抑制函数转成JavaScript，确实因为JS没有现成的maximum_filter这类函数，得自己手动实现核心逻辑。先把你提到的Python代码补全完整，再一步步拆解转成JS版本。

原Python代码（补全完整逻辑） #

你给出的代码没写完，我先补全Numpy版本的完整实现，方便对照：

import numpy as np

# 定义默认阈值
NMS_Threshold = 0.5

def non_max_suppression(plain, window_size=3, threshold=NMS_Threshold):
    # 第一步：把所有低于阈值的元素置0
    under_threshold_indices = plain < threshold
    plain[under_threshold_indices] = 0
    # 第二步：用window_size大小的窗口做最大值滤波，仅保留局部最大值的位置
    return plain * (plain == np.maximum_filter(plain, footprint=np.ones((window_size, window_size))))

核心逻辑拆解 #

要转成JS，得先理清原代码做了三件事：

• 把数组中低于阈值的元素全部置0
• 对整个数组做窗口大小为window_size的最大值滤波：每个位置的值替换成其周围window_size×window_size窗口内的最大值
• 只保留原数组中等于对应窗口最大值的元素（也就是局部最大值），其余元素置0

JavaScript实现 #

因为JS没有内置的最大值滤波函数，我们先实现一个maximumFilter工具函数，再封装完整的NMS函数：

// 定义默认阈值（对应Python里的NMS_Threshold）
const NMS_THRESHOLD = 0.5;

/**
 * 模拟Numpy的maximum_filter，计算二维数组每个位置窗口内的最大值
 * @param {number[][]} arr - 输入的二维浮点数组
 * @param {number} windowSize - 窗口大小（默认3，即3x3窗口）
 * @returns {number[][]} 每个位置窗口内的最大值组成的数组
 */
function maximumFilter(arr, windowSize = 3) {
    const rows = arr.length;
    if (rows === 0) return [];
    const cols = arr[0].length;
    const halfWindow = Math.floor(windowSize / 2);
    const result = Array(rows).fill().map(() => Array(cols).fill(0));

    // 遍历数组的每个元素位置
    for (let i = 0; i < rows; i++) {
        for (let j = 0; j < cols; j++) {
            let maxVal = -Infinity;
            // 遍历当前位置周围的窗口范围
            for (let x = -halfWindow; x <= halfWindow; x++) {
                for (let y = -halfWindow; y <= halfWindow; y++) {
                    const rowIdx = i + x;
                    const colIdx = j + y;
                    // 只处理数组范围内的元素，避免越界
                    if (rowIdx >= 0 && rowIdx < rows && colIdx >= 0 && colIdx < cols) {
                        if (arr[rowIdx][colIdx] > maxVal) {
                            maxVal = arr[rowIdx][colIdx];
                        }
                    }
                }
            }
            result[i][j] = maxVal;
        }
    }
    return result;
}

/**
 * 非极大值抑制函数，移除局部最大值周围的非最大值
 * @param {number[][]} plain - 输入的二维浮点数组
 * @param {number} windowSize - 窗口大小（默认3）
 * @param {number} threshold - 阈值（默认NMS_THRESHOLD）
 * @returns {number[][]} 处理后的数组，仅保留局部最大值（低于阈值的已置0）
 */
function nonMaxSuppression(plain, windowSize = 3, threshold = NMS_THRESHOLD) {
    const rows = plain.length;
    if (rows === 0) return [];
    const cols = plain[0].length;
    // 复制原数组，避免直接修改输入（如果不需要保留原数组，可以跳过这步直接操作）
    const processed = plain.map(row => [...row]);

    // 第一步：把低于阈值的元素置0
    for (let i = 0; i < rows; i++) {
        for (let j = 0; j < cols; j++) {
            if (processed[i][j] < threshold) {
                processed[i][j] = 0;
            }
        }
    }

    // 第二步：计算窗口最大值数组
    const maxFiltered = maximumFilter(processed, windowSize);

    // 第三步：仅保留原数组中等于窗口最大值的元素，其余置0
    for (let i = 0; i < rows; i++) {
        for (let j = 0; j < cols; j++) {
            if (processed[i][j] !== maxFiltered[i][j]) {
                processed[i][j] = 0;
            }
        }
    }

    return processed;
}

使用示例 #

可以用下面的测试代码验证效果：

// 测试用的二维浮点数组
const inputArray = [
    [0.2, 0.8, 0.3],
    [0.6, 0.9, 0.4],
    [0.1, 0.7, 0.5]
];

// 调用NMS函数
const result = nonMaxSuppression(inputArray);
console.log(result);
// 输出结果：
// [[0, 0, 0], [0, 0.9, 0], [0, 0, 0]]
// 只有中心的0.9是3x3窗口内的最大值，其余元素要么低于阈值，要么不是局部最大值，都被置0

补充说明 #

• 上面的maximumFilter处理边界时，会忽略超出数组范围的元素（只取窗口内存在的元素计算最大值），和Numpy的maximum_filter默认的mode='reflect'行为略有不同，如果需要完全对齐Numpy的边界处理逻辑，可以添加边界填充（比如反射填充）的代码，但这个版本已经能满足大部分实际场景。
• 如果你的输入数组很大，可以考虑优化maximumFilter的性能（比如用滑动窗口算法减少重复计算），不过对于一般规模的数组，当前实现已经足够高效。

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