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()