Update app.py

app.py

@@ -1,45 +1,145 @@
-from flask import Flask, request, jsonify
-import openai
 import os
+import numpy as np
+import torch
+from PIL import Image
+import torchxrayvision as xrv
+from transformers import pipeline
+import gradio as gr
 
-app = Flask(__name__)
+# -----------------------------
+# Setup
+# -----------------------------
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load API key from environment variable
-openai.api_key = os.getenv("OPENAI_API_KEY")
+def load_models():
+    chest_model = xrv.models.get_model("densenet121-res224-all").to(DEVICE).eval()
+    clip = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32")
+    return chest_model, clip
 
-@app.route('/')
-def index():
-    return jsonify({"message": "Radiology AI API is running."})
+CHEST_MODEL, CLIP = load_models()
+LABELS = CHEST_MODEL.pathologies
 
-@app.route('/analyze', methods=['POST'])
-def analyze_scan():
-    scan_type = request.form.get('scan_type', 'medical scan')
-    findings_input = request.form.get('findings', '')
+# -----------------------------
+# Helper Functions
+# -----------------------------
+def detect_scan_type(pil_img):
+    candidates = ["chest", "bone", "brain", "abdomen", "dental", "cardiac"]
+    result = CLIP(pil_img, candidate_labels=candidates)
+    return max(result, key=lambda x: x['score'])['label']
 
-    # Prompt focuses on expanding only the Results
-    prompt = f"""
-You are a professional radiology AI assistant.
-A {scan_type} scan was uploaded. Findings extracted: {findings_input}.
+def preprocess_chest(pil_img):
+    if pil_img.mode != "L":
+        pil_img = pil_img.convert("L")
+    img = np.array(pil_img).astype(np.float32)
+    img = xrv.datasets.normalize(img, 255)
+    img = img[None, ...]
+    img = xrv.datasets.XRayCenterCrop()(img)
+    img = xrv.datasets.XRayResizer(224)(img)
+    return torch.from_numpy(img).unsqueeze(0).to(DEVICE)
 
-Write a structured report with two sections:
+def analyse_chest(pil_img):
+    x = preprocess_chest(pil_img)
+    with torch.no_grad():
+        output = CHEST_MODEL(x)
+        probs = torch.sigmoid(output)[0].cpu().numpy() * 100
+    findings = [(LABELS[i], p) for i, p in enumerate(probs)]
+    findings = sorted(findings, key=lambda x: x[1], reverse=True)
+    return [(cond, conf) for cond, conf in findings if conf > 40][:7]
 
-1. Results → Must be elaborated into at least 2–3 sentences (more than 3 lines), written in a formal radiology report style. Expand with detail, clinical context, and interpretation. Avoid bullet points.
+# -----------------------------
+# Clinical Interpretation
+# -----------------------------
+def interpret_findings(findings):
+    strong, moderate, mild = [], [], []
 
-2. Recommendations → Keep concise, 1–2 lines, with direct clinical advice, follow-up, or specialist referral if appropriate.
-"""
+    for cond, conf in findings:
+        if conf >= 75:
+            strong.append((cond, conf))
+        elif conf >= 55:
+            moderate.append((cond, conf))
+        elif conf >= 40:
+            mild.append((cond, conf))
 
-    response = openai.ChatCompletion.create(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": "You are a helpful medical radiology AI assistant."},
-            {"role": "user", "content": prompt}
-        ],
-        max_tokens=600,
-        temperature=0.7
-    )
+    return strong, moderate, mild
 
-    report_text = response['choices'][0]['message']['content'].strip()
-    return jsonify({"report": report_text})
+def recommendations_for_condition(cond):
+    recs = {
+        "Fibrosis": "Pulmonology referral, consider HRCT chest.",
+        "Infiltration": "Suggest clinical correlation (cough, fever, infection signs). Consider antibiotics and follow-up imaging.",
+        "Mass": "Oncology referral; biopsy or CT chest may be warranted.",
+        "Nodule": "Follow Fleischner guidelines; consider repeat imaging or biopsy.",
+        "Pleural_Thickening": "Consider TB or asbestos exposure; recommend pulmonology referral.",
+        "Consolidation": "Likely pneumonia; recommend clinical correlation, antibiotics, and repeat X-ray in 6–8 weeks.",
+        "Effusion": "Evaluate cardiac/renal function; consider thoracentesis if large.",
+        "Atelectasis": "Encourage physiotherapy, bronchoscopy if persistent.",
+        "Cardiomegaly": "Cardiology referral, ECG, and echocardiogram recommended."
+    }
+    return recs.get(cond, "Correlation with symptoms and specialist referral advised.")
 
-if __name__ == '__main__':
-    app.run(host="0.0.0.0", port=5000, debug=True)
+def format_report(scan_type, findings):
+    if scan_type == "chest":
+        if not findings:
+            return "✅ **Chest X-ray Report**\n\nNormal chest. No significant abnormalities detected."
+
+        strong, moderate, mild = interpret_findings(findings)
+
+        report = "## 🫁 Chest X-ray Report\n\n"
+
+        # Summary
+        if strong:
+            report += f"**Overall Impression:** ⚠️ Abnormal – Strong evidence of **{', '.join([c for c, _ in strong])}**.\n\n"
+        elif moderate:
+            report += f"**Overall Impression:** ❗ Suspicious – Possible **{', '.join([c for c, _ in moderate])}**.\n\n"
+        else:
+            report += "**Overall Impression:** ✅ No major abnormalities detected.\n\n"
+
+        # Key Findings
+        report += "### 🔍 Key Findings\n"
+        if strong:
+            report += "\n**Significant Abnormalities:**\n"
+            for cond, conf in strong:
+                report += f"- {cond} ({conf:.1f}%) → {recommendations_for_condition(cond)}\n"
+        if moderate:
+            report += "\n**Possible Abnormalities:**\n"
+            for cond, conf in moderate:
+                report += f"- {cond} ({conf:.1f}%) → {recommendations_for_condition(cond)}\n"
+        if mild:
+            report += "\n**Minor / Non-specific Findings:**\n"
+            for cond, conf in mild:
+                report += f"- {cond} ({conf:.1f}%) → Monitor clinically.\n"
+
+        return report
+
+    # ------------------ Other Scans -------------------
+    elif scan_type == "bone":
+        return "🦴 **Bone X-ray Report**\n\n- Possible fracture or degenerative changes.\n- Recommendation: Orthopedic consultation."
+    elif scan_type == "brain":
+        return "🧠 **Brain MRI/CT Report**\n\n- Possible tumor, lesion, or stroke signs.\n- Recommendation: Neurology referral."
+    elif scan_type == "abdomen":
+        return "🩺 **Abdomen Scan Report**\n\n- Possible mass or opacity detected.\n- Recommendation: Gastroenterology referral."
+    elif scan_type == "dental":
+        return "😁 **Dental X-ray Report**\n\n- Possible cavities or gum disease.\n- Recommendation: Dentist referral."
+    elif scan_type == "cardiac":
+        return "❤️ **Cardiac Scan Report**\n\n- Possible cardiomegaly or valve abnormality.\n- Recommendation: Cardiology referral."
+    else:
+        return "⚠️ Unknown scan type. Please upload a valid medical scan."
+
+# -----------------------------
+# Gradio Interface
+# -----------------------------
+def analyse_scan(image):
+    pil_img = image.convert("RGB")
+    scan_type = detect_scan_type(pil_img)
+    findings = analyse_chest(pil_img) if scan_type == "chest" else None
+    return format_report(scan_type, findings)
+
+demo = gr.Interface(
+    fn=analyse_scan,
+    inputs=gr.Image(type="pil"),
+    outputs="markdown",
+    title="🩻 Universal Radiology AI",
+    description="Upload a scan (Chest, Bone, Brain, Abdomen, Dental, Cardiac) to get an AI-generated structured radiology report. ⚠️ Not a medical device."
+)
+
+if __name__ == "__main__":
+    demo.launch()