Получить трехмерный массив в качестве параметра и вернуть трехмерные массивы

Question

Я должен написать функцию to_red, которая обнуляет зеленый и синий цветовые компоненты и возвращает результат. Я написал приведенный ниже код для изображения (.png), чтобы обнулить зеленый и синий цвета и вернуть красный, и это сработало. Однако, как упоминалось в заголовке, входной параметр должен быть массивом 3-го ряда и возвращать массив 3-го уровня. Как мой код ниже должен быть изменен для этого.

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

Вы можете использовать мощные возможности индексации numpy

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

Answer 2

Вы можете написать свою функцию следующим образом:

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 мерный.