| """CorridorKey Real-Time Edge Refiner: standalone reference implementation. |
| |
| Licensed under the Corridor Key Licence (CC BY-NC-SA 4.0 + additional terms), see LICENSE. |
| Built for use with CorridorKey 2.0 by Corridor Digital LLC. |
| |
| The refiner restores native-resolution edge detail (hair strands, motion blur, soft edges) on top of a |
| CorridorKey 2.0 matte that was computed at a reduced resolution. Inputs per native pixel: |
| |
| I native RGB frame (sRGB, [0, 1]) |
| B smooth background estimate (sRGB), from pixels CorridorKey calls background |
| F CorridorKey foreground transferred onto the native frame (I + upsampled (F_low - I_low)) |
| a_up CorridorKey alpha, bilinear upsampled |
| a_gf alpha of a colour guided filter (He et al.) driven by CorridorKey's low-res alpha |
| |
| It predicts residuals for alpha (on top of a_gf) and foreground (on top of F). The result is only used |
| inside an uncertainty band around CorridorKey's edges; elsewhere the guided-filter result is kept. |
| |
| `refine()` below processes a whole frame at once (simple, a few ms on a modern GPU at 1080p). The real-time |
| engine evaluates exactly the same maths on 32x32 band tiles only. |
| """ |
| from __future__ import annotations |
|
|
| import json |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| IN_CHANNELS = 11 |
|
|
|
|
| class EdgeRefiner(nn.Module): |
| def __init__(self, width: int = 24, dilations=(1, 2, 4, 1)): |
| super().__init__() |
| self.dilations = tuple(int(d) for d in dilations) |
| self.inp = nn.Conv2d(IN_CHANNELS, width, 3, padding=1) |
| self.blocks = nn.ModuleList(nn.Conv2d(width, width, 3, padding=d, dilation=d) for d in self.dilations) |
| self.out = nn.Conv2d(width, 4, 1) |
|
|
| def receptive_radius(self) -> int: |
| return 1 + sum(self.dilations) |
|
|
| def forward(self, x: torch.Tensor): |
| """x: (B, 11, H, W) = [I(3), B(3), F(3), a_up(1), a_gf(1)]. Returns (alpha, fg_srgb) in [0, 1].""" |
| h = F.gelu(self.inp(x)) |
| for blk in self.blocks: |
| h = h + F.gelu(blk(h)) |
| d = self.out(h) |
| alpha = (x[:, 10:11] + d[:, 0:1]).clamp(0.0, 1.0) |
| fg = (x[:, 6:9] + d[:, 1:4]).clamp(0.0, 1.0) |
| return alpha, fg |
|
|
|
|
| def load_refiner(path: str, device="cuda", dtype=torch.float16) -> EdgeRefiner: |
| """Load a .safetensors checkpoint of this repository.""" |
| from safetensors import safe_open |
| from safetensors.torch import load_file |
| with safe_open(path, framework="pt") as f: |
| cfg = json.loads(f.metadata()["config"]) |
| net = EdgeRefiner(width=cfg["width"], dilations=tuple(cfg["dilations"])) |
| net.load_state_dict(load_file(path)) |
| return net.to(device=device, dtype=dtype).eval() |
|
|
|
|
| |
| |
| |
|
|
| def srgb_to_linear(x): |
| return torch.where(x <= 0.04045, x / 12.92, ((x + 0.055) / 1.055) ** 2.4) |
|
|
|
|
| def linear_to_srgb(x): |
| x = x.clamp(min=0.0) |
| return torch.where(x <= 0.0031308, x * 12.92, 1.055 * x.clamp(min=0.0031308) ** (1 / 2.4) - 0.055) |
|
|
|
|
| def guided_coef(rgb_low, alpha_low, radius: int = 2, eps: float = 5e-4): |
| """Colour guided filter (He et al.) at low resolution: smoothed (a_r, a_g, a_b, b) coefficients.""" |
| k = 2 * radius + 1 |
| box = lambda z: F.avg_pool2d(z, k, stride=1, padding=radius, count_include_pad=False) |
| I, p = rgb_low, alpha_low |
| _, _, h, w = p.shape |
| mean_I, mean_p = box(I), box(p) |
| cov_Ip = box(I * p) - mean_I * mean_p |
| mean_II = box((I.unsqueeze(1) * I.unsqueeze(2)).reshape(1, 9, h, w)).reshape(1, 3, 3, h, w) |
| s = mean_II - mean_I.unsqueeze(1) * mean_I.unsqueeze(2) |
| s = s + eps * torch.eye(3, device=s.device, dtype=s.dtype).view(1, 3, 3, 1, 1) |
| a = torch.linalg.solve(s.permute(0, 3, 4, 1, 2), cov_Ip.permute(0, 2, 3, 1).unsqueeze(-1)) |
| a = a.squeeze(-1).permute(0, 3, 1, 2) |
| b = mean_p - (a * mean_I).sum(1, keepdim=True) |
| return torch.cat([box(a), box(b)], 1) |
|
|
|
|
| def push_pull(num, den): |
| """Hole-filling interpolation: weighted average num/den, holes (den ~ 0) filled from coarser levels.""" |
| levels = [(num, den)] |
| while min(levels[-1][0].shape[-2:]) > 4: |
| n, d = levels[-1] |
| levels.append((F.avg_pool2d(n, 2, ceil_mode=True), F.avg_pool2d(d, 2, ceil_mode=True))) |
| n, d = levels[-1] |
| est = (n.sum((-2, -1), keepdim=True) / d.sum((-2, -1), keepdim=True).clamp(min=1e-6)).expand_as(n) |
| for n, d in reversed(levels): |
| est = F.interpolate(est, size=n.shape[-2:], mode="bilinear", align_corners=False) |
| d = d.clamp(0.0, 1.0) |
| est = (n + (1.0 - d) * est) / (d + (1.0 - d)).clamp(min=1e-6) |
| return est |
|
|
|
|
| def background_low(rgb_low, alpha_low): |
| w = (1.0 - alpha_low).clamp(0.0, 1.0) ** 4 |
| return linear_to_srgb(push_pull(srgb_to_linear(rgb_low.clamp(0.0, 1.0)) * w, w).clamp(0.0, 1.0)) |
|
|
|
|
| def refiner_band(alpha_low, gain: float = 10.0, lo: float = 0.004): |
| """Where the refiner may change the guided result: uncertain alpha, dilated, weighted by nearby |
| foreground (gain > 0 keeps it out of noisy/dark screen areas) and feathered.""" |
| unc = ((alpha_low > lo) & (alpha_low < 0.996)).float() |
| unc = F.max_pool2d(unc, 5, stride=1, padding=2) |
| if gain > 0: |
| near_fg = (F.max_pool2d(alpha_low, 5, stride=1, padding=2) * gain).clamp(0.0, 1.0) |
| unc = F.avg_pool2d(unc * near_fg, 3, stride=1, padding=1, count_include_pad=False) |
| return unc |
|
|
|
|
| @torch.inference_mode() |
| def refine(net: EdgeRefiner, rgb, rgb_low, alpha_low, fg_low, band_gain: float = 10.0): |
| """ |
| rgb: (1, 3, H, W) native frame, sRGB in [0, 1] |
| rgb_low: (1, 3, h, w) the frame as CorridorKey saw it (downscaled), sRGB |
| alpha_low: (1, 1, h', w') CorridorKey alpha (any low resolution, e.g. 256x256 from a 512x512 pass) |
| fg_low: (1, 3, h', w') CorridorKey straight foreground, sRGB |
| Returns (alpha, fg) at native resolution, float32. Apply despill / alpha contrast afterwards as needed. |
| """ |
| rgb, alpha_low, fg_low = rgb.float(), alpha_low.float(), fg_low.float() |
| rgb_low = F.interpolate(rgb_low.float(), size=alpha_low.shape[-2:], mode="bilinear", align_corners=False) |
| H, W = rgb.shape[-2:] |
| up = lambda z: F.interpolate(z, size=(H, W), mode="bilinear", align_corners=False) |
| coef = up(guided_coef(rgb_low, alpha_low)) |
| a_gf = ((coef[:, :3] * rgb).sum(1, keepdim=True) + coef[:, 3:4]).clamp(0.0, 1.0) |
| fg_in = (rgb + up(fg_low - rgb_low)).clamp(0.0, 1.0) |
| x = torch.cat([rgb, up(background_low(rgb_low, alpha_low)), fg_in, up(alpha_low), a_gf], 1) |
| p = next(net.parameters()) |
| a_r, f_r = net(x.to(p.dtype)) |
| band = up(refiner_band(alpha_low, band_gain)) |
| alpha = band * a_r.float() + (1.0 - band) * a_gf |
| fg = band * f_r.float() + (1.0 - band) * fg_in |
| return alpha, fg |
|
|
