| """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 |
|
|
|
|
| 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_average(refine_fn, x, flip=False, rot_deg=0.0, progress=None): |
| """Per-pixel mean over the mirror (optional) x rotations (0, -rot, +rot). |
| |
| Each pass: mirror -> rotate the photo (reflected border) -> predict -> |
| rotate the maps back, reorienting the normals -> un-mirror. A pass only |
| counts at the pixels its rotated canvas covers. |
| """ |
| 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)) |
| num = den = None |
| 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) |
| 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) |
|
|
