Spaces:
Running on Zero
Running on Zero
File size: 2,988 Bytes
6ddd1e8 63c46aa 6ddd1e8 63c46aa 6ddd1e8 63c46aa 6ddd1e8 63c46aa 6ddd1e8 63c46aa | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """Test-time averaging for ShadeNet-5, copied from the repository's
`detail_diffuser/inference/core.py` (flip_average / tta_average / unit_normals)
so the Space renders exactly what the model card reports.
`refine_fn(x)` takes a [1, 3, H, W] photo in [-1, 1] and returns the 7 maps
[1, 7, H, W] in [-1, 1]: albedo 0:3, depth 3, normal 4:7.
"""
import itertools
import torch
from rotation import rotate_image, unrotate_maps
NORMAL_X = 4 # first channel of the normal (x right, y up, z to viewer)
def unit_normals(maps, eps=1e-6):
out = maps.clone()
n = out[:, NORMAL_X:NORMAL_X + 3]
out[:, NORMAL_X:NORMAL_X + 3] = n / n.norm(dim=1, keepdim=True).clamp_min(eps)
return out
@torch.no_grad()
def tta_passes(refine_fn, x, flip=False, rot_deg=0.0, progress=None):
"""The passes of the mirror (optional) x rotations (0, -rot, +rot) recipe,
each brought back to the photo's frame, as a list of (maps, valid).
Each pass: mirror -> rotate the photo (reflected border) -> predict ->
rotate the maps back, reorienting the normals -> un-mirror. `valid` marks
the pixels the pass covers (a rotated canvas loses the corners).
"""
angles = (0.0, -float(rot_deg), float(rot_deg)) if rot_deg else (0.0,)
flips = (False, True) if flip else (False,)
passes = list(itertools.product(flips, angles))
out = []
for i, (f, a) in enumerate(passes):
if progress is not None:
progress(i, len(passes))
c = x.flip(-1) if f else x
if a:
c = rotate_image(c, a, "bicubic", "reflection").clamp(-1.0, 1.0)
p = refine_fn(c).float()
valid = torch.ones_like(p[:, :1])
if a:
p, v = unrotate_maps(p, a)
valid = v.to(p.dtype)
if f:
p = p.flip(-1).clone()
p[:, NORMAL_X] = -p[:, NORMAL_X]
valid = valid.flip(-1)
out.append((p, valid))
return out
def mean_of(passes):
"""Per-pixel mean over the passes that cover the pixel (the card's recipe)."""
num = den = None
for p, valid in passes:
num = p * valid if num is None else num + p * valid
den = valid if den is None else den + valid
return num / den.clamp_min(1.0)
def median_of(passes):
"""Per-pixel median over the passes that cover the pixel (the average of the
two middle values when an even number cover it). The unrotated pass covers
everything, so every pixel has at least one value."""
P = torch.stack([p for p, _ in passes]) # [n, B, C, H, W]
V = torch.stack([v for _, v in passes]) > 0.5 # [n, B, 1, H, W]
P = torch.where(V.expand_as(P), P, torch.full_like(P, float("nan")))
return torch.nanquantile(P, 0.5, dim=0)
@torch.no_grad()
def tta_average(refine_fn, x, flip=False, rot_deg=0.0, progress=None):
"""Per-pixel mean over the recipe's passes (see `tta_passes`)."""
return mean_of(tta_passes(refine_fn, x, flip, rot_deg, progress))
|