获取一个三维数组作为参数并返回一个三维数组

Question

我必须编写一个函数 to_red 来将绿色和蓝色分量归零并返回结果。 我为 Image(.png) 编写了以下代码以将绿色和蓝色归零并返回红色并且它起作用了。 但是，正如标题中提到的，输入参数必须是一个 3-d 数组并返回一个 3-d 数组。 我的下面的代码应该如何更改。

import numpy as np

导入 matplotlib.pyplot 作为 plt

from matplotlib.pyplot import imshow

def to_red()

src = plt.imread("C:\src\painting.png")

red_channel = src[:,:,0]

red_img = np.zeros(src.shape)

red_img[:,:,0] = red_channel

plt.imshow(red_img) 

plt.show()

Answer 1

你可以这样写你的函数：

import numpy as np 
import matplotlib.pyplot as plt 
from matplotlib.pyplot import imshow

def to_red(src):
    # Check if the input dimension is 3
    if not src.ndim == 3:
        # Raise exception or do something
        print ("Dimension mismatch")
        return 0

    red_channel = src[:,:,0]
    red_img = np.zeros(src.shape)
    red_img[:,:,0] = red_channel

    return red_img

然后你可以这样称呼它

source_image = plt.imread("C:\src\painting.png")

red_image = to_red(source_image)

plt.imshow(red_image)    
plt.show()

我还添加了一行来检查输入是否实际上是 3 维的。

Answer 2

您可以使用numpy强大的索引功能

def to_red(src):
    ret = a.copy() 
    ret[:,:,1:] = 0 
    return ret