首页 > 解决方案 > 使用numpy中另一个数组的值从一个数组中选择值

问题描述

在 numpy 中,我有一个 3D 数组。沿 0 轴,它存储多个 2D 平面。我需要获取每个平面的梯度,选择这些平面上每个点的中值梯度幅度,从而隔离相应的 x 和 y 梯度分量。但是我很难正确执行此操作。

到目前为止,要获得梯度和中位数,我有:

img_l = #My 3D array of 2D planes
grad = np.gradient(img_l,axis=[1,2]) #Get gradient of each image. This is a list with 2 elements.
mag_grad = np.sqrt(grad[0]**2 + grad[1]**2) #Get magnitude of gradient in each case
med = np.median(mag_grad, axis=0) #Get median value at each point in the planes

然后选择正确的 x & y 渐变分量,我使用:

pos=(mag_grad==med).argmax(axis=0) #This returns the first instance where the median element encountered along axis=0
G = np.stack([np.zeros(med.shape),np.zeros(med.shape)], axis=0) #Will store y and x median components of the gradient, respectively.
for i in range(med.shape[0]):
    for j in range(med.shape[1]):
        G[0,i,j], G[1,i,j] = grad[0][pos[i,j],i,j], grad[1][pos[i,j],i,j] #Manually select the median y and x components of the gradient, and save to G.

我相信第二个代码块可以正常工作。但是,它非常不优雅,并且因为我在 NumPy 中找不到执行此操作的方法,所以我不得不使用 Python 循环,这会增加大量开销。另外,由于这个操作在 NumPy 中经常发生,我怀疑应该有一个内置的方法来做到这一点。

我怎样才能更有效、更优雅地实现这段代码?

标签: pythonnumpymultidimensional-arraynumpy-ndarrayarray-broadcasting

解决方案


使用 itertools 索引您的数组可以使其更高效/优雅。

import itertools

idxs = np.array(list(itertools.product(range(med.shape[0]), range(med.shape[1]))))
G[0,idxs], G[1,idxs] = grad[0][pos[idxs],idxs], grad[1][pos[idxs],idxs]

推荐阅读