exist_ok=True)
newdir_masked = "L1_adv_test_results/random_bbox/{}_images/masked".format(
opt.dataset)
new_root_masked = "/home3/hiya/workspace/inpainting_fft/DnCNN+EC/{}".format(
newdir_masked)
os.makedirs("L1_adv_test_results/random_bbox/{}_images/recon".format(
opt.dataset),
exist_ok=True)
newdir_recon = "L1_adv_test_results/random_bbox/{}_images/recon".format(
opt.dataset)
new_root_recon = "/home3/hiya/workspace/inpainting_fft/DnCNN+EC/{}".format(
newdir_recon)
# loader_train, test_loader = get_data_loader(opt)
model = InpaintingModel(config, gen_in_channel=6).to(device)
Tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
# Initialize generator weights
if use_regular:
PATH = "L1_Adv_results/{}{}_generator.h5f".format(prefix, opt.dataset)
else:
PATH = "L1_Adv_results/{}_generator.h5f".format(opt.dataset)
print('Loading from ', PATH)
ckpt = torch.load(PATH, map_location='cpu')
model.generator.load_state_dict(ckpt)
# Initialize discriminator weights
if use_regular:
PATH = "L1_Adv_results/{}{}_discriminator.h5f".format(prefix, opt.dataset)
else:
os.makedirs("l1_adv_results/{}{}_images".format(prefix, opt.dataset), exist_ok=True)
use_cuda = not opt.no_cuda and torch.cuda.is_available()
torch.manual_seed(1234)
device = torch.device("cuda" if use_cuda else "cpu")
# PSNR metric
psnr_compute = PSNR(255.0).to(config.DEVICE)
# Dataloader
loader_train, test_loader = get_data_loader(opt)
# ----------
# Training
# ----------
model = InpaintingModel(config, gen_in_channel=2, gen_out_channel=1, disc_in_channel=1).to(device)
Tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
all_logs = []
for epoch in range(opt.n_epochs):
for i, (imgs, _) in enumerate(loader_train):
outputs_irregular = get_hirise_masked(imgs.numpy(), 'train')
x_masked, x_masked_parts, all_masks = outputs_irregular
masked_imgs = torch.from_numpy(x_masked).type(Tensor)
masks = torch.from_numpy(all_masks)
masks = masks.type(Tensor)
imgs = imgs.type(Tensor)
masked_imgs_display = masked_imgs.clone()
i_outputs, i_gen_loss, i_dis_loss, i_logs = model.process(imgs, masked_imgs, masks)
if opt.dataset != 'olivetti':
net_stage1 = net_stage1.cuda()
if use_regular:
PATH = '{}/regular_{}_net.pth'.format(opt.stage1_outf, opt.dataset)
else:
PATH = '{}/{}_net.pth'.format(opt.stage1_outf, opt.dataset)
print('Loading from ', PATH)
ckpt = torch.load(PATH, map_location='cpu')
net_stage1.load_state_dict(ckpt)
return net_stage1
net_stage1 = load_stage1_model().to(device)
loader_train, test_loader = get_data_loader(opt)
model = InpaintingModel(config).to(device)
# Optimizers
Tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
loader_train, test_loader = get_data_loader(opt)
# ----------
# Training
# ----------
for epoch in range(opt.n_epochs):
for i, (imgs, _) in enumerate(loader_train):
# Mask the images and compute the FFTs also
if use_regular: