File size: 2,621 Bytes
ebec118 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | """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")
|