Add tested tiled 3D probability Zarr export

README.md

@@ -117,9 +117,25 @@ The script selects the central input depth, slides over XY with half-patch overl
 
 The default is float32 with batch size one. On compatible NVIDIA GPUs, `--precision bf16` enables the mixed precision used for production inference; increase `--batch-size` if memory permits. `--device cpu` also works, but this remains a large 3D network. Temporary disk-backed accumulators use approximately nine bytes per output pixel; set `TMPDIR` to a disk with enough free space for large surfaces.
 
+### Saving a 3D probability volume
+
+Use `--output-3d` to save a compressed uint8 Zarr:
+
+```bash
+python hecate.py --checkpoint hecate_9.6um.pth --input surface.zarr --spacing-um 9.6 --output-3d ink_3d.zarr --reverse --device cuda
+```
+
+Add `--output prediction.png` to also save the 2D map. When only `--output-3d` is requested, attention and the 2D projection are not executed. Requesting both outputs runs two inference passes to keep the implementation and memory use simple.
+
+The 3D Zarr is a **single-resolution array**, not a multiscale OME-Zarr. It stores `round(255 × probability)` with lossless Zstandard compression. Its ZYX shape and voxel order match the input, including after reversed inference: predictions are flipped back before saving. No cropping or per-volume min–max scaling is applied.
+
+Inference tiles over XY and evaluates only the checkpoint's **central depth window**, not sliding windows through the entire Z axis. Unevaluated planes, including the finer model's excluded edge planes, are zero. The Zarr attributes record the half-open `evaluated_z_interval`, spacing, orientation, and inference settings so those zeros can be distinguished from evaluated predictions. Empty CT columns are zeroed; no additional voxel-intensity threshold is applied. Source world-coordinate transforms are not copied; the output is aligned to the input's voxel indices.
+
+Temporary disk use for 3D blending is approximately `(4 × input_patch_depth + 4) × height × width` bytes, in addition to the compressed output. Set `TMPDIR` appropriately for large renders.
+
 ### Accessing the 3D output
 
-The same file provides a small Python API. `model(image)` returns 2D logits; `model.forward_3d(image)` returns 3D logits without executing the learned projection. The CLI writes only the 2D map.
+The same file provides a small Python API. `model(image)` returns 2D logits; `model.forward_3d(image)` returns 3D logits without executing the learned projection. The CLI can write either output or both.
 
 ```python
 import numpy as np
@@ -138,7 +154,7 @@ np.save("ink_3d.npy", ink_3d)
 
 The 2.4 µm model excludes its first and last input planes from the network and pads their 3D logits with −20; these planes are not valid learned predictions. Its interior 3D logits are interpolated to the input grid. The 9.6 µm model uses all sixteen planes. The Python API returns model predictions without applying an additional CT-background threshold.
 
-The standalone runtime was checked against the original inference implementation on real CT crops with both released checkpoints: 2D and 3D patch logits, tiled forward/reverse predictions, padded edges, blank inputs, and lossless PNG decoding. Float32 and CUDA bfloat16 were tested. The checkpoints load strictly, including their saved EMA batch-normalization buffers.
+The standalone runtime was checked against the original inference implementation on real CT crops with both released checkpoints: 2D and 3D patch logits, tiled forward/reverse predictions, padded edges, blank inputs, lossless PNG decoding, and compressed 3D Zarr round trips. Tiled 3D exports matched the original model in both depth directions, and checks confirmed that 3D-only inference bypasses attention. Float32 and CUDA bfloat16 were tested. The checkpoints load strictly, including their saved EMA batch-normalization buffers.
 
 ## Sources and acknowledgements
 

hecate.py

@@ -302,8 +302,8 @@ def open_volume(path, array='0'):
 
 
 @torch.inference_mode()
-def predict(model, volume, output, reverse=False, stride=None, batch_size=1, precision='fp32'):
-    """Tiled 2D uint8 probabilities into a writable HxW array (can be a memmap).
+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.
@@ -313,8 +313,11 @@ def predict(model, volume, output, reverse=False, stride=None, batch_size=1, pre
     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')
-    if tuple(output.shape) != (h, w) or output.dtype != np.uint8:
-        raise ValueError('Output must be an HxW uint8 array')
+    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')
@@ -327,7 +330,7 @@ def predict(model, volume, output, reverse=False, stride=None, batch_size=1, pre
     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=(h,w))
+        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 = [], []
@@ -351,17 +354,28 @@ def predict(model, volume, output, reverse=False, stride=None, batch_size=1, pre
                 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'):
-                    probability = model(image, valid).float().sigmoid()[:,0]
+                    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]
+                    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)
-        for y in range(0,h,256):
-            prob = numerator[y:y+256]/denominator[y:y+256]
-            if not np.isfinite(prob).all():
-                raise FloatingPointError('Non-finite inference output')
-            output[y:y+256] = np.rint(prob.clip(0,1)*255).astype(np.uint8)
+        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
 
 
@@ -389,7 +403,8 @@ def main():
     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', required=True, help='2D probability PNG')
+    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')
@@ -397,31 +412,61 @@ def main():
     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()
-    output = Path(args.output)
-    if output.suffix.lower() != '.png':
-        parser.error('--output must end in .png')
-    if output.exists():
-        parser.error('Output already exists; choose a new filename')
+    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)
-    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
-    output.with_suffix('.json').write_text(json.dumps({
+    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'
-    },indent=2)+'\n')
-    print(output)
+    }
+    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__':

requirements.txt

@@ -1,3 +1,4 @@
 torch>=2.6,<3
 numpy>=1.26,<3
 zarr>=2.18,<4
+numcodecs>=0.12,<1