import os
import torch
import streamlit as st
import hashlib
import io
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
from typing import Union
import time
from config import MODEL_CONFIG, TARGET_LEN, LABEL_MAP
from modules.callbacks import (
on_model_change,
on_input_mode_change,
on_sample_change,
reset_ephemeral_state,
log_message,
clear_batch_results,
)
from core_logic import (
get_sample_files,
load_model,
run_inference,
parse_spectrum_data,
label_file,
)
from modules.callbacks import reset_results
from utils.results_manager import ResultsManager
from utils.confidence import calculate_softmax_confidence
from utils.multifile import process_multiple_files, display_batch_results
from utils.preprocessing import resample_spectrum
def load_css(file_path):
with open(file_path, encoding="utf-8") as f:
st.markdown(f"", unsafe_allow_html=True)
@st.cache_data
def create_spectrum_plot(x_raw, y_raw, x_resampled, y_resampled, _cache_key=None):
"""Create spectrum visualization plot"""
fig, ax = plt.subplots(1, 2, figsize=(13, 5), dpi=100)
# == Raw spectrum ==
ax[0].plot(x_raw, y_raw, label="Raw", color="dimgray", linewidth=1)
ax[0].set_title("Raw Input Spectrum")
ax[0].set_xlabel("Wavenumber (cm⁻¹)")
ax[0].set_ylabel("Intensity")
ax[0].grid(True, alpha=0.3)
ax[0].legend()
# == Resampled spectrum ==
ax[1].plot(
x_resampled, y_resampled, label="Resampled", color="steelblue", linewidth=1
)
ax[1].set_title(f"Resampled ({len(y_resampled)} points)")
ax[1].set_xlabel("Wavenumber (cm⁻¹)")
ax[1].set_ylabel("Intensity")
ax[1].grid(True, alpha=0.3)
ax[1].legend()
fig.tight_layout()
# == Convert to image ==
buf = io.BytesIO()
plt.savefig(buf, format="png", bbox_inches="tight", dpi=100)
buf.seek(0)
plt.close(fig) # Prevent memory leaks
return Image.open(buf)
def render_confidence_progress(
probs: np.ndarray,
labels: list[str] = ["Stable", "Weathered"],
highlight_idx: Union[int, None] = None,
side_by_side: bool = True,
):
"""Render Streamlit native progress bars with scientific formatting."""
p = np.asarray(probs, dtype=float)
p = np.clip(p, 0.0, 1.0)
if side_by_side:
cols = st.columns(len(labels))
for i, (lbl, val, col) in enumerate(zip(labels, p, cols)):
with col:
is_highlighted = highlight_idx is not None and i == highlight_idx
label_text = f"**{lbl}**" if is_highlighted else lbl
st.markdown(f"{label_text}: {val*100:.1f}%")
st.progress(int(round(val * 100)))
else:
# Vertical layout for better readability
for i, (lbl, val) in enumerate(zip(labels, p)):
is_highlighted = highlight_idx is not None and i == highlight_idx
# Create a container for each probability
with st.container():
col1, col2 = st.columns([3, 1])
with col1:
if is_highlighted:
st.markdown(f"**{lbl}** ← Predicted")
else:
st.markdown(f"{lbl}")
with col2:
st.metric(label="", value=f"{val*100:.1f}%", delta=None)
# Progress bar with conditional styling
if is_highlighted:
st.progress(int(round(val * 100)))
st.caption("🎯 **Model Prediction**")
else:
st.progress(int(round(val * 100)))
if i < len(labels) - 1: # Add spacing between items
st.markdown("")
def render_kv_grid(d: dict = {}, ncols: int = 2):
if d is None:
d = {}
if not d:
return
items = list(d.items())
cols = st.columns(ncols)
for i, (k, v) in enumerate(items):
with cols[i % ncols]:
st.caption(f"**{k}:** {v}")
def render_model_meta(model_choice: str):
info = MODEL_CONFIG.get(model_choice, {})
emoji = info.get("emoji", "")
desc = info.get("description", "").strip()
acc = info.get("accuracy", "-")
f1 = info.get("f1", "-")
st.caption(f"{emoji} **Model Snapshot** - {model_choice}")
cols = st.columns(2)
with cols[0]:
st.metric("Accuracy", acc)
with cols[1]:
st.metric("F1 Score", f1)
if desc:
st.caption(desc)
def get_confidence_description(logit_margin):
"""Get human-readable confidence description"""
if logit_margin > 1000:
return "VERY HIGH", "🟢"
elif logit_margin > 250:
return "HIGH", "🟡"
elif logit_margin > 100:
return "MODERATE", "🟠"
else:
return "LOW", "🔴"
def render_sidebar():
with st.sidebar:
# Header
st.header("AI-Driven Polymer Classification")
st.caption(
"Predict polymer degradation (Stable vs Weathered) from Raman spectra using validated CNN models. — v0.1"
)
model_labels = [
f"{MODEL_CONFIG[name]['emoji']} {name}" for name in MODEL_CONFIG.keys()
]
selected_label = st.selectbox(
"Choose AI Model",
model_labels,
key="model_select",
on_change=on_model_change,
)
model_choice = selected_label.split(" ", 1)[1]
# ===Compact metadata directly under dropdown===
render_model_meta(model_choice)
# ===Collapsed info to reduce clutter===
with st.expander("About This App", icon=":material/info:", expanded=False):
st.markdown(
"""
**AI-Driven Polymer Aging Prediction and Classification**
**Purpose**: Classify polymer degradation using AI
**Input**: Raman spectroscopy .txt files
**Models**: CNN architectures for binary classification
**Next**: More trained CNNs in evaluation pipeline
**Contributors**
- Dr. Sanmukh Kuppannagari (Mentor)
- Dr. Metin Karailyan (Mentor)
- Jaser Hasan (Author)
**Links**
[HF Space](https://huggingface.co/spaces/dev-jas/polymer-aging-ml)
[GitHub Repository](https://github.com/KLab-AI3/ml-polymer-recycling)
**Citation Figure2CNN (baseline)**
Neo et al., 2023, *Resour. Conserv. Recycl.*, 188, 106718.
[https://doi.org/10.1016/j.resconrec.2022.106718](https://doi.org/10.1016/j.resconrec.2022.106718)
""",
unsafe_allow_html=True,
)
# col1 goes here
# In modules/ui_components.py
def render_input_column():
st.markdown("##### Data Input")
mode = st.radio(
"Input mode",
["Upload File", "Batch Upload", "Sample Data"],
key="input_mode",
horizontal=True,
on_change=on_input_mode_change,
)
# == Input Mode Logic ==
# ... (The if/elif/else block for Upload, Batch, and Sample modes remains exactly the same) ...
# ==Upload tab==
if mode == "Upload File":
upload_key = st.session_state["current_upload_key"]
up = st.file_uploader(
"Upload Raman spectrum (.txt)",
type="txt",
help="Upload a text file with wavenumber and intensity columns",
key=upload_key, # ← versioned key
)
# ==Process change immediately (no on_change; simpler & reliable)==
if up is not None:
raw = up.read()
text = raw.decode("utf-8") if isinstance(raw, bytes) else raw
# == only reparse if its a different file|source ==
if (
st.session_state.get("filename") != getattr(up, "name", None)
or st.session_state.get("input_source") != "upload"
):
st.session_state["input_text"] = text
st.session_state["filename"] = getattr(up, "name", None)
st.session_state["input_source"] = "upload"
# Ensure single file mode
st.session_state["batch_mode"] = False
st.session_state["status_message"] = (
f"File '{st.session_state['filename']}' ready for analysis"
)
st.session_state["status_type"] = "success"
reset_results("New file uploaded")
# ==Batch Upload tab==
elif mode == "Batch Upload":
st.session_state["batch_mode"] = True
# --- START: BUG 1 & 3 FIX ---
# Use a versioned key to ensure the file uploader resets properly.
batch_upload_key = f"batch_upload_{st.session_state['uploader_version']}"
uploaded_files = st.file_uploader(
"Upload multiple Raman spectrum files (.txt)",
type="txt",
accept_multiple_files=True,
help="Upload one or more text files with wavenumber and intensity columns.",
key=batch_upload_key,
)
# --- END: BUG 1 & 3 FIX ---
if uploaded_files:
# --- START: Bug 1 Fix ---
# Use a dictionary to keep only unique files based on name and size
unique_files = {(file.name, file.size): file for file in uploaded_files}
unique_file_list = list(unique_files.values())
num_uploaded = len(uploaded_files)
num_unique = len(unique_file_list)
# Optionally, inform the user that duplicates were removed
if num_uploaded > num_unique:
st.info(
f"ℹ️ {num_uploaded - num_unique} duplicate file(s) were removed."
)
# Use the unique list
st.session_state["batch_files"] = unique_file_list
st.session_state["status_message"] = (
f"{num_unique} ready for batch analysis"
)
st.session_state["status_type"] = "success"
# --- END: Bug 1 Fix ---
else:
st.session_state["batch_files"] = []
# This check prevents resetting the status if files are already staged
if not st.session_state.get("batch_files"):
st.session_state["status_message"] = (
"No files selected for batch processing"
)
st.session_state["status_type"] = "info"
# ==Sample tab==
elif mode == "Sample Data":
st.session_state["batch_mode"] = False
sample_files = get_sample_files()
if sample_files:
options = ["-- Select Sample --"] + [p.name for p in sample_files]
sel = st.selectbox(
"Choose sample spectrum:",
options,
key="sample_select",
on_change=on_sample_change,
)
if sel != "-- Select Sample --":
st.session_state["status_message"] = (
f"📁 Sample '{sel}' ready for analysis"
)
st.session_state["status_type"] = "success"
else:
st.info("No sample data available")
# == Status box (displays the message) ==
msg = st.session_state.get("status_message", "Ready")
typ = st.session_state.get("status_type", "info")
if typ == "success":
st.success(msg)
elif typ == "error":
st.error(msg)
else:
st.info(msg)
# --- DE-NESTED LOGIC STARTS HERE ---
# This code now runs on EVERY execution, guaranteeing the buttons will appear.
# Safely get model choice from session state
model_choice = st.session_state.get("model_select", " ").split(" ", 1)[1]
model = load_model(model_choice)
# Determine if the app is ready for inference
is_batch_ready = st.session_state.get("batch_mode", False) and st.session_state.get(
"batch_files"
)
is_single_ready = not st.session_state.get(
"batch_mode", False
) and st.session_state.get("input_text")
inference_ready = (is_batch_ready or is_single_ready) and model is not None
# Store for other modules to access
st.session_state["inference_ready"] = inference_ready
# Render buttons
with st.form("analysis_form", clear_on_submit=False):
submitted = st.form_submit_button(
"Run Analysis", type="primary", disabled=not inference_ready
)
st.button(
"Reset All",
on_click=reset_ephemeral_state,
help="Clear all uploaded files and results.",
)
# Handle form submission
if submitted and inference_ready:
if st.session_state.get("batch_mode"):
batch_files = st.session_state.get("batch_files", [])
with st.spinner(f"Processing {len(batch_files)} files ..."):
st.session_state["batch_results"] = process_multiple_files(
uploaded_files=batch_files,
model_choice=model_choice,
load_model_func=load_model,
run_inference_func=run_inference,
label_file_func=label_file,
)
else:
try:
x_raw, y_raw = parse_spectrum_data(st.session_state["input_text"])
x_resampled, y_resampled = resample_spectrum(x_raw, y_raw, TARGET_LEN)
st.session_state.update(
{
"x_raw": x_raw,
"y_raw": y_raw,
"x_resampled": x_resampled,
"y_resampled": y_resampled,
"inference_run_once": True,
}
)
except (ValueError, TypeError) as e:
st.error(f"Error processing spectrum data: {e}")
# col2 goes here
def render_results_column():
# Get the current mode and check for batch results
is_batch_mode = st.session_state.get("batch_mode", False)
has_batch_results = "batch_results" in st.session_state
if is_batch_mode and has_batch_results:
# THEN render the main interactive dashboard from ResultsManager
ResultsManager.display_results_table()
elif st.session_state.get("inference_run_once", False) and not is_batch_mode:
st.markdown("##### Analysis Results")
# Get data from session state
x_raw = st.session_state.get("x_raw")
y_raw = st.session_state.get("y_raw")
x_resampled = st.session_state.get("x_resampled") # ← NEW
y_resampled = st.session_state.get("y_resampled")
filename = st.session_state.get("filename", "Unknown")
if all(v is not None for v in [x_raw, y_raw, y_resampled]):
# ===Run inference===
if y_resampled is None:
raise ValueError(
"y_resampled is None. Ensure spectrum data is properly resampled before proceeding."
)
cache_key = hashlib.md5(
f"{y_resampled.tobytes()}{st.session_state.get('model_select', 'Unknown').split(' ', 1)[1]}".encode()
).hexdigest()
prediction, logits_list, probs, inference_time, logits = run_inference(
y_resampled,
(
st.session_state.get("model_select", "").split(" ", 1)[1]
if "model_select" in st.session_state
else None
),
_cache_key=cache_key,
)
if prediction is None:
st.error(
"❌ Inference failed: Model not loaded. Please check that weights are available."
)
st.stop() # prevents the rest of the code in this block from executing
log_message(
f"Inference completed in {inference_time:.2f}s, prediction: {prediction}"
)
# ===Get ground truth===
true_label_idx = label_file(filename)
true_label_str = (
LABEL_MAP.get(true_label_idx, "Unknown")
if true_label_idx is not None
else "Unknown"
)
# ===Get prediction===
predicted_class = LABEL_MAP.get(int(prediction), f"Class {int(prediction)}")
# Enhanced confidence calculation
if logits is not None:
# Use new softmax-based confidence
probs_np, max_confidence, confidence_level, confidence_emoji = (
calculate_softmax_confidence(logits)
)
confidence_desc = confidence_level
else:
# Fallback to legace method
logit_margin = abs(
(logits_list[0] - logits_list[1])
if logits_list is not None and len(logits_list) >= 2
else 0
)
confidence_desc, confidence_emoji = get_confidence_description(
logit_margin
)
max_confidence = logit_margin / 10.0 # Normalize for display
probs_np = np.array([])
# Store result in results manager for single file too
ResultsManager.add_results(
filename=filename,
model_name=(
st.session_state.get("model_select", "").split(" ", 1)[1]
if "model_select" in st.session_state
else "Unknown"
),
prediction=int(prediction),
predicted_class=predicted_class,
confidence=max_confidence,
logits=logits_list if logits_list else [],
ground_truth=true_label_idx if true_label_idx >= 0 else None,
processing_time=inference_time if inference_time is not None else 0.0,
metadata={
"confidence_level": confidence_desc,
"confidence_emoji": confidence_emoji,
},
)
# ===Precompute Stats===
model_choice = (
st.session_state.get("model_select", "").split(" ", 1)[1]
if "model_select" in st.session_state
else None
)
if not model_choice:
st.error(
"⚠️ Model choice is not defined. Please select a model from the sidebar."
)
st.stop()
model_path = MODEL_CONFIG[model_choice]["path"]
mtime = os.path.getmtime(model_path) if os.path.exists(model_path) else None
file_hash = (
hashlib.md5(open(model_path, "rb").read()).hexdigest()
if os.path.exists(model_path)
else "N/A"
)
# Removed unused variable 'input_tensor'
start_render = time.time()
active_tab = st.selectbox(
"View Results",
["Details", "Technical", "Explanation"],
key="active_tab", # reuse the key you were managing manually
)
if active_tab == "Details":
st.markdown('