Numpy如何计算多通道比值并按阈值赋值对应通道索引

直接用Numpy原生向量化操作即可实现，全程不需要循环，处理大数组效率很高，具体实现如下： #

核心实现代码 #

import numpy as np

def generate_target_array(input_arr: np.ndarray, target_channel: int = 1, threshold: float = 0.3) -> np.ndarray:
    # 校验输入维度为[N, X, Y]
    assert len(input_arr.shape) == 3, "输入数组维度必须是N*X*Y的三通道格式"
    N, X, Y = input_arr.shape
    # 计算每个(X,Y)点位所有通道的数值总和
    point_total = np.sum(input_arr, axis=0)
    # 处理总和为0的极端情况，避免除以0报错，可根据业务需求调整赋值逻辑
    point_total[point_total == 0] = 1e-8
    # 计算指定通道在每个点位的占比
    target_channel_ratio = input_arr[target_channel] / point_total
    # 计算每个点位数值最大的通道索引
    max_channel_idx = np.argmax(input_arr, axis=0)
    # 按条件赋值生成结果数组：占比高于阈值用指定通道索引，否则用最大值通道索引
    result = np.where(target_channel_ratio > threshold, target_channel, max_channel_idx)
    return result

双阈值适配（匹配你举例的场景） #

如果你需要用到高低两个阈值的判断逻辑（比如高于0.3用指定通道，低于0.4用最大通道），可以用下面的变体实现：

def generate_target_array_double_threshold(
    input_arr: np.ndarray, 
    target_channel: int = 1, 
    high_threshold: float = 0.3,
    low_threshold: float = 0.4,
    middle_val: int = 0 # 两个阈值之间的区间默认赋值为0，可自定义
) -> np.ndarray:
    assert len(input_arr.shape) == 3, "输入数组维度必须是N*X*Y的三通道格式"
    N, X, Y = input_arr.shape
    point_total = np.sum(input_arr, axis=0)
    point_total[point_total == 0] = 1e-8
    target_channel_ratio = input_arr[target_channel] / point_total
    max_channel_idx = np.argmax(input_arr, axis=0)
    # 多层条件判断
    result = np.where(
        target_channel_ratio > high_threshold,
        target_channel,
        np.where(target_channel_ratio < low_threshold, max_channel_idx, middle_val)
    )
    return result

测试用例 #

# 测试输入：2通道，2*2大小的数组
test_input = np.array([
    [[1, 2], [3, 4]], # 通道0
    [[1, 1], [1, 5]]  # 通道1
])
res = generate_target_array(test_input, target_channel=1, threshold=0.3)
print(res)

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