Merge patches into image
| 322 | |
| 323 | |
| 324 | class BlockMerging(nn.Module): |
| 325 | """Merge patches into image |
| 326 | """ |
| 327 | |
| 328 | def __init__(self): |
| 329 | super(BlockMerging, self).__init__() |
| 330 | |
| 331 | def forward(self, patches, height, width): |
| 332 | """ |
| 333 | Args: |
| 334 | patches(tensor) batch x height*width/64, height x width |
| 335 | height(int) |
| 336 | width(int) |
| 337 | |
| 338 | Returns: |
| 339 | Tensor: batch x height x width |
| 340 | """ |
| 341 | k = 8 |
| 342 | batch_size = patches.shape[0] |
| 343 | image_reshaped = patches.view(batch_size, height // k, width // k, k, k) |
| 344 | image_transposed = image_reshaped.permute(0, 1, 3, 2, 4) |
| 345 | return image_transposed.contiguous().view(batch_size, height, width) |
| 346 | |
| 347 | |
| 348 | class ChromaUpsampling(nn.Module): |