Numpy - 在多个轴上设置第一个最大值

Numpy - set first max value over multiple axes

我有一个多维数组,最后两个维度对应一个“window”。对于数组中的每个“window”,我的目标是将 window 中的第一个最大值设置为 1,其余值设置为 0。例如,对于一个 window,我们有这个作业:

[[0., 0., 0.],          [[0., 0., 0.],
 [0., 1., 1.],    ->     [0., 1., 0.],
 [0., 1., 1.]],          [0., 0., 0.]]

现在,我想对所有 windows 执行此操作,其中每个 window 都在最后两个维度上,即 axis=(-1,-2)。这是我到目前为止尝试过的:

# windows
x = np.array([[[[[[0., 0., 0.],
                  [0., 1., 1.],
                  [0., 1., 1.]],

                 [[0., 0., 0.],
                  [1., 1., 0.],
                  [1., 1., 0.]]],

  
                [[[0., 1., 1.],
                  [0., 1., 1.],
                  [0., 1., 1.]],

                 [[1., 1., 0.],
                  [1., 1., 0.],
                  [1., 1., 0.]]],


                [[[0., 1., 1.],
                  [0., 1., 1.],
                  [0., 1., 0.]],

                 [[1., 1., 0.],
                  [1., 1., 0.],
                  [1., 0., 1.]]]]]])

max_indices = np.argwhere(x == x.max()) # all indicies of max values in x
mask = np.zeros_like(x) # start of with everything set to 0

# I can set the first window by doing the following
mask[tuple(max_indices[0])] = 1.0

我怎样才能对所有其他 windows 做同样的事情?最好不要循环。

我的想法是在未考虑最后两个维度的 max_indices 中找到唯一值,但我不确定如何去做。

编辑

我有一个案例,上面场景中的 x 是取自另一个 使用 np.lib.stride_tricks.as_strided 的数组,这意味着步幅会有所不同。 情况如下:

# The image or whatever i want to modify
x = np.array([[[[0., 0., 0., 0.],
                [0., 1., 1., 0.],
                [0., 1., 1., 0.],
                [0., 1., 1., 0.],
                [0., 1., 0., 1.]]]])

# the windows of that image
x = np.lib.stride_tricks.as_strided(x, shape=(1, 1, 3, 2, 3, 3), strides=(160, 160, 32, 8, 32, 8))

使用 Felipe 给出的解决方案:

y = x[0][0]  # Just unnesting it a bit
yr = y.reshape(-1, 9)
idx = yr.argmax(1)
y0 = np.zeros_like(yr)
np.put_along_axis(y0, idx[:, None], 1, axis=1)
y = y0.reshape(y.shape)

# plug it back
x[0][0] = y

x

array([[[[[[0., 0., 0.],
           [0., 1., 0.],
           [0., 1., 0.]],

          [[0., 0., 0.],
           [1., 0., 0.],
           [1., 0., 0.]]],


         [[[0., 1., 0.],
           [0., 1., 0.],
           [0., 0., 0.]],

          [[1., 0., 0.],
           [1., 0., 0.],
           [0., 0., 0.]]],


         [[[0., 1., 0.],
           [0., 0., 0.],
           [0., 0., 0.]],

          [[1., 0., 0.],
           [0., 0., 0.],
           [0., 0., 0.]]]]]])

形状相同,但 xy 中的步幅不同,似乎“将其插回”导致 x 和 y 不相同,这是我所期望的。

这是使用 np.put_along_axis 的一种方法(尽管这实际上只是一个更快的循环)。

y = x[0][0]  # Just unnesting it a bit
yr = y.reshape(-1, 9)
idx = yr.argmax(1)
y0 = np.zeros_like(yr)
np.put_along_axis(y0, idx[:, None], 1, axis=1)
y = y0.reshape(y.shape)

输出:

array([[[[0., 0., 0.],
         [0., 1., 0.],
         [0., 0., 0.]],

        [[0., 0., 0.],
         [1., 0., 0.],
         [0., 0., 0.]]],


       [[[0., 1., 0.],
         [0., 0., 0.],
         [0., 0., 0.]],

        [[1., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.]]],


       [[[0., 1., 0.],
         [0., 0., 0.],
         [0., 0., 0.]],

        [[1., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.]]]])

编辑

这不是最聪明的想法,但我的意思是您可以先将其重新分配给它的副本,然后再次获得跨步视图,例如:

x0 = x.copy()  # Get a normal view

# [ Run the previous snippet ]

x0[0][0] = y  # Assign the modified version

# Then if you really want another strided view you can recreate it
x0 = x0.reshape(saved_x.shape)  # Reshape it to the original version pre-stride
x = np.lib.stride_tricks.as_strided(x0, shape=(1, 1, 3, 2, 3, 3), strides=(160, 160, 32, 8, 32, 8))