Spaces:
Sleeping
Sleeping
| """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) | |