def predict(config, args):
gpu_manage(args)
dataset = TestDataset(args.test_dir)
data_loader = DataLoader(dataset=dataset,
num_workers=config.threads,
batch_size=1,
shuffle=False)
### MODELS LOAD ###
print('===> Loading models')
if config.gen_model == 'unet':
gen = UNet(in_ch=config.in_ch,
out_ch=config.out_ch,
gpu_ids=args.gpu_ids)
param = torch.load(args.pretrained)
gen.load_state_dict(param)
if args.cuda:
gen = gen.cuda(0)
with torch.no_grad():
for i, batch in enumerate(tqdm(data_loader)):
x = Variable(batch[0])
filename = batch[1][0]
if args.cuda:
x = x.cuda()
out = gen(x)
h = 1
w = 4
c = 3
p = config.size
allim = np.zeros((h, w, c, p, p))
x_ = x.cpu().numpy()[0]
out_ = out.cpu().numpy()[0]
in_rgb = x_[:3]
in_nir = x_[3]
out_rgb = np.clip(out_[:3], -1, 1)
out_cloud = np.clip(out_[3], -1, 1)
allim[0, 0, :] = np.repeat(in_nir[None, :, :], repeats=3,
axis=0) * 127.5 + 127.5
allim[0, 1, :] = in_rgb * 127.5 + 127.5
allim[0, 2, :] = out_rgb * 127.5 + 127.5
allim[0, 3, :] = np.repeat(
out_cloud[None, :, :], repeats=3, axis=0) * 127.5 + 127.5
allim = allim.transpose(0, 3, 1, 4, 2)
allim = allim.reshape((h * p, w * p, c))
save_image(args.out_dir, allim, i, 1, filename=filename)
def predict(config, args):
gpu_manage(args)
dataset = Dataset(args.test_dir)
data_loader = DataLoader(dataset=dataset,
num_workers=config.threads,
batch_size=1,
shuffle=False)
gen = UNet(in_ch=config.in_ch, out_ch=config.out_ch, gpu_ids=args.gpu_ids)
param = torch.load(args.pretrained)
gen.load_state_dict(param)
criterionMSE = nn.MSELoss()
if args.cuda:
gen = gen.cuda(0)
criterionMSE = criterionMSE.cuda(0)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
avg_mse = 0
avg_psnr = 0
avg_ssim = 0
with torch.no_grad():
for i, batch in enumerate(tqdm(data_loader)):
input_, ground_truth = Variable(batch[0]), Variable(batch[1])
filename = batch[2][0]
input_ = F.interpolate(input_, size=256).to(device)
ground_truth = F.interpolate(ground_truth, size=256).to(device)
output = gen(input_)
save_image_from_tensors(input_, output, ground_truth,
config.out_dir, i, 0, filename)
mse, psnr, ssim = get_metrics(output, ground_truth, criterionMSE)
print(filename)
print('MSE: {:.4f}'.format(mse))
print('PSNR: {:.4f} dB'.format(psnr))
print('SSIM: {:.4f} dB'.format(ssim))
avg_mse += mse
avg_psnr += psnr
avg_ssim += ssim
avg_mse = avg_mse / len(data_loader)
avg_psnr = avg_psnr / len(data_loader)
avg_ssim = avg_ssim / len(data_loader)
print('Average MSE: {:.4f}'.format(avg_mse))
print('Average PSNR: {:.4f} dB'.format(avg_psnr))
print('Average SSIM: {:.4f} dB'.format(avg_ssim))
def predict(config, args):
gpu_manage(args)
dataset = TestDataset(args.test_dir, config.in_ch, config.out_ch)
data_loader = DataLoader(dataset=dataset,
num_workers=config.threads,
batch_size=1,
shuffle=False)
### MODELS LOAD ###
print('===> Loading models')
gen = Generator(gpu_ids=config.gpu_ids)
param = torch.load(args.pretrained)
gen.load_state_dict(param)
if args.cuda:
gen = gen.cuda(0)
with torch.no_grad():
for i, batch in enumerate(tqdm(data_loader)):
x = Variable(batch[0])
filename = batch[1][0]
if args.cuda:
x = x.cuda()
att, out = gen(x)
h = 1
w = 3
c = 3
p = config.width
allim = np.zeros((h, w, c, p, p))
x_ = x.cpu().numpy()[0]
out_ = out.cpu().numpy()[0]
in_rgb = x_[:3]
out_rgb = np.clip(out_[:3], 0, 1)
att_ = att.cpu().numpy()[0] * 255
heat_att = heatmap(att_.astype('uint8'))
allim[0, 0, :] = in_rgb * 255
allim[0, 1, :] = out_rgb * 255
allim[0, 2, :] = heat_att
allim = allim.transpose(0, 3, 1, 4, 2)
allim = allim.reshape((h * p, w * p, c))
save_image(args.out_dir, allim, i, 1, filename=filename)
def predict(args):
gpu_manage(args)
### MODELS LOAD ###
print('===> Loading models')
gen = Generator(gpu_ids=args.gpu_ids)
param = torch.load(args.pretrained)
gen.load_state_dict(param)
if args.cuda:
gen = gen.cuda(0)
print('<=== Model loaded')
print('===> Loading test image')
img = cv2.imread(args.test_filepath, 1).astype(np.float32)
img = img / 255
img = img.transpose(2, 0, 1)
img = img[None]
print('<=== test image loaded')
with torch.no_grad():
x = torch.from_numpy(img)
if args.cuda:
x = x.cuda()
print('===> Removing the cloud...')
start_time = time.time()
att, out = gen(x)
print('<=== finish! %.3fs cost.' % (time.time() - start_time))
x_ = x.cpu().numpy()[0]
x_rgb = x_ * 255
x_rgb = x_rgb.transpose(1, 2, 0).astype('uint8')
out_ = out.cpu().numpy()[0]
out_rgb = np.clip(out_[:3], 0, 1) * 255
out_rgb = out_rgb.transpose(1, 2, 0).astype('uint8')
att_ = att.cpu().numpy()[0] * 255
att_heatmap = heatmap(att_.astype('uint8'))[0]
att_heatmap = att_heatmap.transpose(1, 2, 0)
allim = np.hstack((x_rgb, out_rgb, att_heatmap))
show(allim)
def train(config):
gpu_manage(config)
### DATASET LOAD ###
print('===> Loading datasets')
dataset = Dataset(config)
train_size = int(0.9 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(
dataset, [train_size, test_size])
training_data_loader = DataLoader(dataset=train_dataset,
num_workers=config.threads,
batch_size=config.batchsize,
shuffle=True)
test_data_loader = DataLoader(dataset=test_dataset,
num_workers=config.threads,
batch_size=config.test_batchsize,
shuffle=False)
### MODELS LOAD ###
print('===> Loading models')
if config.gen_model == 'unet':
gen = UNet(in_ch=config.in_ch,
out_ch=config.out_ch,
gpu_ids=config.gpu_ids)
else:
print('The generator model does not exist')
if config.gen_init is not None:
param = torch.load(config.gen_init)
gen.load_state_dict(param)
print('load {} as pretrained model'.format(config.gen_init))
dis = Discriminator(in_ch=config.in_ch,
out_ch=config.out_ch,
gpu_ids=config.gpu_ids)
if config.dis_init is not None:
param = torch.load(config.dis_init)
dis.load_state_dict(param)
print('load {} as pretrained model'.format(config.dis_init))
# setup optimizer
opt_gen = optim.Adam(gen.parameters(),
lr=config.lr,
betas=(config.beta1, 0.999),
weight_decay=0.00001)
opt_dis = optim.Adam(dis.parameters(),
lr=config.lr,
betas=(config.beta1, 0.999),
weight_decay=0.00001)
real_a = torch.FloatTensor(config.batchsize, config.in_ch, 256, 256)
real_b = torch.FloatTensor(config.batchsize, config.out_ch, 256, 256)
criterionL1 = nn.L1Loss()
criterionMSE = nn.MSELoss()
criterionSoftplus = nn.Softplus()
if config.cuda:
gen = gen.cuda(0)
dis = dis.cuda(0)
criterionL1 = criterionL1.cuda(0)
criterionMSE = criterionMSE.cuda(0)
criterionSoftplus = criterionSoftplus.cuda(0)
real_a = real_a.cuda(0)
real_b = real_b.cuda(0)
real_a = Variable(real_a)
real_b = Variable(real_b)
logreport = LogReport(log_dir=config.out_dir)
testreport = TestReport(log_dir=config.out_dir)
# main
for epoch in range(1, config.epoch + 1):
for iteration, batch in enumerate(training_data_loader, 1):
real_a_cpu, real_b_cpu = batch[0], batch[1]
real_a.data.resize_(real_a_cpu.size()).copy_(real_a_cpu)
real_b.data.resize_(real_b_cpu.size()).copy_(real_b_cpu)
fake_b = gen.forward(real_a)
################
### Update D ###
################
opt_dis.zero_grad()
# train with fake
fake_ab = torch.cat((real_a, fake_b), 1)
pred_fake = dis.forward(fake_ab.detach())
batchsize, _, w, h = pred_fake.size()
loss_d_fake = torch.sum(
criterionSoftplus(pred_fake)) / batchsize / w / h
# train with real
real_ab = torch.cat((real_a, real_b), 1)
pred_real = dis.forward(real_ab)
loss_d_real = torch.sum(
criterionSoftplus(-pred_real)) / batchsize / w / h
# Combined loss
loss_d = loss_d_fake + loss_d_real
loss_d.backward()
if epoch % config.minimax == 0:
opt_dis.step()
################
### Update G ###
################
opt_gen.zero_grad()
# First, G(A) should fake the discriminator
fake_ab = torch.cat((real_a, fake_b), 1)
pred_fake = dis.forward(fake_ab)
loss_g_gan = torch.sum(
criterionSoftplus(-pred_fake)) / batchsize / w / h
# Second, G(A) = B
loss_g_l1 = criterionL1(fake_b, real_b) * config.lamb
loss_g = loss_g_gan + loss_g_l1
loss_g.backward()
opt_gen.step()
# log
if iteration % 100 == 0:
print(
"===> Epoch[{}]({}/{}): loss_d_fake: {:.4f} loss_d_real: {:.4f} loss_g_gan: {:.4f} loss_g_l1: {:.4f}"
.format(epoch, iteration, len(training_data_loader),
loss_d_fake.item(), loss_d_real.item(),
loss_g_gan.item(), loss_g_l1.item()))
log = {}
log['epoch'] = epoch
log['iteration'] = len(training_data_loader) * (epoch -
1) + iteration
log['gen/loss'] = loss_g.item()
log['dis/loss'] = loss_d.item()
logreport(log)
with torch.no_grad():
log_test = test(config, test_data_loader, gen, criterionMSE, epoch)
testreport(log_test)
if epoch % config.snapshot_interval == 0:
checkpoint(config, epoch, gen, dis)
logreport.save_lossgraph()
testreport.save_lossgraph()
def train(config):
gpu_manage(config)
### DATASET LOAD ###
print('===> Loading datasets')
dataset = TrainDataset(config)
print('dataset:', len(dataset))
train_size = int((1 - config.validation_size) * len(dataset))
validation_size = len(dataset) - train_size
train_dataset, validation_dataset = torch.utils.data.random_split(
dataset, [train_size, validation_size])
print('train dataset:', len(train_dataset))
print('validation dataset:', len(validation_dataset))
training_data_loader = DataLoader(dataset=train_dataset,
num_workers=config.threads,
batch_size=config.batchsize,
shuffle=True)
validation_data_loader = DataLoader(dataset=validation_dataset,
num_workers=config.threads,
batch_size=config.validation_batchsize,
shuffle=False)
### MODELS LOAD ###
print('===> Loading models')
gen = Generator(gpu_ids=config.gpu_ids)
if config.gen_init is not None:
param = torch.load(config.gen_init)
gen.load_state_dict(param)
print('load {} as pretrained model'.format(config.gen_init))
dis = Discriminator(in_ch=config.in_ch,
out_ch=config.out_ch,
gpu_ids=config.gpu_ids)
if config.dis_init is not None:
param = torch.load(config.dis_init)
dis.load_state_dict(param)
print('load {} as pretrained model'.format(config.dis_init))
# setup optimizer
opt_gen = optim.Adam(gen.parameters(),
lr=config.lr,
betas=(config.beta1, 0.999),
weight_decay=0.00001)
opt_dis = optim.Adam(dis.parameters(),
lr=config.lr,
betas=(config.beta1, 0.999),
weight_decay=0.00001)
real_a = torch.FloatTensor(config.batchsize, config.in_ch, config.width,
config.height)
real_b = torch.FloatTensor(config.batchsize, config.out_ch, config.width,
config.height)
M = torch.FloatTensor(config.batchsize, config.width, config.height)
criterionL1 = nn.L1Loss()
criterionMSE = nn.MSELoss()
criterionSoftplus = nn.Softplus()
if config.cuda:
gen = gen.cuda()
dis = dis.cuda()
criterionL1 = criterionL1.cuda()
criterionMSE = criterionMSE.cuda()
criterionSoftplus = criterionSoftplus.cuda()
real_a = real_a.cuda()
real_b = real_b.cuda()
M = M.cuda()
real_a = Variable(real_a)
real_b = Variable(real_b)
logreport = LogReport(log_dir=config.out_dir)
validationreport = TestReport(log_dir=config.out_dir)
print('===> begin')
start_time = time.time()
# main
for epoch in range(1, config.epoch + 1):
epoch_start_time = time.time()
for iteration, batch in enumerate(training_data_loader, 1):
real_a_cpu, real_b_cpu, M_cpu = batch[0], batch[1], batch[2]
real_a.data.resize_(real_a_cpu.size()).copy_(real_a_cpu)
real_b.data.resize_(real_b_cpu.size()).copy_(real_b_cpu)
M.data.resize_(M_cpu.size()).copy_(M_cpu)
att, fake_b = gen.forward(real_a)
################
### Update D ###
################
opt_dis.zero_grad()
# train with fake
fake_ab = torch.cat((real_a, fake_b), 1)
pred_fake = dis.forward(fake_ab.detach())
batchsize, _, w, h = pred_fake.size()
loss_d_fake = torch.sum(
criterionSoftplus(pred_fake)) / batchsize / w / h
# train with real
real_ab = torch.cat((real_a, real_b), 1)
pred_real = dis.forward(real_ab)
loss_d_real = torch.sum(
criterionSoftplus(-pred_real)) / batchsize / w / h
# Combined loss
loss_d = loss_d_fake + loss_d_real
loss_d.backward()
if epoch % config.minimax == 0:
opt_dis.step()
################
### Update G ###
################
opt_gen.zero_grad()
# First, G(A) should fake the discriminator
fake_ab = torch.cat((real_a, fake_b), 1)
pred_fake = dis.forward(fake_ab)
loss_g_gan = torch.sum(
criterionSoftplus(-pred_fake)) / batchsize / w / h
# Second, G(A) = B
loss_g_l1 = criterionL1(fake_b, real_b) * config.lamb
loss_g_att = criterionMSE(att[:, 0, :, :], M)
loss_g = loss_g_gan + loss_g_l1 + loss_g_att
loss_g.backward()
opt_gen.step()
# log
if iteration % 10 == 0:
print(
"===> Epoch[{}]({}/{}): loss_d_fake: {:.4f} loss_d_real: {:.4f} loss_g_gan: {:.4f} loss_g_l1: {:.4f}"
.format(epoch, iteration, len(training_data_loader),
loss_d_fake.item(), loss_d_real.item(),
loss_g_gan.item(), loss_g_l1.item()))
log = {}
log['epoch'] = epoch
log['iteration'] = len(training_data_loader) * (epoch -
1) + iteration
log['gen/loss'] = loss_g.item()
log['dis/loss'] = loss_d.item()
logreport(log)
print('epoch', epoch, 'finished, use time',
time.time() - epoch_start_time)
with torch.no_grad():
log_validation = test(config, validation_data_loader, gen,
criterionMSE, epoch)
validationreport(log_validation)
print('validation finished')
if epoch % config.snapshot_interval == 0:
checkpoint(config, epoch, gen, dis)
logreport.save_lossgraph()
validationreport.save_lossgraph()
print('training time:', time.time() - start_time)
def train(config):
gpu_manage(config)
train_dataset = Dataset(config.train_dir)
val_dataset = Dataset(config.val_dir)
training_data_loader = DataLoader(dataset=train_dataset,
num_workers=config.threads,
batch_size=config.batchsize,
shuffle=True)
val_data_loader = DataLoader(dataset=val_dataset,
num_workers=config.threads,
batch_size=config.test_batchsize,
shuffle=False)
gen = UNet(in_ch=config.in_ch, out_ch=config.out_ch, gpu_ids=config.gpu_ids)
if config.gen_init is not None:
param = torch.load(config.gen_init)
gen.load_state_dict(param)
print('load {} as pretrained model'.format(config.gen_init))
dis = Discriminator(in_ch=config.in_ch, out_ch=config.out_ch, gpu_ids=config.gpu_ids)
if config.dis_init is not None:
param = torch.load(config.dis_init)
dis.load_state_dict(param)
print('load {} as pretrained model'.format(config.dis_init))
opt_gen = optim.Adam(gen.parameters(), lr=config.lr, betas=(config.beta1, 0.999), weight_decay=0.00001)
opt_dis = optim.Adam(dis.parameters(), lr=config.lr, betas=(config.beta1, 0.999), weight_decay=0.00001)
real_a = torch.FloatTensor(config.batchsize, config.in_ch, 256, 256)
real_b = torch.FloatTensor(config.batchsize, config.out_ch, 256, 256)
criterionL1 = nn.L1Loss()
criterionMSE = nn.MSELoss()
criterionSoftplus = nn.Softplus()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if config.cuda:
gen = gen.cuda(0)
dis = dis.cuda(0)
criterionL1 = criterionL1.cuda(0)
criterionMSE = criterionMSE.cuda(0)
criterionSoftplus = criterionSoftplus.cuda(0)
real_a = real_a.cuda(0)
real_b = real_b.cuda(0)
real_a = Variable(real_a)
real_b = Variable(real_b)
logreport = LogReport(log_dir=config.out_dir)
testreport = TestReport(log_dir=config.out_dir)
for epoch in range(1, config.epoch + 1):
print('Epoch', epoch, datetime.now())
for iteration, batch in enumerate(tqdm(training_data_loader)):
real_a, real_b = batch[0], batch[1]
real_a = F.interpolate(real_a, size=256).to(device)
real_b = F.interpolate(real_b, size=256).to(device)
fake_b = gen.forward(real_a)
# Update D
opt_dis.zero_grad()
fake_ab = torch.cat((real_a, fake_b), 1)
pred_fake = dis.forward(fake_ab.detach())
batchsize, _, w, h = pred_fake.size()
real_ab = torch.cat((real_a, real_b), 1)
pred_real = dis.forward(real_ab)
loss_d_fake = torch.sum(criterionSoftplus(pred_fake)) / batchsize / w / h
loss_d_real = torch.sum(criterionSoftplus(-pred_real)) / batchsize / w / h
loss_d = loss_d_fake + loss_d_real
loss_d.backward()
if epoch % config.minimax == 0:
opt_dis.step()
# Update G
opt_gen.zero_grad()
fake_ab = torch.cat((real_a, fake_b), 1)
pred_fake = dis.forward(fake_ab)
loss_g_gan = torch.sum(criterionSoftplus(-pred_fake)) / batchsize / w / h
loss_g = loss_g_gan + criterionL1(fake_b, real_b) * config.lamb
loss_g.backward()
opt_gen.step()
if iteration % 100 == 0:
logreport({
'epoch': epoch,
'iteration': len(training_data_loader) * (epoch - 1) + iteration,
'gen/loss': loss_g.item(),
'dis/loss': loss_d.item(),
})
with torch.no_grad():
log_test = test(config, val_data_loader, gen, criterionMSE, epoch)
testreport(log_test)
if epoch % config.snapshot_interval == 0:
checkpoint(config, epoch, gen, dis)
logreport.save_lossgraph()
testreport.save_lossgraph()
print('Done', datetime.now())