| |
|
|
| import torch |
| import cv2 |
| from torchvision import transforms |
| import numpy as np |
| import math |
|
|
|
|
| def visual(output, out_path): |
| output = (output + 1)/2 |
| output = torch.clamp(output, 0, 1) |
| if output.shape[1] == 1: |
| output = torch.cat([output, output, output], 1) |
| output = output[0].detach().cpu().permute(1,2,0).numpy() |
| output = (output*255).astype(np.uint8) |
| output = output[:,:,::-1] |
| cv2.imwrite(out_path, output) |
| |
| |
| def get_lr(t, initial_lr, rampdown=0.25, rampup=0.05): |
| |
| lr_ramp = min(1, (1 - t) / rampdown) |
| lr_ramp = 0.5 - 0.5 * math.cos(lr_ramp * math.pi) |
| lr_ramp = lr_ramp * min(1, t / rampup) |
| return initial_lr * lr_ramp |
|
|
|
|
|
|
| def latent_noise(latent, strength): |
| noise = torch.randn_like(latent) * strength |
|
|
| return latent + noise |
|
|
| def noise_regularize_(noises): |
| loss = 0 |
|
|
| for noise in noises: |
| size = noise.shape[2] |
|
|
| while True: |
| loss = ( |
| loss |
| + (noise * torch.roll(noise, shifts=1, dims=3)).mean().pow(2) |
| + (noise * torch.roll(noise, shifts=1, dims=2)).mean().pow(2) |
| ) |
|
|
| if size <= 8: |
| break |
|
|
| noise = noise.reshape([-1, 1, size // 2, 2, size // 2, 2]) |
| noise = noise.mean([3, 5]) |
| size //= 2 |
|
|
| return loss |
|
|
|
|
| def noise_normalize_(noises): |
| for noise in noises: |
| mean = noise.mean() |
| std = noise.std() |
|
|
| noise.data.add_(-mean).div_(std) |
| |
|
|
| def tensor_to_numpy(x): |
| x = x[0].permute(1, 2, 0) |
| x = torch.clamp(x, -1 ,1) |
| x = (x+1) * 127.5 |
| x = x.cpu().detach().numpy().astype(np.uint8) |
| return x |
|
|
| def numpy_to_tensor(x): |
| x = (x / 255 - 0.5) * 2 |
| x = torch.from_numpy(x).unsqueeze(0).permute(0, 3, 1, 2) |
| x = x.cuda().float() |
| return x |
|
|
| def tensor_to_pil(x): |
| x = torch.clamp(x, -1 ,1) |
| x = (x+1) * 127.5 |
| return transforms.ToPILImage()(x.squeeze_(0)) |
|
|
|
|
