图像批处理中的随机补丁

Posted 2023-03-12

【中文标题】图像批处理中的随机补丁

【英文标题】：Shuffle patches in image batch

【发布时间】：2021-07-01 21:27:57

【问题描述】：

我正在尝试创建一个transform，它可以对批次中的每个图像的补丁进行洗牌。 我的目标是以与torchvision 中其他转换相同的方式使用它：

trans = transforms.Compose([
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
            ShufflePatches(patch_size=(16,16)) # our new transform
        ])

更具体地说，输入是一个BxCxHxW 张量。我想将批次中的每个图像拆分为大小为 patch_size 的非重叠补丁，将它们打乱，然后重新组合成单​​个图像。

给定图像（大小为224x224）：

使用ShufflePatches(patch_size=(112,112)) 我想生成输出图像：

我认为解决方案与torch.unfold 和torch.fold 有关，但没有进一步解决。

任何帮助将不胜感激！

【参考方案1】：

确实unfold and fold 在这种情况下似乎很合适。

import torch
import torch.nn.functional as nnf

class ShufflePatches(object):
  def __init__(self, patch_size):
    self.ps = patch_size

  def __call__(self, x):
    # divide the batch of images into non-overlapping patches
    u = nnf.unfold(x, kernel_size=self.ps, stride=self.ps, padding=0)
    # permute the patches of each image in the batch
    pu = torch.cat([b_[:, torch.randperm(b_.shape[-1])][None,...] for b_ in u], dim=0)
    # fold the permuted patches back together
    f = nnf.fold(pu, x.shape[-2:], kernel_size=self.ps, stride=self.ps, padding=0)
    return f

这是一个补丁大小=16 的示例：

【讨论】：

比 sklearn extract_patches_2d 快得多