def test_metrics(model, video_path=None, frames=None, output_folder=None):
if video_path is not None and frames is None:
frames, _ = extract_frames(video_path)
total_ssim = 0
total_psnr = 0
stride = 30
iters = 1 + (len(frames) - 3) // stride
triplets = []
for i in range(iters):
tup = (frames[i * stride], frames[i * stride + 1],
frames[i * stride + 2])
triplets.append(tup)
iters = len(triplets)
for i in range(iters):
x1, gt, x2 = triplets[i]
pred = interpolate(model, x1, x2)
if output_folder is not None:
frame_path = join(output_folder, f'wiz_{i}.jpg')
pred.save(frame_path)
gt = pil_to_tensor(gt)
pred = pil_to_tensor(pred)
total_ssim += ssim(pred, gt).item()
total_psnr += psnr(pred, gt).item()
print(f'#{i+1}/{iters} done')
avg_ssim = total_ssim / iters
avg_psnr = total_psnr / iters
print(f'avg_ssim: {avg_ssim}, avg_psnr: {avg_psnr}')
def test_linear_interp(validation_set=None):
if validation_set is None:
validation_set = get_validation_set()
total_ssim = 0
total_psnr = 0
iters = len(validation_set.tuples)
crop = CenterCrop(config.CROP_SIZE)
for tup in validation_set.tuples:
x1, gt, x2, = [pil_to_tensor(crop(load_img(p))) for p in tup]
pred = torch.mean(torch.stack((x1, x2), dim=0), dim=0)
total_ssim += ssim(pred, gt).item()
total_psnr += psnr(pred, gt).item()
avg_ssim = total_ssim / iters
avg_psnr = total_psnr / iters
print(f'avg_ssim: {avg_ssim}, avg_psnr: {avg_psnr}')
def test_on_validation_set(model, validation_set=None):
if validation_set is None:
validation_set = get_validation_set()
total_ssim = 0
total_psnr = 0
iters = len(validation_set.tuples)
crop = CenterCrop(config.CROP_SIZE)
for i, tup in enumerate(validation_set.tuples):
x1, gt, x2, = [crop(load_img(p)) for p in tup]
pred = interpolate(model, x1, x2)
gt = pil_to_tensor(gt)
pred = pil_to_tensor(pred)
total_ssim += ssim(pred, gt).item()
total_psnr += psnr(pred, gt).item()
print(f'#{i+1} done')
avg_ssim = total_ssim / iters
avg_psnr = total_psnr / iters
print(f'avg_ssim: {avg_ssim}, avg_psnr: {avg_psnr}')