Update app.py

app.py

@@ -11,7 +11,6 @@ import modal
 from PIL import Image
 import numpy as np
 import io
-import base64
 import datetime
 from huggingface_hub import InferenceClient
 from reportlab.lib.pagesizes import A4
@@ -19,20 +18,9 @@ from reportlab.lib import colors
 from reportlab.lib.units import cm
 from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image as RLImage, Table, TableStyle, HRFlowable
 from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
-from reportlab.lib.enums import TA_CENTER, TA_LEFT, TA_RIGHT
-from reportlab.pdfgen import canvas
-from reportlab.platypus import BaseDocTemplate, PageTemplate, Frame
+from reportlab.lib.enums import TA_CENTER, TA_LEFT
 
-os.environ["MODAL_TOKEN_ID"] = os.environ.get("MODAL_TOKEN_ID", "")
-os.environ["MODAL_TOKEN_SECRET"] = os.environ.get("MODAL_TOKEN_SECRET", "")
-
-SurgiSightDetector = modal.Cls.from_name("surgisight", "SurgiSightDetector")
-detector = SurgiSightDetector()
-
-client = InferenceClient(
-    model="meta-llama/Llama-3.1-8B-Instruct",
-    token=os.environ.get("HF_TOKEN")
-)
+print("All imports OK")
 
 CLASS_NAMES = [
     "Black Background", "Abdominal Wall", "Liver", "Gastrointestinal Tract",
@@ -40,9 +28,23 @@ CLASS_NAMES = [
     "L-hook Electrocautery", "Gallbladder", "Hepatic Vein", "Liver Ligament"
 ]
 DANGER_CLASSES = ["Hepatic Vein", "Cystic Duct", "Blood"]
-
-# Store last result globally for PDF generation
 last_result = {}
+chat_context = {}  # stores detected tissues for chat context
+
+try:
+    client = InferenceClient(
+        model="meta-llama/Llama-3.1-8B-Instruct",
+        token=os.environ.get("HF_TOKEN")
+    )
+    print("InferenceClient OK")
+except Exception as e:
+    client = None
+    print(f"InferenceClient failed (non-fatal): {e}")
+
+
+def get_detector():
+    SurgiSightDetector = modal.Cls.from_name("surgisight", "SurgiSightDetector")
+    return SurgiSightDetector()
 
 
 def pil_to_bytes(pil_image):
@@ -51,69 +53,67 @@ def pil_to_bytes(pil_image):
     return buf.getvalue()
 
 
+def generate_suggested_questions(tissue_list):
+    """Generate 3 contextual suggested questions based on detected tissues."""
+    questions = []
+    # Danger-specific questions first
+    for t in tissue_list:
+        if t == "Hepatic Vein":
+            questions.append("Why is the hepatic vein so dangerous to nick?")
+        elif t == "Cystic Duct":
+            questions.append("How do I safely identify the cystic duct?")
+        elif t == "Blood":
+            questions.append("What are the steps to control unexpected bleeding?")
+        elif t == "Gallbladder":
+            questions.append("What is the critical view of safety for gallbladder removal?")
+        elif t == "L-hook Electrocautery":
+            questions.append("What are the risks of using electrocautery near the bile duct?")
+        elif t == "Liver":
+            questions.append("How does liver retraction affect visibility in laparoscopic surgery?")
+        if len(questions) >= 2:
+            break
+    # Always add a general fallback
+    questions.append("What are common complications in laparoscopic cholecystectomy?")
+    return questions[:3]
+
+
 def generate_pdf(original_image, annotated_image, seen, alert, explanation):
     pdf_path = "/tmp/surgisight_report.pdf"
     doc = SimpleDocTemplate(
-        pdf_path,
-        pagesize=A4,
-        rightMargin=2*cm,
-        leftMargin=2*cm,
-        topMargin=2*cm,
-        bottomMargin=2*cm
+        pdf_path, pagesize=A4,
+        rightMargin=2*cm, leftMargin=2*cm,
+        topMargin=2*cm, bottomMargin=2*cm
     )
-
     styles = getSampleStyleSheet()
-
-    # Custom styles
-    title_style = ParagraphStyle(
-        'Title', parent=styles['Title'],
+    title_style = ParagraphStyle('Title', parent=styles['Title'],
         fontSize=22, textColor=colors.HexColor('#0f2027'),
-        spaceAfter=4, fontName='Helvetica-Bold', alignment=TA_CENTER
-    )
-    subtitle_style = ParagraphStyle(
-        'Subtitle', parent=styles['Normal'],
+        spaceAfter=4, fontName='Helvetica-Bold', alignment=TA_CENTER)
+    subtitle_style = ParagraphStyle('Subtitle', parent=styles['Normal'],
         fontSize=10, textColor=colors.HexColor('#888888'),
-        spaceAfter=2, alignment=TA_CENTER
-    )
-    section_style = ParagraphStyle(
-        'Section', parent=styles['Normal'],
+        spaceAfter=2, alignment=TA_CENTER)
+    section_style = ParagraphStyle('Section', parent=styles['Normal'],
         fontSize=13, textColor=colors.HexColor('#203a43'),
-        spaceAfter=6, spaceBefore=12, fontName='Helvetica-Bold'
-    )
-    body_style = ParagraphStyle(
-        'Body', parent=styles['Normal'],
+        spaceAfter=6, spaceBefore=12, fontName='Helvetica-Bold')
+    body_style = ParagraphStyle('Body', parent=styles['Normal'],
         fontSize=10, textColor=colors.HexColor('#1a1a1a'),
-        spaceAfter=4, leading=16
-    )
-    danger_style = ParagraphStyle(
-        'Danger', parent=styles['Normal'],
+        spaceAfter=4, leading=16)
+    danger_style = ParagraphStyle('Danger', parent=styles['Normal'],
         fontSize=10, textColor=colors.HexColor('#cc0000'),
         spaceAfter=4, leading=16, fontName='Helvetica-Bold',
-        backColor=colors.HexColor('#fff0f0'),
-        borderPadding=(6, 8, 6, 8)
-    )
-    safe_style = ParagraphStyle(
-        'Safe', parent=styles['Normal'],
+        backColor=colors.HexColor('#fff0f0'), borderPadding=(6,8,6,8))
+    safe_style = ParagraphStyle('Safe', parent=styles['Normal'],
         fontSize=10, textColor=colors.HexColor('#1a7a40'),
         spaceAfter=4, leading=16, fontName='Helvetica-Bold',
-        backColor=colors.HexColor('#f0fff4'),
-        borderPadding=(6, 8, 6, 8)
-    )
-    caption_style = ParagraphStyle(
-        'Caption', parent=styles['Normal'],
+        backColor=colors.HexColor('#f0fff4'), borderPadding=(6,8,6,8))
+    caption_style = ParagraphStyle('Caption', parent=styles['Normal'],
         fontSize=8, textColor=colors.HexColor('#888888'),
-        alignment=TA_CENTER, spaceAfter=4
-    )
-    footer_style = ParagraphStyle(
-        'Footer', parent=styles['Normal'],
-        fontSize=8, textColor=colors.HexColor('#aaaaaa'),
-        alignment=TA_CENTER
-    )
+        alignment=TA_CENTER, spaceAfter=4)
+    footer_style = ParagraphStyle('Footer', parent=styles['Normal'],
+        fontSize=8, textColor=colors.HexColor('#aaaaaa'), alignment=TA_CENTER)
 
     story = []
     timestamp = datetime.datetime.now().strftime("%B %d, %Y at %H:%M")
 
-    # ── Header ────────────────────────────────────────────────────────────────
     story.append(Spacer(1, 0.3*cm))
     story.append(Paragraph("SurgiSight", title_style))
     story.append(Paragraph("Surgical Anatomy Analysis Report", subtitle_style))
@@ -122,100 +122,80 @@ def generate_pdf(original_image, annotated_image, seen, alert, explanation):
     story.append(HRFlowable(width="100%", thickness=2, color=colors.HexColor('#203a43')))
     story.append(Spacer(1, 0.4*cm))
 
-    # ── Images side by side ───────────────────────────────────────────────────
     story.append(Paragraph("Segmentation Output", section_style))
-
     img_width = 8.5*cm
     img_height = 6.5*cm
-
     orig_buf = io.BytesIO()
     original_image.save(orig_buf, format="PNG")
     orig_buf.seek(0)
-    orig_rl = RLImage(orig_buf, width=img_width, height=img_height)
-
     ann_buf = io.BytesIO()
     annotated_image.save(ann_buf, format="PNG")
     ann_buf.seek(0)
-    ann_rl = RLImage(ann_buf, width=img_width, height=img_height)
-
     img_table = Table(
-        [[orig_rl, ann_rl],
+        [[RLImage(orig_buf, width=img_width, height=img_height),
+          RLImage(ann_buf, width=img_width, height=img_height)],
          [Paragraph("Original Frame", caption_style), Paragraph("AI Segmented Output", caption_style)]],
         colWidths=[img_width + 0.5*cm, img_width + 0.5*cm]
     )
     img_table.setStyle(TableStyle([
-        ('ALIGN', (0,0), (-1,-1), 'CENTER'),
-        ('VALIGN', (0,0), (-1,-1), 'MIDDLE'),
+        ('ALIGN', (0,0), (-1,-1), 'CENTER'), ('VALIGN', (0,0), (-1,-1), 'MIDDLE'),
         ('BACKGROUND', (0,0), (-1,0), colors.HexColor('#f5f5f5')),
-        ('ROUNDEDCORNERS', [6, 6, 6, 6]),
         ('BOX', (0,0), (0,0), 0.5, colors.HexColor('#dddddd')),
         ('BOX', (1,0), (1,0), 0.5, colors.HexColor('#dddddd')),
     ]))
     story.append(img_table)
     story.append(Spacer(1, 0.5*cm))
 
-    # ── Safety Alert ──────────────────────────────────────────────────────────
     story.append(Paragraph("Safety Assessment", section_style))
     if any(d in alert for d in DANGER_CLASSES):
-        story.append(Paragraph(f"⚠  {alert}", danger_style))
+        story.append(Paragraph(f"WARNING: {alert}", danger_style))
     else:
-        story.append(Paragraph(f"✓  {alert}", safe_style))
+        story.append(Paragraph(f"SAFE: {alert}", safe_style))
     story.append(Spacer(1, 0.4*cm))
 
-    # ── Detection Results Table ───────────────────────────────────────────────
     story.append(Paragraph("Detected Tissues & Instruments", section_style))
-
     table_data = [["Structure", "Confidence", "Risk Level"]]
+    rows_danger = []
     for name, conf in sorted(seen.items(), key=lambda x: -x[1]):
         if name == "Black Background":
             continue
-        risk = "⚠ DANGER" if name in DANGER_CLASSES else "Safe"
-        risk_color = colors.HexColor('#cc0000') if name in DANGER_CLASSES else colors.HexColor('#1a7a40')
-        conf_pct = f"{conf:.1%}"
-        bar = "█" * int(conf * 10) + "░" * (10 - int(conf * 10))
-        table_data.append([name, f"{conf_pct}  {bar}", risk])
-
+        risk = "DANGER" if name in DANGER_CLASSES else "Safe"
+        bar = "\u2588" * int(conf * 10) + "\u2591" * (10 - int(conf * 10))
+        table_data.append([name, f"{conf:.1%}  {bar}", risk])
+        rows_danger.append(name in DANGER_CLASSES)
     det_table = Table(table_data, colWidths=[6*cm, 7*cm, 3.5*cm])
-    det_table_style = [
+    det_style = [
         ('BACKGROUND', (0,0), (-1,0), colors.HexColor('#203a43')),
         ('TEXTCOLOR', (0,0), (-1,0), colors.white),
         ('FONTNAME', (0,0), (-1,0), 'Helvetica-Bold'),
         ('FONTSIZE', (0,0), (-1,0), 10),
-        ('ALIGN', (0,0), (-1,-1), 'LEFT'),
-        ('VALIGN', (0,0), (-1,-1), 'MIDDLE'),
+        ('ALIGN', (0,0), (-1,-1), 'LEFT'), ('VALIGN', (0,0), (-1,-1), 'MIDDLE'),
         ('FONTSIZE', (0,1), (-1,-1), 9),
         ('ROWBACKGROUNDS', (0,1), (-1,-1), [colors.HexColor('#f9f9f9'), colors.white]),
         ('GRID', (0,0), (-1,-1), 0.3, colors.HexColor('#dddddd')),
-        ('TOPPADDING', (0,0), (-1,-1), 6),
-        ('BOTTOMPADDING', (0,0), (-1,-1), 6),
+        ('TOPPADDING', (0,0), (-1,-1), 6), ('BOTTOMPADDING', (0,0), (-1,-1), 6),
         ('LEFTPADDING', (0,0), (-1,-1), 8),
     ]
-    # Color danger rows red
-    for i, (name, conf) in enumerate(sorted(seen.items(), key=lambda x: -x[1])):
-        if name == "Black Background":
-            continue
-        if name in DANGER_CLASSES:
-            det_table_style.append(('TEXTCOLOR', (2, i+1), (2, i+1), colors.HexColor('#cc0000')))
-            det_table_style.append(('FONTNAME', (2, i+1), (2, i+1), 'Helvetica-Bold'))
-        else:
-            det_table_style.append(('TEXTCOLOR', (2, i+1), (2, i+1), colors.HexColor('#1a7a40')))
-
-    det_table.setStyle(TableStyle(det_table_style))
+    for i, is_danger in enumerate(rows_danger):
+        r = i + 1
+        col = colors.HexColor('#cc0000') if is_danger else colors.HexColor('#1a7a40')
+        det_style.append(('TEXTCOLOR', (2,r), (2,r), col))
+        if is_danger:
+            det_style.append(('FONTNAME', (2,r), (2,r), 'Helvetica-Bold'))
+    det_table.setStyle(TableStyle(det_style))
     story.append(det_table)
     story.append(Spacer(1, 0.5*cm))
 
-    # ── Anatomy Explanation ───────────────────────────────────────────────────
     story.append(HRFlowable(width="100%", thickness=0.5, color=colors.HexColor('#cccccc')))
     story.append(Spacer(1, 0.3*cm))
     story.append(Paragraph("Anatomy Teaching Note", section_style))
     story.append(Paragraph(explanation, body_style))
     story.append(Spacer(1, 0.5*cm))
 
-    # ── Model Info ────────────────────────────────────────────────────────────
     story.append(HRFlowable(width="100%", thickness=0.5, color=colors.HexColor('#cccccc')))
     story.append(Spacer(1, 0.3*cm))
     model_data = [
-        ["Detection Model", "YOLOv26-seg fine-tuned on CholecSeg8k (MICCAI 2020)"],
+        ["Detection Model", "YOLOv8-seg fine-tuned on CholecSeg8k (MICCAI 2020)"],
         ["LLM", "Meta Llama 3.1 8B Instruct"],
         ["Inference", "Modal GPU (T4)"],
         ["Dataset", "CholecSeg8k — 8,080 frames, 13 classes"],
@@ -223,38 +203,36 @@ def generate_pdf(original_image, annotated_image, seen, alert, explanation):
     ]
     model_table = Table(model_data, colWidths=[4.5*cm, 12*cm])
     model_table.setStyle(TableStyle([
-        ('FONTNAME', (0,0), (0,-1), 'Helvetica-Bold'),
-        ('FONTSIZE', (0,0), (-1,-1), 8),
+        ('FONTNAME', (0,0), (0,-1), 'Helvetica-Bold'), ('FONTSIZE', (0,0), (-1,-1), 8),
         ('TEXTCOLOR', (0,0), (0,-1), colors.HexColor('#203a43')),
         ('TEXTCOLOR', (1,0), (1,-1), colors.HexColor('#555555')),
         ('VALIGN', (0,0), (-1,-1), 'TOP'),
-        ('TOPPADDING', (0,0), (-1,-1), 3),
-        ('BOTTOMPADDING', (0,0), (-1,-1), 3),
+        ('TOPPADDING', (0,0), (-1,-1), 3), ('BOTTOMPADDING', (0,0), (-1,-1), 3),
         ('ROWBACKGROUNDS', (0,0), (-1,-1), [colors.HexColor('#f5f5f5'), colors.white]),
     ]))
     story.append(model_table)
     story.append(Spacer(1, 0.4*cm))
-
-    # ── Disclaimer ────────────────────────────────────────────────────────────
     story.append(HRFlowable(width="100%", thickness=1, color=colors.HexColor('#203a43')))
     story.append(Spacer(1, 0.2*cm))
     story.append(Paragraph(
-        "⚠ DISCLAIMER: This report is generated by an AI research prototype for educational purposes only. "
+        "DISCLAIMER: This report is generated by an AI research prototype for educational purposes only. "
         "It is NOT a medical device and must not be used for clinical decision-making. "
         "All demo footage is from the publicly available CholecSeg8k research dataset — no real patient data involved.",
-        footer_style
-    ))
-    story.append(Paragraph("Built for Build Small Hackathon 2026 · huggingface.co/build-small-hackathon", footer_style))
-
+        footer_style))
+    story.append(Paragraph("Built for Build Small Hackathon 2026", footer_style))
     doc.build(story)
     return pdf_path
 
 
 def segment_image(input_image, conf_threshold=0.25):
-    global last_result
+    global last_result, chat_context
     if input_image is None:
-        return None, "Please upload a surgical frame to begin.", "No image provided.", "No explanation yet.", gr.update(visible=False)
+        return (None, "Please upload a surgical frame to begin.", "No image provided.",
+                "No explanation yet.", gr.update(visible=False),
+                gr.update(visible=False), gr.update(visible=False),
+                gr.update(visible=False), gr.update(visible=False), [])
 
+    detector = get_detector()
     image_bytes = pil_to_bytes(input_image)
     result = detector.run.remote(image_bytes, conf_threshold)
 
@@ -284,50 +262,103 @@ def segment_image(input_image, conf_threshold=0.25):
     else:
         alert = "All clear — no critical structures flagged."
 
-    if seen:
-        tissue_list = [n for n in seen.keys() if n != "Black Background"]
-        if tissue_list:
-            prompt = (
-                f"You are a surgical anatomy teacher helping a junior medical resident. "
-                f"These tissues and instruments were detected in a laparoscopic cholecystectomy frame: {', '.join(tissue_list)}. "
-                f"In 3 sentences, explain what the resident should know about these structures — "
-                f"what they are, why they matter, and what to be careful about."
-            )
-            try:
-                messages = [{"role": "user", "content": prompt}]
-                response = client.chat_completion(messages, max_tokens=180, temperature=0.4)
-                explanation = response.choices[0].message.content.strip()
-            except Exception as e:
-                explanation = f"Explanation unavailable: {str(e)}"
-        else:
-            explanation = "Only background detected — no tissue explanation needed."
-    else:
-        explanation = "No tissues detected to explain."
+    tissue_list = [n for n in seen.keys() if n != "Black Background"]
+    explanation = "No tissues detected to explain."
+    if tissue_list and client:
+        prompt = (
+            f"You are a surgical anatomy teacher helping a junior medical resident. "
+            f"These tissues and instruments were detected in a laparoscopic cholecystectomy frame: {', '.join(tissue_list)}. "
+            f"In 3 sentences, explain what the resident should know about these structures — "
+            f"what they are, why they matter, and what to be careful about."
+        )
+        try:
+            response = client.chat_completion([{"role": "user", "content": prompt}], max_tokens=180, temperature=0.4)
+            explanation = response.choices[0].message.content.strip()
+        except Exception as e:
+            explanation = f"Explanation unavailable: {str(e)}"
+
+    # Store context for chat
+    chat_context = {
+        "tissue_list": tissue_list,
+        "alert": alert,
+        "initial_explanation": explanation
+    }
 
-    # Store for PDF generation
     last_result = {
-        "original": input_image,
-        "annotated": annotated_image,
-        "seen": seen,
-        "alert": alert,
-        "explanation": explanation
+        "original": input_image, "annotated": annotated_image,
+        "seen": seen, "alert": alert, "explanation": explanation
     }
 
-    return annotated_image, summary, alert, explanation, gr.update(visible=True)
+    # Generate suggested questions
+    suggested = generate_suggested_questions(tissue_list) if tissue_list else []
+    q1 = suggested[0] if len(suggested) > 0 else ""
+    q2 = suggested[1] if len(suggested) > 1 else ""
+    q3 = suggested[2] if len(suggested) > 2 else ""
+
+    # Initial chat message from assistant
+    initial_chat = [{"role": "assistant", "content": f"I've analysed the frame. I detected: **{', '.join(tissue_list) if tissue_list else 'no tissues'}**.\n\n{explanation}\n\nFeel free to ask me anything about these structures or general laparoscopic surgery!"}]
+
+    return (annotated_image, summary, alert, explanation,
+            gr.update(visible=True),   # pdf_btn
+            gr.update(visible=True),   # chat section
+            gr.update(value=q1, visible=bool(q1)),
+            gr.update(value=q2, visible=bool(q2)),
+            gr.update(value=q3, visible=bool(q3)),
+            initial_chat)
+
+
+def chat_response(message, history):
+    global chat_context
+    if not message.strip():
+        return history, ""
+
+    tissue_list = chat_context.get("tissue_list", [])
+    alert = chat_context.get("alert", "")
+
+    if tissue_list:
+        system_prompt = (
+            f"You are SurgiSight, an expert surgical anatomy assistant helping medical trainees. "
+            f"The current laparoscopic frame shows these detected tissues/instruments: {', '.join(tissue_list)}. "
+            f"Safety status: {alert}. "
+            f"Answer questions about these specific structures or general laparoscopic/surgical questions. "
+            f"Be concise, educational, and clinically accurate. Use 2-4 sentences max per response."
+        )
+    else:
+        system_prompt = (
+            "You are SurgiSight, an expert surgical anatomy assistant for medical trainees. "
+            "Answer general laparoscopic surgery and anatomy questions concisely and accurately. "
+            "Use 2-4 sentences max per response."
+        )
+
+    messages = [{"role": "system", "content": system_prompt}]
+    for msg in history:
+        messages.append({"role": msg["role"], "content": msg["content"]})
+    messages.append({"role": "user", "content": message})
+
+    try:
+        response = client.chat_completion(messages, max_tokens=200, temperature=0.5)
+        reply = response.choices[0].message.content.strip()
+    except Exception as e:
+        reply = f"Sorry, I couldn't process that: {str(e)}"
+
+    history.append({"role": "user", "content": message})
+    history.append({"role": "assistant", "content": reply})
+    return history, ""
+
+
+def use_suggested(question, history):
+    """Fill chat with a suggested question and immediately get response."""
+    return chat_response(question, history)
 
 
 def export_pdf():
     global last_result
     if not last_result:
         return None
-    pdf_path = generate_pdf(
-        last_result["original"],
-        last_result["annotated"],
-        last_result["seen"],
-        last_result["alert"],
-        last_result["explanation"]
+    return generate_pdf(
+        last_result["original"], last_result["annotated"],
+        last_result["seen"], last_result["alert"], last_result["explanation"]
     )
-    return pdf_path
 
 
 css = """
@@ -342,7 +373,8 @@ body { font-family: 'Inter', sans-serif; }
 #danger-box textarea { font-family: 'Courier New', monospace !important; font-size: 0.9rem !important; font-weight: bold !important; background: #1a0a0a !important; color: #ff6b6b !important; border-radius: 10px !important; border: 1px solid #ff4444 !important; }
 #explain-box textarea { font-family: 'Inter', sans-serif !important; font-size: 0.88rem !important; line-height: 1.7 !important; background: #0a1a0f !important; color: #6ee7b7 !important; border-radius: 10px !important; border: 1px solid #34d399 !important; }
 .footer-note { text-align: center; font-size: 0.78rem; color: #888; margin-top: 16px; padding-bottom: 24px; }
-#pdf-btn { background: linear-gradient(135deg, #203a43, #2c5364) !important; color: white !important; border-radius: 10px !important; margin-top: 8px; }
+#suggested-q button { background: rgba(32,58,67,0.08) !important; border: 1px solid #203a43 !important; color: #203a43 !important; border-radius: 999px !important; font-size: 0.82rem !important; padding: 6px 14px !important; margin: 2px !important; }
+#suggested-q button:hover { background: #203a43 !important; color: white !important; }
 """
 
 header_html = """
@@ -350,7 +382,7 @@ header_html = """
   <h1>🔬 Surgical Tissue Segmentation</h1>
   <p>Upload a laparoscopic surgery frame — or click any example below. The AI identifies tissues, flags danger zones, and explains anatomy for surgical trainees.</p>
   <div class="badge-row">
-    <span class="badge">YOLO26n-seg</span>
+    <span class="badge">YOLOv8-seg</span>
     <span class="badge">Llama 3.1 8B</span>
     <span class="badge">Modal GPU</span>
     <span class="badge">CholecSeg8k Dataset</span>
@@ -369,6 +401,8 @@ footer_html = """
 
 with gr.Blocks(title="SurgiSight — Surgical Tissue Segmentation") as demo:
     gr.HTML(header_html)
+
+    # ── Input / Output ──────────────────────────────────────────────────────
     with gr.Row():
         with gr.Column(scale=1):
             input_img = gr.Image(type="pil", label="Input — Laparoscopic Frame", height=350)
@@ -383,14 +417,40 @@ with gr.Blocks(title="SurgiSight — Surgical Tissue Segmentation") as demo:
                              placeholder="Safety alert will appear here after running segmentation...")
     output_text = gr.Textbox(label="Detected Tissues & Confidence", lines=6, elem_id="output-text",
                               placeholder="Tissue detection results will appear here...")
-    explain_box = gr.Textbox(label="🧠 Anatomy Explanation for Surgical Trainees",
-                              lines=5, elem_id="explain-box",
+    explain_box = gr.Textbox(label="🧠 Initial Anatomy Explanation",
+                              lines=4, elem_id="explain-box",
                               placeholder="AI anatomy explanation will appear here after segmentation...")
 
+    # ── PDF button ──────────────────────────────────────────────────────────
     with gr.Row():
-        pdf_btn = gr.Button("📄  Download Full Report (PDF)", visible=False, elem_id="pdf-btn", size="lg")
+        pdf_btn = gr.Button("📄  Download Full Report (PDF)", visible=False, size="lg")
     pdf_output = gr.File(label="Your Report", visible=False)
 
+    # ── Chat Section ────────────────────────────────────────────────────────
+    with gr.Column(visible=False) as chat_section:
+        gr.HTML("<hr style='margin:20px 0;border-color:#e0e0e0'><h3 style='margin-bottom:8px'>💬 Ask the AI — Follow-up Questions</h3><p style='font-size:0.85rem;color:#666;margin-bottom:12px'>Ask anything about the detected tissues or general laparoscopic surgery</p>")
+
+        # Suggested question chips
+        with gr.Row(elem_id="suggested-q"):
+            sq1 = gr.Button("", visible=False, size="sm")
+            sq2 = gr.Button("", visible=False, size="sm")
+            sq3 = gr.Button("", visible=False, size="sm")
+
+        chatbot = gr.Chatbot(
+            label="SurgiSight Chat",
+            type="messages",
+            height=350,
+            show_label=False,
+            avatar_images=(None, "https://huggingface.co/front/assets/huggingface_logo-noborder.svg")
+        )
+
+        with gr.Row():
+            chat_input = gr.Textbox(
+                placeholder="Ask about detected tissues, surgical anatomy, or any surgery question...",
+                show_label=False, scale=5, container=False
+            )
+            send_btn = gr.Button("Send ➤", variant="primary", scale=1)
+
     gr.Examples(
         examples=[
             ["examples/frame_80_endo.png"],
@@ -402,20 +462,28 @@ with gr.Blocks(title="SurgiSight — Surgical Tissue Segmentation") as demo:
         examples_per_page=3
     )
 
+    gr.HTML(footer_html)
+
+    # ── Event handlers ──────────────────────────────────────────────────────
     run_btn.click(
         fn=segment_image,
         inputs=[input_img, conf_slider],
-        outputs=[output_img, output_text, danger_box, explain_box, pdf_btn],
+        outputs=[output_img, output_text, danger_box, explain_box,
+                 pdf_btn, chat_section, sq1, sq2, sq3, chatbot],
         show_progress="full"
     )
 
-    pdf_btn.click(
-        fn=export_pdf,
-        inputs=[],
-        outputs=[pdf_output]
+    pdf_btn.click(fn=export_pdf, inputs=[], outputs=[pdf_output]
     ).then(fn=lambda: gr.update(visible=True), outputs=[pdf_output])
 
-    gr.HTML(footer_html)
+    # Send button / Enter key
+    send_btn.click(fn=chat_response, inputs=[chat_input, chatbot], outputs=[chatbot, chat_input])
+    chat_input.submit(fn=chat_response, inputs=[chat_input, chatbot], outputs=[chatbot, chat_input])
+
+    # Suggested question chips → send directly into chat
+    sq1.click(fn=use_suggested, inputs=[sq1, chatbot], outputs=[chatbot, chat_input])
+    sq2.click(fn=use_suggested, inputs=[sq2, chatbot], outputs=[chatbot, chat_input])
+    sq3.click(fn=use_suggested, inputs=[sq3, chatbot], outputs=[chatbot, chat_input])
 
 if __name__ == "__main__":
-    demo.launch(css=css, inbrowser=True)
+    demo.launch(css=css)