File size: 4,723 Bytes
dfc2650 | 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 80 81 82 83 84 85 86 87 | """Template-conditioned fingerprint RENDER generator (G_render).
Goal (see design discussion): generate a NEW clean fingerprint image that is forensic and
hallucination-free WITHOUT relying on raw pixels. The trick: decouple *identity information*
(a TEMPLATE extracted once from the observation) from *image rendering* (a generator that only
sees the template, never the raw image). Forensic == the template is faithful; no-hallucination ==
the generator can only render what the template specifies (extracted minutiae + deterministic
ridge flow), and blank/low-coverage regions stay blank.
TEMPLATE conditioning channels (NO raw pixels here — this is the whole point):
0 minutiae heatmap (Gaussian blobs at extracted minutiae; identity anchors)
1-2 orientation field [cos 2θ, sin 2θ] * coherence (deterministic ridge flow)
3 coherence / coverage (where the template HAS evidence; 0 => stay blank)
4 ridge-frequency map (ridge spacing prior for realistic texture)
Self-supervised training: target = a clean rolled print R; cond = template(R). The generator
learns to re-render R from R's own template. At inference, feed a template extracted from a LATENT
=> a clean synthetic print of that identity, with gaps left blank (never invented).
"""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
from g_render.models.structure.heatmap import compute_orientation_field, render_gaussian_points
from g_render.models.frequency.ridge_freq import RidgeFrequencyNormalizer
TEMPLATE_CHANNELS = 5
@torch.no_grad()
def build_template_conditioning(
image: torch.Tensor,
minutiae: list | None,
size: int,
rn: RidgeFrequencyNormalizer | None = None,
minu_sigma: float = 4.0,
) -> torch.Tensor:
"""image: [1,1,H,W] in [0,1] (used ONLY to extract the template, discarded after).
minutiae: list of (x,y[,theta]) in image pixels, or None. Returns cond [1, TEMPLATE_CHANNELS, size, size].
NOTE: the returned tensor contains NO raw pixels — only extracted structure."""
rn = rn or RidgeFrequencyNormalizer()
dev = image.device
H, W = image.shape[-2:]
of = compute_orientation_field(image) # [1,2,H,W] = [cos2,sin2]*coh
coh = (of.pow(2).sum(1, keepdim=True) + 1e-12).sqrt().clamp(0, 1)
fmap, _ = rn.estimate_frequency_map(image) # [1,1,gh,gw]
fmap = F.interpolate(fmap, size=(H, W), mode="bilinear", align_corners=False)
if minutiae:
pts = [(float(m[0]), float(m[1])) for m in minutiae]
minu = render_gaussian_points(pts, out_size=size, src_size=(float(W), float(H)),
sigma=minu_sigma, device=dev)[None] # [1,1,size,size]
else:
minu = torch.zeros(1, 1, size, size, device=dev)
def rs(t):
return F.interpolate(t, size=(size, size), mode="bilinear", align_corners=False)
cond = torch.cat([minu, rs(of), rs(coh), rs(fmap.clamp(0, 1))], dim=1) # [1,5,size,size]
return cond.clamp(-1, 1)
class RenderGenerator(nn.Module):
"""template conditioning -> clean fingerprint image, ATTENTION backbone.
Wraps the in-repo UformerCondGenerator (window self-attention U-Net): attention enforces the
long-range ridge-flow / singular-point coherence that a plain conv U-Net misses, and propagates
the sparse minutiae/OF template globally (esp. across gap regions). No raw pixels in the input.
"""
def __init__(self, in_ch: int = TEMPLATE_CHANNELS, dims=(32, 64, 128, 256),
window: int = 8, embed_dim: int = 128, cover_tau: float = 0.15):
super().__init__()
from g_render.models.generators.uformer_gen import UformerCondGenerator
# skip_mode="lowpass": template has no raw speckle, but keep it for ridge-scale skips.
self.body = UformerCondGenerator(in_channels=in_ch, dims=dims, embed_dim=embed_dim,
window=window, skip_mode="lowpass")
self.head = nn.Conv2d(2, 1, 1) # Uformer emits 2ch -> 1ch image
nn.init.zeros_(self.head.bias)
self.cover_tau = float(cover_tau)
def forward(self, cond: torch.Tensor) -> torch.Tensor:
y = self.body(cond) # [B,2,H,W] window-attention output
img = torch.sigmoid(self.head(y)) # -> [B,1,H,W] in [0,1]
# coverage gate: outside template evidence (coh ch3 ~0) -> white (blank). The generator
# CANNOT paint ridges where the template has no evidence (structural anti-hallucination).
cover = (cond[:, 3:4] / self.cover_tau).clamp(0, 1)
return (img * cover + 1.0 * (1 - cover)).clamp(0, 1)
|