Add optional DirectML inference for Hecate 9.6 µm on Windows

DIRECTML.md

@@ -0,0 +1,67 @@
+# Hecate ink detection on a laptop's integrated graphics
+
+`hecate_directml.py` makes GPU-accelerated Hecate 9.6 µm ink detection usable on a 16 GB Windows laptop with integrated graphics and shared system memory. The tested prediction path requires no discrete graphics card or rented GPU compute.
+
+This lowers the hardware barrier for contributors who already own a laptop but lack a dedicated GPU. They can run the released detector and test rendered regions locally, with faster feedback than the CPU path. The implementation uses ONNX Runtime DirectML to reach the laptop's existing graphics hardware; it preserves the released model and probability-image format.
+
+The demonstrated machine has a Ryzen 5 8540U with **Radeon 740M integrated graphics**, Windows 11 and **16 GB RAM**. A real PHerc0139 crop took **3.54 s**, compared with **13.02 s** for the original CPU predictor: **3.68× faster**, with all 9,216 output pixels identical. These timings cover tiled prediction and blending after model/session loading. They exclude export and initialization.
+
+## Install and run
+
+Use Windows and a DirectX 12 GPU. The tested machine has Windows 11, a Ryzen 5 8540U, Radeon 740M integrated graphics, and 16 GB RAM. Other DirectML devices have not been tested here.
+
+Keep the helper beside the released `hecate.py`, `requirements.txt`, and `hecate_9.6um.pth`. These commands reproduce the tested Python 3.11 / CPU PyTorch environment. Use a separate environment from installations containing `onnxruntime` or `onnxruntime-gpu`:
+
+```powershell
+python -m venv .venv
+.venv\Scripts\python -m pip install torch==2.10.0 --index-url https://download.pytorch.org/whl/cpu
+.venv\Scripts\python -m pip install -r requirements.txt -r requirements-directml.txt
+.venv\Scripts\python hecate_directml.py export --checkpoint hecate_9.6um.pth --output hecate_9.6um.onnx
+.venv\Scripts\python hecate_directml.py predict --model hecate_9.6um.onnx --input surface.zarr --spacing-um 9.6 --output ink.png
+```
+
+Export creates a roughly 513 MB ONNX file and a small `.onnx.json` sidecar. Keep both. The sidecar binds the graph to its checkpoint hash. Export once and reuse it; the measured export took 9.94 s. Allow about 5 GB of free system commit capacity for the verification process. Run export separately from other model jobs.
+
+Inputs must be unnormalized uint8 surface renders in Z,Y,X order, already sampled near 9.6 µm in all axes. A local Zarr array, a Zarr group containing `0`, or `.npy` is accepted. `--array` selects another array within a group. The spacing argument checks the declared spacing; it does not resample a raw scan.
+
+`--reverse` reverses the full depth before selecting the central 16 planes. `--stride` defaults to 32; `--device-id` selects the DirectML adapter, with 0 as the default. `--profile` records executed providers. CPU fallback is disabled; an unsupported graph fails instead of silently running on CPU.
+
+The output is a grayscale PNG with `round(255 * probability)`, plus JSON settings and timings. There is no per-image contrast scaling. The source Hann overlap, empty-column mask and padded-edge behavior are preserved. Accumulators use temporary disk files, approximately eight bytes per output pixel, plus one byte per pixel for the output buffer. Set `TMP` to a drive with adequate space for large surfaces.
+
+This helper supports **9.6 µm, FP32, batch size one, 2D output**. It does not implement training, the 2.4 µm checkpoint or 3D export. DirectML acceleration changes execution, not the detector's accuracy.
+
+## Reproduce the comparison
+
+The source was pinned to Hecate revision `9cb86e500e944b11a06a7020403cde5dffb5bcb2`. The checkpoint SHA-256 is `809f4f10f7cb7afa19b4bee0f7d2ab31edd7e11b664f9f210cc9c17f22fcfe5d`.
+
+The example comes from PHerc0139, segment `20260112000000-w043_2026011217`, native scan `20250728140407`, 9.362 µm surface volume. `fetch_example.py` downloads four hash-checked chunks from the official server and reconstructs the exact comparison crop. CT data and model weights are not bundled with this contribution. The [data terms](https://dl.ash2txt.org/LICENSE.txt) apply to the downloaded data.
+
+```powershell
+.venv\Scripts\python -m pip install -r requirements-verify.txt
+.venv\Scripts\python fetch_example.py --output pherc0139_crop.npy
+.venv\Scripts\python verify_directml.py --checkpoint hecate_9.6um.pth --input pherc0139_crop.npy --output check
+```
+
+The crop is 16×96×96. Its depth spacing is 9.6 µm; its effective XY spacing is 9.6059 µm after rounding the original 512-pixel extent to 499 pixels. Exact depth coordinates, interpolation and source hashes are in `fetch_example.py`. The reconstructed `.npy` SHA-256 is `04224e26f565c4a5c1ff7f018e4904a381e7d0d6a217a6ade09ab6ec313e312a`.
+
+The check runs the original CPU predictor, exports the model, and runs DirectML in separate processes. It requires each comparison to differ by at most one uint8 level, an exact repeated GPU image, exact equality between the array API and Zarr-to-PNG command, and cleanup after injected inference/writer errors. The observed differences were zero in every case:
+
+| Input | CPU prediction | DirectML prediction | Maximum pixel difference |
+|---|---:|---:|---:|
+| Real crop, four overlapping tiles | 13.018 s | 3.538 s | 0 |
+| Reversed depth | 12.973 s | 3.477 s | 0 |
+| Real crop cut to 45×53, with padding | 3.299 s | 0.923 s | 0 |
+| Real crop with nine empty columns | 13.231 s | 3.571 s | 0 |
+| Blank negative control | 0.004 s | 0.009 s | 0 |
+
+The fresh Zarr-to-PNG command took 5.313 s including session initialization. The graph profile contains only DirectML dispatches. The CPU reference used four threads; both paths ran in FP32 on the same laptop. Raw measurements and source/code hashes are in `directml_validation.json`.
+
+![CPU and integrated-GPU pixel comparison](https://huggingface.co/scrollprize/hecate/resolve/468a7bceb71f00edc4a2aca41fcd7ef89d79a870/directml_comparison.png)
+
+The altered and blank inputs are boundary controls, not additional independent scroll examples. This is an exposed backend-equivalence comparison, not held-out accuracy validation or a claim of newly recovered letters.
+
+## Implementation and attribution
+
+The helper reuses `hecate.py` for Zarr opening, tile origins and PNG writing. It uses the same intensity division and floored Hann blend. At 9.6 µm, the attention and output use the same XY grid. The validity mask is constant through depth for each column, so masking empty columns after inference gives the source result. This argument does not extend to the 2.4 µm spatial interpolation path.
+
+Hecate and its released weights are by the Vesuvius Challenge team. This optional backend was developed and tested with Codex during Mohamad Al-Zawahreh's scroll-recovery work. It does not retrain or replace the model. Tested software: Python 3.11.9, PyTorch 2.10.0+cpu, ONNX 1.23.2 and ONNX Runtime DirectML 1.24.4.

directml_validation.json

@@ -0,0 +1,129 @@
+{
+  "status": "PASS",
+  "checkpoint_sha256": "809f4f10f7cb7afa19b4bee0f7d2ab31edd7e11b664f9f210cc9c17f22fcfe5d",
+  "input_sha256": "04224e26f565c4a5c1ff7f018e4904a381e7d0d6a217a6ade09ab6ec313e312a",
+  "hecate_sha256": "c232c18a1a86cfb91257a00db13288202bb8ae33732ba26f7291a9b8adb8da59",
+  "helper_sha256": "69b28298711a90019e2a4e9ec67a81b287391a9c0231e41edf3374817894c152",
+  "export_seconds": 9.94456380000338,
+  "rows": [
+    {
+      "case": "overlap",
+      "shape": [
+        16,
+        96,
+        96
+      ],
+      "max_uint8_difference": 0,
+      "different_pixels": 0,
+      "pixels": 9216,
+      "cpu_seconds": 13.018276199931279,
+      "directml_seconds": 3.5383093000855297,
+      "tiles": 4,
+      "model_forwards": 4,
+      "forward_seconds": [
+        0.9426679000025615,
+        0.8602820000378415,
+        0.8697611999232322,
+        0.8601702000014484
+      ]
+    },
+    {
+      "case": "reverse",
+      "shape": [
+        16,
+        96,
+        96
+      ],
+      "max_uint8_difference": 0,
+      "different_pixels": 0,
+      "pixels": 9216,
+      "cpu_seconds": 12.973447300028056,
+      "directml_seconds": 3.476844499935396,
+      "tiles": 4,
+      "model_forwards": 4,
+      "forward_seconds": [
+        0.8666426000418141,
+        0.8924879000987858,
+        0.8558254999807104,
+        0.856082099955529
+      ]
+    },
+    {
+      "case": "padded_edge",
+      "shape": [
+        16,
+        45,
+        53
+      ],
+      "max_uint8_difference": 0,
+      "different_pixels": 0,
+      "pixels": 2385,
+      "cpu_seconds": 3.299313399940729,
+      "directml_seconds": 0.9228492999682203,
+      "tiles": 1,
+      "model_forwards": 1,
+      "forward_seconds": [
+        0.880084800068289
+      ]
+    },
+    {
+      "case": "empty_columns",
+      "shape": [
+        16,
+        96,
+        96
+      ],
+      "max_uint8_difference": 0,
+      "different_pixels": 0,
+      "pixels": 9216,
+      "cpu_seconds": 13.231093699927442,
+      "directml_seconds": 3.571005600038916,
+      "tiles": 4,
+      "model_forwards": 4,
+      "forward_seconds": [
+        0.8853788999840617,
+        0.8660517999669537,
+        0.9252147999359295,
+        0.8839931000256911
+      ]
+    },
+    {
+      "case": "blank_control",
+      "shape": [
+        16,
+        96,
+        96
+      ],
+      "max_uint8_difference": 0,
+      "different_pixels": 0,
+      "pixels": 9216,
+      "cpu_seconds": 0.0044531000312417746,
+      "directml_seconds": 0.00914159999229014,
+      "tiles": 4,
+      "model_forwards": 0,
+      "forward_seconds": []
+    }
+  ],
+  "exact_repeat": true,
+  "zarr_cli_exact": true,
+  "failure_cleanup": true,
+  "cpu_fallback_disabled": true,
+  "executed_providers": {
+    "DmlExecutionProvider": 17
+  },
+  "environment": {
+    "python": "3.11.9",
+    "torch": "2.10.0+cpu",
+    "os": "Windows 11",
+    "processor": "AMD Ryzen 5 8540U",
+    "gpu": "Radeon 740M integrated",
+    "ram_gib": 16,
+    "onnx": "1.23.2",
+    "onnxruntime": "1.24.4",
+    "cpu_threads": 4
+  },
+  "scope": "Backend equivalence on a real exposed crop and explicit edge controls; no new text or model-accuracy claim.",
+  "cpu_reference_provenance": "CPU references completed in a separate process; exact retained input, checkpoint, source CPU code and images were reused for final DirectML comparison. Export reused after a Windows temporary-mapping cleanup-only change.",
+  "fresh_cli_seconds_including_initialization": 5.313350200071,
+  "publication_note": "No raw CT or model weights included. fetch_example.py reconstructs the hash-checked input from official chunks."
+}

fetch_example.py

@@ -0,0 +1,67 @@
+"""Fetch a small PHerc0139 example from the official server; no CT is redistributed.
+
+Data terms: https://dl.ash2txt.org/LICENSE.txt
+This reconstructs the exposed backend-comparison input, not held-out ink labels.
+"""
+import argparse
+import hashlib
+import json
+from pathlib import Path
+from urllib.request import urlopen
+
+import numpy as np
+from scipy.ndimage import map_coordinates
+
+BASE = ('https://vesuvius-challenge-open-data.s3.us-east-1.amazonaws.com/'
+        'PHerc0139/segments/20260112000000-w043_2026011217/surface-volumes/'
+        '9.362um-1.2m-113keV-volume-20250728140407.zarr/0')
+CHUNKS = {
+    (35, 20): 'e1594169b84799d3556d2907c94be7aa145f53f37741b3fabc674b1b187d29c1',
+    (35, 21): 'f3781ae4d3391ab433457874c2d57901ded5e9ef4a7a81bc4165bf4b6b1b1ca8',
+    (36, 20): '7e103b25b0851a6b340ab03e128c5d006a5ac1bebb7c6ea7084ca2a9772150a4',
+    (36, 21): 'b982a846f80eb78a472d760994446bad9824864e8e42c9d97e69e9ee005776cc',
+}
+DEPTH = [4.4291245739319605, 5.454546492325465, 6.479968410718972,
+         7.505390329112477, 8.530812247505981, 9.55623416589949,
+         10.581656084292995, 11.6070780026865, 12.632499921080004,
+         13.65792183947351, 14.683343757867018, 15.708765676260523,
+         16.734187594654028, 17.759609513047533, 18.785031431441038,
+         19.810453349834543]
+EXPECTED = '04224e26f565c4a5c1ff7f018e4904a381e7d0d6a217a6ade09ab6ec313e312a'
+
+
+def fetch(output: Path) -> None:
+    if output.exists():
+        raise FileExistsError(output)
+    with urlopen(BASE + '/.zarray', timeout=30) as response:
+        meta = json.load(response)
+    expected = {'shape': [28, 6120, 8120], 'chunks': [28, 128, 128], 'dtype': '|u1',
+                'compressor': None, 'filters': None, 'order': 'C', 'dimension_separator': '/'}
+    if any(meta.get(k) != v for k, v in expected.items()):
+        raise ValueError('The public source metadata changed')
+    block = np.empty((28, 256, 256), dtype=np.uint8)
+    for (y, x), digest in CHUNKS.items():
+        with urlopen(f'{BASE}/0/{y}/{x}', timeout=30) as response:
+            raw = response.read(28 * 128 * 128 + 1)
+        if hashlib.sha256(raw).hexdigest() != digest:
+            raise ValueError(f'Source chunk changed: {y}/{x}')
+        block[:, (y-35)*128:(y-34)*128, (x-20)*128:(x-19)*128] = np.frombuffer(raw, np.uint8).reshape(28, 128, 128)
+    # The original 512-pixel crop began at Y=4320, X=2440. Its resampled
+    # extent was 499 pixels; this example takes [192:288,192:288] from it.
+    xy = (np.arange(192, 288) + .5) * 512 / 499 - .5
+    coordinates = np.array(np.broadcast_arrays(np.array(DEPTH)[:, None, None],
+        (xy - 160)[None, :, None], (xy - 120)[None, None, :]))
+    image = map_coordinates(block.astype(np.float32), coordinates, order=1, prefilter=False)
+    image = image.clip(0, 255).astype(np.uint8)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    with output.open('xb') as stream:
+        np.save(stream, image, allow_pickle=False)
+    if hashlib.sha256(output.read_bytes()).hexdigest() != EXPECTED:
+        raise ValueError('The reconstructed example differs from the recorded input')
+    print(f'{output}: {EXPECTED}')
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--output', type=Path, default=Path('pherc0139_crop.npy'))
+    fetch(parser.parse_args().output)

hecate_directml.py

@@ -0,0 +1,228 @@
+"""Optional Windows DirectML inference for Hecate's 9.6 um 2D ink map.
+
+Keep this file beside the released hecate.py. Export once, then reuse the graph.
+CPU/CUDA inference and the released weights are unchanged.
+"""
+from __future__ import annotations
+
+import argparse
+from collections import Counter
+from contextlib import contextmanager
+import hashlib
+import json
+from pathlib import Path
+import tempfile
+import time
+
+import numpy as np
+
+
+@contextmanager
+def mapped_array(path: Path, dtype: str, shape: tuple):
+    """Close the Windows file mapping even when inference or writing raises."""
+    array = np.memmap(path, mode='w+', dtype=dtype, shape=shape)
+    try:
+        yield array
+    finally:
+        array._mmap.close()
+
+
+def sha256(path: Path) -> str:
+    with path.open('rb') as stream:
+        return hashlib.file_digest(stream, 'sha256').hexdigest()
+
+
+def export_model(checkpoint: Path, output: Path) -> None:
+    import onnx
+    import torch
+    from hecate import load_model
+
+    sidecar = output.with_suffix('.onnx.json')
+    if output.suffix != '.onnx' or output.exists() or sidecar.exists():
+        raise ValueError('Choose new .onnx and .onnx.json output paths')
+    model = load_model(checkpoint, 'cpu')
+    if model.sampling_um != 9.6 or tuple(model.patch_size) != (16, 64, 64):
+        raise ValueError('This backend supports only Hecate 9.6 um / 16x64x64')
+    output.parent.mkdir(parents=True, exist_ok=True)
+    checkpoint_hash = sha256(checkpoint)
+    with tempfile.TemporaryDirectory(dir=output.parent) as tmp:
+        graph_path = Path(tmp) / 'model.onnx'
+        torch.onnx.export(
+            model, (torch.zeros(1, 1, 16, 64, 64),), str(graph_path),
+            input_names=['volume'], output_names=['logits'], opset_version=18,
+            dynamo=False, external_data=False, do_constant_folding=True,
+        )
+        # A small sidecar avoids loading and serializing another 513 MB protobuf
+        # while the traced PyTorch model is still resident.
+        onnx.checker.check_model(str(graph_path))
+        metadata = {
+            'hecate_backend_contract': '9.6um-2d-fp32-v1',
+            'checkpoint_sha256': checkpoint_hash,
+            'graph_sha256': sha256(graph_path),
+        }
+        metadata_path = Path(tmp) / 'model.onnx.json'
+        metadata_path.write_text(json.dumps(metadata, indent=2) + '\n', encoding='utf-8')
+        graph_path.replace(output)
+        metadata_path.replace(sidecar)
+
+
+def load_session(graph: Path, device_id: int = 0, profile_prefix: str | None = None):
+    import onnxruntime as ort
+
+    if device_id < 0:
+        raise ValueError('device_id must be nonnegative')
+    if 'DmlExecutionProvider' not in ort.get_available_providers():
+        raise RuntimeError('Install onnxruntime-directml in a Windows environment')
+    metadata = json.loads(graph.with_suffix('.onnx.json').read_text(encoding='utf-8'))
+    if metadata.get('hecate_backend_contract') != '9.6um-2d-fp32-v1' or metadata.get('graph_sha256') != sha256(graph):
+        raise ValueError('Use the unchanged graph and sidecar exported by this helper')
+    options = ort.SessionOptions()
+    options.enable_mem_pattern = False
+    options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+    options.intra_op_num_threads = 1
+    options.inter_op_num_threads = 1
+    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    options.add_session_config_entry('session.disable_cpu_ep_fallback', '1')
+    if profile_prefix:
+        options.enable_profiling = True
+        options.profile_file_prefix = profile_prefix
+    session = ort.InferenceSession(str(graph), sess_options=options,
+        providers=[('DmlExecutionProvider', {'device_id': device_id})])
+    inputs, outputs = session.get_inputs(), session.get_outputs()
+    if (len(inputs) != 1 or inputs[0].name != 'volume'
+            or inputs[0].shape != [1, 1, 16, 64, 64] or inputs[0].type != 'tensor(float)'
+            or len(outputs) != 1 or outputs[0].name != 'logits'
+            or outputs[0].shape != [1, 1, 64, 64] or outputs[0].type != 'tensor(float)'):
+        raise ValueError('Unexpected model input/output contract')
+    return session
+
+
+def predict_patch(session, raw: np.ndarray) -> np.ndarray:
+    if raw.shape != (16, 64, 64) or raw.dtype != np.uint8:
+        raise ValueError('Expected uint8 16x64x64 CT patch')
+    image = (raw.astype(np.float32) / 255)[None, None]
+    logits = session.run(['logits'], {'volume': image})[0]
+    if logits.shape != (1, 1, 64, 64) or not np.isfinite(logits).all():
+        raise FloatingPointError('Invalid model output')
+    return (1 / (1 + np.exp(-logits)))[0, 0]
+
+
+def predict(session, volume, output, reverse: bool = False, stride: int = 32) -> dict:
+    """Stream uint8 ZYX input into an HxW uint8 probability array.
+
+    At 9.6 um, attention and output use the same XY grid. Its validity mask is
+    constant through depth for each column. Nonempty columns therefore match
+    unmasked attention exactly; empty columns are zeroed before blending.
+    This equivalence does not justify using this helper for the 2.4 um model.
+    """
+    import torch
+    from hecate import axis_starts
+
+    if len(volume.shape) != 3 or np.dtype(volume.dtype) != np.dtype('uint8'):
+        raise ValueError('Input must be an unnormalized uint8 ZYX surface volume')
+    depth, height, width = volume.shape
+    if depth < 16 or height < 1 or width < 1 or not 0 < stride <= 64:
+        raise ValueError('Need depth >= 16, positive XY dimensions and stride 1..64')
+    if output.shape != (height, width) or output.dtype != np.uint8:
+        raise ValueError('Output must be an HxW uint8 array')
+    z0 = depth // 2 - 8
+    zs = slice(depth - z0 - 16, depth - z0) if reverse else slice(z0, z0 + 16)
+    axis = torch.hann_window(64, periodic=False)
+    window = axis[:, None] * axis[None, :]
+    window = (window / window.max()).clamp_min(.001).numpy()
+    ys, xs = axis_starts(height, 64, stride), axis_starts(width, 64, stride)
+    forward_seconds = []
+    with (tempfile.TemporaryDirectory(prefix='hecate-dml-') as tmp,
+          mapped_array(Path(tmp) / 'sum', 'float32', (height, width)) as numerator,
+          mapped_array(Path(tmp) / 'weight', 'float32', (height, width)) as denominator):
+        for y in ys:
+            for x in xs:
+                ph, pw = min(64, height - y), min(64, width - x)
+                region = np.s_[y:y+ph, x:x+pw]
+                denominator[region] += 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((16, 64, 64), dtype=np.uint8)
+                patch[:, :ph, :pw] = raw
+                tick = time.perf_counter()
+                probability = predict_patch(session, patch)
+                forward_seconds.append(time.perf_counter() - tick)
+                weighted = (probability * patch.any(0)) * window
+                numerator[region] += weighted[:ph, :pw]
+        for y in range(0, height, 128):
+            for x in range(0, width, 256):
+                region = np.s_[y:y+128, x:x+256]
+                probability = numerator[region] / denominator[region]
+                if not np.isfinite(probability).all():
+                    raise FloatingPointError('Non-finite blended output')
+                output[region] = np.rint(probability.clip(0, 1) * 255).astype(np.uint8)
+    return {'tiles': len(ys) * len(xs), 'model_forwards': len(forward_seconds),
+            'forward_seconds': forward_seconds}
+
+
+def finish_profile(session) -> dict:
+    path = Path(session.end_profiling())
+    providers = Counter(event.get('args', {}).get('provider')
+        for event in json.loads(path.read_text(encoding='utf-8'))
+        if event.get('args', {}).get('provider'))
+    if providers and set(providers) != {'DmlExecutionProvider'}:
+        raise RuntimeError(f'Unexpected executed providers: {dict(providers)}')
+    return {'path': str(path), 'executed_providers': dict(providers)}
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    sub = parser.add_subparsers(dest='command', required=True)
+    export = sub.add_parser('export', help='Export the unchanged released checkpoint once')
+    export.add_argument('--checkpoint', type=Path, required=True)
+    export.add_argument('--output', type=Path, required=True)
+    run = sub.add_parser('predict', help='Run the exported model on a Windows GPU')
+    run.add_argument('--model', type=Path, required=True)
+    run.add_argument('--input', required=True, help='uint8 ZYX surface .npy or local Zarr')
+    run.add_argument('--array', default='0')
+    run.add_argument('--output', type=Path, required=True, help='New probability .png')
+    run.add_argument('--spacing-um', type=float, required=True)
+    run.add_argument('--reverse', action='store_true')
+    run.add_argument('--stride', type=int, default=32)
+    run.add_argument('--device-id', type=int, default=0)
+    run.add_argument('--profile', action='store_true')
+    args = parser.parse_args()
+    if args.command == 'export':
+        export_model(args.checkpoint, args.output)
+        return
+    if not np.isfinite(args.spacing_um) or abs(args.spacing_um / 9.6 - 1) > .02:
+        parser.error('Resample the surface render to 9.6 um first; this script does not resample')
+    if args.output.suffix.lower() != '.png' or args.output.exists() or args.output.with_suffix('.json').exists():
+        parser.error('Choose new .png and matching .json output paths')
+    from hecate import open_volume, write_png
+    volume = open_volume(args.input, args.array)
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    start = time.perf_counter()
+    session = load_session(args.model, args.device_id,
+        str(args.output.with_suffix('.profile')) if args.profile else None)
+    with tempfile.TemporaryDirectory(dir=args.output.parent, prefix='hecate-dml-') as tmp:
+        partial = Path(tmp) / 'prediction.png'
+        with mapped_array(Path(tmp) / 'prediction', 'uint8', volume.shape[-2:]) as result:
+            stats = predict(session, volume, result, args.reverse, args.stride)
+            write_png(result, partial)
+        partial.replace(args.output)
+    metadata = {
+        'sampling_um': args.spacing_um, 'shape_yx': list(volume.shape[-2:]),
+        'reverse_z': args.reverse, 'central_depth': 16, 'stride': args.stride,
+        'precision': 'fp32', 'provider': 'DmlExecutionProvider', 'device_id': args.device_id,
+        'cpu_fallback_disabled': True,
+        'checkpoint_sha256': json.loads(args.model.with_suffix('.onnx.json').read_text(encoding='utf-8'))['checkpoint_sha256'],
+        'probabilities': 'uint8 = round(255 * probability); no min-max scaling',
+        'elapsed_seconds': time.perf_counter() - start, **stats,
+    }
+    if args.profile:
+        metadata['profile'] = finish_profile(session)
+    args.output.with_suffix('.json').write_text(json.dumps(metadata, indent=2) + '\n', encoding='utf-8')
+    print(args.output)
+
+
+if __name__ == '__main__':
+    main()

requirements-directml.txt

@@ -0,0 +1,4 @@
+# Install in the same environment as Hecate's requirements.txt.
+# Use a separate environment from onnxruntime / onnxruntime-gpu installations.
+onnxruntime-directml==1.24.4
+onnx==1.23.2

requirements-verify.txt

@@ -0,0 +1,3 @@
+-r requirements-directml.txt
+Pillow>=10,<13
+scipy>=1.11,<2

verify_directml.py

@@ -0,0 +1,163 @@
+"""Compare the optional backend with released Hecate on an actual CT crop."""
+from __future__ import annotations
+
+import argparse
+import gc
+import hashlib
+import json
+from pathlib import Path
+import subprocess
+import sys
+import time
+import tempfile
+from unittest.mock import patch
+
+import numpy as np
+from PIL import Image
+
+import hecate_directml as dml
+
+
+def sha(path: Path) -> str:
+    with path.open('rb') as stream:
+        return hashlib.file_digest(stream, 'sha256').hexdigest()
+
+
+def check_failure_cleanup(source: np.ndarray, folder: Path) -> None:
+    class BrokenSession:
+        def run(self, *args, **kwargs):
+            raise RuntimeError('injected inference failure')
+
+    real_temporary_directory = tempfile.TemporaryDirectory
+    with real_temporary_directory(dir=folder) as parent:
+        def directory(*args, **kwargs):
+            return real_temporary_directory(dir=parent)
+        with patch.object(dml.tempfile, 'TemporaryDirectory', side_effect=directory):
+            try:
+                dml.predict(BrokenSession(), source, np.zeros(source.shape[-2:], np.uint8))
+            except RuntimeError as exc:
+                assert str(exc) == 'injected inference failure'
+            else:
+                raise AssertionError('The injected failure did not propagate')
+        assert not list(Path(parent).iterdir()), 'Temporary mappings survived the error'
+        # The outer PNG buffer must close on errors too.
+        try:
+            with real_temporary_directory(dir=parent) as tmp:
+                with dml.mapped_array(Path(tmp) / 'prediction', 'uint8', source.shape[-2:]):
+                    raise RuntimeError('injected writer failure')
+        except RuntimeError as exc:
+            assert str(exc) == 'injected writer failure'
+        assert not list(Path(parent).iterdir()), 'Output mapping survived the error'
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--checkpoint', type=Path, required=True)
+    parser.add_argument('--input', type=Path, required=True, help='Real uint8 ZYX crop, at least 16x96x96')
+    parser.add_argument('--output', type=Path, required=True, help='New verification directory')
+    parser.add_argument('--stage', choices=['all', 'cpu', 'export', 'directml'], default='all')
+    args = parser.parse_args()
+    if args.stage == 'all':
+        args.output.mkdir(parents=True, exist_ok=False)
+        # End each process to release the CPU convolution/export allocator before
+        # loading the next model. gc.collect() does not release those native arenas.
+        for stage in ('cpu', 'export', 'directml'):
+            subprocess.run([sys.executable, str(Path(__file__)), '--checkpoint', str(args.checkpoint),
+                '--input', str(args.input), '--output', str(args.output), '--stage', stage], check=True)
+        return
+    import torch
+    import hecate
+    torch.set_num_threads(4)
+    torch.set_num_interop_threads(1)
+    graph = args.output / 'hecate.onnx'
+    if args.stage == 'export':
+        start = time.perf_counter()
+        dml.export_model(args.checkpoint, graph)
+        (args.output / 'export.json').write_text(json.dumps({'seconds': time.perf_counter() - start}), encoding='utf-8')
+        return
+    import zarr
+    source = np.load(args.input, allow_pickle=False)
+    assert source.dtype == np.uint8 and source.ndim == 3 and min(source.shape[-2:]) >= 96 and source.shape[0] >= 16
+    source = source[:, :96, :96]
+    assert source.any(), 'A blank array is not a real-data demonstration'
+    cases = [('overlap', source, False), ('reverse', source, True),
+             ('padded_edge', source[:, :45, :53], False)]
+    masked = source.copy()
+    masked[:, :, :9] = 0
+    cases.append(('empty_columns', masked, False))
+    cases.append(('blank_control', np.zeros_like(source), False))
+    reference, cpu_seconds = {}, {}
+    if args.stage == 'cpu':
+        model = hecate.load_model(args.checkpoint, 'cpu')
+        for name, volume, reverse in cases:
+            value = np.zeros(volume.shape[-2:], np.uint8)
+            start = time.perf_counter()
+            hecate.predict(model, volume, value, reverse=reverse)
+            cpu_seconds[name] = time.perf_counter() - start
+            Image.fromarray(value).save(args.output / f'{name}_cpu.png')
+        (args.output / 'cpu.json').write_text(json.dumps(cpu_seconds), encoding='utf-8')
+        return
+    cpu_seconds = json.loads((args.output / 'cpu.json').read_text())
+    export_seconds = json.loads((args.output / 'export.json').read_text())['seconds']
+    reference = {name: np.asarray(Image.open(args.output / f'{name}_cpu.png')) for name, _, _ in cases}
+    check_failure_cleanup(source, args.output)
+    session = dml.load_session(graph, profile_prefix=str(args.output / 'dispatch'))
+    rows = []
+    for name, volume, reverse in cases:
+        value = np.zeros(volume.shape[-2:], np.uint8)
+        start = time.perf_counter()
+        stats = dml.predict(session, volume, value, reverse=reverse)
+        elapsed = time.perf_counter() - start
+        difference = np.abs(value.astype(np.int16) - reference[name].astype(np.int16))
+        row = {'case': name, 'shape': list(volume.shape), 'max_uint8_difference': int(difference.max()),
+               'different_pixels': int(np.count_nonzero(difference)), 'pixels': int(value.size),
+               'cpu_seconds': cpu_seconds[name], 'directml_seconds': elapsed, **stats}
+        Image.fromarray(value).save(args.output / f'{name}_directml.png')
+        rows.append(row)
+        print(json.dumps(row), flush=True)
+        assert row['max_uint8_difference'] <= 1, row
+        if name == 'empty_columns':
+            assert not value[:, :9].any()
+        if name == 'blank_control':
+            assert not value.any() and stats['model_forwards'] == 0
+    repeated = np.zeros(source.shape[-2:], np.uint8)
+    dml.predict(session, source, repeated)
+    assert np.array_equal(repeated, np.asarray(Image.open(args.output / 'overlap_directml.png')))
+    profile = dml.finish_profile(session)
+    assert set(profile['executed_providers']) == {'DmlExecutionProvider'}
+    del session
+    gc.collect()
+    volume_path = args.output / 'surface.zarr'
+    options = {'zarr_format': 2} if int(zarr.__version__.split('.')[0]) >= 3 else {}
+    volume = zarr.open_array(str(volume_path), mode='w', shape=source.shape,
+                             chunks=(16, 32, 32), dtype='uint8', **options)
+    volume[:] = source
+    cli_png = args.output / 'cli.png'
+    command = [sys.executable, str(Path(dml.__file__)), 'predict', '--model', str(graph),
+               '--input', str(volume_path), '--spacing-um', '9.6', '--output', str(cli_png), '--profile']
+    subprocess.run(command, check=True)
+    assert np.array_equal(np.asarray(Image.open(cli_png)), repeated)
+    cli = json.loads(cli_png.with_suffix('.json').read_text())
+    assert set(cli['profile']['executed_providers']) == {'DmlExecutionProvider'}
+    assert cli['checkpoint_sha256'] == sha(args.checkpoint)
+    # Refuse accidental output replacement and a wrong physical spacing.
+    for change in ('existing_output', 'wrong_spacing'):
+        bad = command.copy()
+        if change == 'wrong_spacing':
+            bad[bad.index('--spacing-um') + 1] = '2.4'
+            bad[bad.index('--output') + 1] = str(args.output / 'wrong_spacing.png')
+        result = subprocess.run(bad, capture_output=True, text=True)
+        assert result.returncode != 0, change
+    result = {'status': 'PASS', 'checkpoint_sha256': sha(args.checkpoint), 'input_sha256': sha(args.input),
+              'hecate_sha256': sha(Path(hecate.__file__)), 'helper_sha256': sha(Path(dml.__file__)),
+              'export_seconds': export_seconds, 'rows': rows, 'exact_repeat': True,
+              'zarr_cli_exact': True, 'failure_cleanup': True, 'cpu_fallback_disabled': True,
+              'executed_providers': profile['executed_providers'],
+              'environment': {'python': sys.version.split()[0], 'torch': torch.__version__},
+              'scope': 'Backend equivalence on a real exposed crop and explicit edge controls; no new text or model-accuracy claim.'}
+    (args.output / 'result.json').write_text(json.dumps(result, indent=2) + '\n', encoding='utf-8')
+    print('PASS: real-data parity, reversed depth, padding, empty columns, replay and Zarr CLI', flush=True)
+
+
+if __name__ == '__main__':
+    main()