def main(data_dir):
# 0) Tensoboard Writer.
writer = SummaryWriter(FLAGS['summary_path'])
origin_img, uv_map_gt, uv_map_predicted = None, None, None
if not os.path.exists(FLAGS['images']):
os.mkdir(FLAGS['images'])
# 1) Create Dataset of 300_WLP & Dataloader.
wlp300 = PRNetDataset(root_dir=data_dir,
transform=transforms.Compose([
ToTensor(),
ToResize((416, 416)),
ToNormalize(FLAGS["normalize_mean"],
FLAGS["normalize_std"])
]))
wlp300_dataloader = DataLoader(dataset=wlp300,
batch_size=FLAGS['batch_size'],
shuffle=True,
num_workers=1)
# 2) Intermediate Processing.
transform_img = transforms.Compose([
transforms.Normalize(FLAGS["normalize_mean"], FLAGS["normalize_std"])
])
# 3) Create PRNet model.
start_epoch, target_epoch = FLAGS['start_epoch'], FLAGS['target_epoch']
model = ResFCN256(resolution_input=416,
resolution_output=416,
channel=3,
size=16)
discriminator = Discriminator()
# Load the pre-trained weight
if FLAGS['resume'] != "" and os.path.exists(
os.path.join(FLAGS['pretrained'], FLAGS['resume'])):
state = torch.load(os.path.join(FLAGS['pretrained'], FLAGS['resume']))
model.load_state_dict(state['prnet'])
start_epoch = state['start_epoch']
INFO("Load the pre-trained weight! Start from Epoch", start_epoch)
else:
start_epoch = 0
INFO(
"Pre-trained weight cannot load successfully, train from scratch!")
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(FLAGS["device"])
discriminator.to(FLAGS["device"])
optimizer = torch.optim.Adam(model.parameters(),
lr=FLAGS["lr"],
betas=(0.5, 0.999))
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=FLAGS["lr"])
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.99)
stat_loss = SSIM(mask_path=FLAGS["mask_path"], gauss=FLAGS["gauss_kernel"])
loss = WeightMaskLoss(mask_path=FLAGS["mask_path"])
bce_loss = torch.nn.BCEWithLogitsLoss()
bce_loss.to(FLAGS["device"])
#Loss function for adversarial
for ep in range(start_epoch, target_epoch):
bar = tqdm(wlp300_dataloader)
loss_list_G, stat_list = [], []
loss_list_D = []
for i, sample in enumerate(bar):
uv_map, origin = sample['uv_map'].to(
FLAGS['device']), sample['origin'].to(FLAGS['device'])
# Inference.
optimizer.zero_grad()
uv_map_result = model(origin)
# Update D
optimizer_D.zero_grad()
fake_detach = uv_map_result.detach()
d_fake = discriminator(fake_detach)
d_real = discriminator(uv_map)
retain_graph = False
if FLAGS['gan_type'] == 'GAN':
loss_d = bce_loss(d_real, d_fake)
elif FLAGS['gan_type'].find('WGAN') >= 0:
loss_d = (d_fake - d_real).mean()
if FLAGS['gan_type'].find('GP') >= 0:
epsilon = torch.rand(fake_detach.shape[0]).view(
-1, 1, 1, 1)
epsilon = epsilon.to(fake_detach.device)
hat = fake_detach.mul(1 - epsilon) + uv_map.mul(epsilon)
hat.requires_grad = True
d_hat = discriminator(hat)
gradients = torch.autograd.grad(outputs=d_hat.sum(),
inputs=hat,
retain_graph=True,
create_graph=True,
only_inputs=True)[0]
gradients = gradients.view(gradients.size(0), -1)
gradient_norm = gradients.norm(2, dim=1)
gradient_penalty = 10 * gradient_norm.sub(1).pow(2).mean()
loss_d += gradient_penalty
# from ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks
elif FLAGS['gan_type'] == 'RGAN':
better_real = d_real - d_fake.mean(dim=0, keepdim=True)
better_fake = d_fake - d_real.mean(dim=0, keepdim=True)
loss_d = bce_loss(better_real, better_fake)
retain_graph = True
if discriminator.training:
loss_list_D.append(loss_d.item())
loss_d.backward(retain_graph=retain_graph)
optimizer_D.step()
if 'WGAN' in FLAGS['gan_type']:
for p in discriminator.parameters():
p.data.clamp_(-1, 1)
# Update G
d_fake_bp = discriminator(
uv_map_result) # for backpropagation, use fake as it is
if FLAGS['gan_type'] == 'GAN':
label_real = torch.ones_like(d_fake_bp)
loss_g = bce_loss(d_fake_bp, label_real)
elif FLAGS['gan_type'].find('WGAN') >= 0:
loss_g = -d_fake_bp.mean()
elif FLAGS['gan_type'] == 'RGAN':
better_real = d_real - d_fake_bp.mean(dim=0, keepdim=True)
better_fake = d_fake_bp - d_real.mean(dim=0, keepdim=True)
loss_g = bce_loss(better_fake, better_real)
loss_g.backward()
loss_list_G.append(loss_g.item())
optimizer.step()
stat_logit = stat_loss(uv_map_result, uv_map)
stat_list.append(stat_logit.item())
#bar.set_description(" {} [Loss(Paper)] {} [Loss(D)] {} [SSIM({})] {}".format(ep, loss_list_G[-1], loss_list_D[-1],FLAGS["gauss_kernel"], stat_list[-1]))
# Record Training information in Tensorboard.
"""
if origin_img is None and uv_map_gt is None:
origin_img, uv_map_gt = origin, uv_map
uv_map_predicted = uv_map_result
writer.add_scalar("Original Loss", loss_list_G[-1], FLAGS["summary_step"])
writer.add_scalar("D Loss", loss_list_D[-1], FLAGS["summary_step"])
writer.add_scalar("SSIM Loss", stat_list[-1], FLAGS["summary_step"])
grid_1, grid_2, grid_3 = make_grid(origin_img, normalize=True), make_grid(uv_map_gt), make_grid(uv_map_predicted)
writer.add_image('original', grid_1, FLAGS["summary_step"])
writer.add_image('gt_uv_map', grid_2, FLAGS["summary_step"])
writer.add_image('predicted_uv_map', grid_3, FLAGS["summary_step"])
writer.add_graph(model, uv_map)
"""
if ep % FLAGS["save_interval"] == 0:
with torch.no_grad():
print(" {} [Loss(Paper)] {} [Loss(D)] {} [SSIM({})] {}".format(
ep, loss_list_G[-1], loss_list_D[-1],
FLAGS["gauss_kernel"], stat_list[-1]))
origin = cv2.imread("./test_data/obama_origin.jpg")
gt_uv_map = np.load("./test_data/test_obama.npy")
origin, gt_uv_map = test_data_preprocess(
origin), test_data_preprocess(gt_uv_map)
origin, gt_uv_map = transform_img(origin), transform_img(
gt_uv_map)
origin_in = origin.unsqueeze_(0).cuda()
pred_uv_map = model(origin_in).detach().cpu()
save_image(
[origin.cpu(),
gt_uv_map.unsqueeze_(0).cpu(), pred_uv_map],
os.path.join(FLAGS['images'],
str(ep) + '.png'),
nrow=1,
normalize=True)
# Save model
print("Save model")
state = {
'prnet': model.state_dict(),
'Loss': loss_list_G,
'start_epoch': ep,
'Loss_D': loss_list_D,
}
torch.save(state, os.path.join(FLAGS['images'],
'{}.pth'.format(ep)))
scheduler.step()
writer.close()
def main(data_dir):
origin_img, uv_map_gt, uv_map_predicted = None, None, None
if not os.path.exists(FLAGS['images']):
os.mkdir(FLAGS['images'])
# 1) Create Dataset of 300_WLP & Dataloader.
wlp300 = PRNetDataset(root_dir=data_dir,
transform=transforms.Compose([
ToTensor(),
ToResize((416, 416)),
ToNormalize(FLAGS["normalize_mean"],
FLAGS["normalize_std"])
]))
wlp300_dataloader = DataLoader(dataset=wlp300,
batch_size=FLAGS['batch_size'],
shuffle=True,
num_workers=1)
# 2) Intermediate Processing.
transform_img = transforms.Compose([
#transforms.ToTensor(),
transforms.Normalize(FLAGS["normalize_mean"], FLAGS["normalize_std"])
])
# 3) Create PRNet model.
start_epoch, target_epoch = FLAGS['start_epoch'], FLAGS['target_epoch']
g_x = ResFCN256(resolution_input=416,
resolution_output=416,
channel=3,
size=16)
g_y = ResFCN256(resolution_input=416,
resolution_output=416,
channel=3,
size=16)
d_x = Discriminator()
d_y = Discriminator()
# Load the pre-trained weight
if FLAGS['resume'] != "" and os.path.exists(
os.path.join(FLAGS['pretrained'], FLAGS['resume'])):
state = torch.load(os.path.join(FLAGS['pretrained'], FLAGS['resume']))
try:
g_x.load_state_dict(state['g_x'])
g_y.load_state_dict(state['g_y'])
d_x.load_state_dict(state['d_x'])
d_y.load_state_dict(state['d_y'])
except Exception:
g_x.load_state_dict(state['prnet'])
start_epoch = state['start_epoch']
INFO("Load the pre-trained weight! Start from Epoch", start_epoch)
else:
start_epoch = 0
INFO(
"Pre-trained weight cannot load successfully, train from scratch!")
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
g_x.to(FLAGS["device"])
g_y.to(FLAGS["device"])
d_x.to(FLAGS["device"])
d_y.to(FLAGS["device"])
optimizer_g = torch.optim.Adam(itertools.chain(g_x.parameters(),
g_y.parameters()),
lr=FLAGS["lr"],
betas=(0.5, 0.999))
optimizer_d = torch.optim.Adam(itertools.chain(d_x.parameters(),
d_y.parameters()),
lr=FLAGS["lr"])
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer_g, gamma=0.99)
stat_loss = SSIM(mask_path=FLAGS["mask_path"], gauss=FLAGS["gauss_kernel"])
loss = WeightMaskLoss(mask_path=FLAGS["mask_path"])
bce_loss = torch.nn.BCEWithLogitsLoss()
bce_loss.to(FLAGS["device"])
l1_loss = nn.L1Loss().to(FLAGS["device"])
lambda_X = 10
lambda_Y = 10
#Loss function for adversarial
for ep in range(start_epoch, target_epoch):
bar = tqdm(wlp300_dataloader)
loss_list_cycle_x = []
loss_list_cycle_y = []
loss_list_d_x = []
loss_list_d_y = []
real_label = torch.ones(FLAGS['batch_size'])
fake_label = torch.zeros(FLAGS['batch_size'])
for i, sample in enumerate(bar):
real_y, real_x = sample['uv_map'].to(
FLAGS['device']), sample['origin'].to(FLAGS['device'])
# x -> y' -> x^
optimizer_g.zero_grad()
fake_y = g_x(real_x)
prediction = d_x(fake_y)
loss_g_x = bce_loss(prediction, real_label)
x_hat = g_y(fake_y)
loss_cycle_x = l1_loss(x_hat, real_x) * lambda_X
loss_x = loss_g_x + loss_cycle_x
loss_x.backward(retain_graph=True)
optimizer_g.step()
loss_list_cycle_x.append(loss_x.item())
# y -> x' -> y^
optimizer_g.zero_grad()
fake_x = g_y(real_y)
prediction = d_y(fake_x)
loss_g_y = bce_loss(prediction, real_label)
y_hat = g_x(fake_x)
loss_cycle_y = l1_loss(y_hat, real_y) * lambda_Y
loss_y = loss_g_y + loss_cycle_y
loss_y.backward(retain_graph=True)
optimizer_g.step()
loss_list_cycle_y.append(loss_y.item())
# d_x
optimizer_d.zero_grad()
pred_real = d_x(real_y)
loss_d_x_real = bce_loss(pred_real, real_label)
pred_fake = d_x(fake_y)
loss_d_x_fake = bce_loss(pred_fake, fake_label)
loss_d_x = (loss_d_x_real + loss_d_x_fake) * 0.5
loss_d_x.backward()
loss_list_d_x.append(loss_d_x.item())
optimizer_d.step()
if 'WGAN' in FLAGS['gan_type']:
for p in d_x.parameters():
p.data.clamp_(-1, 1)
# d_y
optimizer_d.zero_grad()
pred_real = d_y(real_x)
loss_d_y_real = bce_loss(pred_real, real_label)
pred_fake = d_y(fake_x)
loss_d_y_fake = bce_loss(pred_fake, fake_label)
loss_d_y = (loss_d_y_real + loss_d_y_fake) * 0.5
loss_d_y.backward()
loss_list_d_y.append(loss_d_y.item())
optimizer_d.step()
if 'WGAN' in FLAGS['gan_type']:
for p in d_y.parameters():
p.data.clamp_(-1, 1)
if ep % FLAGS["save_interval"] == 0:
with torch.no_grad():
print(
" {} [Loss_G_X] {} [Loss_G_Y] {} [Loss_D_X] {} [Loss_D_Y] {}"
.format(ep, loss_list_g_x[-1], loss_list_g_y[-1],
loss_list_d_x[-1], loss_list_d_y[-1]))
origin = cv2.imread("./test_data/obama_origin.jpg")
gt_uv_map = np.load("./test_data/test_obama.npy")
origin, gt_uv_map = test_data_preprocess(
origin), test_data_preprocess(gt_uv_map)
origin, gt_uv_map = transform_img(origin), transform_img(
gt_uv_map)
origin_in = origin.unsqueeze_(0).cuda()
pred_uv_map = g_x(origin_in).detach().cpu()
save_image(
[origin.cpu(),
gt_uv_map.unsqueeze_(0).cpu(), pred_uv_map],
os.path.join(FLAGS['images'],
str(ep) + '.png'),
nrow=1,
normalize=True)
# Save model
print("Save model")
state = {
'g_x': g_x.state_dict(),
'g_y': g_y.state_dict(),
'd_x': d_x.state_dict(),
'd_y': d_y.state_dict(),
'start_epoch': ep,
}
torch.save(state, os.path.join(FLAGS['images'],
'{}.pth'.format(ep)))
scheduler.step()
for each_hazy in hazy_lst:
hazy_list.append(hazy_img_path + each_hazy)
clear_list.append(clear_img_path + each_hazy.split('_')[0] + '.png')
#divide the dataset into train/val
train_loader, test_loader = split_train_val(train_size, image_size, hazy_list,
clear_list, batch_size)
#inititialize the Generator,Discriminator and move them to GPU if available
gen = Generator()
gen.apply(init_weights)
gen.to(device)
disc = Discriminator()
disc.apply(init_weights)
disc.to(device)
#initialize the gan,content,perceptual and brightness loss
gan_criterion = nn.BCELoss().to(device)
content_criterion = nn.L1Loss().to(device)
perceptual_criterion = nn.MSELoss().to(device)
brightness_criterion = nn.L1Loss().to(device)
#set the feature extractor to evaluation mode as it will be used only to calculate perceptual loss
feature_extractor = FeatureExtractor()
feature_extractor.eval()
feature_extractor.to(device)
#initialize the optimizers for Generator and Discriminator
optimizerD = optim.Adam(disc.parameters(), lr=0.0003, betas=(0.5, 0.999))
optimizerG = optim.Adam(gen.parameters(), lr=0.0001, betas=(0.5, 0.999))
def train(train_sources, eval_source):
path = sys.argv[1]
dr = DataReader(path, train_sources)
dr.read()
print(len(dr.train.x))
batch_size = 8
device = torch.device('cpu')
if torch.cuda.is_available():
device = torch.device('cuda')
dataset_s_train = MultiDomainDataset(dr.train.x, dr.train.y, dr.train.vendor, device, DomainAugmentation())
dataset_s_dev = MultiDomainDataset(dr.dev.x, dr.dev.y, dr.dev.vendor, device)
dataset_s_test = MultiDomainDataset(dr.test.x, dr.test.y, dr.test.vendor, device)
loader_s_train = DataLoader(dataset_s_train, batch_size, shuffle=True)
dr_eval = DataReader(path, [eval_source])
dr_eval.read()
dataset_eval_dev = MultiDomainDataset(dr_eval.dev.x, dr_eval.dev.y, dr_eval.dev.vendor, device)
dataset_eval_test = MultiDomainDataset(dr_eval.test.x, dr_eval.test.y, dr_eval.test.vendor, device)
dataset_da_train = MultiDomainDataset(dr.train.x+dr_eval.train.x, dr.train.y+dr_eval.train.y, dr.train.vendor+dr_eval.train.vendor, device, DomainAugmentation())
loader_da_train = DataLoader(dataset_da_train, batch_size, shuffle=True)
segmentator = UNet()
discriminator = Discriminator(n_domains=len(train_sources))
discriminator.to(device)
segmentator.to(device)
sigmoid = nn.Sigmoid()
selector = Selector()
s_criterion = nn.BCELoss()
d_criterion = nn.CrossEntropyLoss()
s_optimizer = optim.AdamW(segmentator.parameters(), lr=0.0001, weight_decay=0.01)
d_optimizer = optim.AdamW(discriminator.parameters(), lr=0.001, weight_decay=0.01)
a_optimizer = optim.AdamW(segmentator.encoder.parameters(), lr=0.001, weight_decay=0.01)
lmbd = 1/150
s_train_losses = []
s_dev_losses = []
d_train_losses = []
eval_domain_losses = []
train_dices = []
dev_dices = []
eval_dices = []
epochs = 3
da_loader_iter = iter(loader_da_train)
for epoch in tqdm(range(epochs)):
s_train_loss = 0.0
d_train_loss = 0.0
for index, sample in enumerate(loader_s_train):
img = sample['image']
target_mask = sample['target']
da_sample = next(da_loader_iter, None)
if epoch == 100:
s_optimizer.defaults['lr'] = 0.001
d_optimizer.defaults['lr'] = 0.0001
if da_sample is None:
da_loader_iter = iter(loader_da_train)
da_sample = next(da_loader_iter, None)
if epoch < 50 or epoch >= 100:
# Training step of segmentator
predicted_activations, inner_repr = segmentator(img)
predicted_mask = sigmoid(predicted_activations)
s_loss = s_criterion(predicted_mask, target_mask)
s_optimizer.zero_grad()
s_loss.backward()
s_optimizer.step()
s_train_loss += s_loss.cpu().detach().numpy()
if epoch >= 50:
# Training step of discriminator
predicted_activations, inner_repr = segmentator(da_sample['image'])
predicted_activations = predicted_activations.clone().detach()
inner_repr = inner_repr.clone().detach()
predicted_vendor = discriminator(predicted_activations, inner_repr)
d_loss = d_criterion(predicted_vendor, da_sample['vendor'])
d_optimizer.zero_grad()
d_loss.backward()
d_optimizer.step()
d_train_loss += d_loss.cpu().detach().numpy()
if epoch >= 100:
# adversarial training step
predicted_mask, inner_repr = segmentator(da_sample['image'])
predicted_vendor = discriminator(predicted_mask, inner_repr)
a_loss = -1 * lmbd * d_criterion(predicted_vendor, da_sample['vendor'])
a_optimizer.zero_grad()
a_loss.backward()
a_optimizer.step()
lmbd += 1/150
inference_model = nn.Sequential(segmentator, selector, sigmoid)
inference_model.to(device)
inference_model.eval()
d_train_losses.append(d_train_loss / len(loader_s_train))
s_train_losses.append(s_train_loss / len(loader_s_train))
s_dev_losses.append(calculate_loss(dataset_s_dev, inference_model, s_criterion, batch_size))
eval_domain_losses.append(calculate_loss(dataset_eval_dev, inference_model, s_criterion, batch_size))
train_dices.append(calculate_dice(inference_model, dataset_s_train))
dev_dices.append(calculate_dice(inference_model, dataset_s_dev))
eval_dices.append(calculate_dice(inference_model, dataset_eval_dev))
segmentator.train()
date_time = datetime.now().strftime("%m%d%Y_%H%M%S")
model_path = os.path.join(pathlib.Path(__file__).parent.absolute(), "model", "weights", "segmentator"+str(date_time)+".pth")
torch.save(segmentator.state_dict(), model_path)
util.plot_data([(s_train_losses, 'train_losses'), (s_dev_losses, 'dev_losses'), (d_train_losses, 'discriminator_losses'),
(eval_domain_losses, 'eval_domain_losses')],
'losses.png')
util.plot_dice([(train_dices, 'train_dice'), (dev_dices, 'dev_dice'), (eval_dices, 'eval_dice')],
'dices.png')
inference_model = nn.Sequential(segmentator, selector, sigmoid)
inference_model.to(device)
inference_model.eval()
print('Dice on annotated: ', calculate_dice(inference_model, dataset_s_test))
print('Dice on unannotated: ', calculate_dice(inference_model, dataset_eval_test))
