Add --resample so Hecate can read the 8.64 µm eligible volumes

hecate.py

@@ -411,6 +411,9 @@ def main():
     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')
+    parser.add_argument('--resample', action='store_true',
+        help='Rescale the render from --spacing-um to the model spacing before inference (linear), and map the 2D '
+             'prediction back to the input YX grid. Not supported with --output-3d.')
     args = parser.parse_args()
     if not args.output and not args.output_3d:
         parser.error('Specify --output and/or --output-3d')
@@ -420,21 +423,37 @@ def main():
         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)')
+    off_scale = abs(args.spacing_um/model.sampling_um-1) > .02
+    if off_scale and not args.resample:
+        parser.error(f'Resample the render to {model.sampling_um} um first, or pass --resample')
+    if args.resample and output3d is not None:
+        parser.error('--resample is only supported for the 2D --output')
     volume = open_volume(args.input,args.array)
+    input_yx = volume.shape[-2:]
+    if args.resample and off_scale:
+        from scipy import ndimage
+        factor = args.spacing_um/model.sampling_um
+        volume = ndimage.zoom(np.asarray(volume),factor,order=1)
+        print(f'resampled {args.spacing_um} um -> {model.sampling_um} um: {input_yx} -> {volume.shape[-2:]}',flush=True)
     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)
+            if tuple(result.shape)!=tuple(input_yx):
+                from scipy import ndimage
+                back = ndimage.zoom(np.asarray(result,dtype=np.float32),
+                    (input_yx[0]/result.shape[0],input_yx[1]/result.shape[1]),order=1)
+                back = np.clip(np.round(back),0,255).astype('uint8')[:input_yx[0],:input_yx[1]]
+                result = np.zeros(input_yx,dtype='uint8'); result[:back.shape[0],:back.shape[1]] = back
             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,
+        'shape_yx':list(input_yx),'reverse_z':args.reverse,
+        'resampled_to_um':model.sampling_um if (args.resample and off_scale) else None,
         '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'

requirements.txt

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