File size: 2,620 Bytes
97fd768 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """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()
|