| """Inference on raw (height, width, time-bin) HDF5 histograms.""" |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import h5py |
| import numpy as np |
| import torch |
| from model import HistCoarseMultiScaleMLP |
|
|
|
|
| def preprocess(histogram, stats): |
| data = np.asarray(histogram, dtype=np.float32) |
| if data.shape != (32, 32, 64): |
| raise ValueError(f'Expected (32, 32, 64), got {data.shape}') |
| if not np.isfinite(data).all() or (data < 0).any(): |
| raise ValueError('Histogram counts must be finite and nonnegative') |
| if stats['mode'] != 'photon_count_10000_log1p_standardize': |
| raise ValueError('Unsupported preprocessing mode') |
| total = float(data.sum(dtype=np.float64)) |
| if total > 0: |
| data = data * (10000.0 / total) |
| data = np.log1p(data) |
| data = (data - stats['mean']) / max(stats['std'], 1e-6) |
| return torch.from_numpy(np.transpose(data, (2, 0, 1)).copy()).unsqueeze(0).unsqueeze(0) |
|
|
|
|
| def load_model(model_dir=None, device='cpu'): |
| root = Path(model_dir) if model_dir else Path(__file__).resolve().parent |
| config = json.loads((root / 'config.json').read_text()) |
| stats = json.loads((root / 'preprocess_stats.json').read_text()) |
| checkpoint = torch.load(root / 'best.pt', map_location='cpu', weights_only=True) |
| if stats != checkpoint['preprocess_stats']: |
| raise ValueError('Preprocessing statistics differ from checkpoint') |
| model = HistCoarseMultiScaleMLP(config['hidden_features'], config['dropout']) |
| model.load_state_dict(checkpoint['model_state'], strict=True) |
| model.to(device).eval() |
| return model, config, stats |
|
|
|
|
| def predict(histogram, model, config, stats): |
| device = next(model.parameters()).device |
| with torch.inference_mode(): |
| logit = model(preprocess(histogram, stats).to(device)).item() |
| probability = torch.sigmoid(torch.tensor(logit)).item() |
| label = int(probability >= config['threshold']) |
| return {'logit': logit, 'probability_positive': probability, |
| 'label': label, 'label_name': config['id2label'][str(label)]} |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('histogram', type=Path) |
| parser.add_argument('--model-dir', type=Path, default=Path(__file__).resolve().parent) |
| parser.add_argument('--device', default='cpu') |
| args = parser.parse_args() |
| model, config, stats = load_model(args.model_dir, args.device) |
| with h5py.File(args.histogram, 'r') as f: |
| histogram = f[config['h5_key']][()] |
| print(json.dumps(predict(histogram, model, config, stats), indent=2)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
