| """Score in-betweens where things move: PSNR over pixels whose two endpoint |
| frames differ (dilated), on held-out triplets, for a blend, RIFE and ours.""" |
| import json, random, sys |
| from pathlib import Path |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| sys.path.insert(0, str(Path(__file__).parent)) |
| from liquid_model import IFNet, load_teacher |
|
|
| eval_dir, teacher_path, ours_path = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3]) |
| device = torch.device("mps" if torch.backends.mps.is_available() else "cpu") |
| index = json.loads((eval_dir / "index.json").read_text()) |
| videos = [(np.load(eval_dir / f"{v['name']}.npy", mmap_mode="r"), set(v["cuts"])) for v in index] |
| teacher = load_teacher(teacher_path).to(device).eval() |
| ours = IFNet(timestep=True).to(device).eval() |
| ours.load_state_dict(torch.load(ours_path, map_location=device)) |
|
|
| random.seed(99) |
| samples = [] |
| while len(samples) < 90: |
| frames, cuts = random.choice(videos) |
| gap = random.choice([6, 8, 10, 12]) |
| i = random.randrange(0, len(frames) - gap) |
| if any(i < c <= i + gap for c in cuts): |
| continue |
| j = random.choice([gap // 4, gap // 2, 3 * gap // 4]) |
| samples.append((frames[i], frames[i + j], frames[i + gap], j / gap)) |
|
|
| def tensor(f): |
| return torch.from_numpy(np.ascontiguousarray(f)).permute(2, 0, 1)[None].float().div(255).to(device) |
|
|
| def rife_at(t, x0, x1, lo=0.0, hi=1.0, depth=0): |
| mid = teacher(x0, x1) |
| c = (lo + hi) / 2 |
| if abs(c - t) < 1e-6 or depth > 4: |
| return mid |
| return rife_at(t, x0, mid, lo, c, depth + 1) if t < c else rife_at(t, mid, x1, c, hi, depth + 1) |
|
|
| scores = {"blend": [], "rife": [], "ours": []} |
| with torch.no_grad(): |
| for a, m, b, t in samples: |
| A, M, B = tensor(a), tensor(m), tensor(b) |
| h, w = A.shape[-2:] |
| H, W = (h + 31) // 32 * 32, (w + 31) // 32 * 32 |
| pad = lambda x: F.pad(x, (0, W - w, 0, H - h), mode="replicate") |
| moving = ((A - B).abs().max(1, keepdim=True)[0] > 0.08).float() |
| moving = (F.max_pool2d(moving, 9, 1, 4) > 0).float() |
| if moving.mean() < 0.01: |
| continue |
| outs = { |
| "blend": A * (1 - t) + B * t, |
| "rife": rife_at(t, pad(A), pad(B))[..., :h, :w], |
| "ours": ours(pad(A), pad(B), torch.tensor([[[[t]]]], device=device))[..., :h, :w], |
| } |
| for k, o in outs.items(): |
| mse = (((o.clamp(0, 1) - M) ** 2) * moving).sum() / (moving.sum() * 3) |
| scores[k].append(-10 * torch.log10(mse + 1e-10).item()) |
| print({k: round(float(np.mean(v)), 2) for k, v in scores.items()}, "on", len(scores["ours"]), "moving samples") |
|
|
