def __init__(self, args): self.args = args self.model, self.device = configure(args) # Image preprocessing operation self.tensor2pil = transforms.ToPILImage() self.video_capture = cv2.VideoCapture(args.file) # Prepare to write the processed image into the video. self.fps = self.video_capture.get(cv2.CAP_PROP_FPS) self.total_frames = int( self.video_capture.get(cv2.CAP_PROP_FRAME_COUNT)) # Set video size self.size = (int(self.video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(self.video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))) self.sr_size = (self.size[0] * args.upscale_factor, self.size[1] * args.upscale_factor) self.pare_size = (self.sr_size[0] * 2 + 10, self.sr_size[1] + 10 + self.sr_size[0] // 5 - 9) # Video write loader. self.sr_writer = cv2.VideoWriter( os.path.join( "video", f"sr_{args.upscale_factor}x_{os.path.basename(args.file)}"), cv2.VideoWriter_fourcc(*"MPEG"), self.fps, self.sr_size) self.compare_writer = cv2.VideoWriter( os.path.join( "video", f"compare_{args.upscale_factor}x_{os.path.basename(args.file)}" ), cv2.VideoWriter_fourcc(*"MPEG"), self.fps, self.pare_size)
def main_worker(gpu, args): args.gpu = gpu if args.gpu is not None: logger.info(f"Use GPU: {args.gpu} for training.") model = configure(args) if not torch.cuda.is_available(): logger.warning("Using CPU, this will be slow.") elif args.gpu is not None: torch.cuda.set_device(args.gpu) model = model.cuda(args.gpu) # Set eval mode. model.eval() cudnn.benchmark = True # Get image filename. filename = os.path.basename(args.lr) # Read all pictures. lr = Image.open(args.lr) bicubic = transforms.Resize( (lr.size[1] * args.upscale_factor, lr.size[0] * args.upscale_factor), Mode.BICUBIC)(lr) lr = process_image(lr, args.gpu) bicubic = process_image(bicubic, args.gpu) with torch.no_grad(): sr = model(lr) if args.hr: hr = process_image(Image.open(args.hr), args.gpu) vutils.save_image(hr, os.path.join("test", f"hr_{filename}")) images = torch.cat([bicubic, sr, hr], dim=-1) value = iqa(sr, hr, args.gpu) print(f"Performance avg results:\n") print(f"indicator Score\n") print(f"--------- -----\n") print(f"MSE {value[0]:6.4f}\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)
def main_worker(gpu, args): global best_mse_value, best_rmse_value, best_psnr_value, best_ssim_value, best_lpips_value, best_gmsd_value args.gpu = gpu if args.gpu is not None: logger.info(f"Use GPU: {args.gpu} for training.") cudnn.benchmark = True model = configure(args) if not torch.cuda.is_available(): logger.warning("Using CPU, this will be slow.") # Read all pictures. lr = process_image(Image.open(args.lr), args.gpu) hr = process_image(Image.open(args.hr), args.gpu) model_paths = glob( os.path.join(f"{args.model_dir}", "Generator_epoch*.pth")) best_model = model_paths[0] for model_path in model_paths: print(f"Process `{model_path}`") value = inference(lr, hr, model, model_path, gpu) is_best = value[2] > best_psnr_value if is_best: best_model = os.path.basename(model_path) best_mse_value = value[0] best_rmse_value = value[1] best_psnr_value = value[2] best_ssim_value = value[3] best_lpips_value = value[4] best_gmsd_value = value[5] print("\n##################################################") print(f"Best model: `{best_model}`.") print(f"indicator Score") print(f"--------- -----") print(f"MSE {best_mse_value:6.4f}" f"RMSE {best_rmse_value:6.2f}" f"PSNR {best_psnr_value:6.2f}\n" f"SSIM {best_ssim_value:6.4f}\n" f"LPIPS {best_lpips_value:6.4f}\n" f"GMSD {best_gmsd_value:6.4f}") print(f"--------- -----") print("##################################################\n")
def __init__(self, args): self.args = args self.model, self.device = configure(args) logger.info("Load testing dataset") dataset = BaseTestDataset(root=os.path.join(args.data, "test"), image_size=args.image_size) self.dataloader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, pin_memory=True, num_workers=int(args.workers)) logger.info(f"Dataset information\n" f"\tDataset dir is `{os.getcwd()}/{args.data}/test`\n" f"\tBatch size is {args.batch_size}\n" f"\tWorkers is {int(args.workers)}\n" f"\tLoad dataset to CUDA")
def main_worker(gpu, args): args.gpu = gpu if args.gpu is not None: logger.info(f"Use GPU: {args.gpu} for testing.") model = configure(args) if not torch.cuda.is_available(): logger.warning("Using CPU, this will be slow.") if args.gpu is not None: torch.cuda.set_device(args.gpu) model = model.cuda(args.gpu) cudnn.benchmark = True # Set eval mode. model.eval() # Get video filename. filename = os.path.basename(args.file) # Image preprocessing operation tensor2pil = transforms.ToPILImage() video_capture = cv2.VideoCapture(args.file) # Prepare to write the processed image into the video. fps = video_capture.get(cv2.CAP_PROP_FPS) total_frames = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT)) # Set video size size = (int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))) sr_size = (size[0] * args.upscale_factor, size[1] * args.upscale_factor) pare_size = (sr_size[0] * 2 + 10, sr_size[1] + 10 + sr_size[0] // 5 - 9) # Video write loader. sr_writer = cv2.VideoWriter( os.path.join("videos", f"sr_{args.upscale_factor}x_{filename}"), cv2.VideoWriter_fourcc(*"MPEG"), fps, sr_size) compare_writer = cv2.VideoWriter( os.path.join("videos", f"compare_{args.upscale_factor}x_{filename}"), cv2.VideoWriter_fourcc(*"MPEG"), fps, pare_size) # read frame. with torch.no_grad(): success, raw_frame = video_capture.read() progress_bar = tqdm( range(total_frames), desc="[processing video and saving/view result videos]") for _ in progress_bar: if success: # Read image to tensor and transfer to the specified device for processing. lr = process_image(raw_frame, args.gpu) sr = model(lr) sr = sr.cpu() sr = sr.data[0].numpy() sr *= 255.0 sr = (np.uint8(sr)).transpose((1, 2, 0)) # save sr video sr_writer.write(sr) # make compared video and crop shot of left top\right top\center\left bottom\right bottom sr = tensor2pil(sr) # Five areas are selected as the bottom contrast map. crop_sr_images = transforms.FiveCrop(size=sr.width // 5 - 9)(sr) crop_sr_images = [ np.asarray(transforms.Pad(padding=(10, 5, 0, 0))(image)) for image in crop_sr_images ] sr = transforms.Pad(padding=(5, 0, 0, 5))(sr) # Five areas in the contrast map are selected as the bottom contrast map compare_image = transforms.Resize( (sr_size[1], sr_size[0]), interpolation=Mode.BICUBIC)(tensor2pil(raw_frame)) crop_compare_images = transforms.FiveCrop( size=compare_image.width // 5 - 9)(compare_image) crop_compare_images = [ np.asarray(transforms.Pad((0, 5, 10, 0))(image)) for image in crop_compare_images ] compare_image = transforms.Pad(padding=(0, 0, 5, 5))(compare_image) # concatenate all the pictures to one single picture # 1. Mosaic the left and right images of the video. top_image = np.concatenate( (np.asarray(compare_image), np.asarray(sr)), axis=1) # 2. Mosaic the bottom left and bottom right images of the video. bottom_image = np.concatenate(crop_compare_images + crop_sr_images, axis=1) bottom_image_height = int(top_image.shape[1] / bottom_image.shape[1] * bottom_image.shape[0]) bottom_image_width = top_image.shape[1] # 3. Adjust to the right size. bottom_image = np.asarray( transforms.Resize( (bottom_image_height, 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()
def main_worker(gpu, ngpus_per_node, args): global best_psnr args.gpu = gpu if args.gpu is not None: logger.info(f"Use GPU: {args.gpu} for training.") if args.distributed: if args.dist_url == "env://" and args.rank == -1: args.rank = int(os.environ["RANK"]) if args.multiprocessing_distributed: # For multiprocessing distributed training, rank needs to be the # global rank among all the processes args.rank = args.rank * ngpus_per_node + gpu dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank) # create model generator = configure(args) discriminator = discriminator_for_vgg(args.image_size) if not torch.cuda.is_available(): logger.warning("Using CPU, this will be slow.") elif args.distributed: # For multiprocessing distributed, DistributedDataParallel constructor # should always set the single device scope, otherwise, # DistributedDataParallel will use all available devices. if args.gpu is not None: torch.cuda.set_device(args.gpu) discriminator.cuda(args.gpu) generator.cuda(args.gpu) # When using a single GPU per process and per # DistributedDataParallel, we need to divide the batch size # ourselves based on the total number of GPUs we have args.batch_size = int(args.batch_size / ngpus_per_node) args.workers = int( (args.workers + ngpus_per_node - 1) / ngpus_per_node) discriminator = nn.parallel.DistributedDataParallel( module=discriminator, device_ids=[args.gpu]) generator = nn.parallel.DistributedDataParallel( module=generator, device_ids=[args.gpu]) else: discriminator.cuda() generator.cuda() # DistributedDataParallel will divide and allocate batch_size to all # available GPUs if device_ids are not set discriminator = nn.parallel.DistributedDataParallel(discriminator) generator = nn.parallel.DistributedDataParallel(generator) elif args.gpu is not None: torch.cuda.set_device(args.gpu) discriminator = discriminator.cuda(args.gpu) generator = generator.cuda(args.gpu) else: # DataParallel will divide and allocate batch_size to all available GPUs if args.arch.startswith("alexnet") or args.arch.startswith("vgg"): discriminator.features = torch.nn.DataParallel( discriminator.features) generator.features = torch.nn.DataParallel(generator.features) discriminator.cuda() generator.cuda() else: discriminator = torch.nn.DataParallel(discriminator).cuda() generator = torch.nn.DataParallel(generator).cuda() # Loss = content loss + 0.001 * adversarial loss pixel_criterion = nn.MSELoss().cuda(args.gpu) # We use VGG5.4 as our feature extraction method by default. content_criterion = VGGLoss().cuda(args.gpu) adversarial_criterion = nn.BCELoss().cuda(args.gpu) logger.info(f"Losses function information:\n" f"\tPixel: MSELoss\n" f"\tContent: VGG19_36th\n" f"\tAdversarial: BCELoss") # All optimizer function and scheduler function. psnr_optimizer = torch.optim.Adam(generator.parameters(), lr=args.psnr_lr, betas=(0.9, 0.999)) discriminator_optimizer = torch.optim.Adam(discriminator.parameters(), lr=args.gan_lr, betas=(0.9, 0.999)) generator_optimizer = torch.optim.Adam(generator.parameters(), lr=args.gan_lr, betas=(0.9, 0.999)) discriminator_scheduler = torch.optim.lr_scheduler.StepLR( discriminator_optimizer, args.gan_epochs // 2, 0.1) generator_scheduler = torch.optim.lr_scheduler.StepLR( generator_optimizer, args.gan_epochs // 2, 0.1) logger.info( f"Optimizer information:\n" f"\tPSNR learning rate: {args.psnr_lr}\n" f"\tDiscriminator learning rate: {args.gan_lr}\n" f"\tGenerator learning rate: {args.gan_lr}\n" f"\tPSNR optimizer: Adam, [betas=(0.9,0.999)]\n" f"\tDiscriminator optimizer: Adam, [betas=(0.9,0.999)]\n" f"\tGenerator optimizer: Adam, [betas=(0.9,0.999)]\n" f"\tPSNR scheduler: None\n" f"\tDiscriminator scheduler: StepLR, [step_size=self.gan_epochs // 2, gamma=0.1]\n" f"\tGenerator scheduler: StepLR, [step_size=self.gan_epochs // 2, gamma=0.1]" ) logger.info("Load training dataset") # Selection of appropriate treatment equipment. train_dataset = BaseTrainDataset(root=os.path.join(args.data, "train"), image_size=args.image_size, upscale_factor=args.upscale_factor) test_dataset = BaseTestDataset(root=os.path.join(args.data, "test"), image_size=args.image_size, upscale_factor=args.upscale_factor) if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset) else: train_sampler = None train_dataloader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), pin_memory=True, sampler=train_sampler, num_workers=args.workers) test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=args.workers) logger.info(f"Dataset information:\n" f"\tTrain Path: {os.getcwd()}/{args.data}/train\n" f"\tTest Path: {os.getcwd()}/{args.data}/test\n" f"\tNumber of train samples: {len(train_dataset)}\n" f"\tNumber of test samples: {len(test_dataset)}\n" f"\tNumber of train batches: {len(train_dataloader)}\n" f"\tNumber of test batches: {len(test_dataloader)}\n" f"\tShuffle of train: True\n" f"\tShuffle of test: False\n" f"\tSampler of train: {bool(train_sampler)}\n" f"\tSampler of test: None\n" f"\tWorkers of train: {args.workers}\n" f"\tWorkers of test: {args.workers}") # optionally resume from a checkpoint if args.resume_psnr: if os.path.isfile(args.resume_psnr): logger.info(f"Loading checkpoint '{args.resume_psnr}'.") if args.gpu is None: checkpoint = torch.load(args.resume_psnr) else: # Map model to be loaded to specified single gpu. checkpoint = torch.load(args.resume_psnr, map_location=f"cuda:{args.gpu}") args.start_psnr_epoch = checkpoint["epoch"] best_psnr = checkpoint["best_psnr"] if args.gpu is not None: # best_psnr may be from a checkpoint from a different GPU best_psnr = best_psnr.to(args.gpu) generator.load_state_dict(checkpoint["state_dict"]) psnr_optimizer.load_state_dict(checkpoint["optimizer"]) logger.info( f"Loaded checkpoint '{args.resume_psnr}' (epoch {checkpoint['epoch']})." ) else: logger.info(f"No checkpoint found at '{args.resume_psnr}'.") if args.resume_d or args.resume_g: if os.path.isfile(args.resume_d) or os.path.isfile(args.resume_g): logger.info(f"Loading checkpoint '{args.resume_d}'.") logger.info(f"Loading checkpoint '{args.resume_g}'.") if args.gpu is None: checkpoint_d = torch.load(args.resume_d) checkpoint_g = torch.load(args.resume_g) else: # Map model to be loaded to specified single gpu. checkpoint_d = torch.load(args.resume_d, map_location=f"cuda:{args.gpu}") checkpoint_g = torch.load(args.resume_g, map_location=f"cuda:{args.gpu}") args.start_gan_epoch = checkpoint_g["epoch"] best_psnr = checkpoint_g["best_psnr"] if args.gpu is not None: # best_psnr may be from a checkpoint from a different GPU best_psnr = best_psnr.to(args.gpu) discriminator.load_state_dict(checkpoint_d["state_dict"]) discriminator_optimizer.load_state_dict(checkpoint_d["optimizer"]) generator.load_state_dict(checkpoint_g["state_dict"]) generator_optimizer.load_state_dict(checkpoint_g["optimizer"]) logger.info( f"Loaded checkpoint '{args.resume_d}' (epoch {checkpoint_d['epoch']})." ) logger.info( f"Loaded checkpoint '{args.resume_g}' (epoch {checkpoint_g['epoch']})." ) else: logger.info( f"No checkpoint found at '{args.resume_d}' or '{args.resume_g}'." ) cudnn.benchmark = True # Create a SummaryWriter at the beginning of training. psnr_writer = SummaryWriter(f"runs/{args.arch}_psnr_logs") gan_writer = SummaryWriter(f"runs/{args.arch}_gan_logs") logger.info(f"Train information:\n" f"\tPSNR-oral epochs: {args.psnr_epochs}\n" f"\tGAN-oral epochs: {args.gan_epochs}") for epoch in range(args.start_psnr_epoch, args.psnr_epochs): if args.distributed: train_sampler.set_epoch(epoch) # train for one epoch train_psnr(train_dataloader=train_dataloader, generator=generator, pixel_criterion=pixel_criterion, psnr_optimizer=psnr_optimizer, epoch=epoch, writer=psnr_writer, args=args) # Test for every epoch. psnr, ssim, lpips, gmsd = test(generator, test_dataloader, args.gpu) gan_writer.add_scalar("Test/PSNR", psnr, epoch + 1) gan_writer.add_scalar("Test/SSIM", ssim, epoch + 1) gan_writer.add_scalar("Test/LPIPS", lpips, epoch + 1) gan_writer.add_scalar("Test/GMSD", gmsd, epoch + 1) if not args.multiprocessing_distributed or ( args.multiprocessing_distributed and args.rank % ngpus_per_node == 0): torch.save( { "epoch": epoch + 1, "arch": args.arch, "state_dict": generator.module.state_dict() if args.multiprocessing_distributed else generator.state_dict(), "optimizer": psnr_optimizer.state_dict(), }, os.path.join("weights", f"PSNR_epoch{epoch}.pth")) if psnr > best_psnr: best_psnr = max(psnr, best_psnr) torch.save(generator.state_dict(), os.path.join("weights", f"PSNR.pth")) # Load best model weight. best_psnr = 0.0 generator.load_state_dict( torch.load(os.path.join("weights", f"PSNR.pth"), map_location=f"cuda:{args.gpu}")) for epoch in range(args.start_gan_epoch, args.gan_epochs): if args.distributed: train_sampler.set_epoch(epoch) # train for one epoch train_gan(train_dataloader=train_dataloader, discriminator=discriminator, generator=generator, content_criterion=content_criterion, adversarial_criterion=adversarial_criterion, discriminator_optimizer=discriminator_optimizer, generator_optimizer=generator_optimizer, epoch=epoch, writer=gan_writer, args=args) discriminator_scheduler.step() generator_scheduler.step() # Test for every epoch. psnr, ssim, lpips, gmsd = test(generator, test_dataloader, args.gpu) gan_writer.add_scalar("Test/PSNR", psnr, epoch + 1) gan_writer.add_scalar("Test/SSIM", ssim, epoch + 1) gan_writer.add_scalar("Test/LPIPS", lpips, epoch + 1) gan_writer.add_scalar("Test/GMSD", gmsd, epoch + 1) if not args.multiprocessing_distributed or ( args.multiprocessing_distributed and args.rank % ngpus_per_node == 0): torch.save( { "epoch": epoch + 1, "arch": "vgg", "state_dict": discriminator.module.state_dict() if args.multiprocessing_distributed else discriminator.state_dict(), "optimizer": discriminator_optimizer.state_dict() }, os.path.join("weights", f"Discriminator_epoch{epoch}.pth")) torch.save( { "epoch": epoch + 1, "arch": args.arch, "state_dict": generator.module.state_dict() if args.multiprocessing_distributed else generator.state_dict(), "optimizer": generator_optimizer.state_dict() }, os.path.join("weights", f"Generator_epoch{epoch}.pth")) if psnr > best_psnr: best_psnr = max(psnr, best_psnr) torch.save(generator.state_dict(), os.path.join("weights", f"GAN.pth"))
# See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== import argparse import logging import torch import srgan_pytorch.models as models from srgan_pytorch.utils.common import configure model_names = sorted(name for name in models.__dict__ if name.islower() and not name.startswith("__") and callable(models.__dict__[name])) logger = logging.getLogger(__name__) logging.basicConfig(format="[ %(levelname)s ] %(message)s", level=logging.INFO) parser = argparse.ArgumentParser("Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network.") parser.add_argument("-a", "--arch", metavar="ARCH", default="srgan", choices=model_names, help="Model architecture: " + " | ".join(model_names) + ". (Default: srgan)") parser.add_argument("--model-path", type=str, metavar="PATH", required=True, help="Path to latest checkpoint for model.") args = parser.parse_args() model = configure(args) model.load_state_dict(torch.load(args.model_path)["state_dict"]) torch.save(model.state_dict(), "Generator.pth") logger.info("Model convert done.")
def __init__(self, args): self.args = args self.model, self.device = configure(args)
def main_worker(gpu, ngpus_per_node, args): global total_mse_value, total_rmse_value, total_psnr_value global total_ssim_value, total_lpips_value, total_gmsd_value args.gpu = gpu if args.gpu is not None: logger.info(f"Use GPU: {args.gpu} for training.") if args.distributed: if args.dist_url == "env://" and args.rank == -1: args.rank = int(os.environ["RANK"]) if args.multiprocessing_distributed: # For multiprocessing distributed training, rank needs to be the # global rank among all the processes args.rank = args.rank * ngpus_per_node + gpu dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank) model = configure(args) if not torch.cuda.is_available(): logger.warning("Using CPU, this will be slow.") elif args.distributed: # For multiprocessing distributed, DistributedDataParallel constructor # should always set the single device scope, otherwise, # DistributedDataParallel will use all available devices. if args.gpu is not None: torch.cuda.set_device(args.gpu) model.cuda(args.gpu) # When using a single GPU per process and per # DistributedDataParallel, we need to divide the batch size # ourselves based on the total number of GPUs we have args.batch_size = int(args.batch_size / ngpus_per_node) args.workers = int( (args.workers + ngpus_per_node - 1) / ngpus_per_node) model = nn.parallel.DistributedDataParallel(module=model, device_ids=[args.gpu]) else: model.cuda() # DistributedDataParallel will divide and allocate batch_size to all # available GPUs if device_ids are not set model = nn.parallel.DistributedDataParallel(model) elif args.gpu is not None: torch.cuda.set_device(args.gpu) model = model.cuda(args.gpu) else: # DataParallel will divide and allocate batch_size to all available GPUs if args.arch.startswith("alexnet") or args.arch.startswith("vgg"): model.features = torch.nn.DataParallel(model.features) model.cuda() else: model = torch.nn.DataParallel(model).cuda() logger.info("Load testing dataset") # Selection of appropriate treatment equipment. dataset = BaseTestDataset(root=os.path.join(args.data, "test"), image_size=args.image_size, upscale_factor=args.upscale_factor) dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, pin_memory=True, num_workers=args.workers) logger.info(f"Dataset information:\n" f"\tPath: {os.getcwd()}/{args.data}/test\n" f"\tNumber of samples: {len(dataset)}\n" f"\tNumber of batches: {len(dataloader)}\n" f"\tShuffle: False\n" f"\tSampler: None\n" f"\tWorkers: {args.workers}") cudnn.benchmark = True # Set eval mode. model.eval() # Start evaluate model performance. progress_bar = tqdm(enumerate(dataloader), total=len(dataloader)) for i, (lr, bicubic, hr) in progress_bar: # Move data to special device. if args.gpu is not None: lr = lr.cuda(args.gpu, non_blocking=True) bicubic = bicubic.cuda(args.gpu, non_blocking=True) hr = hr.cuda(args.gpu, non_blocking=True) with torch.no_grad(): sr = model(lr) # Evaluate performance value = iqa(sr, hr, args.gpu) total_mse_value += value[0] total_rmse_value += value[1] total_psnr_value += value[2] total_ssim_value += value[3] total_lpips_value += value[4] total_gmsd_value += value[5] progress_bar.set_description( f"[{i + 1}/{len(dataloader)}] " f"PSNR: {total_psnr_value / (i + 1):6.2f} " f"SSIM: {total_ssim_value / (i + 1):6.4f}") 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}")
def main_worker(gpu, args): global total_mse_value, total_rmse_value, total_psnr_value, total_ssim_value, total_lpips_value, total_gmsd_value args.gpu = gpu if args.gpu is not None: logger.info(f"Use GPU: {args.gpu} for testing.") model = configure(args) if not torch.cuda.is_available(): logger.warning("Using CPU, this will be slow.") if args.gpu is not None: torch.cuda.set_device(args.gpu) model = model.cuda(args.gpu) logger.info("Load testing dataset.") # Selection of appropriate treatment equipment. dataset = BaseTestDataset(os.path.join(args.data, "test"), args.image_size, args.upscale_factor) dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, pin_memory=True, num_workers=args.workers) logger.info(f"Dataset information:\n" f"\tPath: {os.getcwd()}/{args.data}/test\n" f"\tNumber of samples: {len(dataset)}\n" f"\tNumber of batches: {len(dataloader)}\n" f"\tShuffle: False\n" f"\tSampler: None\n" f"\tWorkers: {args.workers}") cudnn.benchmark = True # Set eval mode. model.eval() with torch.no_grad(): # Start evaluate model performance. progress_bar = tqdm(enumerate(dataloader), total=len(dataloader)) for i, (lr, bicubic, hr) in progress_bar: # Move data to special device. if args.gpu is not None: lr = lr.cuda(args.gpu, non_blocking=True) bicubic = bicubic.cuda(args.gpu, non_blocking=True) hr = hr.cuda(args.gpu, non_blocking=True) sr = model(lr) # Evaluate performance value = iqa(sr, hr, args.gpu) total_mse_value += value[0] total_rmse_value += value[1] total_psnr_value += value[2] total_ssim_value += value[3] total_lpips_value += value[4] total_gmsd_value += value[5] progress_bar.set_description( f"[{i + 1}/{len(dataloader)}] " f"PSNR: {total_psnr_value / (i + 1):6.2f} " f"SSIM: {total_ssim_value / (i + 1):6.4f}") images = torch.cat([bicubic, sr, hr], dim=-1) vutils.save_image(images, os.path.join("benchmarks", 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}")