使用Numpy apply_along_axis时维度异常的原因咨询

使用环境 #

Python 3.10.2，Numpy 1.23.0

核心代码 #

# Setup
import numpy as np

# functions
def get_RGBA_from_Int(num):
    b = num & 255
    g = (num >> 8) & 255
    r = (num >> 16) & 255
    a = (num >> 24) & 255
    return (r,g,b,a)

def get_INT_from_RGBA(ls):
    r,g,b,a = ls
    num = (a << 24) + (r << 16) + (g << 8) + b
    return num

测试用例1：将3D数组的整数转换为RGBA #

# Convert a 3d array value into a rgba
# 4294967295 == [255,255,255,255] for a RGBA color of white
# Int to RGBA array of (4,)
x = np.zeros((5,5,1), dtype='int64') + 4294967295 # a 5x5 image colored white in array format
output = np.apply_along_axis(get_RGBA_from_Int, 2, x) # Testing conversion
print(output.shape) # prints to screen :: (5,5,4,1)
# This adds 1 dimension to 3d array

测试用例2：将3D数组的RGBA转换为整数 #

# Convert a 3d array rgba into a int
# 4294967295 == [255,255,255,255] for a RGBA color of white
# RGBA array of (4,) to Int of (1,)
x = np.zeros((5,5,4), dtype='int64') + 255 # a 5x5 image colored white in array format
output = np.apply_along_axis(get_INT_from_RGBA, 2, x)
print(output.shape) # prints to screen :: (5,5)
# This removes 1 dimension from 3d array

问题：为何测试用例1增加1个维度，测试用例2却减少1个维度？ #

原因解析 #

这种维度变化的核心是np.apply_along_axis的输出维度完全由函数返回值的维度决定，结合输入数组在指定轴上的元素维度共同作用：

1. 测试用例1维度增加的原因

   • 输入x的形状是(5,5,1)，指定axis=2后，每次传给get_RGBA_from_Int的是一个长度为1的Numpy数组（比如array([4294967295])）。
   • 函数里的位运算都是对数组操作的，返回的b、g、r、a都是长度为1的数组，最终函数返回的是包含4个长度为1数组的元组，等价于(4,1)形状的数组。
   • np.apply_along_axis会把这个结果替换原数组的axis=2位置，因此最终输出形状变成(5,5,4,1)，比原数组多了一个维度。

2. 测试用例2维度减少的原因

   • 输入x的形状是(5,5,4)，指定axis=2后，每次传给get_INT_from_RGBA的是一个长度为4的数组（比如array([255,255,255,255])）。
   • 函数将数组解包为r、g、b、a四个标量，计算后返回的num是一个标量值。
   • np.apply_along_axis遇到函数返回标量时，会直接移除原来的axis=2维度，因此最终输出形状变成(5,5)，比原数组少了一个维度。

修正方案 #

如果想让输出维度符合预期，可以调整函数的返回值类型：

• 对于测试用例1，取出数组中的标量值再计算，确保返回的是普通数值组成的元组：

  def get_RGBA_from_Int(num):
      num_val = num.item()  # 提取数组中的标量
      b = num_val & 255
      g = (num_val >> 8) & 255
      r = (num_val >> 16) & 255
      a = (num_val >> 24) & 255
      return (r,g,b,a)
  

  修改后输出形状会变为(5,5,4)，符合预期。

• 对于测试用例2，如果想保留最后一个维度（比如(5,5,1)），让函数返回长度为1的数组：

  def get_INT_from_RGBA(ls):
      r,g,b,a = ls
      num = (a << 24) + (r << 16) + (g << 8) + b
      return [num]  # 返回长度为1的列表
  

  修改后输出形状会变为(5,5,1)。

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