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_results,
reset_ephemeral_state,
log_message,
)
from core_logic import (
get_sample_files,
load_model,
run_inference,
parse_spectrum_data,
label_file,
)
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("")
from typing import Optional
def render_kv_grid(d: Optional[dict] = None, 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/FTIR spectra using validated CNN models. — v0.01"
)
# Modality Selection
st.markdown("##### Spectroscopy Modality")
modality = st.selectbox(
"Choose Modality",
["raman", "ftir"],
index=0,
key="modality_select",
format_func=lambda x: f"{'Raman' if x == 'raman' else 'FTIR'}",
)
# Display modality info
if modality == "ftir":
st.info("FTIR mode: 400-4000 cm-1 range with atmospheric correction")
else:
st.info("Raman mode: 200-4000 cm-1 range with standard preprocessing")
# Model selection
st.markdown("##### AI Model Selection")
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 classification
**Modalities**: Raman and FTIR spectroscopy support
**Features**: Multi-model comparison and analysis
**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,
)
# //////////////////////////////////////////
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 ==
if mode == "Upload File":
upload_key = st.session_state["current_upload_key"]
up = st.file_uploader(
"Upload spectrum file (.txt, .csv, .json)",
type=["txt", "csv", "json"],
help="Upload spectroscopy data: TXT (2-column), CSV (with headers), or JSON format",
key=upload_key, # ← versioned key
)
# Process change immediately
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
# 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 spectrum files (.txt, .csv, .json)",
type=["txt", "csv", "json"],
accept_multiple_files=True,
help="Upload spectroscopy files in TXT, CSV, or JSON format.",
key=batch_upload_key,
)
if uploaded_files:
# 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"
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)
# 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}")
# //////////////////////////////////////////
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 legacy 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"
)
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('