def main(*args):
if not opt.input_dir:
raise ValueError("--input_dir is required")
if not opt.dataset.upper() in DATASETS.keys():
raise ValueError("--dataset is missing, or can't be found")
data_ref = DATASETS.get(opt.dataset.upper())
data = load_folder(opt.input_dir)
skip = opt.offset
metric_config = Config(depth=opt.clip, batch=1, scale=1, modcrop=False)
loader = BasicLoader(data, 'test', metric_config, False)
ref_loader = BasicLoader(data_ref, 'test', metric_config, False)
# make sure len(ref_loader) == len(loader)
loader_iter = loader.make_one_shot_iterator()
ref_iter = ref_loader.make_one_shot_iterator()
for ref, _, name in ref_iter:
name = str(name)
img, _, _ = next(loader_iter)
# reduce the batch dimension for video clips
if img.ndim == 5: img = img[0]
if ref.ndim == 5: ref = ref[0]
if opt.shave:
img = shave(img, opt.shave)
ref = shave(ref, opt.shave)
if opt.l_only:
img = rgb_to_yuv(img, max_val=255, standard=opt.l_standard)[...,
0:1]
ref = rgb_to_yuv(ref, max_val=255, standard=opt.l_standard)[...,
0:1]
if ref.shape[0] - skip != img.shape[0]:
b_min = np.minimum(ref.shape[0] - skip, img.shape[0])
ref = ref[:b_min + skip, ...]
img = img[:b_min, ...]
img = tf.constant(img.astype(np.float32))
ref = tf.constant(ref.astype(np.float32))
psnr = tf.reduce_mean(tf.image.psnr(
ref[skip:], img, 255)).eval() if not opt.no_psnr else 0
ssim = tf.reduce_mean(tf.image.ssim(
ref[skip:], img, 255)).eval() if not opt.no_ssim else 0
tf.logging.info(f'[{name}] PSNR = {psnr}, SSIM = {ssim}')
tf.add_to_collection('PSNR', psnr)
tf.add_to_collection('SSIM', ssim)
for key in ('PSNR', 'SSIM'):
mp = np.mean(tf.get_collection(key))
tf.logging.info(f'Mean {key}: {mp}')
def action(x, y):
xname = f'{x.parent.name}/{x.stem}'
yname = f'{y.parent.name}/{y.stem}'
x = Image.open(x)
y = Image.open(y)
assert x.width == y.width and x.height == y.height, "Image size mismatch!"
xx = np.asarray(x, dtype=np.float) / 255.0
yy = np.asarray(y, dtype=np.float) / 255.0
if FLAGS.l_only:
xx = rgb_to_yuv(xx, standard='matlab')[..., :1]
yy = rgb_to_yuv(yy, standard='matlab')[..., :1]
if FLAGS.shave:
xx = xx[..., FLAGS.shave:-FLAGS.shave, FLAGS.shave:-FLAGS.shave, :]
yy = yy[..., FLAGS.shave:-FLAGS.shave, FLAGS.shave:-FLAGS.shave, :]
mse = np.mean((xx - yy)**2)
psnr = np.log10(1.0 / mse) * 10.0
info = {"x": xname, "y": yname}
if FLAGS.ssim:
ssim = compare_ssim(xx, yy, multichannel=True)
info.update(SSIM=ssim)
info.update(PSNR=psnr)
info.update(MSE=mse)
return info
def test_rgb2yuv():
img = imread(URL)
img = img.astype('float32')
yuv = rgb_to_yuv(img, 255, 'matlab')
array_to_img(yuv).show()