是否有在PyTorch中提取图像补丁的功能?

时间:2017-08-22 23:08:36

标签: python tensorflow pytorch

鉴于一批图像,我想提取所有可能的图像补丁,类似于卷积。在TensorFlow中,我们可以使用tf.extract_image_patches来实现这一目标。 PyTorch中是否有等效函数?

谢谢。

3 个答案:

答案 0 :(得分:3)

不幸的是,可能没有直接的方法可以实现您的目标。
但是Tensor.unfold函数可能是一个解决方案。
https://discuss.pytorch.org/t/how-to-extract-smaller-image-patches-3d/16837/2
该网站可能会为您提供帮助。

答案 1 :(得分:1)

也许这个代码示例将有助于了解如何使用unfold,这受@gasoon链接的this线程的启发,但更为冗长:

batch_size, n_channels, n_rows, n_cols = 32, 3, 64, 64
kernel_h, kernel_w = 7, 9
step = 5

x = torch.arange(batch_size*n_channels*n_rows*n_cols).view(batch_size, n_channels, n_rows, n_cols)

# unfold(dimension, size, step)
windows = x.unfold(2, kernel_h, step).unfold(3, kernel_w, step).permute(2, 3, 0, 1, 4, 5).reshape(-1, n_channels, kernel_h, kernel_w)
print(windows.shape)
# result: torch.Size([4608, 3, 7, 9]) = [n_windows, n_channels, krenel_h, kernel_w]

答案 2 :(得分:0)

我也花了一些时间研究这个,我发现 this pytorch thread 对我很有用,PyTorch dev ptrblck(保佑这个家伙)提供了一个等效的 pytorch 版本的 tensorflow 函数。< /p>

为了简单起见,我将在这里重新发布代码(来自用户 FloCF)。

import math
import torch.nn.functional as F

def extract_image_patches(x, kernel, stride=1, dilation=1):
    # Do TF 'SAME' Padding
    b,c,h,w = x.shape
    h2 = math.ceil(h / stride)
    w2 = math.ceil(w / stride)
    pad_row = (h2 - 1) * stride + (kernel - 1) * dilation + 1 - h
    pad_col = (w2 - 1) * stride + (kernel - 1) * dilation + 1 - w
    x = F.pad(x, (pad_row//2, pad_row - pad_row//2, pad_col//2, pad_col - pad_col//2))
    
    # Extract patches
    patches = x.unfold(2, kernel, stride).unfold(3, kernel, stride)
    patches = patches.permute(0,4,5,1,2,3).contiguous()
    
    return patches.view(b,-1,patches.shape[-2], patches.shape[-1])

在 PyTorch 论坛上给这些人点赞 :)