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