numpy 使用PyTorchTensor从索引图像中索引一批图像

zpjtge22  于 9个月前  发布在  其他
关注(0)|答案(1)|浏览(123)

假设我有一批图像M,其形式为 Torch Tensor(B, W, H),还有一个大小为(W, H)的图像I,其像素是索引。
我想得到一个图像(W, H),其中每个像素来自图像批中的相应图像(在I的索引之后)。

示例

给定(3, 4, 8)的形状M

tensor([[[ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
         [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
         [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
         [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.]],

        [[-1., -1., -1., -1., -1., -1., -1., -1.],
         [-1., -1., -1., -1., -1., -1., -1., -1.],
         [-1., -1., -1., -1., -1., -1., -1., -1.],
         [-1., -1., -1., -1., -1., -1., -1., -1.]],

        [[-2., -2., -2., -2., -2., -2., -2., -2.],
         [-2., -2., -2., -2., -2., -2., -2., -2.],
         [-2., -2., -2., -2., -2., -2., -2., -2.],
         [-2., -2., -2., -2., -2., -2., -2., -2.]]])

字符串
(4, 8)形状的I

tensor([[2, 0, 2, 0, 1, 0, 1, 0],
        [2, 2, 1, 0, 0, 2, 1, 0],
        [2, 0, 0, 2, 1, 1, 0, 0],
        [0, 1, 0, 0, 2, 0, 2, 1]], dtype=torch.int32)


所得到的图像将是:

tensor([[-2.,  0., -2.,  0., -1.,  0., -1.,  0.],
        [-2., -2., -1.,  0.,  0., -2., -1.,  0.],
        [-2.,  0.,  0., -2., -1., -1.,  0.,  0.],
        [ 0., -1.,  0.,  0., -2.,  0., -2., -1.]])

注1

我不关心M维度的顺序,如果它提供了一个更简单的解决方案,它也可以是(W, H, B)

注2

我对NumPy解决方案也很感兴趣。

c9qzyr3d

c9qzyr3d1#

一种解决办法是:

indices = torch.meshgrid(torch.arange(I.shape[0]), torch.arange(I.shape[1]))
result = M[I, *indices]

字符串
使用numpy:

indices = np.indices(I)
result = M[I, indices[0], indices[1]]

相关问题