Splitting image into patches
| 96 | |
| 97 | |
| 98 | class BlockSplitting(nn.Module): |
| 99 | """ Splitting image into patches |
| 100 | """ |
| 101 | |
| 102 | def __init__(self): |
| 103 | super(BlockSplitting, self).__init__() |
| 104 | self.k = 8 |
| 105 | |
| 106 | def forward(self, image): |
| 107 | """ |
| 108 | Args: |
| 109 | image(tensor): batch x height x width |
| 110 | |
| 111 | Returns: |
| 112 | Tensor: batch x h*w/64 x h x w |
| 113 | """ |
| 114 | height, _ = image.shape[1:3] |
| 115 | batch_size = image.shape[0] |
| 116 | image_reshaped = image.view(batch_size, height // self.k, self.k, -1, self.k) |
| 117 | image_transposed = image_reshaped.permute(0, 1, 3, 2, 4) |
| 118 | return image_transposed.contiguous().view(batch_size, -1, self.k, self.k) |
| 119 | |
| 120 | |
| 121 | class DCT8x8(nn.Module): |