Pytorch 展开和折叠:如何将这个图像张量重新组合在一起?
Posted
技术标签:
【中文标题】Pytorch 展开和折叠:如何将这个图像张量重新组合在一起?【英文标题】:Pytorch Unfold and Fold: How do I put this image tensor back together again? 【发布时间】:2021-01-10 21:08:00 【问题描述】:我正在尝试使用展开过滤大小为 256x256 的单通道 2D 图像,以创建重叠为 8 的 16x16 块。如下所示:
# I = [256, 256] image
kernel_size = 16
stride = bx/2
patches = I.unfold(1, kernel_size,
int(stride)).unfold(0, kernel_size, int(stride)) # size = [31, 31, 16, 16]
我已经开始尝试用 fold 将图像重新组合在一起,但我还没有完全做到。我尝试使用视图让图像“适合”它应该的方式,但我不明白这将如何保留原始图像。可能是我想多了。
# patches.shape = [31, 31, 16, 16]
patches = = filt_data_block.contiguous().view(-1, kernel_size*kernel_size) # [961, 256]
patches = patches.permute(1, 0) # size = [951, 256]
任何帮助将不胜感激。非常感谢。
【问题讨论】:
【参考方案1】:我相信在这种情况下使用 torch.nn.functional.fold
和 torch.nn.functional.unfold
会让您受益,因为这些函数是专门为图像(或任何 4D 张量,即形状为 B X C X H X W)而构建的。
让我们从展开图像开始:
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
from sklearn.datasets import load_sample_image #Used to load a sample image
dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
#Load a flower image from sklearn.datasets, crop it to shape 1 X 3 X 256 X 256:
I = torch.from_numpy(load_sample_image('flower.jpg')).permute(2,0,1).unsqueeze(0).type(dtype)[...,128:128+256,256:256+256]
kernel_size = 16
stride = kernel_size//2
I_unf = F.unfold(I, kernel_size, stride=stride)
在这里,我们使用F.unfold
函数获得所有步幅为 8 的 16x16 图像块。这将产生一个形状为torch.Size([1, 768, 961])
的 3D 张量。即 - 961 个补丁,每个补丁内有 768 = 16 X 16 X 3 像素。
现在,假设我们希望将其折叠回 I:
I_f = F.fold(I_unf,I.shape[-2:],kernel_size,stride=stride)
norm_map = F.fold(F.unfold(torch.ones(I.shape).type(dtype),kernel_size,stride=stride),I.shape[-2:],kernel_size,stride=stride)
I_f /= norm_map
我们使用F.fold
告诉它I
的原始形状,我们用来展开的kernel_size
和使用的stride
。折叠I_unf
后,我们将获得一个有重叠的总和。这意味着生成的图像将显得饱和。因此,我们需要计算一个归一化映射,它将归一化由于重叠导致的多个像素总和。一种有效地做到这一点的方法是使用一个张量并使用unfold
,然后使用fold
- 来模拟重叠的求和。这为我们提供了归一化映射,我们通过它对 I_f
进行归一化以恢复 I
。
现在,我们希望绘制I_f
和I
来证明内容被保留:
#Plot I:
plt.imshow(I[0,...].permute(1,2,0).cpu()/255)
#Plot I_f:
plt.imshow(I_f[0,...].permute(1,2,0).cpu()/255)
整个过程也适用于单通道图像。需要注意的一件事是,如果图像的空间维度不能被步幅整除,由于某些像素无法到达,您将得到带有零(在边缘)的 norm_map
,但您也可以轻松处理这种情况。
【讨论】:
非常感谢这个 Gil,我很感激。我认为您的回答对以后做同样事情的其他人非常有用。 我的荣幸,布莱德。【参考方案2】:比 Gil 提出的方案稍微不优雅:
I took inspiration from this post on the Pytorch forums,将我的图像张量格式化为标准形状 B x C x H x W (1 x 1 x 256 x 256)。展开:
# CREATE THE UNFOLDED IMAGE SLICES
I = image # shape [256, 256]
kernel_size = bx #shape [16]
stride = int(bx/2) #shape [8]
I2 = I.unsqueeze(0).unsqueeze(0) #shape [1, 1, 256, 256]
patches2 = I2.unfold(2, kernel_size, stride).unfold(3, kernel_size, stride)
#shape [1, 1, 31, 31, 16, 16]
在此之后,我对我的张量堆栈进行了一些转换和过滤。在此之前,我应用余弦窗口并进行标准化:
# NORMALISE AND WINDOW
Pvv = torch.mean(torch.pow(win, 2))*torch.numel(win)*(noise_std**2)
Pvv = Pvv.double()
mean_patches = torch.mean(patches2, (4, 5), keepdim=True)
mean_patches = mean_patches.repeat(1, 1, 1, 1, 16, 16)
window_patches = win.unsqueeze(0).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(1, 1, 31, 31, 1, 1)
zero_mean = patches2 - mean_patches
windowed_patches = zero_mean * window_patches
#SOME FILTERING ....
#ADD MEAN AND WINDOW BEFORE FOLDING BACK TOGETHER.
filt_data_block = (filt_data_block + mean_patches*window_patches) * window_patches
上面的代码对我有用,但是掩码会更简单。接下来,我准备将 [1, 1, 31, 31, 16, 16] 形式的张量转换回原来的 [1, 1, 256, 256]:
# REASSEMBLE THE IMAGE USING FOLD
patches = filt_data_block.contiguous().view(1, 1, -1, kernel_size*kernel_size)
patches = patches.permute(0, 1, 3, 2)
patches = patches.contiguous().view(1, kernel_size*kernel_size, -1)
IR = F.fold(patches, output_size=(256, 256), kernel_size=kernel_size, stride=stride)
IR = IR.squeeze()
这让我可以创建一个重叠的滑动窗口并将图像无缝缝合在一起。去掉过滤可以得到相同的图像。
【讨论】:
以上是关于Pytorch 展开和折叠:如何将这个图像张量重新组合在一起?的主要内容,如果未能解决你的问题,请参考以下文章
如何在 PyTorch 中将 RGB 图像编码为 n_class One 热张量