Spaces:
Running on Zero
Running on Zero
| """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 | |
| 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) | |
| 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)) | |