"""Local profile refinement and ordinary-coordinate-zero mixture models.""" import torch from torch import nn from torch.nn import functional as F class ProfileRefiner(nn.Module): def __init__(self, boundary_keys, radius=16, channels=16): super().__init__() if not 1 <= radius <= 64 or channels % 4: raise ValueError("Unsupported profile-refiner geometry") self.radius = radius self.boundary_keys = boundary_keys self.register_buffer( "offsets", torch.arange(-radius, radius + 1, dtype=torch.float32) ) self.features = nn.Sequential( nn.Conv1d(3, channels, 5, padding=2), nn.GroupNorm(4, channels), nn.SiLU(), nn.Conv1d(channels, channels, 5, padding=2), nn.GroupNorm(4, channels), nn.SiLU(), nn.Conv1d(channels, channels, 3, padding=1), nn.GroupNorm(4, channels), nn.SiLU(), ) self.heads = nn.ParameterDict( { d: nn.Parameter(torch.zeros(len(keys), channels)) for d, keys in boundary_keys.items() } ) self.register_buffer("prior", -0.5 * (self.offsets / max(radius / 3, 1.0)) ** 2) def profiles(self, image, center): n, one, h, a = image.shape K = center.shape[1] D = len(self.offsets) if one != 1 or center.shape != (n, K, a) or h < 2 or a < 2: raise ValueError("Profile extraction geometry mismatch") y = (center[:, :, None, :] + self.offsets[None, None, :, None]) * ( 2 / (h - 1) ) - 1 x = torch.arange(a, device=image.device, dtype=torch.float32)[ None, None, None, : ] parts = [] for shift in [-1, 0, 1]: gx = ((x + shift).clamp(0, a - 1) * (2 / (a - 1)) - 1).expand(n, K, D, a) grid = torch.stack([gx, y], dim=-1).reshape(n, K * D, a, 2) sampled = ( F.grid_sample( image, grid, mode="bilinear", padding_mode="border", align_corners=True, )[:, 0] .reshape(n, K, D, a) .permute(0, 1, 3, 2) ) parts.append(sampled[:, :, :, None, :]) return torch.cat(parts, dim=3).reshape(n * K * a, 3, D) def forward(self, image, center, device): n, K, a = center.shape if device not in self.boundary_keys or K != len(self.boundary_keys[device]): raise ValueError("Profile-refiner boundary mismatch") features = self.features(self.profiles(image, center)) features = features.reshape(n, K, a, features.shape[1], features.shape[2]) return ( torch.einsum("nkacd,kc->nkda", features, self.heads[device]) + self.prior[None, None, :, None] ) def residual(self, logits): return (logits.float().softmax(dim=2) * self.offsets[None, None, :, None]).sum( dim=2 ) class ProfileCascade(nn.Module): def __init__(self, config): super().__init__() from src.models import make_model self.coarse = make_model(config["coarse_config"]) self.refiner = ProfileRefiner( config["boundary_keys"], config["refine_radius"], config["refine_channels"] ) def forward(self, *args, **kwargs): raise ValueError("A profile cascade requires coarse and fine OCT grids") def refinement_loss( logits, target_residual, valid, offsets, sigma=0.8, coordinate_weight=0.1 ): """All valid targets contribute; CE projects to the representable interval. Expected absolute error uses the original, unprojected residual target, so out-of-window cases are still supervised and never discarded. """ from src.models.unet import boundary_loss radius = (len(offsets) - 1) // 2 ce = boundary_loss( logits, (target_residual + radius).clamp(0, 2 * radius), valid, sigma, 0.0 ) distance = ( logits.float().softmax(dim=2) * (offsets[None, None, :, None] - target_residual[:, :, None, :]).abs() ).sum(dim=2) counts = valid.sum(-1) terms = (distance * valid).sum(-1) / counts.clamp_min(1) risk = (terms * (counts > 0)).sum() / (counts > 0).sum().clamp_min(1) return ce + coordinate_weight * risk def fine_preparation_shape(config): if config.get("architecture") == "boundary_profile_refine_v1": return config["fine_height"], config["fine_width"] if config.get("architecture") == "boundary_ensemble_v1": shapes = {fine_preparation_shape(c["config"]) for c in config["components"]} shapes.discard(None) if len(shapes) > 1: raise ValueError("Embedded refiners need the same fine preprocessing grid") return next(iter(shapes)) if shapes else None return None class ZeroAtomNet(nn.Module): def __init__(self, config): super().__init__() from src.models import make_model child = config["coarse_config"] if child["architecture"] not in [ "boundary_unet_v1", "boundary_unet_axial_v1", "boundary_unet_global_v1", ]: raise ValueError("Zero atom needs an image U-Net backbone") self.coarse = make_model(child) self.boundary_keys = config["boundary_keys"] base = child["base"] hidden = config.get("zero_hidden_channels", 64) self.zero_heads = nn.ModuleDict( { d: nn.Sequential( nn.Conv1d(26 * base, hidden, 3, padding=1), nn.GroupNorm(8, hidden), nn.SiLU(), nn.Conv1d(hidden, hidden, 3, padding=1), nn.SiLU(), nn.Conv1d(hidden, len(keys), 1), ) for d, keys in self.boundary_keys.items() } ) for head in self.zero_heads.values(): nn.init.zeros_(head[-1].weight) nn.init.constant_(head[-1].bias, -8.0) def forward(self, x, device): x, skips = self.coarse.encode(x) for context in getattr(self.coarse, "axial", []): x = context(x) bottom = x x = self.coarse.decode(x, skips) logits = self.coarse.heads[device](x) if hasattr(self.coarse, "edge_heads"): edge = self.coarse.edge_heads[device]( torch.cat([bottom.mean(2), bottom.amax(2)], dim=1) ) edge = F.interpolate( edge, size=logits.shape[-1], mode="linear", align_corners=True ) logits = torch.cat( [logits[:, :, :1] + edge[:, :, None, :], logits[:, :, 1:]], dim=2 ) global_features = F.interpolate( torch.cat([bottom.mean(2), bottom.amax(2)], dim=1), size=x.shape[-1], mode="linear", align_corners=True, ) features = torch.cat([x.mean(2), x.amax(2), global_features], dim=1) return logits, self.zero_heads[device](features) def zero_atom_loss( logits, zero_logits, positive_positions, zero_fraction, valid, sigma=0.8, coordinate_weight=0.01, ): """All valid reference values contribute to a mixture likelihood and risk.""" h = logits.shape[-2] width = positive_positions.shape[-1] if logits.shape[-1] != width: logits = F.interpolate( logits.float(), size=(h, width), mode="bilinear", align_corners=True ) zero_logits = F.interpolate( zero_logits.float(), size=width, mode="linear", align_corners=True ) grid = torch.arange(h, device=logits.device, dtype=torch.float32)[ None, None, :, None ] zt = zero_fraction.float() yp = positive_positions.float() positive_weight = 1 - zt target = torch.exp(-0.5 * ((grid - yp[:, :, None, :]) / sigma) ** 2) target /= target.sum(2, keepdim=True).clamp_min(1e-8) logp = logits.float().log_softmax(2) probabilities = logp.exp() ce = -(target * logp).sum(2) likelihood = ( F.binary_cross_entropy_with_logits(zero_logits, zt, reduction="none") + positive_weight * ce ) pzero = zero_logits.sigmoid() positive_distance = (probabilities * (grid - yp[:, :, None, :]).abs()).sum(2) distance_from_zero = (probabilities * grid).sum(2) risk = pzero * positive_weight * yp + (1 - pzero) * ( zt * distance_from_zero + positive_weight * positive_distance ) point = likelihood + coordinate_weight * risk count = valid.sum(-1) term = (point * valid).sum(-1) / count.clamp_min(1) return (term * (count > 0)).sum() / (count > 0).sum().clamp_min(1)