xie_x1andClaude Sonnet 5
2683fd3576
MC generation -> pile-up assembly -> CNN vs. classical eta-interpolation (scored against MC ground truth) -> optional training -> qualitative inference on real measurement data, self-contained and runnable from a single conda env (see README.md). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
32 lines
1.3 KiB
Python
32 lines
1.3 KiB
Python
"""Small registry mapping photon count -> which pretrained 12 keV checkpoint
|
|
and model class to use, so the demo scripts don't repeat this three times.
|
|
"""
|
|
from pathlib import Path
|
|
import torch
|
|
from models import get_model_class, get_double_photon_model_class, get_triple_photon_model_class
|
|
|
|
MODELS_DIR = Path(__file__).resolve().parent.parent / 'models'
|
|
|
|
REGISTRY = {
|
|
1: dict(version='260511', checkpoint=MODELS_DIR / 'singlePhoton_12keV.pth',
|
|
sample_size=3, get_class=get_model_class),
|
|
2: dict(version='260610', checkpoint=MODELS_DIR / 'doublePhoton_12keV.pth',
|
|
sample_size=6, get_class=get_double_photon_model_class),
|
|
3: dict(version='260611', checkpoint=MODELS_DIR / 'triplePhoton_12keV.pth',
|
|
sample_size=9, get_class=get_triple_photon_model_class),
|
|
}
|
|
|
|
|
|
def load_pretrained(n_photons, device='cuda'):
|
|
info = REGISTRY[n_photons]
|
|
model = info['get_class'](info['version'])().to(device)
|
|
model.load_state_dict(torch.load(info['checkpoint'], map_location=device, weights_only=True))
|
|
model.eval()
|
|
return model
|
|
|
|
|
|
def new_model(n_photons, device='cuda'):
|
|
"""Same architecture as load_pretrained(), but randomly initialized (for training from scratch)."""
|
|
info = REGISTRY[n_photons]
|
|
return info['get_class'](info['version'])().to(device)
|