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 run(self): args = self.args # Get file name. filename = os.path.basename(args.lr) # Read all pictures. input = Image.open(args.lr) bicubic = transforms.Resize((args.image_size, args.image_size), interpolation=Image.BICUBIC)(input) target = Image.open(args.hr) # Convert image to tensor format. lr = process_image(input, self.device) bicubic = process_image(bicubic, self.device) hr = process_image(target, self.device) if args.eval: sr, use_time = inference(self.model, lr, statistical_time=True) images = torch.cat([bicubic, sr, hr], dim=-1) vutils.save_image(sr, os.path.join("test", f"sr_{filename}"), padding=10) vutils.save_image(images, os.path.join("test", f"compare_{filename}"), padding=10) value = image_quality_evaluation(sr_filename=os.path.join( "test", f"sr_{filename}"), hr_filename=args.hr, device=self.device) print(f"Performance avg results:\n") print(f"indicator Score\n") print(f"--------- -----\n") if self.args.detail: print(f"MSE {value[0]:.2f}\n" f"RMSE {value[1]:.2f}\n" f"PSNR {value[2]:.2f}\n" f"SSIM {value[3][0]:.4f}\n" f"MS-SSIM {value[4].real:.4f}\n" f"NIQE {value[5]:.2f}\n" f"SAM {value[6]:.4f}\n" f"VIF {value[7]:.4f}\n" f"LPIPS {value[8].item():.4f}\n" f"Use time: {use_time * 1000:.2f}ms | {use_time:.4f}s") else: print(f"PSNR {value[2]:.2f}\n" f"SSIM {value[3][0]:.2f}\n" f"Use time: {use_time * 1000:.2f}ms | {use_time:.4f}s") else: sr = inference(self.model, lr) images = torch.cat([bicubic, sr, hr], dim=-1) vutils.save_image(sr, os.path.join("test", f"sr_{filename}")) vutils.save_image(images, os.path.join("test", f"compare_{filename}"))
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 # Configure model model = models.__dict__[args.arch]() 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 run(self): args = self.args # Read img to tensor and transfer to the specified device for processing. img = Image.open(args.lr) lr = process_image(img, self.device) sr, use_time = inference(self.model, lr, statistical_time=True) vutils.save_image(sr, os.path.join( "test", f"{args.lr.split('/')[-1]}")) # Save super resolution image. value = image_quality_evaluation( os.path.join("test", f"{args.lr.split('/')[-1]}"), args.hr, self.device) print(f"Performance avg results:\n") print(f"indicator Score\n") print(f"--------- -----\n") if self.args.detail: print(f"MSE {value[0]:.2f}\n" f"RMSE {value[1]:.2f}\n" f"PSNR {value[2]:.2f}\n" f"SSIM {value[3][0]:.4f}\n" f"MS-SSIM {value[4].real:.4f}\n" f"NIQE {value[5]:.2f}\n" f"SAM {value[6]:.4f}\n" f"VIF {value[7]:.4f}\n" f"LPIPS {value[8].item():.4f}\n" f"Use time: {use_time * 1000:.2f}ms | {use_time:.4f}s") else: print(f"PSNR {value[0]:.2f}\n" f"SSIM {value[1][0]:.2f}\n" f"Use time: {use_time * 1000:.2f}ms | {use_time:.4f}s")
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 run(self): args = self.args # Set eval model. self.model.eval() # read frame success, raw_frame = self.video_capture.read() progress_bar = tqdm( range(self.total_frames), desc="[processing video and saving/view result videos]") for _ in progress_bar: if success: # Read img to tensor and transfer to the specified device for processing. img = Image.open(args.lr) lr = process_image(img, self.device) sr = inference(self.model, lr) sr = sr.cpu() sr = sr.data[0].numpy() sr *= 255.0 sr = (np.uint8(sr)).transpose((1, 2, 0)) # save sr video self.sr_writer.write(sr) # make compared video and crop shot of left top\right top\center\left bottom\right bottom sr = self.tensor2pil(sr) # Five areas are selected as the bottom contrast map. crop_sr_imgs = transforms.FiveCrop(size=sr.width // 5 - 9)(sr) crop_sr_imgs = [ np.asarray(transforms.Pad(padding=(10, 5, 0, 0))(img)) for img in crop_sr_imgs ] sr = transforms.Pad(padding=(5, 0, 0, 5))(sr) # Five areas in the contrast map are selected as the bottom contrast map compare_img = transforms.Resize( (self.sr_size[1], self.sr_size[0]), interpolation=Image.BICUBIC)(self.tensor2pil(raw_frame)) crop_compare_imgs = transforms.FiveCrop( size=compare_img.width // 5 - 9)(compare_img) crop_compare_imgs = [ np.asarray(transforms.Pad((0, 5, 10, 0))(img)) for img in crop_compare_imgs ] compare_img = transforms.Pad(padding=(0, 0, 5, 5))(compare_img) # concatenate all the pictures to one single picture # 1. Mosaic the left and right images of the video. top_img = np.concatenate( (np.asarray(compare_img), np.asarray(sr)), axis=1) # 2. Mosaic the bottom left and bottom right images of the video. bottom_img = np.concatenate(crop_compare_imgs + crop_sr_imgs, axis=1) bottom_img_height = int(top_img.shape[1] / bottom_img.shape[1] * bottom_img.shape[0]) bottom_img_width = top_img.shape[1] # 3. Adjust to the right size. bottom_img = np.asarray( transforms.Resize( (bottom_img_height, bottom_img_width))(self.tensor2pil(bottom_img))) # 4. Combine the bottom zone with the upper zone. final_image = np.concatenate((top_img, bottom_img)) # save compare video self.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 = self.video_capture.read()