Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,131 +1,327 @@
-
 import streamlit as st
-import matplotlib.pyplot as plt
-import joblib
 import pandas as pd
-import json
 import numpy as np
-import pandas as pd
 import joblib
+import json
+import plotly.express as px
+import plotly.graph_objects as go
+from datetime import datetime
 
-DEFAULT_MATERIALS = ["Local Anaesthesia","Dry Needle","Botox","Saline","Magnesium","PRF"]
-
-def load_materials(path="material_list.json"):
-    try:
-        with open(path, "r", encoding="utf-8") as f:
-            mats = json.load(f).get("materials", [])
-        return mats if mats else DEFAULT_MATERIALS
-    except FileNotFoundError:
-        return DEFAULT_MATERIALS
-
-def build_baseline_from_form(row_dict):
-    # Convert to single-row DataFrame
-    df = pd.DataFrame([row_dict])
-    # Ensure presence of all baseline columns model might expect.
-    # (The pipeline has imputation & one-hot; we just pass what's available.)
-    # Guarantee numeric types for core measures
-    for c in ["age","pain_m0","mmo_m0","ohip_14_m0"]:
-        if c in df.columns:
-            df[c] = pd.to_numeric(df[c], errors="coerce")
-    # previous_injection as 0/1
-    if "previous_injection" in df.columns:
-        df["previous_injection"] = pd.to_numeric(df["previous_injection"], errors="coerce")
-    # include placeholder for time2_days if training had it
-    if "time2_days" not in df.columns:
-        df["time2_days"] = np.nan
-    # Fill optional text columns if missing
-    for tcol in ["location","muscle_injected","adjunctive_tretment","material_in_previous_injection"]:
-        if tcol not in df.columns:
-            df[tcol] = ""
-    return df
-
-def score_material(model, row_dict, material):
-    base = build_baseline_from_form(row_dict)
-    base["material_injected"] = material
-    proba = float(model.predict_proba(base)[0,1])
-    return proba
-
-def rank_materials(model, row_dict, materials):
-    rows = []
-    for m in materials:
-        p = score_material(model, row_dict, m)
-        rows.append((m, p))
-    rows.sort(key=lambda x: x[1], reverse=True)
-    return rows
-
-st.set_page_config(page_title="TMJ Success Predictor (Material-specific)", page_icon="🦷", layout="wide")
-st.title("🦷 TMJ Success Predictor — Material-specific")
-st.caption("Predicts 3-month success probability **conditioned on the selected injection material**. Also compares across all materials.")
-
-MODEL_PATH = "src/best_tmj_success_classifier.pkl"
-MATS_PATH = "src/material_list.json"
+# Page config
+st.set_page_config(
+    page_title="TMJ Injection Success Predictor",
+    page_icon="💉",
+    layout="wide"
+)
 
+# Load model and materials
 @st.cache_resource
 def load_artifacts():
-    model = joblib.load(MODEL_PATH)
-    mats = load_materials(MATS_PATH)
-    return model, mats
+    """Load the trained model and materials list"""
+    try:
+        # Load model
+        model = joblib.load('src/best_tmj_success_classifier_without_fe.pkl')
+        
+        # Load materials list
+        try:
+            with open('src/material_list.json', 'r') as f:
+                materials_data = json.load(f)
+                materials = materials_data.get('materials', [])
+        except FileNotFoundError:
+            # Fallback to default materials
+            materials = ['Local Anaesthesia', 'Dry Needle', 'Botox', 
+                        'Saline', 'Magnesium', 'PRF']
+            st.warning("Using default materials list. Train the model to generate actual materials from your data.")
+        
+        # Load metadata if available
+        metadata = {}
+        try:
+            with open('model_metadata.json', 'r') as f:
+                metadata = json.load(f)
+        except FileNotFoundError:
+            pass
+        
+        return model, materials, metadata
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        st.stop()
+
+# Initialize
+model, materials, metadata = load_artifacts()
 
-try:
-    model, materials = load_artifacts()
-except Exception as e:
-    st.error(f"Failed to load model: {e}")
-    st.stop()
+# Title and description
+st.title("🦷 TMJ Injection Success Predictor")
+st.markdown("""
+This tool predicts the 3-month treatment success probability for TMJ injections based on patient baseline characteristics.
+Enter the patient information below to see predictions for different injection materials.
+""")
 
-with st.form("frm"):
-    st.subheader("Baseline inputs")
-    c1,c2,c3 = st.columns(3)
-    with c1:
-        sex = st.selectbox("Sex", ["male","female","unknown"], index=1)
-        age = st.number_input("Age", 10, 100, value=30)
-    with c2:
-        pain_m0 = st.number_input("Pain M0 (0–10)", min_value=0.0, max_value=10.0, value=7.0, step=0.1)
-        mmo_m0 = st.number_input("MMO M0 (mm)", min_value=0.0, max_value=80.0, value=35.0, step=1.0)
-    with c3:
-        ohip_14_m0 = st.number_input("OHIP-14 M0 (0–56)", min_value=0.0, max_value=56.0, value=20.0, step=1.0)
-        prev = st.selectbox("Previous injection?", ["no","yes"], index=0)
+# Display model info if available
+if metadata:
+    with st.expander("ℹ️ Model Information"):
+        col1, col2, col3 = st.columns(3)
+        with col1:
+            st.metric("Model Type", metadata.get('model_type', 'Unknown'))
+        with col2:
+            st.metric("Test ROC-AUC", f"{metadata.get('test_roc_auc', 0):.3f}")
+        with col3:
+            st.metric("Training Date", metadata.get('training_date', 'Unknown')[:10])
+        
+        st.write(f"**Success Definition:** {metadata.get('success_definition', 'Unknown')}")
+        
+        if metadata.get('simplified_version', False):
+            st.info("This model uses the simplified feature set without text analysis.")
 
-    st.subheader("Context (optional)")
-    c4,c5 = st.columns(2)
-    with c4:
-        location = st.text_input("Location", value="TMJ Right")
-        muscle_injected = st.text_input("Muscle injected", value="masseter")
-        adjunctive_treatment = st.text_input("Adjunctive treatment", value="physiotherapy")
-    with c5:
-        material_prev = st.text_input("Material in previous injection", value="")
-        # free text fields are optional; model pipeline is robust
-    material_choice = st.selectbox("Material injected (to score)", materials)
-    do_compare = st.checkbox("Compare all materials", value=True)
+st.divider()
 
-    submitted = st.form_submit_button("Predict")
+# Create form
+with st.form("patient_form"):
+    st.subheader("Patient Information")
+    
+    # Required fields
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        st.markdown("**Required Fields**")
+        sex = st.selectbox("Sex", options=['Male', 'Female'], help="Patient's biological sex")
+        age = st.number_input("Age", min_value=10, max_value=100, value=45, help="Patient age in years")
+        pain_m0 = st.slider("Baseline Pain (M0)", min_value=0, max_value=10, value=7, 
+                           help="Pain score at baseline (0-10 scale)")
+    
+    with col2:
+        st.markdown("** **")  # Empty space to align with "Required Fields"
+        mmo_m0 = st.slider("Baseline MMO (M0)", min_value=0, max_value=80, value=35,
+                          help="Maximum mouth opening at baseline (mm)")
+        ohip_14_m0 = st.slider("Baseline OHIP-14 (M0)", min_value=0, max_value=56, value=28,
+                              help="Oral Health Impact Profile score at baseline (0-56)")
+    
+    st.divider()
+    
+    # Optional fields
+    st.markdown("**Optional Fields**")
+    col3, col4 = st.columns(2)
+    
+    with col3:
+        location = st.text_input("Location", placeholder="e.g., Right TMJ",
+                                help="Injection location")
+        muscle_injected = st.text_input("Muscle Injected", placeholder="e.g., Masseter",
+                                      help="Specific muscle targeted")
+        adjunctive_treatment = st.text_input("Adjunctive Treatment", placeholder="e.g., Physical therapy",
+                                           help="Additional treatments")
+    
+    with col4:
+        previous_injection = st.selectbox("Previous Injection", options=['No', 'Yes'],
+                                        help="Has the patient had previous TMJ injections?")
+        if previous_injection == 'Yes':
+            material_in_previous_injection = st.selectbox("Previous Material", 
+                                                         options=[''] + materials,
+                                                         help="Material used in previous injection")
+        else:
+            material_in_previous_injection = ''
+    
+    st.divider()
+    
+    # Material selection for primary prediction
+    st.markdown("**Primary Prediction**")
+    selected_material = st.selectbox("Select Material for Prediction", 
+                                   options=materials,
+                                   help="Choose the material you're considering for this patient")
+    
+    # Compare all materials option
+    compare_all = st.checkbox("Compare all available materials", value=True,
+                            help="Show predictions for all materials to help with decision making")
+    
+    # Submit button
+    submitted = st.form_submit_button("🔮 Predict Success", use_container_width=True, type="primary")
 
+# Process form submission
 if submitted:
-    row = {
-        "sex": sex,
-        "age": age,
-        "pain_m0": pain_m0,
-        "mmo_m0": mmo_m0,
-        "ohip_14_m0": ohip_14_m0,
-        "previous_injection": 1 if prev=="yes" else 0,
-        "location": location,
-        "muscle_injected": muscle_injected,
-        "adjunctive_treatment": adjunctive_treatment,
-        "material_in_previous_injection": material_prev
-    }
-    # Score selected material
-    proba = score_material(model, row, material_choice)
-    st.metric("Success probability (selected material)", f"{proba:.3f}")
+    # Create input dataframe
+    input_data = pd.DataFrame({
+        'sex': [sex],
+        'age': [age],
+        'pain_m0': [pain_m0],
+        'mmo_m0': [mmo_m0],
+        'ohip_14_m0': [ohip_14_m0],
+        'location': [location if location else np.nan],
+        'muscle_injected': [muscle_injected if muscle_injected else np.nan],
+        'adjunctive_treatment': [adjunctive_treatment if adjunctive_treatment else np.nan],
+        'previous_injection': [1 if previous_injection == 'Yes' else 0],
+        'material_in_previous_injection': [material_in_previous_injection if material_in_previous_injection else np.nan],
+        'material_injected': [selected_material]
+    })
+    
+    # Make prediction for selected material
+    try:
+        prediction_proba = model.predict_proba(input_data)[0, 1]
+        
+        # Display primary prediction
+        st.divider()
+        st.subheader("Prediction Results")
+        
+        # Create a visual indicator
+        col1, col2, col3 = st.columns([1, 2, 1])
+        with col2:
+            # Success probability gauge
+            fig = go.Figure(go.Indicator(
+                mode = "gauge+number+delta",
+                value = prediction_proba * 100,
+                domain = {'x': [0, 1], 'y': [0, 1]},
+                title = {'text': f"Success Probability with {selected_material}"},
+                number = {'suffix': "%", 'font': {'size': 40}},
+                gauge = {
+                    'axis': {'range': [None, 100]},
+                    'bar': {'color': "darkblue"},
+                    'steps': [
+                        {'range': [0, 30], 'color': "lightgray"},
+                        {'range': [30, 70], 'color': "gray"},
+                        {'range': [70, 100], 'color': "lightgreen"}
+                    ],
+                    'threshold': {
+                        'line': {'color': "red", 'width': 4},
+                        'thickness': 0.75,
+                        'value': 50
+                    }
+                }
+            ))
+            fig.update_layout(height=400)
+            st.plotly_chart(fig, use_container_width=True)
+        
+        # Interpretation
+        if prediction_proba >= 0.7:
+            st.success(f"✅ High likelihood of success ({prediction_proba:.1%}) with {selected_material}")
+        elif prediction_proba >= 0.5:
+            st.warning(f"⚠️ Moderate likelihood of success ({prediction_proba:.1%}) with {selected_material}")
+        else:
+            st.error(f"❌ Low likelihood of success ({prediction_proba:.1%}) with {selected_material}")
+        
+        # Compare all materials if requested
+        if compare_all:
+            st.divider()
+            st.subheader("📊 Material Comparison")
+            
+            # Predict for all materials
+            material_results = []
+            for material in materials:
+                temp_data = input_data.copy()
+                temp_data['material_injected'] = material
+                prob = model.predict_proba(temp_data)[0, 1]
+                material_results.append({
+                    'Material': material,
+                    'Success Probability': prob,
+                    'Success %': f"{prob:.1%}"
+                })
+            
+            # Sort by probability
+            material_df = pd.DataFrame(material_results)
+            material_df = material_df.sort_values('Success Probability', ascending=False)
+            
+            # Display results
+            col1, col2 = st.columns([1, 1])
+            
+            with col1:
+                # Table view
+                st.markdown("**Ranked Materials**")
+                display_df = material_df[['Material', 'Success %']].reset_index(drop=True)
+                display_df.index += 1  # Start index at 1
+                st.dataframe(display_df, use_container_width=True)
+                
+                # Highlight best option
+                best_material = material_df.iloc[0]['Material']
+                best_prob = material_df.iloc[0]['Success Probability']
+                if best_material != selected_material:
+                    st.info(f"💡 Consider using **{best_material}** for potentially better outcomes ({best_prob:.1%} vs {prediction_proba:.1%})")
+            
+            with col2:
+                # Bar chart
+                st.markdown("**Visual Comparison**")
+                fig = px.bar(material_df, 
+                           x='Success Probability', 
+                           y='Material',
+                           orientation='h',
+                           color='Success Probability',
+                           color_continuous_scale='RdYlGn',
+                           range_color=[0, 1],
+                           text='Success %')
+                
+                fig.update_traces(textposition='outside')
+                fig.update_layout(
+                    xaxis_title="Success Probability",
+                    yaxis_title="",
+                    showlegend=False,
+                    xaxis=dict(range=[0, 1.1]),
+                    height=400
+                )
+                
+                # Add vertical line at 50%
+                fig.add_vline(x=0.5, line_dash="dash", line_color="gray", 
+                            annotation_text="50% threshold")
+                
+                st.plotly_chart(fig, use_container_width=True)
+        
+        # Additional insights
+        st.divider()
+        with st.expander("📋 Patient Summary"):
+            st.write("**Baseline Characteristics:**")
+            summary_cols = st.columns(3)
+            with summary_cols[0]:
+                st.write(f"- Age: {age} years")
+                st.write(f"- Sex: {sex}")
+                st.write(f"- Previous injection: {previous_injection}")
+            with summary_cols[1]:
+                st.write(f"- Pain score: {pain_m0}/10")
+                st.write(f"- MMO: {mmo_m0} mm")
+                st.write(f"- OHIP-14: {ohip_14_m0}/56")
+            with summary_cols[2]:
+                if location:
+                    st.write(f"- Location: {location}")
+                if muscle_injected:
+                    st.write(f"- Muscle: {muscle_injected}")
+                if adjunctive_treatment:
+                    st.write(f"- Adjunctive: {adjunctive_treatment}")
+        
+    except Exception as e:
+        st.error(f"Error making prediction: {str(e)}")
+        st.info("Please ensure the model was trained with all the necessary features.")
+
+# Footer
+st.divider()
+st.markdown("""
+<div style='text-align: center; color: gray;'>
+    <small>
+    TMJ Injection Success Predictor | 
+    Model trained on historical patient data | 
+    Predictions are probabilistic and should be used alongside clinical judgment
+    </small>
+</div>
+""", unsafe_allow_html=True)
 
-    if do_compare:
-        st.markdown("---")
-        st.subheader("Compare all materials (ranked)")
-        ranked = rank_materials(model, row, materials)
-        df_rank = pd.DataFrame(ranked, columns=["material","success_proba"])
-        st.dataframe(df_rank, width=True)
-        # Bar plot (matplotlib; single plot; default colors)
-        fig = plt.figure(figsize=(6,4))
-        plt.bar(df_rank["material"], df_rank["success_proba"])
-        plt.xticks(rotation=45, ha="right")
-        plt.ylabel("Predicted success probability")
-        plt.title("Material comparison — success probability")
-        st.pyplot(fig)
+# Sidebar with instructions
+with st.sidebar:
+    st.header("📖 Instructions")
+    st.markdown("""
+    1. **Enter patient baseline data** in the form
+    2. **Select the material** you're considering
+    3. **Click Predict** to see the success probability
+    4. **Compare materials** to find the optimal choice
+    
+    ---
+    
+    ### 🎯 Success Definition
+    Treatment success is typically defined as:
+    - Pain reduction > 2 points
+    - MMO increase > 5 mm  
+    - OHIP-14 reduction > 5 points
+    
+    ---
+    
+    ### 📊 Interpretation Guide
+    - **70%+**: High success likelihood ✅
+    - **50-70%**: Moderate success ⚠️
+    - **<50%**: Low success likelihood ❌
+    
+    ---
+    
+    ### ⚕️ Clinical Note
+    These predictions are based on statistical models and should complement, not replace, clinical expertise and patient-specific considerations.
+    """)