По какой-то причине я пишу функцию, которая создает массив для хранения данного массива (обычно изображения), который будет иметь свойство ndim
, равное 2 или 3, я не могу знать заранее. В настоящее время я делаю что-то вроде:
import numpy as np
def create_greater_array(im, scale):
shape = (im.shape[0] * scale + 2, im.shape[1] * scale + 2)
if im.ndim==3: #if image is RGB
shape = *shape, 3
greater_array = np.fill(shape, np.nan)
return greater_array
Что меня бесит, так это если я смогу избежать дополнительной проверки if
. На самом деле мой код на самом деле является определением класса и делает немного больше, поэтому я в итоге добавил атрибут ._is_RBG
к классу.
Я думаю, что, поскольку последнее измерение всегда либо 3, либо не существует, оно может быть автоматизированным.
РЕДАКТИРОВАТЬ: Другой случай, когда мне нужно проверить в той же программе.
Позже у меня есть другая функция, которая вставляет изображения в массив:
def insert_image(im):
big = create_greater_array(im, 2)
if im.ndims == 2:
big[1:im.shape[0]+1, 1:im.shape[1]+1] = im
else:
big[1:im.shape[0]+1, 1:im.shape[1]+1, :] = im
return big