def __init__(self, ref_im, loss_str, eps):
super(LossBuilder, self).__init__()
assert ref_im.shape[2]==ref_im.shape[3]
im_size = ref_im.shape[2]
factor = 1024//im_size
assert im_size*factor==1024
self.D = BicubicDownSample(factor=factor)
self.ref_im = ref_im
self.parsed_loss = [loss_term.split('*') for loss_term in loss_str.split('+')]
self.eps = eps
def __init__(self, ref_im, loss_str, eps):
super().__init__()
assert ref_im.shape[2] == ref_im.shape[3]
im_size = ref_im.shape[2]
factor = 1024 // im_size
assert im_size * factor == 1024
self.D = BicubicDownSample(factor=factor)
self.ref_im = ref_im
self.parsed_loss = [
loss_term.split("*") for loss_term in loss_str.split("+")
]
self.eps = eps
if "DISC" in set([loss_type for _, loss_type in self.parsed_loss]):
self.d_basic = stylegan.D_basic()
self.d_basic.load_state_dict(
torch.load("karras2019stylegan-ffhq-1024x1024.for_d_basic.pt"))
def __init__(self, ref_im, loss_str, eps):
super(LossBuilder, self).__init__()
if ref_im.shape[2]!=ref_im.shape[3]:
print('ensure that your image dimensions are divisible by 4')
assert ref_im.shape[2]==ref_im.shape[3]
im_size = ref_im.shape[2]
factor=1024//im_size
if im_size*factor!=1024:
print('ensure that your image dimensions are divisible by 4')
assert im_size*factor==1024
self.D = BicubicDownSample(factor=factor)
self.ref_im = ref_im
self.parsed_loss = [loss_term.split('*') for loss_term in loss_str.split('+')]
self.eps = eps
args = parser.parse_args()
cache_dir = Path(args.cache_dir)
cache_dir.mkdir(parents=True, exist_ok=True)
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
print("Downloading Shape Predictor")
f = open_url(
"https://drive.google.com/uc?id=1xtxqSWYHADTEO9ptPG5174DbV3FYJPev",
cache_dir=cache_dir,
return_path=True)
predictor = dlib.shape_predictor(f)
for im in Path(args.input_dir).glob("*.*"):
faces = align_face(str(im), predictor)
for i, face in enumerate(faces):
if (args.output_size):
factor = 1024 // args.output_size
assert args.output_size * factor == 1024
D = BicubicDownSample(factor=factor)
face_tensor = torchvision.transforms.ToTensor()(face).unsqueeze(
0).cuda()
face_tensor_lr = D(face_tensor)[0].cpu().detach().clamp(0, 1)
face = torchvision.transforms.ToPILImage()(face_tensor_lr)
face.save(Path(args.output_dir) / (im.stem + f"_{i}.png"))
phase_base_dir = '../edsr_bases/'
phase_test_dir = '../demo/'
hr_dir = '../hr/'
lr_dir = '../lr/'
img_name = '0833'
hr_img = imread(os.path.join(hr_dir, img_name+ '.png')).astype(np.float32)
lr_img = imread(os.path.join(lr_dir, img_name+ '.png')).astype(np.float32)
test_img = imread(os.path.join(phase_test_dir, img_name+ '_hr_145.png')).astype(np.floa\
t32)
hr_tensor = torch.tensor(hr_img.reshape(1, *hr_img.shape)).type('torch.DoubleTensor')
lr_tensor = torch.tensor(lr_img.reshape(1, *lr_img.shape)).type('torch.DoubleTensor')
test_tensor = torch.tensor(test_img.reshape(1, *test_img.shape)).type('torch.DoubleTens\
or')
bds = BicubicDownSample()
ds_hr_tensor = bds(hr_tensor, nhwc=True)
l = nn.MSELoss()
ls_loss = l(ds_hr_tensor, lr_tensor)
# hr_loss = l(base_tensor, hr_tensor)
# print(ls_loss): 0.1046
diff = test_tensor - hr_tensor
diff_div = diff.div(255)
#print(diff_div)
mse = diff_div.pow(2).mean()
print(mse)
psnr = -10 * math.log10(mse)
def __init__(self, config):
super().__init__()
self.config = config
if config['image_size'][0] != config['image_size'][1]:
raise Exception('Non-square images are not supported yet.')
self.reconstruction = 'invert'
#self.project = config["project"]
self.steps = 1000
self.layer_in = None
self.best = None
self.skip = None
self.lr = 0.1
self.lr_record = []
self.current_step = 0
# prepare images
device = 'cuda'
resized_imgs = []
original_imgs = []
transform_lpips = get_transformation(256)
transform = get_transformation(256)
for imgfile in '/content/dcgan/alex.png':
resized_imgs.append(
transform_lpips(Image.open(imgfile).convert("RGB")))
original_imgs.append(transform(Image.open(imgfile).convert("RGB")))
self.resized_imgs = torch.stack(resized_imgs, 0).to(device)
self.original_imgs = torch.stack(original_imgs, 0).to(device)
self.downsampler_1024_image = BicubicDownSample(4)
# Load models and pre-trained weights
gen = Generator(1024, 512, 8)
gen.load_state_dict(torch.load(config["ckpt"])["g_ema"], strict=False)
gen.eval()
self.gen = gen.to(device)
self.gen.start_layer = 0
self.gen.end_layer = 4
self.mpl = MappingProxy(torch.load('gaussian_fit.pt'))
self.percept = lpips.PerceptualLoss(model="net-lin",
net="vgg",
use_gpu=device.startswith("cuda"))
self.cls = imagenet_models.resnet50()
state_dict = torch.load('imagenet_l2_3_0.pt')['model']
new_dict = OrderedDict()
for key in state_dict.keys():
if 'module.model' in key:
new_dict[key[13:]] = state_dict[key]
self.cls.load_state_dict(new_dict)
self.cls.to(config['device'])
bs = self.original_imgs.shape[0]
# initialization
if self.gen.start_layer == 0:
noises_single = self.gen.make_noise(bs)
self.noises = []
for noise in noises_single:
self.noises.append(noise.normal_())
self.latent_z = torch.randn((bs, 18, 512),
dtype=torch.float,
requires_grad=True,
device='cuda')
self.gen_outs = [None]
else:
# restore noises
self.noises = torch.load(config['saved_noises'][0])
self.latent_z = torch.load(config['saved_noises'][1]).to(
config['device'])
self.gen_outs = torch.load(config['saved_noises'][2])
self.latent_z.requires_grad = True