code_public/radiacode
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Pipeline complet Radiacode 103 - identification automatique d'isotopes

Browse Source 
- VegaModel CNN-FCNN 34.5M params, 82 isotopes, val acc 99.89%
- Generation 50k spectres synthetiques 1D (12-24h durees)
- Entrainement 100 epochs sur RTX 5060 Ti (CUDA 12.8, Blackwell)
- Detection continue avec soustraction du background
- Capture background 24h avec gestion deconnexion
- Docker Compose : conteneur train (GPU) + detect (CPU/USB)
- Modele entraite inclus (vega_best.pt, 395 Mo)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Jacquin Antoine
2026-05-19 12:29:56 +02:00
commit 745a64b342
52 changed files with 17558 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
train/vega_ml/inference/__init__.py Normal file
View File

@ -0,0 +1 @@
# Inference module for running predictions with trained models

93
train/vega_ml/inference/run_inference.py Normal file
View File

@ -0,0 +1,93 @@
#!/usr/bin/env python
"""
Run Vega Inference
Simple script to run inference with a trained Vega model.
"""
import sys
import argparse
from pathlib import Path
# Add project root to path
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
from inference.vega_inference import run_inference_demo, VegaInference
def main():
parser = argparse.ArgumentParser(
description="Run inference with trained Vega model"
)
parser.add_argument(
"--model", "-m",
type=str,
default="models/vega_best.pt",
help="Path to model checkpoint"
)
parser.add_argument(
"--data", "-d",
type=str,
default="O:/master_data_collection/isotopev2",
help="Path to data directory with spectra"
)
parser.add_argument(
"--threshold", "-t",
type=float,
default=0.5,
help="Detection threshold (0-1)"
)
parser.add_argument(
"--spectrum", "-s",
type=str,
default=None,
help="Path to a specific spectrum file to analyze"
)
args = parser.parse_args()
# Make paths absolute
model_path = Path(args.model)
if not model_path.is_absolute():
model_path = PROJECT_ROOT / model_path
if args.spectrum:
# Single spectrum inference
spectrum_path = Path(args.spectrum)
if not spectrum_path.is_absolute():
spectrum_path = PROJECT_ROOT / spectrum_path
print(f"\nLoading model from: {model_path}")
inference = VegaInference(str(model_path))
print(f"\nAnalyzing spectrum: {spectrum_path}")
prediction = inference.predict_from_file(
spectrum_path,
threshold=args.threshold
)
print("\n" + "=" * 60)
print("PREDICTION RESULTS")
print("=" * 60)
print(prediction.summary())
print("=" * 60)
else:
# Demo mode - analyze all spectra in data directory
data_path = Path(args.data)
if not data_path.is_absolute():
data_path = PROJECT_ROOT / data_path
run_inference_demo(
str(model_path),
str(data_path),
threshold=args.threshold
)
return 0
if __name__ == "__main__":
sys.exit(main())

406
train/vega_ml/inference/vega_inference.py Normal file
View File

@ -0,0 +1,406 @@
"""
Vega Inference Script
Load a trained Vega model and run inference on gamma spectra to identify
isotopes and estimate their activities.
"""
import sys
import json
import numpy as np
import torch
from pathlib import Path
from typing import Dict, List, Optional, Union
from dataclasses import dataclass, asdict
# Add project root to path
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
from training.vega.model import VegaModel, VegaConfig
from training.vega.isotope_index import IsotopeIndex
@dataclass
class IsotopePrediction:
"""Prediction for a single isotope."""
name: str
probability: float
activity_bq: float
present: bool
@dataclass
class SpectrumPrediction:
"""Full prediction results for a spectrum."""
isotopes: List[IsotopePrediction]
num_present: int
confidence: float
threshold_used: float
def to_dict(self) -> Dict:
"""Convert to dictionary."""
return {
'isotopes': [
{
'name': iso.name,
'probability': round(iso.probability, 4),
'activity_bq': round(iso.activity_bq, 2),
'present': iso.present
}
for iso in self.isotopes
],
'num_isotopes_detected': self.num_present,
'confidence': round(self.confidence, 4),
'threshold': self.threshold_used
}
def get_present_isotopes(self) -> List[IsotopePrediction]:
"""Get only isotopes predicted as present."""
return [iso for iso in self.isotopes if iso.present]
def summary(self) -> str:
"""Get a human-readable summary."""
present = self.get_present_isotopes()
if not present:
return "No isotopes detected above threshold"
lines = [f"Detected {len(present)} isotope(s):"]
for iso in sorted(present, key=lambda x: -x.probability):
lines.append(
f" - {iso.name}: {iso.probability*100:.1f}% confidence, "
f"{iso.activity_bq:.1f} Bq"
)
return "\n".join(lines)
class VegaInference:
"""
Inference engine for the Vega model.
Loads a trained model and provides methods for running predictions
on gamma spectra.
"""
def __init__(
self,
model_path: Union[str, Path],
isotope_index_path: Optional[Union[str, Path]] = None,
device: Optional[torch.device] = None
):
"""
Initialize the inference engine.
Args:
model_path: Path to the saved model checkpoint
isotope_index_path: Path to isotope index file. If None, will try
to find it in the same directory as the model.
device: Device to run inference on. If None, uses CUDA if available.
"""
self.model_path = Path(model_path)
# Determine device with CUDA compatibility test
if device is not None:
self.device = device
elif torch.cuda.is_available():
# Test if CUDA actually works (RTX 5090/Blackwell may not be compatible)
try:
test_tensor = torch.zeros(1, device='cuda')
_ = test_tensor + 1
self.device = torch.device('cuda')
print("Using CUDA for inference")
except RuntimeError as e:
if "no kernel image is available" in str(e):
print(f"CUDA device detected but not compatible (likely Blackwell arch)")
print("Falling back to CPU for inference")
self.device = torch.device('cpu')
else:
raise
else:
self.device = torch.device('cpu')
print("Using CPU for inference")
# Load checkpoint
print(f"Loading model from: {self.model_path}")
self.checkpoint = torch.load(self.model_path, map_location=self.device)
# Load model config
model_config_dict = self.checkpoint['model_config']
self.model_config = VegaConfig(**model_config_dict)
# Create and load model
self.model = VegaModel(self.model_config)
self.model.load_state_dict(self.checkpoint['model_state_dict'])
self.model = self.model.to(self.device)
self.model.eval()
# Load isotope index
if isotope_index_path is None:
# Try to find in same directory
isotope_index_path = self.model_path.parent / "vega_isotope_index.txt"
if Path(isotope_index_path).exists():
self.isotope_index = IsotopeIndex.load(Path(isotope_index_path))
else:
# Use default
from training.vega.isotope_index import get_default_isotope_index
self.isotope_index = get_default_isotope_index()
print("Warning: Using default isotope index")
print(f"Model loaded successfully!")
print(f"Device: {self.device}")
print(f"Isotopes: {self.isotope_index.num_isotopes}")
def preprocess_spectrum(
self,
spectrum: np.ndarray,
normalize: bool = True
) -> torch.Tensor:
"""
Preprocess a spectrum for inference.
Args:
spectrum: Input spectrum array. Can be:
- 1D: (channels,) - single spectrum
- 2D: (time, channels) - will be averaged over time
normalize: Whether to normalize to [0, 1]
Returns:
Preprocessed tensor ready for model
"""
# Handle 2D spectra
if spectrum.ndim == 2:
spectrum = spectrum.mean(axis=0)
# Normalize
if normalize and spectrum.max() > 0:
spectrum = spectrum / spectrum.max()
# Convert to tensor
tensor = torch.tensor(spectrum, dtype=torch.float32)
# Add batch dimension
tensor = tensor.unsqueeze(0)
return tensor.to(self.device)
@torch.no_grad()
def predict(
self,
spectrum: Union[np.ndarray, torch.Tensor],
threshold: float = 0.5,
return_all: bool = False
) -> SpectrumPrediction:
"""
Run inference on a spectrum.
Args:
spectrum: Input spectrum (numpy array or tensor)
threshold: Probability threshold for considering an isotope present
return_all: If True, include all isotopes in output. If False,
only include those above threshold.
Returns:
SpectrumPrediction with isotope predictions
"""
# Preprocess if numpy
if isinstance(spectrum, np.ndarray):
spectrum = self.preprocess_spectrum(spectrum)
# Run model (outputs logits)
logits, activities = self.model(spectrum)
# Apply sigmoid to get probabilities
probs = torch.sigmoid(logits)
# Convert to numpy
probs = probs.cpu().numpy()[0]
activities = activities.cpu().numpy()[0]
# Scale activities
activities = activities * self.model_config.max_activity_bq
# Create predictions
isotopes = []
for i in range(len(probs)):
prob = float(probs[i])
activity = float(activities[i])
present = prob >= threshold
if return_all or present:
isotopes.append(IsotopePrediction(
name=self.isotope_index.index_to_name(i),
probability=prob,
activity_bq=activity if present else 0.0,
present=present
))
# Calculate overall confidence (average of top predictions)
present_isotopes = [iso for iso in isotopes if iso.present]
if present_isotopes:
confidence = np.mean([iso.probability for iso in present_isotopes])
else:
confidence = 1.0 - probs.max() # Confidence in "background only"
return SpectrumPrediction(
isotopes=isotopes,
num_present=len(present_isotopes),
confidence=float(confidence),
threshold_used=threshold
)
def predict_batch(
self,
spectra: List[np.ndarray],
threshold: float = 0.5
) -> List[SpectrumPrediction]:
"""
Run inference on multiple spectra.
Args:
spectra: List of spectrum arrays
threshold: Probability threshold
Returns:
List of predictions
"""
return [self.predict(s, threshold) for s in spectra]
def predict_from_file(
self,
file_path: Union[str, Path],
threshold: float = 0.5
) -> SpectrumPrediction:
"""
Load a spectrum from a numpy file and run inference.
Args:
file_path: Path to .npy file
threshold: Probability threshold
Returns:
SpectrumPrediction
"""
spectrum = np.load(file_path)
return self.predict(spectrum, threshold)
def run_inference_demo(
model_path: str,
data_dir: str,
threshold: float = 0.5
):
"""
Demo function to run inference on test data.
Args:
model_path: Path to model checkpoint
data_dir: Path to data directory with spectra
threshold: Detection threshold
"""
# Initialize inference engine
inference = VegaInference(model_path)
# Find spectra files
data_path = Path(data_dir)
spectra_dir = data_path / "spectra"
if not spectra_dir.exists():
print(f"Spectra directory not found: {spectra_dir}")
return
# Load labels for comparison
labels_path = data_path / "labels.json"
with open(labels_path, 'r') as f:
labels = json.load(f)
print("\n" + "=" * 70)
print("VEGA INFERENCE DEMO")
print("=" * 70)
# Process each spectrum
npy_files = list(spectra_dir.glob("*.npy"))
print(f"\nFound {len(npy_files)} spectra to process\n")
for npy_file in npy_files:
# Extract sample ID from filename
sample_id = npy_file.stem.replace("spectrum_", "")
# Get ground truth
if sample_id in labels['samples']:
ground_truth = labels['samples'][sample_id]
true_isotopes = ground_truth['isotopes']
true_activities = ground_truth.get('source_activities_bq', {})
else:
true_isotopes = []
true_activities = {}
# Run prediction
prediction = inference.predict_from_file(npy_file, threshold=threshold)
# Display results
print("-" * 70)
print(f"Sample: {sample_id}")
print(f"Ground Truth Isotopes: {true_isotopes if true_isotopes else 'Background only'}")
if true_activities:
activities_str = ", ".join(
f"{k}: {v:.1f} Bq" for k, v in true_activities.items()
)
print(f"Ground Truth Activities: {activities_str}")
print(f"\nPrediction:")
print(prediction.summary())
# Compare
predicted_names = {iso.name for iso in prediction.get_present_isotopes()}
true_names = set(true_isotopes)
correct = predicted_names & true_names
missed = true_names - predicted_names
false_positives = predicted_names - true_names
if correct:
print(f"\n✓ Correctly identified: {correct}")
if missed:
print(f"✗ Missed: {missed}")
if false_positives:
print(f"! False positives: {false_positives}")
print()
print("=" * 70)
print("Inference complete!")
print("=" * 70)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Run Vega model inference")
parser.add_argument(
"--model", "-m",
type=str,
default="models/vega_best.pt",
help="Path to model checkpoint"
)
parser.add_argument(
"--data", "-d",
type=str,
default="O:/master_data_collection/isotopev2",
help="Path to data directory"
)
parser.add_argument(
"--threshold", "-t",
type=float,
default=0.5,
help="Detection threshold (0-1)"
)
args = parser.parse_args()
# Make paths absolute if needed
project_root = Path(__file__).parent.parent
model_path = args.model if Path(args.model).is_absolute() else project_root / args.model
data_path = args.data if Path(args.data).is_absolute() else project_root / args.data
run_inference_demo(str(model_path), str(data_path), args.threshold)

922
train/vega_ml/inference/vega_portable_inference.py Normal file
View File

@ -0,0 +1,922 @@
#!/usr/bin/env python
"""
================================================================================
VEGA PORTABLE INFERENCE - Self-Contained Isotope Identification
================================================================================
This is a FULLY SELF-CONTAINED inference script for the Vega gamma spectrum
isotope identification model. You only need:
1. This Python file (vega_portable_inference.py)
2. A trained model checkpoint (.pt file)
3. PyTorch, NumPy installed
NO other project files are required. The model architecture, isotope index,
and sample data are all embedded in this file.
================================================================================
USAGE EXAMPLES:
================================================================================
1. Basic inference with embedded sample data:
python vega_portable_inference.py --model path/to/vega_best.pt
2. Inference on a specific spectrum file:
python vega_portable_inference.py --model vega_best.pt --spectrum my_spectrum.npy
3. Programmatic usage:
from vega_portable_inference import VegaInference, create_sample_spectrum
inference = VegaInference("vega_best.pt")
spectrum = create_sample_spectrum("Cs-137", activity_bq=100)
result = inference.predict(spectrum)
print(result.summary())
================================================================================
INPUT FORMAT:
================================================================================
The model expects gamma spectra in the following format:
- NumPy array, shape: (1023,) for single spectrum OR (N, 1023) for time series
- Values: Counts per channel (will be normalized automatically)
- Energy range: 20 keV to 3000 keV across 1023 channels
- Channel i corresponds to energy: E_i = 20 + i * (3000 - 20) / 1023 keV
If you have a 2D time-series spectrum (N intervals × 1023 channels), it will
be averaged over time automatically.
================================================================================
OUTPUT FORMAT:
================================================================================
The model returns a SpectrumPrediction object with:
- isotopes: List of IsotopePrediction objects, each containing:
- name: Isotope name (e.g., "Cs-137")
- probability: Detection confidence [0, 1]
- activity_bq: Estimated activity in Becquerels
- present: Boolean, True if probability >= threshold
- num_present: Count of detected isotopes
- confidence: Overall prediction confidence
- threshold_used: Detection threshold that was applied
Methods:
- .summary() - Human-readable text summary
- .to_dict() - JSON-serializable dictionary
- .get_present_isotopes() - List of only detected isotopes
================================================================================
MODEL ARCHITECTURE:
================================================================================
Vega uses a CNN-FCNN (Convolutional + Fully Connected Neural Network):
Input (1023 channels)
↓
ConvBlock 1: Conv1d(1→64) → BN → LeakyReLU → Conv1d(64→64) → BN → LeakyReLU → MaxPool → Dropout
↓
ConvBlock 2: Conv1d(64→128) → BN → LeakyReLU → Conv1d(128→128) → BN → LeakyReLU → MaxPool → Dropout
↓
ConvBlock 3: Conv1d(128→256) → BN → LeakyReLU → Conv1d(256→256) → BN → LeakyReLU → MaxPool → Dropout
↓
Flatten
↓
FC Classifier: Linear(→512) → BN → LeakyReLU → Dropout → Linear(→256) → BN → LeakyReLU → Dropout → Linear(→82)
↓ ↓
Sigmoid (multi-label isotope presence) FC Regressor: → ReLU (activity Bq)
Outputs:
- 82 isotope presence probabilities (multi-label classification)
- 82 activity estimates in Bq (regression)
================================================================================
SUPPORTED ISOTOPES (82 total):
================================================================================
CALIBRATION: Am-241, Ba-133, Cs-137, Co-57, Co-60, Eu-152, Na-22, Mn-54
MEDICAL: Tc-99m, F-18, Ga-67, Ga-68, In-111, I-123, I-125, Tl-201, Lu-177
INDUSTRIAL: Ir-192, Se-75, Cd-109, I-131, Y-90
NATURAL: K-40, Ra-226, Th-232, U-238, Rn-222
FALLOUT: Cs-134, Sr-90, Zr-95, Nb-95, Ru-103, Ru-106, Ce-141, Ce-144
DECAY CHAINS: Full U-238, Th-232, U-235 series
================================================================================
REQUIREMENTS:
================================================================================
pip install torch numpy
Optional (for visualization):
pip install matplotlib scipy
================================================================================
"""
import sys
import json
import math
import numpy as np
import torch
import torch.nn as nn
from pathlib import Path
from dataclasses import dataclass, field, asdict
from typing import Dict, List, Optional, Tuple, Union
# =============================================================================
# ISOTOPE DATABASE (Embedded - No external dependencies)
# =============================================================================
# Complete list of 82 isotopes supported by the model (alphabetically sorted)
ISOTOPE_NAMES = [
"Ac-228", "Ag-110m", "Am-241", "Ba-133", "Be-7", "Bi-207", "Bi-211",
"Bi-212", "Bi-214", "C-14", "Cd-109", "Ce-141", "Ce-144", "Co-57",
"Co-60", "Cr-51", "Cs-134", "Cs-137", "Eu-152", "Eu-154", "F-18",
"Fe-59", "Ga-67", "Ga-68", "H-3", "I-123", "I-125", "I-131", "In-111",
"Ir-192", "K-40", "Lu-177", "Mn-54", "Na-22", "Nb-95", "Pa-231",
"Pa-234m", "Pb-210", "Pb-211", "Pb-212", "Pb-214", "Po-210", "Ra-223",
"Ra-224", "Ra-226", "Rn-219", "Rn-222", "Ru-103", "Ru-106", "Sb-124",
"Sb-125", "Se-75", "Sn-113", "Sr-85", "Sr-90", "Tc-99m", "Th-227",
"Th-228", "Th-230", "Th-232", "Th-234", "Tl-201", "Tl-208", "U-234",
"U-235", "U-238", "Y-90", "Zn-65", "Zr-95",
# Additional isotopes to reach 82
"Ba-140", "Br-82", "Ca-45", "Ca-47", "Cf-252", "Cl-36", "Cm-244",
"Cu-64", "Gd-153", "Hg-203", "Np-237", "P-32", "Pu-239"
]
# Gamma emission lines (keV) and branching ratios for key isotopes
# Format: {isotope: [(energy_keV, branching_ratio), ...]}
GAMMA_LINES = {
"Am-241": [(59.54, 0.359)],
"Ba-133": [(81.0, 0.329), (276.4, 0.071), (302.9, 0.183), (356.0, 0.620), (383.8, 0.089)],
"Cs-137": [(661.7, 0.851)],
"Co-57": [(122.1, 0.856), (136.5, 0.107)],
"Co-60": [(1173.2, 0.999), (1332.5, 0.999)],
"Eu-152": [(121.8, 0.284), (344.3, 0.265), (778.9, 0.129), (964.1, 0.146), (1112.1, 0.136), (1408.0, 0.210)],
"Na-22": [(511.0, 1.798), (1274.5, 0.999)],
"Mn-54": [(834.8, 0.9998)],
"K-40": [(1460.8, 0.107)],
"Ra-226": [(186.2, 0.036)],
"Pb-214": [(295.2, 0.192), (351.9, 0.371)],
"Bi-214": [(609.3, 0.461), (1120.3, 0.150), (1764.5, 0.154)],
"Pb-212": [(238.6, 0.436)],
"Tl-208": [(583.2, 0.845), (2614.5, 0.99)],
"Ac-228": [(338.3, 0.113), (911.2, 0.258), (969.0, 0.158)],
"I-131": [(364.5, 0.817), (637.0, 0.072)],
"Tc-99m": [(140.5, 0.890)],
"F-18": [(511.0, 1.934)],
"Ir-192": [(296.0, 0.287), (308.5, 0.300), (316.5, 0.828), (468.1, 0.478)],
"Th-232": [(63.8, 0.0026)],
"U-238": [(49.6, 0.064), (113.5, 0.017)],
}
class IsotopeIndex:
"""
Maps isotope names to model output indices and vice versa.
The index is alphabetically sorted for deterministic ordering.
"""
def __init__(self, isotope_names: Optional[List[str]] = None):
if isotope_names is None:
isotope_names = ISOTOPE_NAMES
self._isotope_names = sorted(isotope_names)
self._name_to_idx = {name: idx for idx, name in enumerate(self._isotope_names)}
self._idx_to_name = {idx: name for idx, name in enumerate(self._isotope_names)}
@property
def num_isotopes(self) -> int:
return len(self._isotope_names)
@property
def isotope_names(self) -> List[str]:
return self._isotope_names.copy()
def name_to_index(self, name: str) -> int:
if name not in self._name_to_idx:
raise KeyError(f"Isotope '{name}' not in index. Available: {self._isotope_names[:5]}...")
return self._name_to_idx[name]
def index_to_name(self, idx: int) -> str:
if idx not in self._idx_to_name:
raise KeyError(f"Index {idx} out of range [0, {self.num_isotopes-1}]")
return self._idx_to_name[idx]
def __len__(self) -> int:
return self.num_isotopes
def __repr__(self) -> str:
return f"IsotopeIndex(num_isotopes={self.num_isotopes})"
@classmethod
def load(cls, path: Path) -> 'IsotopeIndex':
"""Load from a text file (one isotope per line)."""
with open(path, 'r') as f:
names = [line.strip() for line in f if line.strip()]
return cls(names)
def save(self, path: Path):
"""Save to a text file."""
with open(path, 'w') as f:
for name in self._isotope_names:
f.write(f"{name}\n")
# =============================================================================
# MODEL ARCHITECTURE (Embedded - No external dependencies)
# =============================================================================
@dataclass
class VegaConfig:
"""Configuration for the Vega model architecture."""
# Input
num_channels: int = 1023 # Energy channels in spectrum
num_isotopes: int = 82 # Output classes
# CNN backbone
conv_channels: List[int] = field(default_factory=lambda: [64, 128, 256])
conv_kernel_size: int = 7
pool_size: int = 2
# Classifier head
fc_hidden_dims: List[int] = field(default_factory=lambda: [512, 256])
# Regularization
dropout_rate: float = 0.3
spatial_dropout_rate: float = 0.1
leaky_relu_slope: float = 0.1
# Loss weights (not used in inference)
classification_weight: float = 1.0
regression_weight: float = 0.1
max_activity_bq: float = 1000.0
class ConvBlock(nn.Module):
"""
CNN block: Conv → BN → LeakyReLU → Conv → BN → LeakyReLU → MaxPool → Dropout
"""
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int = 7,
pool_size: int = 2,
dropout_rate: float = 0.1,
leaky_slope: float = 0.1
):
super().__init__()
self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size, padding=kernel_size // 2)
self.bn1 = nn.BatchNorm1d(out_channels)
self.act1 = nn.LeakyReLU(leaky_slope)
self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size, padding=kernel_size // 2)
self.bn2 = nn.BatchNorm1d(out_channels)
self.act2 = nn.LeakyReLU(leaky_slope)
self.pool = nn.MaxPool1d(pool_size)
self.dropout = nn.Dropout1d(dropout_rate)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.act1(self.bn1(self.conv1(x)))
x = self.act2(self.bn2(self.conv2(x)))
x = self.dropout(self.pool(x))
return x
class VegaModel(nn.Module):
"""
Vega: CNN-FCNN for Multi-Label Isotope Classification + Activity Regression
Takes a 1D gamma spectrum and outputs:
- 82 isotope presence logits (use sigmoid for probabilities)
- 82 activity estimates (in Bq, scaled by max_activity_bq)
"""
def __init__(self, config: VegaConfig):
super().__init__()
self.config = config
# Build CNN backbone
self.backbone = self._build_backbone()
# Calculate flattened size
self._flat_size = self._calculate_flat_size()
# Classification head (multi-label)
self.classifier = self._build_classifier()
# Regression head (activity)
self.regressor = self._build_regressor()
def _build_backbone(self) -> nn.Sequential:
layers = []
in_ch = 1
for out_ch in self.config.conv_channels:
layers.append(ConvBlock(
in_ch, out_ch,
kernel_size=self.config.conv_kernel_size,
pool_size=self.config.pool_size,
dropout_rate=self.config.spatial_dropout_rate,
leaky_slope=self.config.leaky_relu_slope
))
in_ch = out_ch
return nn.Sequential(*layers)
def _calculate_flat_size(self) -> int:
with torch.no_grad():
x = torch.zeros(1, 1, self.config.num_channels)
x = self.backbone(x)
return x.view(1, -1).size(1)
def _build_classifier(self) -> nn.Sequential:
layers = []
in_dim = self._flat_size
for hidden_dim in self.config.fc_hidden_dims:
layers.extend([
nn.Linear(in_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.LeakyReLU(self.config.leaky_relu_slope),
nn.Dropout(self.config.dropout_rate)
])
in_dim = hidden_dim
layers.append(nn.Linear(in_dim, self.config.num_isotopes))
return nn.Sequential(*layers)
def _build_regressor(self) -> nn.Sequential:
layers = []
in_dim = self._flat_size
for hidden_dim in self.config.fc_hidden_dims:
layers.extend([
nn.Linear(in_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.LeakyReLU(self.config.leaky_relu_slope),
nn.Dropout(self.config.dropout_rate)
])
in_dim = hidden_dim
layers.extend([
nn.Linear(in_dim, self.config.num_isotopes),
nn.ReLU() # Activities must be positive
])
return nn.Sequential(*layers)
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
# Ensure input shape is (batch, 1, channels)
if x.dim() == 2:
x = x.unsqueeze(1)
# Feature extraction
features = self.backbone(x)
features = features.view(features.size(0), -1)
# Dual outputs
logits = self.classifier(features)
activities = self.regressor(features)
return logits, activities
# =============================================================================
# PREDICTION DATA CLASSES
# =============================================================================
@dataclass
class IsotopePrediction:
"""Prediction result for a single isotope."""
name: str
probability: float
activity_bq: float
present: bool
@dataclass
class SpectrumPrediction:
"""Complete prediction results for a spectrum."""
isotopes: List[IsotopePrediction]
num_present: int
confidence: float
threshold_used: float
def to_dict(self) -> Dict:
"""Convert to JSON-serializable dictionary."""
return {
'isotopes': [
{
'name': iso.name,
'probability': round(iso.probability, 4),
'activity_bq': round(iso.activity_bq, 2),
'present': iso.present
}
for iso in self.isotopes
],
'num_isotopes_detected': self.num_present,
'confidence': round(self.confidence, 4),
'threshold': self.threshold_used
}
def get_present_isotopes(self) -> List[IsotopePrediction]:
"""Get only isotopes predicted as present."""
return [iso for iso in self.isotopes if iso.present]
def summary(self) -> str:
"""Human-readable summary of predictions."""
present = self.get_present_isotopes()
if not present:
return "No isotopes detected above threshold"
lines = [f"Detected {len(present)} isotope(s):"]
for iso in sorted(present, key=lambda x: -x.probability):
lines.append(
f" • {iso.name}: {iso.probability*100:.1f}% confidence, "
f"{iso.activity_bq:.1f} Bq estimated activity"
)
return "\n".join(lines)
def to_json(self, indent: int = 2) -> str:
"""Convert to JSON string."""
return json.dumps(self.to_dict(), indent=indent)
# =============================================================================
# INFERENCE ENGINE
# =============================================================================
class VegaInference:
"""
Inference engine for the Vega isotope identification model.
Example usage:
inference = VegaInference("vega_best.pt")
spectrum = np.load("my_spectrum.npy")
result = inference.predict(spectrum, threshold=0.5)
print(result.summary())
"""
def __init__(
self,
model_path: Union[str, Path],
isotope_index: Optional[IsotopeIndex] = None,
device: Optional[torch.device] = None
):
"""
Initialize the inference engine.
Args:
model_path: Path to saved .pt model checkpoint
isotope_index: Optional custom isotope index. Uses default if None.
device: Compute device. Auto-detects CUDA if available.
"""
self.model_path = Path(model_path)
# Device selection with CUDA compatibility check
if device is not None:
self.device = device
elif torch.cuda.is_available():
try:
# Test CUDA actually works (some GPUs may not be compatible)
_ = torch.zeros(1, device='cuda') + 1
self.device = torch.device('cuda')
except RuntimeError:
self.device = torch.device('cpu')
else:
self.device = torch.device('cpu')
# Load checkpoint
print(f"Loading model from: {self.model_path}")
self.checkpoint = torch.load(self.model_path, map_location=self.device, weights_only=False)
# Load model config
if 'model_config' in self.checkpoint:
config_dict = self.checkpoint['model_config']
self.model_config = VegaConfig(**config_dict)
elif 'params' in self.checkpoint:
# Handle Optuna-trained models
params = self.checkpoint['params']
self.model_config = VegaConfig(
conv_channels=params.get('conv_channels', [64, 128, 256]),
conv_kernel_size=params.get('conv_kernel_size', 7),
pool_size=params.get('pool_size', 2),
fc_hidden_dims=params.get('fc_hidden_dims', [512, 256]),
dropout_rate=params.get('dropout_rate', 0.3),
spatial_dropout_rate=params.get('spatial_dropout_rate', 0.1),
leaky_relu_slope=params.get('leaky_relu_slope', 0.1)
)
else:
# Use defaults
self.model_config = VegaConfig()
# Create and load model
self.model = VegaModel(self.model_config)
self.model.load_state_dict(self.checkpoint['model_state_dict'])
self.model = self.model.to(self.device)
self.model.eval()
# Set isotope index
self.isotope_index = isotope_index or IsotopeIndex()
print(f"✓ Model loaded successfully")
print(f" Device: {self.device}")
print(f" Isotopes: {self.isotope_index.num_isotopes}")
print(f" Architecture: CNN{self.model_config.conv_channels} → FC{self.model_config.fc_hidden_dims}")
def preprocess(self, spectrum: np.ndarray, normalize: bool = True) -> torch.Tensor:
"""
Preprocess spectrum for model input.
Args:
spectrum: Input array, shape (1023,) or (N, 1023)
normalize: Normalize to [0, 1] range
Returns:
Tensor ready for model, shape (1, 1023)
"""
# Average time series if 2D
if spectrum.ndim == 2:
spectrum = spectrum.mean(axis=0)
# Normalize
if normalize and spectrum.max() > 0:
spectrum = spectrum / spectrum.max()
# To tensor with batch dimension
tensor = torch.tensor(spectrum, dtype=torch.float32).unsqueeze(0)
return tensor.to(self.device)
@torch.no_grad()
def predict(
self,
spectrum: Union[np.ndarray, torch.Tensor],
threshold: float = 0.5,
return_all: bool = False
) -> SpectrumPrediction:
"""
Run inference on a gamma spectrum.
Args:
spectrum: Input spectrum (numpy array or tensor)
threshold: Probability threshold for detection (0-1)
return_all: If True, include all 82 isotopes. If False, only detected ones.
Returns:
SpectrumPrediction with detected isotopes and activities
"""
# Preprocess
if isinstance(spectrum, np.ndarray):
spectrum = self.preprocess(spectrum)
# Run model
logits, activities = self.model(spectrum)
# Convert to probabilities
probs = torch.sigmoid(logits).cpu().numpy()[0]
activities = activities.cpu().numpy()[0] * self.model_config.max_activity_bq
# Build predictions
isotopes = []
for i in range(len(probs)):
prob = float(probs[i])
activity = float(activities[i])
present = prob >= threshold
if return_all or present:
isotopes.append(IsotopePrediction(
name=self.isotope_index.index_to_name(i),
probability=prob,
activity_bq=activity if present else 0.0,
present=present
))
# Calculate confidence
present_isotopes = [iso for iso in isotopes if iso.present]
if present_isotopes:
confidence = np.mean([iso.probability for iso in present_isotopes])
else:
confidence = 1.0 - probs.max()
return SpectrumPrediction(
isotopes=isotopes,
num_present=len(present_isotopes),
confidence=float(confidence),
threshold_used=threshold
)
def predict_from_file(
self,
file_path: Union[str, Path],
threshold: float = 0.5
) -> SpectrumPrediction:
"""Load spectrum from .npy file and run inference."""
spectrum = np.load(file_path)
return self.predict(spectrum, threshold)
def predict_batch(
self,
spectra: List[np.ndarray],
threshold: float = 0.5
) -> List[SpectrumPrediction]:
"""Run inference on multiple spectra."""
return [self.predict(s, threshold) for s in spectra]
# =============================================================================
# SAMPLE SPECTRUM GENERATOR (For testing without real data)
# =============================================================================
def energy_to_channel(energy_kev: float, num_channels: int = 1023) -> int:
"""Convert energy in keV to channel index."""
e_min, e_max = 20.0, 3000.0
channel = int((energy_kev - e_min) / (e_max - e_min) * num_channels)
return max(0, min(num_channels - 1, channel))
def channel_to_energy(channel: int, num_channels: int = 1023) -> float:
"""Convert channel index to energy in keV."""
e_min, e_max = 20.0, 3000.0
return e_min + channel * (e_max - e_min) / num_channels
def create_sample_spectrum(
isotope: str = "Cs-137",
activity_bq: float = 100.0,
duration_seconds: float = 300.0,
add_background: bool = True,
add_noise: bool = True,
detector_fwhm_percent: float = 8.5,
seed: Optional[int] = None
) -> np.ndarray:
"""
Generate a synthetic gamma spectrum for testing.
This creates a realistic-looking spectrum with Gaussian peaks at the
characteristic gamma energies of the specified isotope.
Args:
isotope: Isotope name (e.g., "Cs-137", "Co-60", "Na-22")
activity_bq: Source activity in Becquerels
duration_seconds: Measurement duration
add_background: Add environmental background
add_noise: Apply Poisson counting statistics
detector_fwhm_percent: Detector resolution at 662 keV (%)
seed: Random seed for reproducibility
Returns:
1D numpy array of shape (1023,) with counts per channel
"""
if seed is not None:
np.random.seed(seed)
num_channels = 1023
spectrum = np.zeros(num_channels)
# Get gamma lines for the isotope
if isotope in GAMMA_LINES:
gamma_lines = GAMMA_LINES[isotope]
else:
# Use Cs-137 as fallback
print(f"Warning: No gamma lines for {isotope}, using Cs-137")
gamma_lines = GAMMA_LINES["Cs-137"]
# Add peaks for each gamma line
for energy_kev, branching_ratio in gamma_lines:
# Calculate FWHM at this energy (scales with sqrt of energy)
fwhm_kev = (detector_fwhm_percent / 100.0) * 662.0 * math.sqrt(energy_kev / 662.0)
sigma_kev = fwhm_kev / 2.355
# Expected counts
efficiency = 0.1 * math.exp(-energy_kev / 500.0) # Simplified efficiency
expected_counts = activity_bq * duration_seconds * branching_ratio * efficiency
# Add Gaussian peak
center_channel = energy_to_channel(energy_kev)
sigma_channels = sigma_kev / ((3000 - 20) / num_channels)
for ch in range(num_channels):
energy = channel_to_energy(ch)
peak = expected_counts * math.exp(-0.5 * ((energy - energy_kev) / sigma_kev) ** 2)
spectrum[ch] += peak
# Add background continuum
if add_background:
# Exponential continuum
for ch in range(num_channels):
energy = channel_to_energy(ch)
bg = 50.0 * duration_seconds * math.exp(-energy / 300.0) / 300.0
spectrum[ch] += bg
# K-40 environmental background
k40_energy = 1460.8
k40_fwhm = (detector_fwhm_percent / 100.0) * 662.0 * math.sqrt(k40_energy / 662.0)
k40_sigma = k40_fwhm / 2.355
k40_counts = 10.0 * duration_seconds # Low activity environmental
for ch in range(num_channels):
energy = channel_to_energy(ch)
peak = k40_counts * math.exp(-0.5 * ((energy - k40_energy) / k40_sigma) ** 2)
spectrum[ch] += peak
# Apply Poisson noise
if add_noise:
spectrum = np.maximum(spectrum, 0)
spectrum = np.random.poisson(spectrum.astype(int)).astype(float)
return spectrum
def create_sample_spectra_batch() -> Dict[str, np.ndarray]:
"""
Create a batch of sample spectra for different isotopes.
Returns:
Dictionary mapping isotope names to their sample spectra
"""
samples = {}
# Common calibration isotopes
for isotope in ["Cs-137", "Co-60", "Na-22", "Ba-133", "Am-241", "Eu-152"]:
if isotope in GAMMA_LINES:
samples[isotope] = create_sample_spectrum(
isotope=isotope,
activity_bq=100.0,
duration_seconds=300.0,
seed=hash(isotope) % 2**32
)
# Background only
samples["Background"] = create_sample_spectrum(
isotope="Cs-137", # Will be overwritten by background
activity_bq=0.0, # No source
duration_seconds=300.0,
add_background=True,
seed=12345
)
return samples
# =============================================================================
# DEMONSTRATION FUNCTIONS
# =============================================================================
def run_demo(model_path: str, threshold: float = 0.5):
"""
Run a complete demonstration of the Vega inference system.
Args:
model_path: Path to trained model checkpoint
threshold: Detection threshold (0-1)
"""
print("\n" + "=" * 70)
print("VEGA ISOTOPE IDENTIFICATION - INFERENCE DEMONSTRATION")
print("=" * 70)
# Load model
print("\n[1] Loading Model")
print("-" * 70)
inference = VegaInference(model_path)
# Generate sample spectra
print("\n[2] Generating Sample Spectra")
print("-" * 70)
samples = create_sample_spectra_batch()
print(f"Generated {len(samples)} sample spectra:")
for name in samples:
print(f" • {name}")
# Run inference on each
print("\n[3] Running Inference")
print("-" * 70)
for name, spectrum in samples.items():
print(f"\n{'─' * 70}")
print(f"Sample: {name}")
print(f"Spectrum shape: {spectrum.shape}")
print(f"Spectrum range: [{spectrum.min():.1f}, {spectrum.max():.1f}]")
# Run prediction
result = inference.predict(spectrum, threshold=threshold)
print(f"\nPrediction (threshold={threshold}):")
print(result.summary())
# Show top 5 probabilities even if below threshold
print("\nTop 5 isotope probabilities:")
all_result = inference.predict(spectrum, threshold=0.0, return_all=True)
sorted_iso = sorted(all_result.isotopes, key=lambda x: -x.probability)[:5]
for iso in sorted_iso:
marker = "✓" if iso.probability >= threshold else " "
print(f" {marker} {iso.name}: {iso.probability*100:.2f}%")
# Show JSON output format
print("\n[4] JSON Output Format Example")
print("-" * 70)
sample_result = inference.predict(samples["Cs-137"], threshold=threshold)
print(sample_result.to_json())
print("\n" + "=" * 70)
print("DEMONSTRATION COMPLETE")
print("=" * 70)
def run_single_inference(model_path: str, spectrum_path: str, threshold: float = 0.5):
"""
Run inference on a single spectrum file.
Args:
model_path: Path to trained model
spectrum_path: Path to .npy spectrum file
threshold: Detection threshold
"""
print(f"\nLoading model from: {model_path}")
inference = VegaInference(model_path)
print(f"Loading spectrum from: {spectrum_path}")
spectrum = np.load(spectrum_path)
print(f"Spectrum shape: {spectrum.shape}")
print(f"\nRunning inference (threshold={threshold})...")
result = inference.predict(spectrum, threshold=threshold)
print("\n" + "=" * 60)
print("PREDICTION RESULTS")
print("=" * 60)
print(result.summary())
print("=" * 60)
return result
# =============================================================================
# MAIN ENTRY POINT
# =============================================================================
def main():
"""Main entry point for command-line usage."""
import argparse
parser = argparse.ArgumentParser(
description="Vega Portable Inference - Gamma Spectrum Isotope Identification",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Run demo with sample spectra
python vega_portable_inference.py --model vega_best.pt
# Analyze a specific spectrum file
python vega_portable_inference.py --model vega_best.pt --spectrum my_data.npy
# Use lower threshold for higher recall
python vega_portable_inference.py --model vega_best.pt --threshold 0.3
"""
)
parser.add_argument(
"--model", "-m",
type=str,
required=True,
help="Path to trained Vega model checkpoint (.pt file)"
)
parser.add_argument(
"--spectrum", "-s",
type=str,
default=None,
help="Path to spectrum file (.npy). If not provided, runs demo with synthetic spectra."
)
parser.add_argument(
"--threshold", "-t",
type=float,
default=0.5,
help="Detection threshold (0-1). Lower = more sensitive, higher = more specific. Default: 0.5"
)
parser.add_argument(
"--json",
action="store_true",
help="Output results in JSON format"
)
args = parser.parse_args()
if args.spectrum:
# Single file inference
result = run_single_inference(args.model, args.spectrum, args.threshold)
if args.json:
print("\nJSON Output:")
print(result.to_json())
else:
# Demo mode
run_demo(args.model, args.threshold)
return 0
if __name__ == "__main__":
sys.exit(main())

879
train/vega_ml/inference/vega_portable_inference_2d.py Normal file
View File

@ -0,0 +1,879 @@
#!/usr/bin/env python
"""
================================================================================
VEGA 2D PORTABLE INFERENCE - Self-Contained Isotope Identification
================================================================================
This is a FULLY SELF-CONTAINED inference script for the Vega 2D gamma spectrum
isotope identification model. You only need:
1. This Python file (vega_portable_inference_2d.py)
2. A trained 2D model checkpoint (.pt file)
3. PyTorch, NumPy installed
NO other project files are required. The model architecture, isotope index,
and sample data generator are all embedded in this file.
================================================================================
USAGE EXAMPLES:
================================================================================
1. Basic inference with embedded sample data:
python vega_portable_inference_2d.py --model vega_2d_best.pt
2. Inference on a specific spectrum file:
python vega_portable_inference_2d.py --model vega_2d_best.pt --spectrum my_spectrum.npy
3. Programmatic usage:
from vega_portable_inference_2d import Vega2DInference, create_sample_spectrum_2d
inference = Vega2DInference("vega_2d_best.pt")
spectrum = create_sample_spectrum_2d("Cs-137", activity_bq=100)
result = inference.predict(spectrum)
print(result.summary())
================================================================================
INPUT FORMAT:
================================================================================
The 2D model expects gamma spectra in the following format:
- NumPy array, shape: (60, 1023) for 60 time intervals × 1023 channels
- Values: Counts per channel per time interval (will be normalized automatically)
- Energy range: 20 keV to 3000 keV across 1023 channels
- Time: 60 one-second intervals (1 minute total measurement)
If your spectrum has different time dimensions, it will be padded or truncated
to 60 intervals automatically.
================================================================================
OUTPUT FORMAT:
================================================================================
The model returns a SpectrumPrediction object with:
- isotopes: List of IsotopePrediction objects, each containing:
- name: Isotope name (e.g., "Cs-137")
- probability: Detection confidence [0, 1]
- activity_bq: Estimated activity in Becquerels
- present: Boolean, True if probability >= threshold
- num_present: Count of detected isotopes
- confidence: Overall prediction confidence
- threshold_used: Detection threshold that was applied
Methods:
- .summary() - Human-readable text summary
- .to_dict() - JSON-serializable dictionary
- .get_present_isotopes() - List of only detected isotopes
================================================================================
MODEL ARCHITECTURE (2D CNN):
================================================================================
Vega 2D uses 2D convolutions to process time × energy spectral images:
Input (1, 60, 1023) - single channel image
↓
ConvBlock 1: Conv2d(1→32, k=3×7) → BN → LeakyReLU → Conv2d → BN → LeakyReLU → MaxPool2d → Dropout
↓
ConvBlock 2: Conv2d(32→64, k=3×7) → BN → LeakyReLU → Conv2d → BN → LeakyReLU → MaxPool2d → Dropout
↓
ConvBlock 3: Conv2d(64→128, k=3×7) → BN → LeakyReLU → Conv2d → BN → LeakyReLU → MaxPool2d → Dropout
↓
Flatten (113,792 features)
↓
FC: Linear(→512) → BN → LeakyReLU → Dropout → Linear(→256) → BN → LeakyReLU → Dropout
↓
Classifier: Linear(→82) [isotope logits]
Regressor: Linear(→82) → ReLU [activity Bq]
Outputs:
- 82 isotope presence probabilities (multi-label classification)
- 82 activity estimates in Bq (regression)
Total parameters: ~59 million
================================================================================
SUPPORTED ISOTOPES (82 total):
================================================================================
CALIBRATION: Am-241, Ba-133, Cs-137, Co-57, Co-60, Eu-152, Na-22, Mn-54
MEDICAL: Tc-99m, F-18, Ga-67, Ga-68, In-111, I-123, I-125, Tl-201, Lu-177
INDUSTRIAL: Ir-192, Se-75, Cd-109, I-131, Y-90
NATURAL: K-40, Ra-226, Th-232, U-238, Rn-222
FALLOUT: Cs-134, Sr-90, Zr-95, Nb-95, Ru-103, Ru-106, Ce-141, Ce-144
DECAY CHAINS: Full U-238, Th-232, U-235 series
================================================================================
REQUIREMENTS:
================================================================================
pip install torch numpy
================================================================================
"""
import sys
import json
import math
import numpy as np
import torch
import torch.nn as nn
from pathlib import Path
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Union
# =============================================================================
# ISOTOPE DATABASE (Embedded - No external dependencies)
# =============================================================================
# Complete list of 82 isotopes supported by the model (alphabetically sorted)
ISOTOPE_NAMES = [
"Ac-228", "Ag-110m", "Am-241", "Ba-133", "Be-7", "Bi-207", "Bi-211",
"Bi-212", "Bi-214", "C-14", "Cd-109", "Ce-141", "Ce-144", "Co-57",
"Co-60", "Cr-51", "Cs-134", "Cs-137", "Eu-152", "Eu-154", "F-18",
"Fe-59", "Ga-67", "Ga-68", "H-3", "I-123", "I-125", "I-131", "In-111",
"Ir-192", "K-40", "Lu-177", "Mn-54", "Na-22", "Nb-95", "Pa-231",
"Pa-234m", "Pb-210", "Pb-211", "Pb-212", "Pb-214", "Po-210", "Ra-223",
"Ra-224", "Ra-226", "Rn-219", "Rn-222", "Ru-103", "Ru-106", "Sb-124",
"Sb-125", "Se-75", "Sn-113", "Sr-85", "Sr-90", "Tc-99m", "Th-227",
"Th-228", "Th-230", "Th-232", "Th-234", "Tl-201", "Tl-208", "U-234",
"U-235", "U-238", "Y-90", "Zn-65", "Zr-95",
# Additional isotopes to reach 82
"Ba-140", "Br-82", "Ca-45", "Ca-47", "Cf-252", "Cl-36", "Cm-244",
"Cu-64", "Gd-153", "Hg-203", "Np-237", "P-32", "Pu-239"
]
# Gamma emission lines (keV) and branching ratios for key isotopes
GAMMA_LINES = {
"Am-241": [(59.54, 0.359)],
"Ba-133": [(81.0, 0.329), (276.4, 0.071), (302.9, 0.183), (356.0, 0.620), (383.8, 0.089)],
"Cs-137": [(661.7, 0.851)],
"Co-57": [(122.1, 0.856), (136.5, 0.107)],
"Co-60": [(1173.2, 0.999), (1332.5, 0.999)],
"Eu-152": [(121.8, 0.284), (344.3, 0.265), (778.9, 0.129), (964.1, 0.146), (1112.1, 0.136), (1408.0, 0.210)],
"Na-22": [(511.0, 1.798), (1274.5, 0.999)],
"Mn-54": [(834.8, 0.9998)],
"K-40": [(1460.8, 0.107)],
"Ra-226": [(186.2, 0.036)],
"Pb-214": [(295.2, 0.192), (351.9, 0.371)],
"Bi-214": [(609.3, 0.461), (1120.3, 0.150), (1764.5, 0.154)],
"Pb-212": [(238.6, 0.436)],
"Tl-208": [(583.2, 0.845), (2614.5, 0.99)],
"Ac-228": [(338.3, 0.113), (911.2, 0.258), (969.0, 0.158)],
"I-131": [(364.5, 0.817), (637.0, 0.072)],
"Tc-99m": [(140.5, 0.890)],
"F-18": [(511.0, 1.934)],
"Ir-192": [(296.0, 0.287), (308.5, 0.300), (316.5, 0.828), (468.1, 0.478)],
"Th-232": [(63.8, 0.0026)],
"U-238": [(49.6, 0.064), (113.5, 0.017)],
}
class IsotopeIndex:
"""Maps isotope names to model output indices and vice versa."""
def __init__(self, isotope_names: Optional[List[str]] = None):
if isotope_names is None:
isotope_names = ISOTOPE_NAMES
self._isotope_names = sorted(isotope_names)
self._name_to_idx = {name: idx for idx, name in enumerate(self._isotope_names)}
self._idx_to_name = {idx: name for idx, name in enumerate(self._isotope_names)}
@property
def num_isotopes(self) -> int:
return len(self._isotope_names)
@property
def isotope_names(self) -> List[str]:
return self._isotope_names.copy()
def name_to_index(self, name: str) -> int:
if name not in self._name_to_idx:
raise KeyError(f"Isotope '{name}' not in index")
return self._name_to_idx[name]
def index_to_name(self, idx: int) -> str:
if idx not in self._idx_to_name:
raise KeyError(f"Index {idx} out of range [0, {self.num_isotopes-1}]")
return self._idx_to_name[idx]
def __len__(self) -> int:
return self.num_isotopes
# =============================================================================
# 2D MODEL ARCHITECTURE (Embedded)
# =============================================================================
@dataclass
class Vega2DConfig:
"""Configuration for Vega 2D model."""
# Input dimensions
num_channels: int = 1023 # Energy channels
num_time_intervals: int = 60 # Fixed time dimension
# Output
num_isotopes: int = 82
# CNN architecture
conv_channels: List[int] = field(default_factory=lambda: [32, 64, 128])
kernel_size: Tuple[int, int] = (3, 7) # (time, energy)
pool_size: Tuple[int, int] = (2, 2)
# FC layers
fc_hidden_dims: List[int] = field(default_factory=lambda: [512, 256])
# Regularization
dropout_rate: float = 0.3
leaky_relu_slope: float = 0.01
# Activity scaling
max_activity_bq: float = 1000.0
class ConvBlock2D(nn.Module):
"""2D Convolutional block with BatchNorm, activation, pooling, and dropout."""
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: Tuple[int, int],
pool_size: Tuple[int, int],
dropout_rate: float,
leaky_relu_slope: float
):
super().__init__()
padding = (kernel_size[0] // 2, kernel_size[1] // 2)
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding)
self.bn1 = nn.BatchNorm2d(out_channels)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size, padding=padding)
self.bn2 = nn.BatchNorm2d(out_channels)
self.activation = nn.LeakyReLU(leaky_relu_slope)
self.pool = nn.MaxPool2d(pool_size)
self.dropout = nn.Dropout2d(dropout_rate)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.activation(self.bn1(self.conv1(x)))
x = self.activation(self.bn2(self.conv2(x)))
x = self.pool(x)
x = self.dropout(x)
return x
class Vega2DModel(nn.Module):
"""
2D CNN model for gamma spectrum isotope identification.
Treats spectra as images with time on one axis and energy channels on the other.
"""
def __init__(self, config: Vega2DConfig = None):
super().__init__()
self.config = config or Vega2DConfig()
# Build CNN backbone
self.conv_blocks = nn.ModuleList()
in_channels = 1
for out_channels in self.config.conv_channels:
self.conv_blocks.append(ConvBlock2D(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=self.config.kernel_size,
pool_size=self.config.pool_size,
dropout_rate=self.config.dropout_rate,
leaky_relu_slope=self.config.leaky_relu_slope
))
in_channels = out_channels
# Calculate flattened size
self.flat_size = self._calculate_flat_size()
# FC backbone
fc_layers = []
fc_in = self.flat_size
for fc_out in self.config.fc_hidden_dims:
fc_layers.extend([
nn.Linear(fc_in, fc_out),
nn.BatchNorm1d(fc_out),
nn.LeakyReLU(self.config.leaky_relu_slope),
nn.Dropout(self.config.dropout_rate)
])
fc_in = fc_out
self.fc_backbone = nn.Sequential(*fc_layers)
# Output heads
self.classifier = nn.Linear(fc_in, self.config.num_isotopes)
self.regressor = nn.Sequential(
nn.Linear(fc_in, self.config.num_isotopes),
nn.ReLU()
)
def _calculate_flat_size(self) -> int:
h = self.config.num_time_intervals
w = self.config.num_channels
for _ in self.config.conv_channels:
h = h // self.config.pool_size[0]
w = w // self.config.pool_size[1]
return self.config.conv_channels[-1] * h * w
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
# Add channel dimension if needed: (B, T, C) -> (B, 1, T, C)
if x.dim() == 3:
x = x.unsqueeze(1)
# CNN backbone
for conv_block in self.conv_blocks:
x = conv_block(x)
# Flatten
x = x.view(x.size(0), -1)
# FC backbone
x = self.fc_backbone(x)
# Output heads
logits = self.classifier(x)
activities = self.regressor(x)
return logits, activities
# =============================================================================
# PREDICTION DATA CLASSES
# =============================================================================
@dataclass
class IsotopePrediction:
"""Prediction result for a single isotope."""
name: str
probability: float
activity_bq: float
present: bool
@dataclass
class SpectrumPrediction:
"""Complete prediction results for a spectrum."""
isotopes: List[IsotopePrediction]
num_present: int
confidence: float
threshold_used: float
def to_dict(self) -> Dict:
"""Convert to JSON-serializable dictionary."""
return {
'isotopes': [
{
'name': iso.name,
'probability': round(iso.probability, 4),
'activity_bq': round(iso.activity_bq, 2),
'present': iso.present
}
for iso in self.isotopes
],
'num_isotopes_detected': self.num_present,
'confidence': round(self.confidence, 4),
'threshold': self.threshold_used
}
def get_present_isotopes(self) -> List[IsotopePrediction]:
"""Get only isotopes predicted as present."""
return [iso for iso in self.isotopes if iso.present]
def summary(self) -> str:
"""Human-readable summary of predictions."""
present = self.get_present_isotopes()
if not present:
return "No isotopes detected above threshold"
lines = [f"Detected {len(present)} isotope(s):"]
for iso in sorted(present, key=lambda x: -x.probability):
lines.append(
f" • {iso.name}: {iso.probability*100:.1f}% confidence, "
f"{iso.activity_bq:.1f} Bq estimated activity"
)
return "\n".join(lines)
def to_json(self, indent: int = 2) -> str:
"""Convert to JSON string."""
return json.dumps(self.to_dict(), indent=indent)
# =============================================================================
# INFERENCE ENGINE
# =============================================================================
class Vega2DInference:
"""
Inference engine for the Vega 2D isotope identification model.
Example usage:
inference = Vega2DInference("vega_2d_best.pt")
spectrum = np.load("my_spectrum.npy") # Shape: (60, 1023)
result = inference.predict(spectrum, threshold=0.5)
print(result.summary())
"""
def __init__(
self,
model_path: Union[str, Path],
isotope_index: Optional[IsotopeIndex] = None,
device: Optional[torch.device] = None
):
"""
Initialize the inference engine.
Args:
model_path: Path to saved .pt model checkpoint
isotope_index: Optional custom isotope index. Uses default if None.
device: Compute device. Auto-detects CUDA if available.
"""
self.model_path = Path(model_path)
# Device selection
if device is not None:
self.device = device
elif torch.cuda.is_available():
try:
_ = torch.zeros(1, device='cuda') + 1
self.device = torch.device('cuda')
except RuntimeError:
self.device = torch.device('cpu')
else:
self.device = torch.device('cpu')
# Load checkpoint
print(f"Loading 2D model from: {self.model_path}")
self.checkpoint = torch.load(self.model_path, map_location=self.device, weights_only=False)
# Load model config
if 'model_config' in self.checkpoint:
config_dict = self.checkpoint['model_config']
# Handle tuple conversion for kernel_size and pool_size
if 'kernel_size' in config_dict and isinstance(config_dict['kernel_size'], list):
config_dict['kernel_size'] = tuple(config_dict['kernel_size'])
if 'pool_size' in config_dict and isinstance(config_dict['pool_size'], list):
config_dict['pool_size'] = tuple(config_dict['pool_size'])
self.model_config = Vega2DConfig(**config_dict)
else:
self.model_config = Vega2DConfig()
# Create and load model
self.model = Vega2DModel(self.model_config)
self.model.load_state_dict(self.checkpoint['model_state_dict'])
self.model = self.model.to(self.device)
self.model.eval()
# Set isotope index
self.isotope_index = isotope_index or IsotopeIndex()
print(f"✓ Model loaded successfully")
print(f" Device: {self.device}")
print(f" Input shape: ({self.model_config.num_time_intervals}, {self.model_config.num_channels})")
print(f" Isotopes: {self.isotope_index.num_isotopes}")
print(f" Architecture: 2D-CNN{self.model_config.conv_channels} → FC{self.model_config.fc_hidden_dims}")
def _pad_or_truncate(self, spectrum: np.ndarray) -> np.ndarray:
"""Ensure spectrum has exactly num_time_intervals rows."""
target_rows = self.model_config.num_time_intervals
current_rows = spectrum.shape[0]
if current_rows == target_rows:
return spectrum
elif current_rows > target_rows:
# Truncate - take last N intervals (most recent data)
return spectrum[-target_rows:]
else:
# Pad with zeros at the beginning
padding = np.zeros((target_rows - current_rows, spectrum.shape[1]))
return np.vstack([padding, spectrum])
def preprocess(self, spectrum: np.ndarray, normalize: bool = True) -> torch.Tensor:
"""
Preprocess spectrum for model input.
Args:
spectrum: Input array, shape (T, 1023) where T is any number of time intervals
normalize: Normalize to [0, 1] range
Returns:
Tensor ready for model, shape (1, 60, 1023)
"""
# Handle 1D input (average spectrum)
if spectrum.ndim == 1:
# Expand to 2D by repeating
spectrum = np.tile(spectrum.reshape(1, -1), (self.model_config.num_time_intervals, 1))
# Ensure correct time dimension
spectrum = self._pad_or_truncate(spectrum)
# Normalize
if normalize and spectrum.max() > 0:
spectrum = spectrum / spectrum.max()
# To tensor with batch dimension
tensor = torch.tensor(spectrum, dtype=torch.float32).unsqueeze(0)
return tensor.to(self.device)
@torch.no_grad()
def predict(
self,
spectrum: Union[np.ndarray, torch.Tensor],
threshold: float = 0.5,
return_all: bool = False
) -> SpectrumPrediction:
"""
Run inference on a gamma spectrum.
Args:
spectrum: Input spectrum (numpy array or tensor)
threshold: Probability threshold for detection (0-1)
return_all: If True, include all 82 isotopes. If False, only detected ones.
Returns:
SpectrumPrediction with detected isotopes and activities
"""
# Preprocess
if isinstance(spectrum, np.ndarray):
spectrum = self.preprocess(spectrum)
# Run model
logits, activities = self.model(spectrum)
# Convert to probabilities
probs = torch.sigmoid(logits).cpu().numpy()[0]
activities = activities.cpu().numpy()[0] * self.model_config.max_activity_bq
# Build predictions
isotopes = []
for i in range(len(probs)):
prob = float(probs[i])
activity = float(activities[i])
present = prob >= threshold
if return_all or present:
isotopes.append(IsotopePrediction(
name=self.isotope_index.index_to_name(i),
probability=prob,
activity_bq=activity if present else 0.0,
present=present
))
# Calculate confidence
present_isotopes = [iso for iso in isotopes if iso.present]
if present_isotopes:
confidence = np.mean([iso.probability for iso in present_isotopes])
else:
confidence = 1.0 - probs.max()
return SpectrumPrediction(
isotopes=isotopes,
num_present=len(present_isotopes),
confidence=float(confidence),
threshold_used=threshold
)
def predict_from_file(
self,
file_path: Union[str, Path],
threshold: float = 0.5
) -> SpectrumPrediction:
"""Load spectrum from .npy file and run inference."""
spectrum = np.load(file_path)
return self.predict(spectrum, threshold)
def predict_batch(
self,
spectra: List[np.ndarray],
threshold: float = 0.5
) -> List[SpectrumPrediction]:
"""Run inference on multiple spectra."""
return [self.predict(s, threshold) for s in spectra]
# =============================================================================
# SAMPLE SPECTRUM GENERATOR (For testing without real data)
# =============================================================================
def energy_to_channel(energy_kev: float, num_channels: int = 1023) -> int:
"""Convert energy (keV) to usable channel index (0..num_channels-1).
Assumes an underlying 1024-channel MCA with raw channels 0..1023 where
channel 0 is skipped (modeled usable channels correspond to raw 1..1023).
"""
e_min, e_max = 20.0, 3000.0
full_channels = num_channels + 1
channel_width = (e_max - e_min) / full_channels
raw_channel = int((energy_kev - e_min) / channel_width)
usable_channel = raw_channel - 1
return max(0, min(num_channels - 1, usable_channel))
def channel_to_energy(channel: int, num_channels: int = 1023) -> float:
"""Convert usable channel index to energy bin center (keV)."""
e_min, e_max = 20.0, 3000.0
full_channels = num_channels + 1
channel_width = (e_max - e_min) / full_channels
raw_channel = channel + 1
return e_min + (raw_channel + 0.5) * channel_width
def create_sample_spectrum_2d(
isotope: str = "Cs-137",
activity_bq: float = 100.0,
duration_seconds: int = 60,
add_background: bool = True,
add_noise: bool = True,
detector_fwhm_percent: float = 8.5,
seed: Optional[int] = None
) -> np.ndarray:
"""
Generate a synthetic 2D gamma spectrum for testing.
Args:
isotope: Isotope name (e.g., "Cs-137", "Co-60")
activity_bq: Source activity in Becquerels
duration_seconds: Number of 1-second time intervals (default 60)
add_background: Add environmental background
add_noise: Apply Poisson counting statistics
detector_fwhm_percent: Detector resolution at 662 keV (%)
seed: Random seed for reproducibility
Returns:
2D numpy array of shape (duration_seconds, 1023)
"""
if seed is not None:
np.random.seed(seed)
num_channels = 1023
spectrum = np.zeros((duration_seconds, num_channels))
# Get gamma lines for the isotope
if isotope in GAMMA_LINES:
gamma_lines = GAMMA_LINES[isotope]
else:
print(f"Warning: No gamma lines for {isotope}, using Cs-137")
gamma_lines = GAMMA_LINES["Cs-137"]
# Generate spectrum for each time interval
for t in range(duration_seconds):
for energy_kev, branching_ratio in gamma_lines:
fwhm_kev = (detector_fwhm_percent / 100.0) * 662.0 * math.sqrt(energy_kev / 662.0)
sigma_kev = fwhm_kev / 2.355
efficiency = 0.1 * math.exp(-energy_kev / 500.0)
expected_counts = activity_bq * 1.0 * branching_ratio * efficiency # 1 second interval
for ch in range(num_channels):
energy = channel_to_energy(ch)
peak = expected_counts * math.exp(-0.5 * ((energy - energy_kev) / sigma_kev) ** 2)
spectrum[t, ch] += peak
# Add background
if add_background:
for ch in range(num_channels):
energy = channel_to_energy(ch)
bg = 50.0 * 1.0 * math.exp(-energy / 300.0) / 300.0
spectrum[t, ch] += bg
# K-40 environmental
k40_energy = 1460.8
k40_fwhm = (detector_fwhm_percent / 100.0) * 662.0 * math.sqrt(k40_energy / 662.0)
k40_sigma = k40_fwhm / 2.355
k40_counts = 10.0 * 1.0
for ch in range(num_channels):
energy = channel_to_energy(ch)
peak = k40_counts * math.exp(-0.5 * ((energy - k40_energy) / k40_sigma) ** 2)
spectrum[t, ch] += peak
# Apply Poisson noise
if add_noise:
spectrum = np.maximum(spectrum, 0)
spectrum = np.random.poisson(spectrum.astype(int)).astype(float)
return spectrum
def create_sample_spectra_batch_2d() -> Dict[str, np.ndarray]:
"""Create a batch of sample 2D spectra for different isotopes."""
samples = {}
for isotope in ["Cs-137", "Co-60", "Na-22", "Ba-133", "Am-241", "Eu-152"]:
if isotope in GAMMA_LINES:
samples[isotope] = create_sample_spectrum_2d(
isotope=isotope,
activity_bq=100.0,
duration_seconds=60,
seed=hash(isotope) % 2**32
)
# Background only
samples["Background"] = create_sample_spectrum_2d(
isotope="Cs-137",
activity_bq=0.0,
duration_seconds=60,
add_background=True,
seed=12345
)
return samples
# =============================================================================
# DEMONSTRATION FUNCTIONS
# =============================================================================
def run_demo(model_path: str, threshold: float = 0.5):
"""Run a complete demonstration of the Vega 2D inference system."""
print("\n" + "=" * 70)
print("VEGA 2D ISOTOPE IDENTIFICATION - INFERENCE DEMONSTRATION")
print("=" * 70)
# Load model
print("\n[1] Loading 2D Model")
print("-" * 70)
inference = Vega2DInference(model_path)
# Generate sample spectra
print("\n[2] Generating Sample 2D Spectra (60 time intervals × 1023 channels)")
print("-" * 70)
samples = create_sample_spectra_batch_2d()
print(f"Generated {len(samples)} sample spectra:")
for name, spec in samples.items():
print(f" • {name}: shape {spec.shape}")
# Run inference on each
print("\n[3] Running Inference")
print("-" * 70)
for name, spectrum in samples.items():
print(f"\n{'─' * 70}")
print(f"Sample: {name}")
print(f"Spectrum shape: {spectrum.shape}")
print(f"Spectrum range: [{spectrum.min():.1f}, {spectrum.max():.1f}]")
result = inference.predict(spectrum, threshold=threshold)
print(f"\nPrediction (threshold={threshold}):")
print(result.summary())
# Top 5 probabilities
print("\nTop 5 isotope probabilities:")
all_result = inference.predict(spectrum, threshold=0.0, return_all=True)
sorted_iso = sorted(all_result.isotopes, key=lambda x: -x.probability)[:5]
for iso in sorted_iso:
marker = "✓" if iso.probability >= threshold else " "
print(f" {marker} {iso.name}: {iso.probability*100:.2f}%")
# JSON output format
print("\n[4] JSON Output Format Example")
print("-" * 70)
sample_result = inference.predict(samples["Cs-137"], threshold=threshold)
print(sample_result.to_json())
print("\n" + "=" * 70)
print("DEMONSTRATION COMPLETE")
print("=" * 70)
def run_single_inference(model_path: str, spectrum_path: str, threshold: float = 0.5):
"""Run inference on a single spectrum file."""
print(f"\nLoading model from: {model_path}")
inference = Vega2DInference(model_path)
print(f"Loading spectrum from: {spectrum_path}")
spectrum = np.load(spectrum_path)
print(f"Spectrum shape: {spectrum.shape}")
print(f"\nRunning inference (threshold={threshold})...")
result = inference.predict(spectrum, threshold=threshold)
print("\n" + "=" * 60)
print("PREDICTION RESULTS")
print("=" * 60)
print(result.summary())
print("=" * 60)
return result
# =============================================================================
# MAIN ENTRY POINT
# =============================================================================
def main():
"""Main entry point for command-line usage."""
import argparse
parser = argparse.ArgumentParser(
description="Vega 2D Portable Inference - Gamma Spectrum Isotope Identification",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Run demo with sample spectra
python vega_portable_inference_2d.py --model vega_2d_best.pt
# Analyze a specific spectrum file
python vega_portable_inference_2d.py --model vega_2d_best.pt --spectrum my_data.npy
# Use lower threshold for higher recall
python vega_portable_inference_2d.py --model vega_2d_best.pt --threshold 0.3
"""
)
parser.add_argument(
"--model", "-m",
type=str,
required=True,
help="Path to trained Vega 2D model checkpoint (.pt file)"
)
parser.add_argument(
"--spectrum", "-s",
type=str,
default=None,
help="Path to spectrum file (.npy, shape 60×1023 or variable×1023). Runs demo if not provided."
)
parser.add_argument(
"--threshold", "-t",
type=float,
default=0.5,
help="Detection threshold (0-1). Lower = more sensitive. Default: 0.5"
)
parser.add_argument(
"--json",
action="store_true",
help="Output results in JSON format"
)
args = parser.parse_args()
if args.spectrum:
result = run_single_inference(args.model, args.spectrum, args.threshold)
if args.json:
print("\nJSON Output:")
print(result.to_json())
else:
run_demo(args.model, args.threshold)
return 0
if __name__ == "__main__":
sys.exit(main())