如何用Numpy实现遍历图像像素识别红色像素的函数？

原代码存在的问题 #

• 语法错误：map_filename = 这行多余，导致后续赋值语句非法
• 形状解包错误：PNG图像是三维数组（行、列、3通道），row, col = rgb.shape会抛出解包错误，需改为row, col, _ = rgb.shape
• 参数未复用：硬编码图像路径./data/map.png，忽略了传入的map_filename参数
• 逻辑偏差：找到红色像素后仅修改原RGB数组为黑色，未生成要求的二维二值数组，也未返回该数组
• 二值图像缺失：未创建独立的二值标记数组，而是直接修改原图像数据

修正后的完整代码 #

import numpy as np
from matplotlib import pyplot as plt

def find_red_pixels(map_filename, upper_threshold=100, lower_threshold=50):
    """
    Returns a 2D numpy array representing the binary image (1 for red pixels, 0 otherwise)
    and saves the binary image as a JPG file.
    """
    # 读取输入图像（matplotlib读入的是[0,1]范围的float数组）
    img = plt.imread(map_filename)
    # 转换为[0,255]范围的uint8类型
    img = (img * 255).astype(np.uint8)
    # 获取图像的行、列尺寸
    rows, cols, _ = img.shape
    # 初始化二维二值数组，0代表非红色像素
    binary_img = np.zeros((rows, cols), dtype=np.uint8)
    
    # 双重循环遍历所有像素
    for i in range(rows):
        for j in range(cols):
            r, g, b = img[i, j]
            # 严格按照判定条件标记红色像素
            if r > upper_threshold and g < lower_threshold and b < lower_threshold:
                binary_img[i, j] = 1  # 1代表红色像素
    
    # 将二值数组转换为[0,255]范围的灰度图，便于保存为JPG
    save_img = (binary_img * 255).astype(np.uint8)
    plt.imsave('map-red-pixels.jpg', save_img, cmap='gray')
    
    # 返回要求的二维二值数组
    return binary_img

代码关键说明 #

• 复用函数传入的map_filename参数，支持动态指定图像路径
• 独立创建二维二值数组存储检测结果，避免修改原图像数据
• 严格遵循红色像素判定规则：R值大于upper_threshold，G、B值均小于lower_threshold
• 保存JPG时将二值数组转换为灰度格式（255表示红色区域，0表示非红色区域），保证可视化效果
• 函数正确返回要求的二维二值numpy数组

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