Update app.py

app.py

@@ -6,7 +6,7 @@ from groq import Groq
 
 from reportlab.platypus import (
     SimpleDocTemplate, Paragraph, Spacer,
-    Table, TableStyle
+    Table, TableStyle, Image
 )
 from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
 from reportlab.lib import colors
@@ -21,6 +21,8 @@ client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 last_llm_report = ""
 last_basic_data = {}
 
+LOGO_PATH = "logo.png"   # 🔹 Place your logo file in HF repo
+
 # ===============================
 # LLM FUNCTION
 # ===============================
@@ -124,7 +126,8 @@ def predict_basic(name, age, gender, platelets, hb, duration,
         "prob_pct": prob_pct,
         "risk_label": risk_label,
         "recommendation": recommendation,
-        "factors": factors
+        "factors": factors,
+        "color": color
     }
 
     urdu = f"""
@@ -138,8 +141,13 @@ def predict_basic(name, age, gender, platelets, hb, duration,
     return 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: {risk_label}</h3>
+
+    <h3>
+    Risk of Treatment Failure:
+    <span style="background:{color};color:white;padding:6px 12px;border-radius:6px;">
+    {prob_pct}% - {risk_label}
+    </span>
+    </h3>
 
     <h4>Clinical Risk Factors:</h4>
     <ul>
@@ -161,7 +169,6 @@ def generate_detailed(*inputs):
     global last_basic_data, last_llm_report
 
     basic_html = predict_basic(*inputs)
-
     d = last_basic_data
 
     patient_info = f"""
@@ -195,7 +202,7 @@ Medication: {d['medication']}
 
 
 # ===============================
-# PDF GENERATION (PROFESSIONAL)
+# PDF GENERATION
 # ===============================
 def generate_pdf(*inputs):
 
@@ -207,11 +214,33 @@ def generate_pdf(*inputs):
     elements = []
     styles = getSampleStyleSheet()
 
-    # Title
+    # ---- LOGO ----
+    if os.path.exists(LOGO_PATH):
+        elements.append(Image(LOGO_PATH, width=1.5*inch, height=1.5*inch))
+        elements.append(Spacer(1,0.2*inch))
+
     elements.append(Paragraph("<b>TyphoidGuard AI Clinical Report</b>", styles["Title"]))
     elements.append(Spacer(1, 0.3 * inch))
 
-    # ---- TABLE ----
+    # ---- RISK BADGE ----
+    risk_color = colors.HexColor(d["color"])
+    badge_style = ParagraphStyle(
+        name="Badge",
+        parent=styles["Normal"],
+        backColor=risk_color,
+        textColor=colors.white,
+        fontSize=12,
+        leftIndent=5,
+        rightIndent=5,
+        spaceAfter=10
+    )
+
+    elements.append(Paragraph(
+        f"Risk: {d['prob_pct']}% - {d['risk_label']}",
+        badge_style
+    ))
+    elements.append(Spacer(1,0.3*inch))
+
     table_data = [
         ["Patient Name", d["name"]],
         ["Age", d["age"]],
@@ -222,13 +251,10 @@ def generate_pdf(*inputs):
         ["Bacteria", d["bacteria"]],
         ["Severity", d["severity"]],
         ["Medication", d["medication"]],
-        ["Predicted Risk %", d["prob_pct"]],
-        ["Risk Level", d["risk_label"]],
     ]
 
     table = Table(table_data, colWidths=[2.5*inch, 3*inch])
     table.setStyle(TableStyle([
-        ('BACKGROUND',(0,0),(-1,0),colors.lightgrey),
         ('GRID',(0,0),(-1,-1),1,colors.grey),
         ('FONTNAME',(0,0),(-1,-1),'Helvetica')
     ]))
@@ -236,30 +262,18 @@ def generate_pdf(*inputs):
     elements.append(table)
     elements.append(Spacer(1,0.4*inch))
 
-    # ---- BASIC RECOMMENDATION ----
     elements.append(Paragraph("<b>Agent Recommendation</b>", styles["Heading2"]))
-    elements.append(Spacer(1,0.2*inch))
     elements.append(Paragraph(d["recommendation"], styles["Normal"]))
     elements.append(Spacer(1,0.4*inch))
 
-    # ---- LLM SECTION ----
     if last_llm_report:
-        elements.append(Paragraph("<b>Detailed AI Report & Recommendations</b>", styles["Heading2"]))
-        elements.append(Spacer(1,0.2*inch))
+        elements.append(Paragraph("<b>Detailed AI Report</b>", styles["Heading2"]))
         elements.append(Paragraph(last_llm_report.replace("\n","<br/>"), styles["Normal"]))
-        elements.append(Spacer(1,0.4*inch))
-
-    # Disclaimer
-    elements.append(Paragraph(
-        "<i>This report is AI-generated. Final medical decisions must be made by a qualified doctor.</i>",
-        styles["Normal"]
-    ))
 
-    # Team Footer
     elements.append(Spacer(1,0.5*inch))
     elements.append(Paragraph("<b>Team Chinar AI (AJK)</b>", styles["Normal"]))
     elements.append(Paragraph(
-        "Anees Qumar Abbasi (Team Lead)| Sharafat Hussain | Salma Asghar | "
+        "Anees Qumar Abbasi | Sharafat Hussain | Salma Asghar | "
         "Kaleem Hussain | Munazza Zahra | Kokub Khurishid",
         styles["Normal"]
     ))
@@ -274,9 +288,14 @@ def generate_pdf(*inputs):
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     gr.Markdown("""
-    # 🏥 TyphoidGuard AI  
-    ### HEC Generative AI Hackathon – Pakistan  
-    #### Team Chinar AI (AJK)
+    <div style="text-align:center;">
+    <img src="logo.png" width="120">
+    <h1>🏥 TyphoidGuard AI</h1>
+    <h4>HEC Generative AI Hackathon – Pakistan</h4>
+    <h4>Team Chinar AI (AJK)</h4>
+    <p><b>Anees Qumar Abbasi (Lead)</b> | Sharafat Hussain | 
+    Salma Asghar | Kaleem Hussain | Munazza Zahra | Kokub Khurishid</p>
+    </div>
     """)
 
     with gr.Row():
@@ -303,10 +322,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             label="Current Medication"
         )
 
-    # ---- BUTTONS ABOVE OUTPUT ----
     with gr.Row():
-        predict_btn = gr.Button("Basic Risk Assessment", variant="primary")
-        detailed_btn = gr.Button("Detailed Report & Recommendations", variant="secondary")
+        predict_btn = gr.Button("Basic Risk Assessment")
+        detailed_btn = gr.Button("Detailed Report & Recommendations")
         pdf_btn = gr.Button("Generate PDF Report")
 
     output = gr.HTML()