| 300 | |
| 301 | @torch.no_grad() |
| 302 | def estimate_augmented(self, pixel, latent): |
| 303 | args = [ |
| 304 | [False, 0], [False, 1], [False, 2], [False, 3], [True, 0], [True, 1], [True, 2], [True, 3], |
| 305 | ] |
| 306 | |
| 307 | result = [] |
| 308 | |
| 309 | for flip, rok in tqdm(args): |
| 310 | feed_pixel = pixel.clone() |
| 311 | feed_latent = latent.clone() |
| 312 | |
| 313 | if flip: |
| 314 | feed_pixel = torch.flip(feed_pixel, dims=(3,)) |
| 315 | feed_latent = torch.flip(feed_latent, dims=(3,)) |
| 316 | |
| 317 | feed_pixel = torch.rot90(feed_pixel, k=rok, dims=(2, 3)) |
| 318 | feed_latent = torch.rot90(feed_latent, k=rok, dims=(2, 3)) |
| 319 | |
| 320 | eps = self.decoder(feed_pixel, feed_latent).clip(0, 1) |
| 321 | eps = torch.rot90(eps, k=-rok, dims=(2, 3)) |
| 322 | |
| 323 | if flip: |
| 324 | eps = torch.flip(eps, dims=(3,)) |
| 325 | |
| 326 | result += [eps] |
| 327 | |
| 328 | result = torch.stack(result, dim=0) |
| 329 | median = torch.median(result, dim=0).values |
| 330 | return median |
| 331 | |
| 332 | |
| 333 | |