bottom_image_width))(tensor2pil(bottom_image)))
# 4. Combine the bottom zone with the upper zone.
final_image = np.concatenate((top_image, bottom_image))
# save compare video
compare_writer.write(final_image)
if args.view:
# display video
cv2.imshow("LR video convert SR video ", final_image)
if cv2.waitKey(1) & 0xFF == ord("q"):
break
# next frame
success, raw_frame = video_capture.read()
if __name__ == "__main__":
print("##################################################\n")
print("Run SR Engine.\n")
create_folder("videos")
logger.info("SREngine:")
print("\tAPI version .......... 0.2.1")
print("\tBuild ................ 2021.04.09")
print("##################################################\n")
main()
logger.info("Super-resolution video completed successfully.\n")
f"RMSE {value[1]:6.4f}\n"
f"PSNR {value[2]:6.2f}\n"
f"SSIM {value[3]:6.4f}\n"
f"LPIPS {value[4]:6.4f}\n"
f"GMSD {value[5]:6.4f}\n")
else:
images = torch.cat([bicubic, sr], dim=-1)
vutils.save_image(lr, os.path.join("test", f"lr_{filename}"))
vutils.save_image(bicubic, os.path.join("test", f"bicubic_{filename}"))
vutils.save_image(sr, os.path.join("test", f"sr_{filename}"))
vutils.save_image(images,
os.path.join("test", f"compare_{filename}"),
padding=10)
if __name__ == "__main__":
print("##################################################\n")
print("Run Testing Engine.\n")
create_folder("test")
logger.info("TestingEngine:")
print("\tAPI version .......... 0.1.0")
print("\tBuild ................ 2021.03.23")
print("##################################################\n")
main()
logger.info(
"Test single image performance evaluation completed successfully.\n")
writer.add_scalar("Train/D(SR2)", d_sr2, iters)
# Output results every 100 batches.
if i % 100 == 0:
progress.display(i)
# Save image every 1000 batches.
if iters % 1000 == 0:
vutils.save_image(hr, os.path.join("runs", "hr",
f"GAN_{iters}.bmp"))
sr = generator(lr)
vutils.save_image(sr.detach(),
os.path.join("runs", "sr", f"GAN_{iters}.bmp"))
if __name__ == "__main__":
print("##################################################\n")
print("Run Training Engine.\n")
create_folder("runs")
create_folder("runs/hr")
create_folder("runs/sr")
create_folder("weights")
logger.info("TrainingEngine:")
print("\tAPI version .......... 0.1.0")
print("\tBuild ................ 2021.03.25")
print("##################################################\n")
main()
logger.info("All training has been completed successfully.\n")
f"RMSE {value[1]:6.4f}\n"
f"PSNR {value[2]:6.2f}\n"
f"SSIM {value[3]:6.4f}\n"
f"LPIPS {value[4]:6.4f}\n"
f"GMSD {value[5]:6.4f}\n")
else:
images = torch.cat([bicubic, sr], dim=-1)
vutils.save_image(lr, os.path.join("tests", f"lr_{filename}"))
vutils.save_image(bicubic, os.path.join("tests", f"bicubic_{filename}"))
vutils.save_image(sr, os.path.join("tests", f"sr_{filename}"))
vutils.save_image(images,
os.path.join("tests", f"compare_{filename}"),
padding=10)
if __name__ == "__main__":
print("##################################################\n")
print("Run Testing Engine.\n")
create_folder("tests")
logger.info("TestingEngine:")
print("\tAPI version .......... 0.2.2")
print("\tBuild ................ 2021.04.23")
print("##################################################\n")
main()
logger.info(
"Test single image performance evaluation completed successfully.\n")
help="Path to latest generator checkpoint. (default: ````).")
parser.add_argument("--manualSeed",
type=int,
default=1111,
help="Seed for initializing training. (default:1111)")
parser.add_argument(
"--device",
default="",
help="device id i.e. `0` or `0,1` or `cpu`. (default: ````).")
args = parser.parse_args()
print("##################################################\n")
print("Run Training Engine.\n")
print(args)
create_folder("run")
create_folder("run/hr")
create_folder("run/sr")
create_folder("weights")
logger.info("TrainingEngine:")
print("\tAPI version .......... 0.1.1")
print("\tBuild ................ 2020.11.30-1116-0c5adc7e")
logger.info("Creating Training Engine")
trainer = Trainer(args)
logger.info("Staring training model")
trainer.run()
print("##################################################\n")
images = torch.cat([bicubic, sr, hr], dim=-1)
vutils.save_image(images,
os.path.join("benchmark", f"{i + 1}.bmp"),
padding=10)
print(f"Performance average results:\n")
print(f"indicator Score\n")
print(f"--------- -----\n")
print(f"MSE {total_mse_value / len(dataloader):6.4f}\n"
f"RMSE {total_rmse_value / len(dataloader):6.4f}\n"
f"PSNR {total_psnr_value / len(dataloader):6.2f}\n"
f"SSIM {total_ssim_value / len(dataloader):6.4f}\n"
f"LPIPS {total_lpips_value / len(dataloader):6.4f}\n"
f"GMSD {total_gmsd_value / len(dataloader):6.4f}")
if __name__ == "__main__":
print("##################################################\n")
print("Run Testing Engine.\n")
create_folder("benchmark")
logger.info("TestingEngine:")
print("\tAPI version .......... 0.1.0")
print("\tBuild ................ 2021.03.23")
print("##################################################\n")
main()
logger.info("Test dataset performance evaluation completed successfully.")
adversarial_losses.update(adversarial_loss.item(), lr.size(0))
iters = i + epoch * len(dataloader) + 1
writer.add_scalar("Train/D_Loss", d_loss.item(), iters)
writer.add_scalar("Train/G_Loss", g_loss.item(), iters)
writer.add_scalar("Train/Content_Loss", content_loss.item(), iters)
writer.add_scalar("Train/Adversarial_Loss", adversarial_loss.item(), iters)
# Output results every 100 batches.
if i % 100 == 0:
progress.display(i)
# Each Epoch validates the model once.
sr = generator(base_image)
vutils.save_image(sr.detach(), os.path.join("runs", f"GAN_epoch_{epoch}.png"))
if __name__ == "__main__":
print("##################################################\n")
print("Run Training Engine.\n")
create_folder("runs")
create_folder("weights")
logger.info("TrainingEngine:")
print("\tAPI version .......... 0.2.2")
print("\tBuild ................ 2021.04.25")
print("##################################################\n")
main()
logger.info("All training has been completed successfully.\n")
def train_gan(train_dataloader: torch.utils.data.DataLoader,
discriminator: nn.Module,
generator: nn.Module,
content_criterion: VGGLoss,
adversarial_criterion: nn.BCELoss,
discriminator_optimizer: torch.optim.Adam,
generator_optimizer: torch.optim.Adam,
epoch: int,
writer: SummaryWriter,
args: argparse.ArgumentParser.parse_args):
batch_time = AverageMeter("Time", ":.4f")
d_losses = AverageMeter("D Loss", ":.6f")
g_losses = AverageMeter("G Loss", ":.6f")
content_losses = AverageMeter("Content Loss", ":.4f")
adversarial_losses = AverageMeter("Adversarial Loss", ":.4f")
d_hr_values = AverageMeter("D(x)", ":.4f")
d_sr1_values = AverageMeter("D(SR1)", ":.4f")
d_sr2_values = AverageMeter("D(SR2)", ":.4f")
progress = ProgressMeter(
len(train_dataloader),
[batch_time,
d_losses, g_losses,
content_losses, adversarial_losses,
d_hr_values, d_sr1_values, d_sr2_values],
prefix=f"Epoch: [{epoch}]")
# switch to train mode
discriminator.train()
generator.train()
end = time.time()
for i, (lr, hr) in enumerate(train_dataloader):
# Move data to special device.
if args.gpu is not None:
lr = lr.cuda(args.gpu, non_blocking=True)
hr = hr.cuda(args.gpu, non_blocking=True)
batch_size = lr.size(0)
# The real sample label is 1, and the generated sample label is 0.
real_label = torch.full((batch_size, 1), 1, dtype=lr.dtype).cuda(args.gpu, non_blocking=True)
fake_label = torch.full((batch_size, 1), 0, dtype=lr.dtype).cuda(args.gpu, non_blocking=True)
##############################################
# (1) Update D network: maximize - E(hr)[log(D(hr))] + E(lr)[log(1- D(G(lr))]
##############################################
# Set discriminator gradients to zero.
discriminator.zero_grad()
real_output = discriminator(hr)
# Let the discriminator realize that the sample is real.
d_loss_real = adversarial_criterion(real_output, real_label)
# Generating fake high resolution images from real low resolution images.
sr = generator(lr)
fake_output = discriminator(sr.detach())
# Let the discriminator realize that the sample is false.
d_loss_fake = adversarial_criterion(fake_output, fake_label)
# Count all discriminator losses.
d_loss = (d_loss_real + d_loss_fake) / 2
d_loss.backward()
d_hr = real_output.mean().item()
d_sr1 = fake_output.mean().item()
# Update discriminator optimizer gradient information.
discriminator_optimizer.step()
##############################################
# (2) Update G network: content loss + 0.001 * adversarial loss
##############################################
# Set discriminator gradients to zero.
generator.zero_grad()
# Based on VGG19_36th pre training model to find the maximum square error between feature maps.
content_loss = content_criterion(sr, hr.detach())
fake_output = discriminator(sr)
# Let the discriminator realize that the sample is true.
adversarial_loss = adversarial_criterion(fake_output, real_label)
g_loss = content_loss + 0.001 * adversarial_loss
g_loss.backward()
d_sr2 = fake_output.mean().item()
# Update generator optimizer gradient information.
generator_optimizer.step()
# measure accuracy and record loss
d_losses.update(d_loss.item(), lr.size(0))
g_losses.update(g_loss.item(), lr.size(0))
content_losses.update(content_loss.item(), lr.size(0))
adversarial_losses.update(adversarial_loss.item(), lr.size(0))
d_hr_values.update(d_hr, lr.size(0))
d_sr1_values.update(d_sr1, lr.size(0))
d_sr2_values.update(d_sr2, lr.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
iters = i + epoch * len(train_dataloader) + 1
writer.add_scalar("Train/D Loss", d_loss.item(), iters)
writer.add_scalar("Train/G Loss", d_loss.item(), iters)
writer.add_scalar("Train/Content Loss", content_loss.item(), iters)
writer.add_scalar("Train/Adversarial Loss", adversarial_loss.item(), iters)
writer.add_scalar("Train/D(LR)", d_hr, iters)
writer.add_scalar("Train/D(SR1)", d_sr1, iters)
writer.add_scalar("Train/D(SR2)", d_sr2, iters)
# Output results every 100 batches.
if i % 100 == 0:
progress.display(i)
# Save image every 1000 batches.
if iters % 1000 == 0:
vutils.save_image(hr, os.path.join("runs", "hr", f"GAN_{iters}.bmp"))
sr = generator(lr)
vutils.save_image(sr.detach(), os.path.join("runs", "sr", f"GAN_{iters}.bmp"))
if __name__ == "__main__":
print("##################################################\n")
print("Run Training Engine.\n")
create_folder("runs")
create_folder("runs/hr")
create_folder("runs/sr")
create_folder("weights")
logger.info("TrainingEngine:")
print("\tAPI version .......... 0.1.0")
print("\tBuild ................ 2021.03.25")
print("##################################################\n")
main()
logger.info("All training has been completed successfully.\n")
