Update app.py

app.py

@@ -1,217 +1,244 @@
 import gradio as gr
 import joblib
 import pandas as pd
+import numpy as np
 import os
 from groq import Groq
 
-from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
-from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table
+from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
 from reportlab.lib import colors
 from reportlab.lib.units import inch
+from reportlab.platypus import TableStyle
 
 # ===============================
-# Load ML Model
+# Load Model
 # ===============================
-model = joblib.load("typhoidguard_model_reduced.pkl")
+model = joblib.load("typhoidguard_model_reduced2.pkl")
 
-# ===============================
 # Groq Client
-# ===============================
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 
-# ===============================
-# GLOBAL CACHE (avoid double LLM call)
-# ===============================
-LLM_CACHE = {
-    "english": "",
-    "urdu": ""
-}
-
-# ===============================
-# FORMAT LLM TEXT TO HTML
-# ===============================
-def format_llm_text(text):
-    lines = text.split("\n")
-    html = ""
-
-    for line in lines:
-        line = line.strip()
-
-        if line.lower().startswith("1.") or line.lower().startswith("2.") or line.lower().startswith("3.") \
-           or line.lower().startswith("4.") or line.lower().startswith("5."):
-            html += f"<h3 style='color:#003366'>{line}</h3>"
-
-        elif line.startswith("*"):
-            html += f"<li>{line[1:].strip()}</li>"
-
-        elif line:
-            html += f"<p>{line}</p>"
-
-    return f"<ul>{html}</ul>"
+# Cache LLM result (avoid double call)
+last_llm_report = ""
+last_basic_output = ""
 
 
 # ===============================
-# LLM REPORT (English + Urdu)
+# LLM FUNCTION
 # ===============================
 def generate_llm_report(patient_info, prob_pct, risk_label, recommendation, factors):
 
     prompt = f"""
-You are a medical AI assistant helping doctors in rural Pakistan.
+You are a responsible medical AI assistant helping rural healthcare units in Pakistan.
 
 Patient Details:
 {patient_info}
 
-Risk of Treatment Failure: {prob_pct}%
+Predicted Risk of Treatment Failure: {prob_pct}%
 Risk Category: {risk_label}
 
-Risk Factors:
+Clinical Risk Factors:
 {", ".join(factors) if factors else "None"}
 
 Agent Recommendation:
 {recommendation}
 
-TASK:
-1. Write a detailed medical report with headings:
-   - Clinical Interpretation
-   - Possible Complications
-   - Immediate Precautions
-   - Preventive Measures
-   - Follow-up Recommendations
-
-2. Then write a SHORT summary in URDU for healthcare staff.
+Generate a structured detailed report including:
+1. Clinical Interpretation
+2. Possible Complications
+3. Immediate Precautions
+4. Preventative Measures
+5. Follow-up Recommendations
 
-Rules:
-- Use bullet points
-- No markdown (** or ###)
-- Simple clinical language
+Use professional formatting.
+End with reminder that final decision must be taken by qualified doctor.
 """
 
     response = client.chat.completions.create(
         model="llama-3.3-70b-versatile",
         messages=[{"role": "user", "content": prompt}],
-        temperature=0.3,
+        temperature=0.4,
     )
 
-    full_text = response.choices[0].message.content
-
-    # Split English & Urdu
-    if "اردو" in full_text:
-        eng, urdu = full_text.split("اردو", 1)
-        urdu = "اردو" + urdu
-    else:
-        eng, urdu = full_text, ""
-
-    return eng.strip(), urdu.strip()
+    return response.choices[0].message.content
 
 
 # ===============================
-# BASIC ANALYSIS (NO LLM)
+# PREDICT FUNCTION (UPDATED FEATURES)
 # ===============================
-def analyze_patient(name, age, gender, platelets, hb, duration, calcium, potassium):
+def predict_basic(name, age, gender, platelets, hb, duration,
+                  blood_bacteria, severity, medication):
+
+    global last_llm_report, last_basic_output
+    last_llm_report = ""  # reset
 
     data = pd.DataFrame([{
-        "Platelet Count": platelets,
         "Age": age,
+        "Platelet Count": platelets,
         "Hemoglobin (g/dL)": hb,
         "Treatment Duration": duration,
-        "Calcium (mg/dL)": calcium,
-        "Potassium (mmol/L)": potassium,
-        "Gender": gender
+        "Gender": gender,
+        "Blood Culture Bacteria": blood_bacteria,
+        "Symptoms Severity": severity,
+        "Current Medication": medication
     }])
 
-    X = model.named_steps["preprocessor"].transform(data)
-    prob = model.named_steps["classifier"].predict_proba(X)[0][1]
+    prob = model.predict_proba(data)[0][1]
+    prob_pct = round(prob * 100, 2)
 
+    # Agentic Override Logic
     low_platelets = platelets < 150000
     low_hb = hb < 10
-
-    prob_pct = 90.0 if (low_platelets or low_hb) else round(prob * 100, 2)
+    prolonged_treatment = duration > 7
 
     if low_platelets or low_hb:
-        status = "🔴 HIGH RISK"
+        prob_pct = 90.0
+        risk_label = "🔴 HIGH RISK"
         color = "#d62828"
         recommendation = "Immediate referral to secondary care facility recommended."
+    elif prob >= 0.65:
+        risk_label = "🔴 HIGH RISK"
+        color = "#d62828"
+        recommendation = "Urgent clinical evaluation required."
     elif prob >= 0.35:
-        status = "🟡 MODERATE RISK"
+        risk_label = "🟡 MODERATE RISK"
         color = "#f4a261"
-        recommendation = "Close monitoring advised."
+        recommendation = "Close monitoring advised. Repeat tests in 48 hours."
     else:
-        status = "🟢 LOW RISK"
+        risk_label = "🟢 LOW RISK"
         color = "#2a9d8f"
-        recommendation = "Continue standard treatment."
+        recommendation = "Continue standard treatment protocol."
 
     factors = []
-    if low_platelets: factors.append("Low Platelet Count")
-    if low_hb: factors.append("Low Hemoglobin")
+    if low_platelets:
+        factors.append("Low Platelet Count")
+    if low_hb:
+        factors.append("Low Hemoglobin")
+    if prolonged_treatment:
+        factors.append("Prolonged Treatment Duration")
+    if severity == "High":
+        factors.append("High Symptom Severity")
+
+    urdu_summary = f"""
+    <div style='background:#eef7ff;padding:10px;border-radius:8px;'>
+    <b>اردو خلاصہ:</b><br>
+    مریض میں علاج ناکامی کا خطرہ <b>{prob_pct}%</b> ہے۔ موجودہ کیفیت: <b>{risk_label}</b>۔<br>
+    براہ کرم ڈاکٹر سے مشورہ ضرور کریں۔
+    </div>
+    """
 
-    html = f"""
-    <div style="padding:15px;font-family:Arial">
-        <h2 style="color:{color}">Patient: {name}</h2>
+    basic_output = f"""
+    <div style="font-family:Arial; padding:15px;">
+        <h2 style="color:{color};">Patient: {name}</h2>
         <h3>Risk of Treatment Failure: {prob_pct}%</h3>
-        <h3>Status: {status}</h3>
-        <h4>Agent Recommendation</h4>
+        <h3>Status: {risk_label}</h3>
+
+        <h4>Clinical Risk Factors:</h4>
+        <ul>
+            {''.join([f"<li>{f}</li>" for f in factors]) if factors else "<li>None</li>"}
+        </ul>
+
+        <h4>Agent Recommendation:</h4>
         <p><b>{recommendation}</b></p>
+
+        {urdu_summary}
     </div>
     """
 
-    return html, prob_pct, status, recommendation, factors
+    last_basic_output = basic_output
+
+    return basic_output
 
 
 # ===============================
 # DETAILED REPORT BUTTON
 # ===============================
-def detailed_report(name, age, gender, platelets, hb, duration, calcium, potassium):
+def generate_detailed(name, age, gender, platelets, hb, duration,
+                      blood_bacteria, severity, medication):
 
-    html, prob_pct, status, rec, factors = analyze_patient(
-        name, age, gender, platelets, hb, duration, calcium, potassium
-    )
+    global last_llm_report, last_basic_output
+
+    basic_html = predict_basic(name, age, gender, platelets, hb,
+                               duration, blood_bacteria, severity, medication)
+
+    data = pd.DataFrame([{
+        "Age": age,
+        "Platelet Count": platelets,
+        "Hemoglobin (g/dL)": hb,
+        "Treatment Duration": duration,
+        "Gender": gender,
+        "Blood Culture Bacteria": blood_bacteria,
+        "Symptoms Severity": severity,
+        "Current Medication": medication
+    }])
+
+    prob = model.predict_proba(data)[0][1]
+    prob_pct = round(prob * 100, 2)
+
+    factors = []
+    if platelets < 150000:
+        factors.append("Low Platelet Count")
+    if hb < 10:
+        factors.append("Low Hemoglobin")
 
     patient_info = f"""
-Age: {age}, Gender: {gender}
-Platelets: {platelets}, Hb: {hb}
-Duration: {duration} days
-Calcium: {calcium}, Potassium: {potassium}
+Age: {age}
+Gender: {gender}
+Platelets: {platelets}
+Hemoglobin: {hb}
+Treatment Duration: {duration}
+Blood Bacteria: {blood_bacteria}
+Severity: {severity}
+Medication: {medication}
 """
 
-    eng, urdu = generate_llm_report(patient_info, prob_pct, status, rec, factors)
+    llm_report = generate_llm_report(
+        patient_info, prob_pct, "See Above", "See Above", factors
+    )
 
-    LLM_CACHE["english"] = eng
-    LLM_CACHE["urdu"] = urdu
+    last_llm_report = llm_report
 
-    formatted_eng = format_llm_text(eng)
+    formatted_llm = llm_report.replace("\n", "<br>")
 
-    return html + f"""
+    return basic_html + f"""
     <hr>
-    <h2 style="color:#003366">🧠 Detailed AI Report</h2>
-    {formatted_eng}
-    <h3 style="color:#006400">اردو خلاصہ</h3>
-    <p>{urdu}</p>
+    <h3 style="color:#003366;">🧠 Detailed AI Report & Recommendations</h3>
+    <div style="background:#f8f9fa;padding:15px;border-radius:10px;line-height:1.6;">
+    {formatted_llm}
+    </div>
     """
 
 
 # ===============================
-# PDF GENERATION (USES CACHE)
+# PDF GENERATION
 # ===============================
-def generate_pdf(name, age, gender, platelets, hb, duration, calcium, potassium):
+def generate_pdf(name, age, gender, platelets, hb, duration,
+                 blood_bacteria, severity, medication):
+
+    global last_llm_report, last_basic_output
 
     file_path = f"{name}_Typhoid_Report.pdf"
     doc = SimpleDocTemplate(file_path)
-    styles = getSampleStyleSheet()
     elements = []
+    styles = getSampleStyleSheet()
 
     elements.append(Paragraph("<b>TyphoidGuard AI Clinical Report</b>", styles["Title"]))
     elements.append(Spacer(1, 0.3 * inch))
 
-    elements.append(Paragraph("<b>Detailed AI Clinical Report</b>", styles["Heading2"]))
-    elements.append(Paragraph(LLM_CACHE["english"].replace("\n", "<br/>"), styles["Normal"]))
-    elements.append(Spacer(1, 0.2 * inch))
+    # Basic Output
+    elements.append(Paragraph(last_basic_output.replace("<br>", "<br/>"), styles["Normal"]))
+    elements.append(Spacer(1, 0.3 * inch))
 
-    elements.append(Paragraph("<b>اردو خلاصہ</b>", styles["Heading2"]))
-    elements.append(Paragraph(LLM_CACHE["urdu"], styles["Normal"]))
+    # Detailed LLM Output (if exists)
+    if last_llm_report:
+        elements.append(Paragraph("<b>Detailed AI Report</b>", styles["Heading2"]))
+        elements.append(Spacer(1, 0.2 * inch))
+        elements.append(Paragraph(last_llm_report.replace("\n", "<br/>"), styles["Normal"]))
+        elements.append(Spacer(1, 0.3 * inch))
 
-    elements.append(Spacer(1, 0.3 * inch))
     elements.append(Paragraph(
-        "<i>Disclaimer: This report is AI-generated. Consult a qualified doctor.</i>",
+        "<i>Disclaimer: This report is AI-generated. Please consult a qualified doctor before making clinical decisions.</i>",
         styles["Normal"]
     ))
 
@@ -220,11 +247,12 @@ def generate_pdf(name, age, gender, platelets, hb, duration, calcium, potassium)
 
 
 # ===============================
-# UI
+# INTERFACE
 # ===============================
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     gr.Markdown("# 🏥 TyphoidGuard AI")
+    gr.Markdown("### HEC Generative AI Training Hackathon – Pakistan")
 
     with gr.Row():
         name = gr.Textbox(label="Patient Name")
@@ -233,29 +261,48 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     with gr.Row():
         platelets = gr.Number(label="Platelet Count")
-        hb = gr.Number(label="Hemoglobin")
-        duration = gr.Number(label="Treatment Duration")
+        hb = gr.Number(label="Hemoglobin (g/dL)")
+        duration = gr.Number(label="Treatment Duration (Days)")
 
     with gr.Row():
-        calcium = gr.Number(label="Calcium")
-        potassium = gr.Number(label="Potassium")
+        blood_bacteria = gr.Dropdown(
+            ["Escherichia coli", "Salmonella typhi", "Unknown", "Staphylococcus"],
+            label="Blood Culture Bacteria"
+        )
+        severity = gr.Dropdown(
+            ["Low", "Moderate", "High", "Unknown"],
+            label="Symptoms Severity"
+        )
+        medication = gr.Dropdown(
+            ["Amoxicillin", "Ceftriaxone", "Azithromycin"],
+            label="Current Medication"
+        )
 
     output = gr.HTML()
 
-    btn1 = gr.Button("🔍 Analyze Patient")
-    btn2 = gr.Button("🧠 Detailed Report and Recommendations")
+    predict_btn = gr.Button("🔍 Basic Risk Assessment")
+    detailed_btn = gr.Button("🧠 Detailed Report and Recommendations")
     pdf_btn = gr.Button("📄 Generate PDF Report")
 
-    btn1.click(analyze_patient, 
-               inputs=[name, age, gender, platelets, hb, duration, calcium, potassium],
-               outputs=[output, gr.State(), gr.State(), gr.State(), gr.State()])
+    predict_btn.click(
+        predict_basic,
+        inputs=[name, age, gender, platelets, hb, duration,
+                blood_bacteria, severity, medication],
+        outputs=output
+    )
 
-    btn2.click(detailed_report,
-               inputs=[name, age, gender, platelets, hb, duration, calcium, potassium],
-               outputs=output)
+    detailed_btn.click(
+        generate_detailed,
+        inputs=[name, age, gender, platelets, hb, duration,
+                blood_bacteria, severity, medication],
+        outputs=output
+    )
 
-    pdf_btn.click(generate_pdf,
-                  inputs=[name, age, gender, platelets, hb, duration, calcium, potassium],
-                  outputs=gr.File())
+    pdf_btn.click(
+        generate_pdf,
+        inputs=[name, age, gender, platelets, hb, duration,
+                blood_bacteria, severity, medication],
+        outputs=gr.File()
+    )
 
 demo.launch()