self-supervised-attack/transform.py at master · chenxi52/self-supervised-attack · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import torch

from torch.utils.data import Dataset, DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import transforms


class overlap:
    def __init__(self, size):
        self.h , self.w = size, size

    def __call__(self, img):
        assert img.size(2) % self.h == 0, 'Image height must be divided by size[0]'
        assert img.size(1) % self.w == 0, 'Image height must be divided by size[1]'

        imgs = []
        for i in range(img.size(2)//self.h + 1):
            for j in range(img.size(1)//self.w + 1):
                imgs.append(img[:,i * self.h//2: i * self.h//2 + self.h , j * self.w//2 :self.w + j * self.w//2])
        imgs = torch.stack(imgs, dim=0)

        return imgs