| """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 |
| 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() |
|
|
