Spaces:
Sleeping
(REFAC)(core): Enhance core logic for model loading, inference, and spectrum parsing
Browse files- Centralized imports and configuration for spectrum preprocessing and sample data handling.
- Improved `label_file` function for robust filename-based label extraction, supporting unknown patterns.
- Refactored model loading with registry support, multiple weight path fallbacks, and legacy compatibility.
- Added cache decorators (`@st.cache_data`, `@st.cache_resource`) for efficient state management in Streamlit.
- Enhanced memory cleanup with explicit garbage collection and CUDA cache clearing.
- Updated `run_inference` to include performance tracking, confidence calculation, and error-tolerant logging.
- Implemented `_get_memory_usage` for accurate resource monitoring using `psutil`.
- Improved sample file discovery and spectrum data parsing with validation for format, monotonicity, and range.
- Strengthened error handling throughout to prevent UI crashes and ensure robust operation.
- Streamlined code structure for
- core_logic.py +70 -44
@@ -1,7 +1,7 @@
|
|
1 |
import os
|
2 |
|
3 |
# --- New Imports ---
|
4 |
-
from config import
|
5 |
import time
|
6 |
import gc
|
7 |
import torch
|
@@ -11,6 +11,7 @@ import streamlit as st
|
|
11 |
from pathlib import Path
|
12 |
from config import SAMPLE_DATA_DIR
|
13 |
from datetime import datetime
|
|
|
14 |
|
15 |
|
16 |
def label_file(filename: str) -> int:
|
@@ -37,48 +38,74 @@ def load_state_dict(_mtime, model_path):
|
|
37 |
|
38 |
@st.cache_resource
|
39 |
def load_model(model_name):
|
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 |
-
|
66 |
-
|
67 |
-
|
68 |
-
|
69 |
-
|
70 |
-
|
71 |
-
|
72 |
-
|
73 |
-
|
74 |
-
model.eval()
|
75 |
-
return model, True
|
76 |
-
else:
|
77 |
-
return model, False
|
78 |
|
79 |
-
|
80 |
-
|
81 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
82 |
|
83 |
|
84 |
def cleanup_memory():
|
@@ -89,7 +116,7 @@ def cleanup_memory():
|
|
89 |
|
90 |
|
91 |
@st.cache_data
|
92 |
-
def run_inference(y_resampled, model_choice, _cache_key=None):
|
93 |
"""Run model inference and cache results with performance tracking"""
|
94 |
from utils.performance_tracker import get_performance_tracker, PerformanceMetrics
|
95 |
from datetime import datetime
|
@@ -126,7 +153,6 @@ def run_inference(y_resampled, model_choice, _cache_key=None):
|
|
126 |
|
127 |
# Log performance metrics
|
128 |
try:
|
129 |
-
modality = st.session_state.get("modality_select", "raman")
|
130 |
confidence = float(max(probs)) if probs is not None and len(probs) > 0 else 0.0
|
131 |
|
132 |
metrics = PerformanceMetrics(
|
|
|
1 |
import os
|
2 |
|
3 |
# --- New Imports ---
|
4 |
+
from config import TARGET_LEN
|
5 |
import time
|
6 |
import gc
|
7 |
import torch
|
|
|
11 |
from pathlib import Path
|
12 |
from config import SAMPLE_DATA_DIR
|
13 |
from datetime import datetime
|
14 |
+
from models.registry import build, choices
|
15 |
|
16 |
|
17 |
def label_file(filename: str) -> int:
|
|
|
38 |
|
39 |
@st.cache_resource
|
40 |
def load_model(model_name):
|
41 |
+
# First try registry system (new approach)
|
42 |
+
if model_name in choices():
|
43 |
+
# Use registry system
|
44 |
+
model = build(model_name, TARGET_LEN)
|
45 |
+
|
46 |
+
# Try to load weights from standard locations
|
47 |
+
weight_paths = [
|
48 |
+
f"model_weights/{model_name}_model.pth",
|
49 |
+
f"outputs/{model_name}_model.pth",
|
50 |
+
f"model_weights/{model_name}.pth",
|
51 |
+
f"outputs/{model_name}.pth",
|
52 |
+
]
|
53 |
+
|
54 |
+
weights_loaded = False
|
55 |
+
for weight_path in weight_paths:
|
56 |
+
if os.path.exists(weight_path):
|
57 |
+
try:
|
58 |
+
mtime = os.path.getmtime(weight_path)
|
59 |
+
state_dict = load_state_dict(mtime, weight_path)
|
60 |
+
if state_dict:
|
61 |
+
model.load_state_dict(state_dict, strict=True)
|
62 |
+
model.eval()
|
63 |
+
weights_loaded = True
|
64 |
+
|
65 |
+
except Exception:
|
66 |
+
continue
|
67 |
+
|
68 |
+
if not weights_loaded:
|
69 |
+
st.warning(
|
70 |
+
f"⚠️ Model weights not found for '{model_name}'. Using randomly initialized model."
|
71 |
+
)
|
72 |
+
st.info(
|
73 |
+
"This model will provide random predictions for demonstration purposes."
|
74 |
+
)
|
|
|
|
|
|
|
|
|
75 |
|
76 |
+
return model, weights_loaded
|
77 |
+
|
78 |
+
# Legacy system fallback (for backward compatibility)
|
79 |
+
if model_name in MODEL_CONFIG:
|
80 |
+
config = MODEL_CONFIG[model_name]
|
81 |
+
model_class = config["class"]
|
82 |
+
model_path = config["path"]
|
83 |
+
|
84 |
+
# Initialize model
|
85 |
+
model = model_class(input_length=TARGET_LEN)
|
86 |
+
|
87 |
+
# Check if model file exists
|
88 |
+
if not os.path.exists(model_path):
|
89 |
+
st.warning(f"⚠️ Model weights not found: {model_path}")
|
90 |
+
st.info("Using randomly initialized model for demonstration purposes.")
|
91 |
+
return model, False
|
92 |
+
|
93 |
+
# Get mtime for cache invalidation
|
94 |
+
mtime = os.path.getmtime(model_path)
|
95 |
+
|
96 |
+
# Load weights
|
97 |
+
state_dict = load_state_dict(mtime, model_path)
|
98 |
+
if state_dict:
|
99 |
+
model.load_state_dict(state_dict, strict=True)
|
100 |
+
model.eval()
|
101 |
+
return model, True
|
102 |
+
else:
|
103 |
+
return model, False
|
104 |
+
else:
|
105 |
+
st.error(
|
106 |
+
f"❌ Unknown model '{model_name}'. Available models: {list(MODEL_CONFIG.keys())}"
|
107 |
+
)
|
108 |
+
return None, False
|
109 |
|
110 |
|
111 |
def cleanup_memory():
|
|
|
116 |
|
117 |
|
118 |
@st.cache_data
|
119 |
+
def run_inference(y_resampled, model_choice, modality: str, _cache_key=None):
|
120 |
"""Run model inference and cache results with performance tracking"""
|
121 |
from utils.performance_tracker import get_performance_tracker, PerformanceMetrics
|
122 |
from datetime import datetime
|
|
|
153 |
|
154 |
# Log performance metrics
|
155 |
try:
|
|
|
156 |
confidence = float(max(probs)) if probs is not None and len(probs) > 0 else 0.0
|
157 |
|
158 |
metrics = PerformanceMetrics(
|