"""Standalone Hecate inference; architecture derived from scrollprize/villa's canonical ResNet-152 3D decoder and CanonicalLogitProjection. Only the two released Hecate EMA architectures are supported. """ import argparse import json from pathlib import Path import tempfile import struct import zlib import numpy as np import torch from torch import nn import torch.nn.functional as F def conv3x3x3(in_planes, out_planes, stride=1): return nn.Conv3d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) def conv1x1x1(in_planes, out_planes, stride=1): return nn.Conv3d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) class Bottleneck(nn.Module): expansion = 4 def __init__(self, in_planes, planes, stride=1, downsample=None): super().__init__() self.conv1 = conv1x1x1(in_planes, planes) self.bn1 = nn.BatchNorm3d(planes) self.conv2 = conv3x3x3(planes, planes, stride) self.bn2 = nn.BatchNorm3d(planes) self.conv3 = conv1x1x1(planes, planes * self.expansion) self.bn3 = nn.BatchNorm3d(planes * self.expansion) self.relu = nn.ReLU(inplace=True) self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out = self.relu(out) out = self.conv3(out) out = self.bn3(out) if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out class ResNet(nn.Module): def __init__(self, block, layers, block_inplanes, n_input_channels=3, conv1_t_size=7, conv1_t_stride=1, no_max_pool=False, shortcut_type='B', widen_factor=1.0, n_classes=400, forward_features=False): super().__init__() self.forward_features = forward_features block_inplanes = [int(x * widen_factor) for x in block_inplanes] self.in_planes = block_inplanes[0] self.no_max_pool = no_max_pool self.conv1 = nn.Conv3d(n_input_channels, self.in_planes, kernel_size=(conv1_t_size, 7, 7), stride=(conv1_t_stride, 2, 2), padding=(conv1_t_size // 2, 3, 3), bias=False) self.bn1 = nn.BatchNorm3d(self.in_planes) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)) self.layer1 = self._make_layer(block, block_inplanes[0], layers[0], shortcut_type) self.layer2 = self._make_layer(block, block_inplanes[1], layers[1], shortcut_type, stride=2) self.layer3 = self._make_layer(block, block_inplanes[2], layers[2], shortcut_type, stride=2) self.layer4 = self._make_layer(block, block_inplanes[3], layers[3], shortcut_type, stride=2) self.avgpool = nn.AdaptiveMaxPool3d((1, 1, 1)) self.fc = nn.Linear(block_inplanes[3] * block.expansion, n_classes) for m in self.modules(): if isinstance(m, nn.Conv3d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') elif isinstance(m, nn.BatchNorm3d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) def _make_layer(self, block, planes, blocks, shortcut_type, stride=1): downsample = None if stride != 1 or self.in_planes != planes * block.expansion: if False: raise ValueError('Only shortcut B is supported') else: downsample = nn.Sequential(conv1x1x1(self.in_planes, planes * block.expansion, stride), nn.BatchNorm3d(planes * block.expansion)) layers = [] layers.append(block(in_planes=self.in_planes, planes=planes, stride=stride, downsample=downsample)) self.in_planes = planes * block.expansion for i in range(1, blocks): layers.append(block(self.in_planes, planes)) return nn.Sequential(*layers) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu(x) if not self.no_max_pool: x = self.maxpool(x) x1 = self.layer1(x) x2 = self.layer2(x1) x3 = self.layer3(x2) x4 = self.layer4(x3) if self.forward_features: return [x1, x2, x3, x4] else: x = self.avgpool(x4) x = x.view(x.size(0), -1) x = self.fc(x) return x class ResConvBlock3D(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv1 = nn.Conv3d(in_ch, out_ch, 3, padding=1, bias=False) self.gn1 = nn.GroupNorm(min(32, out_ch), out_ch) self.conv2 = nn.Conv3d(out_ch, out_ch, 3, padding=1, bias=False) self.gn2 = nn.GroupNorm(min(32, out_ch), out_ch) self.shortcut = nn.Conv3d(in_ch, out_ch, 1, bias=False) if in_ch != out_ch else nn.Identity() def forward(self, x): identity = self.shortcut(x) x = F.relu(self.gn1(self.conv1(x)), inplace=True) x = self.gn2(self.conv2(x)) x = F.relu(x + identity, inplace=True) return x class DepthAttentionCollapse(nn.Module): def __init__(self, in_ch): super().__init__() self.attn_conv = nn.Conv3d(in_ch, 1, 1) class AuxHead(nn.Module): def __init__(self, in_ch, target_size=64): super().__init__() self.pool = nn.AdaptiveAvgPool3d((1, None, None)) self.conv = nn.Conv2d(in_ch, 1, 1) self.target_size = target_size class Decoder3DUNet(nn.Module): def __init__(self, encoder_dims=(256, 512, 1024, 2048), decoder_dims=(64, 128, 256, 512), deep_supervision=True): super().__init__() self.deep_supervision = deep_supervision self.channel_reduce = nn.ModuleList([nn.Sequential(nn.Conv3d(enc_d, dec_d, 1, bias=False), nn.GroupNorm(min(32, dec_d), dec_d), nn.ReLU(inplace=True)) for enc_d, dec_d in zip(encoder_dims, decoder_dims)]) self.decoder_blocks = nn.ModuleList([ResConvBlock3D(decoder_dims[i] + decoder_dims[i - 1], decoder_dims[i - 1]) for i in range(len(decoder_dims) - 1, 0, -1)]) self.depth_collapse = DepthAttentionCollapse(decoder_dims[0]) self.logit = nn.Conv2d(decoder_dims[0], 1, 1) if deep_supervision: self.aux_head_s2 = AuxHead(decoder_dims[2]) self.aux_head_s1 = AuxHead(decoder_dims[1]) def forward_features(self, feat_maps): """Decode once, retaining the feature volume before depth collapse.""" feats = [self.channel_reduce[i](feat_maps[i]) for i in range(4)] aux_outputs = [] x = feats[3] x = F.interpolate(x, size=feats[2].shape[2:], mode='trilinear', align_corners=False) x = torch.cat([x, feats[2]], dim=1) x = self.decoder_blocks[0](x) if self.deep_supervision and self.training: aux_outputs.append(self.aux_head_s2(x)) x = F.interpolate(x, size=feats[1].shape[2:], mode='trilinear', align_corners=False) x = torch.cat([x, feats[1]], dim=1) x = self.decoder_blocks[1](x) if self.deep_supervision and self.training: aux_outputs.append(self.aux_head_s1(x)) x = F.interpolate(x, size=feats[0].shape[2:], mode='trilinear', align_corners=False) x = torch.cat([x, feats[0]], dim=1) x = self.decoder_blocks[2](x) return (x, aux_outputs) class RegressionModel(nn.Module): def __init__(self, with_norm=False): super().__init__() self.backbone = ResNet(Bottleneck, [3, 8, 36, 3], [64, 128, 256, 512], n_input_channels=1, forward_features=True, n_classes=1039) self.decoder = Decoder3DUNet(encoder_dims=(256, 512, 1024, 2048), decoder_dims=(64, 128, 256, 512), deep_supervision=True) self.normalization = nn.BatchNorm3d(num_features=1) if with_norm else None def encode(self, x: torch.Tensor): if x.ndim == 4: x = x[:, None] if self.normalization is not None: x = self.normalization(x) return self.backbone(x) def forward_features(self, x: torch.Tensor): return self.decoder.forward_features(self.encode(x)) class SpatialMaxPool3d(nn.Module): def forward(self, x): b, c, d, h, w = x.shape out = F.max_pool2d(x.reshape(b*c*d, 1, h, w), 3, 2, 1) return out.reshape(b, c, d, *out.shape[-2:]) class Hecate(nn.Module): def __init__(self, patch_size, sampling_um): super().__init__() self.patch_size = tuple(patch_size) if (sampling_um, self.patch_size) not in ( (2.4, (64, 256, 256)), (9.6, (16, 64, 64))): raise ValueError('Unsupported Hecate sampling/patch geometry') self.sampling_um = sampling_um self.margin = 1 if sampling_um == 2.4 else 0 self.xy_stride = 4 if sampling_um == 2.4 else 1 self.divisor = 200. if sampling_um == 2.4 else 255. self.canonical = RegressionModel(with_norm=True) self.canonical.decoder.deep_supervision = False self.canonical.backbone.maxpool = SpatialMaxPool3d() if self.xy_stride == 1: self.canonical.backbone.conv1 = nn.Conv3d(1, 64, 3, padding=1, bias=False) self.canonical.backbone.maxpool = nn.Identity() self.depth_coordinate_scale = nn.Parameter(torch.zeros(())) def features_logits(self, image): if tuple(image.shape[1:]) != (1, *self.patch_size): raise ValueError(f'Expected B,1,{self.patch_size}, got {tuple(image.shape)}') d = self.patch_size[0] features, _ = self.canonical.forward_features(image[:, :, self.margin:d-self.margin]) head = self.canonical.decoder.logit logits = F.conv3d(features.float(), head.weight.float().unsqueeze(2), head.bias.float()) return features, logits.float() def forward_3d(self, image): """3D logits on the input grid; image is already normalized, B1ZYX. This path does not execute attention or the 2D projection. The 2.4um model's outermost planes are untrained and padded with -20. """ _, logits = self.features_logits(image) full = F.interpolate(logits, size=(self.patch_size[0]-2*self.margin, *image.shape[-2:]), mode='trilinear', align_corners=False) return F.pad(full, (0, 0, 0, 0, self.margin, self.margin), value=-20.) def forward(self, image, valid=None): """2D logits; image is already normalized, B1ZYX; valid is optional B1ZYX.""" features, logits = self.features_logits(image) scores = self.canonical.decoder.depth_collapse.attn_conv(features).float() half = scores.shape[2] // 2 z = (torch.arange(scores.shape[2], device=scores.device, dtype=torch.float32)-half)/half scores = scores + self.depth_coordinate_scale*z.view(1, 1, -1, 1, 1) if valid is None: support = torch.ones_like(logits, dtype=torch.bool) else: support = F.max_pool3d(valid[:, :, self.margin:self.patch_size[0]-self.margin].float(), (1, self.xy_stride, self.xy_stride)).bool() weights = scores.masked_fill(~support, -1e4).softmax(2)*support weights = weights/weights.sum(2, keepdim=True).clamp_min(1e-8) probability = F.interpolate((weights*logits).sum(2).sigmoid(), size=image.shape[-2:], mode='bilinear', align_corners=False) eps = torch.finfo(torch.float32).eps return torch.logit(probability.clamp(eps, 1-eps)) def load_model(checkpoint, device='cpu'): """Strictly load a released inference-only EMA checkpoint.""" payload = torch.load(checkpoint, map_location='cpu', weights_only=True, mmap=True) if payload.get('architecture_family') != 'hecate' or payload.get('weight_source') != 'ema_model': raise ValueError('Use a released Hecate inference checkpoint') model = Hecate(payload['config']['patch_size'], payload['sampling_um']) config = payload['config'] canonical = config['model_config']['canonical'] if (config['image_normalization'] != {'mode': 'divide', 'divisor': model.divisor} or canonical.get('native_refinement', False) or canonical['native_xy_stride'] != model.xy_stride or canonical['depth_margin'] != model.margin or not canonical['with_norm']): raise ValueError('Checkpoint configuration does not match the released architecture') model.load_state_dict(payload['ema_model'], strict=True) return model.requires_grad_(False).eval().to(device) def axis_starts(length, patch, stride): return list(range(0, max(0, length-patch), stride)) + [max(0, length-patch)] def open_volume(path, array='0'): if str(path).endswith('.npy'): volume = np.load(path, mmap_mode='r', allow_pickle=False) else: import zarr volume = zarr.open(str(path), mode='r') if not hasattr(volume, 'shape'): volume = volume[array] if len(volume.shape) != 3 or np.dtype(volume.dtype) != np.dtype('uint8'): raise ValueError('Input must be an unnormalized uint8 Z,Y,X surface volume') return volume @torch.inference_mode() def predict(model, volume, output, reverse=False, stride=None, batch_size=1, precision='fp32', three_d=False): """Tiled 2D or 3D uint8 probabilities into a writable HxW array (can be a memmap). Reads only individual CT patches. Accumulators use temporary disk-backed arrays, so a full surface volume never needs to fit in RAM or GPU memory. """ d, p, _ = model.patch_size depth, h, w = volume.shape stride = p//2 if stride is None else stride if depth < d or not 0 < stride <= p or batch_size < 1: raise ValueError('Insufficient depth, invalid stride, or invalid batch size') expected_shape = (depth,h,w) if three_d else (h,w) if tuple(output.shape) != expected_shape or output.dtype != np.uint8: raise ValueError(f'Output must be a uint8 array of shape {expected_shape}') if three_d: output[:] = 0 device = next(model.parameters()).device if precision not in ('fp32', 'bf16') or (precision == 'bf16' and device.type != 'cuda'): raise ValueError('bf16 requires CUDA; otherwise use fp32') z0 = depth//2-d//2 zs = slice(depth-z0-d, depth-z0) if reverse else slice(z0, z0+d) win = torch.hann_window(p, periodic=False) win = win[:, None]*win[None, :] win = (win/win.max().clamp_min(torch.finfo(torch.float32).eps)).clamp_min(.001) window = win.numpy() win = win.to(device) origins = [(y, x) for y in axis_starts(h,p,stride) for x in axis_starts(w,p,stride)] with tempfile.TemporaryDirectory(prefix='hecate-') as tmp: numerator = np.memmap(Path(tmp)/'sum', mode='w+', dtype='float32', shape=(d,h,w) if three_d else (h,w)) denominator = np.memmap(Path(tmp)/'weight', mode='w+', dtype='float32', shape=(h,w)) for start in range(0, len(origins), batch_size): active, patches = [], [] for y, x in origins[start:start+batch_size]: ph, pw = min(p,h-y), min(p,w-x) denominator[y:y+ph,x:x+pw] += window[:ph,:pw] raw = np.asarray(volume[zs,y:y+ph,x:x+pw]) if reverse: raw = raw[::-1] if not raw.any(): continue patch = np.zeros((d,p,p), dtype=np.uint8) patch[:,:ph,:pw] = raw patches.append(patch) active.append((y,x,ph,pw)) if patches: raw = torch.from_numpy(np.stack(patches)).to(device) support = raw.any(1) normalized = np.stack(patches).astype(np.float32) normalized /= model.divisor # Match original CPU division; no clipping. image = torch.from_numpy(normalized[:,None]).to(device) valid = support[:,None,None].expand(-1,1,d,-1,-1) with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=precision=='bf16'): logits = model.forward_3d(image) if three_d else model(image, valid) probability = logits.float().sigmoid()[:,0] if three_d: support = support[:,None] weighted = (probability*support*win).cpu().numpy() for pred,(y,x,ph,pw) in zip(weighted,active): numerator[...,y:y+ph,x:x+pw] += pred[..., :ph,:pw] if start == 0 or start+batch_size >= len(origins) or start % 100 == 0: print(f'{min(start+batch_size,len(origins))}/{len(origins)} patches', flush=True) if three_d and model.margin: numerator[:model.margin] = 0 numerator[-model.margin:] = 0 for y in range(0,h,128): for x in range(0,w,256): prob = numerator[...,y:y+128,x:x+256]/denominator[y:y+128,x:x+256] if not np.isfinite(prob).all(): raise FloatingPointError('Non-finite inference output') values = np.rint(prob.clip(0,1)*255).astype(np.uint8) if three_d: output[zs,y:y+128,x:x+256] = values[::-1] if reverse else values else: output[y:y+128,x:x+256] = values del numerator, denominator def write_png(array, path): """Losslessly compressed grayscale PNG; streamed to avoid full-image copies.""" with Path(path).open('wb') as stream: def chunk(kind, data): stream.write(struct.pack('!I',len(data))+kind+data+struct.pack('!I',zlib.crc32(kind+data)&0xffffffff)) stream.write(b'\x89PNG\r\n\x1a\n') h,w = array.shape chunk(b'IHDR',struct.pack('!2I5B',w,h,8,0,0,0,0)) compressor = zlib.compressobj(9) for row in array: filtered = row.copy() filtered[1:] = row[1:]-row[:-1] data = compressor.compress(b'\x01'+filtered.tobytes()) if data: chunk(b'IDAT',data) chunk(b'IDAT',compressor.flush()) chunk(b'IEND',b'') def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--checkpoint', required=True) parser.add_argument('--input', required=True, help='uint8 ZYX .npy or local Zarr') parser.add_argument('--array', default='0', help='Array within a Zarr group') parser.add_argument('--output', help='2D probability PNG (optional with --output-3d)') parser.add_argument('--output-3d', help='Compressed uint8 3D probability Zarr') parser.add_argument('--spacing-um', type=float, required=True, help='Actual isotropic render spacing, not raw scan spacing') parser.add_argument('--reverse', action='store_true', help='Reverse full Z before selecting the central input planes') parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu') parser.add_argument('--batch-size', type=int, default=1) parser.add_argument('--stride', type=int, help='XY stride; default half the trained patch width') parser.add_argument('--precision', choices=('fp32','bf16'), default='fp32') args = parser.parse_args() if not args.output and not args.output_3d: parser.error('Specify --output and/or --output-3d') output = Path(args.output) if args.output else None output3d = Path(args.output_3d) if args.output_3d else None for path,suffix in ((output,'.png'),(output3d,'.zarr')): if path is not None and (path.suffix.lower()!=suffix or path.exists()): parser.error(f'Choose a new output filename ending in {suffix}: {path}') model = load_model(args.checkpoint,args.device) if abs(args.spacing_um/model.sampling_um-1) > .02: parser.error(f'Resample the render to {model.sampling_um} um first (no automatic resampling)') volume = open_volume(args.input,args.array) if output is not None: output.parent.mkdir(parents=True,exist_ok=True) with tempfile.TemporaryDirectory(prefix='hecate-') as tmp: result = np.memmap(Path(tmp)/'prediction',mode='w+',dtype='uint8',shape=volume.shape[-2:]) predict(model,volume,result,args.reverse,args.stride,args.batch_size,args.precision) partial = output.with_suffix('.partial') write_png(result,partial) partial.replace(output) del result metadata = { 'checkpoint':Path(args.checkpoint).name,'sampling_um':args.spacing_um, 'shape_yx':list(volume.shape[-2:]),'reverse_z':args.reverse, 'central_depth':model.patch_size[0],'precision':args.precision, 'stride':args.stride or model.patch_size[1]//2, 'probabilities':'uint8 = round(255 * probability); no min-max scaling' } if output is not None: output.with_suffix('.json').write_text(json.dumps(metadata,indent=2)+'\n') print(output) if output3d is not None: import zarr from numcodecs import Blosc output3d.parent.mkdir(parents=True,exist_ok=True) partial = output3d.with_name(output3d.name+'.partial') if partial.exists(): parser.error(f'Incomplete output exists: {partial}; remove it or choose another output') options = {'zarr_format':2} if int(zarr.__version__.split('.')[0])>=3 else {} result = zarr.open_array(str(partial),mode='w',shape=volume.shape, chunks=(1,256,256),dtype='uint8',fill_value=0, compressor=Blosc(cname='zstd',clevel=3,shuffle=Blosc.BITSHUFFLE),**options) predict(model,volume,result,args.reverse,args.stride,args.batch_size,args.precision,three_d=True) depth = volume.shape[0] d = model.patch_size[0] start = depth//2-d//2 if args.reverse: start = depth-start-d result.attrs.update({**metadata,'axes':['z','y','x'], 'spacing_um':[args.spacing_um]*3,'z_order':'same as input', 'evaluated_z_interval':[start+model.margin,start+d-model.margin], 'outside_evaluated_z':'zero; not evaluated', 'background':'empty CT columns zeroed; no intensity threshold'}) partial.rename(output3d) print(output3d) if __name__ == '__main__': main()