我有一批图像和一批针对每幅图像的索引 (x, y)。每幅图像的索引都不同,因此我无法使用简单的索引。获取另一批具有每幅图像所选像素颜色的批次的最佳或最快方法是什么?
n_images = 4
width = 100
height = 100
channels = 3
n_samples = 30
images = torch.rand((n_images, height, width, channels))
indices = (torch.rand((n_images, n_samples, 2)) * width).to(torch.int32)
# preferred function
# result = images[indices]
# with result.shape = (n_images, n_samples, 3)
# I just found this solution but I would rather like to call a general torch function
xs = indices.reshape((-1, 2))[:, 0]
ys = indices.reshape((-1, 2))[:, 1]
ix = torch.arange(n_images, dtype=torch.int32)
ix = ix[..., None].expand((-1, n_samples)).flatten()
result = images[ix, ys, xs].reshape((n_images, n_samples, 3))