如何独立于维度数量选择第一个项目

Question

我有一个多维numpy.array，称为my_imgs。我想将数组存储为图像。数组的维度不是恒定的。目前我有四五个维度；稍后我将添加更多维度。对于从第4个diedimension开始的所有diedimensions，我总是想选择第一个元素。是否有一个函数可以始终选择第一个元素，而不考虑维度的数量？

目前，我对每个维度都有一个单独的if语句。

import matplotlib.pyplot as plt
# shape (number, row, col, channel, stack, ...)
my_imgs = '' # Multidimensional image array

fig, axes = plt.subplots(nrows=3, ncols=3)
for i, ax in enumerate(axes.flat):
ax.axis('off')
if len(my_imgs.shape) == 4:
ax.imshow(my_imgs[i, :, :, 0], cmap='gray_r')
elif len(my_imgs.shape) == 5:
ax.imshow(my_imgs[i, :, :, 0, 0], cmap='gray_r')
elif len(my_imgs.shape) == 6:
ax.imshow(my_imgs[i, :, :, 0, 0, 0], cmap='gray_r')

Answer 1

numpy用于指定"整个维度"的切片:等于slice(None)。

这意味着你可以进行

import numpy as np

arr_5d = np.zeros((10, ) * 5)
arr_6d = np.zeros((10, ) * 6)

def my_slicer(arr, i=0):
indexer = (i, slice(None), slice(None)) + (0,) * (len(arr.shape)-3)
return arr[indexer]

# will print (10, 10) (10, 10)
print(
my_slicer(arr_5d, i=2).shape,
my_slicer(arr_6d, i=-3).shape,
)