Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from PIL import Image
 import torch
 from torchvision import models, transforms
+import fitz  # PyMuPDF
 import logging
 import os
 
@@ -9,37 +10,68 @@ import os
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# ChestX-ray14 condition labels
+# Condition list
 conditions = [
-    "Atelectasis", "Cardiomegaly", "Effusion", "Infiltration", "Mass", "Nodule",
-    "Pneumonia", "Pneumothorax", "Consolidation", "Edema", "Emphysema", "Fibrosis",
-    "Pleural Thickening", "Hernia"
+    "Normal", "Pneumonia", "Cancer", "TB", "Other",
+    "Coronary Artery Disease", "Aortic Aneurysm", "Stroke", "Peripheral Artery Disease",
+    "Brain Tumor", "Alzheimer's Disease", "Multiple Sclerosis", "Epilepsy",
+    "COPD", "Lung Cancer", "Pulmonary Embolism",
+    "Fractures", "Arthritis", "Osteoporosis",
+    "Appendicitis", "Gallstones", "Kidney Stones", "Infections", "Abdominal Aortic Aneurysm", "Diverticulitis"
 ]
 
-# Load DenseNet121 base model
-model = models.densenet121(pretrained=False)
-model.classifier = torch.nn.Linear(model.classifier.in_features, len(conditions))
+# Condition details
+condition_details = {
+    "Normal": {"description": "No abnormal signs detected.", "recommendation": "Routine check-ups recommended."},
+    "Pneumonia": {"description": "Lung inflammation detected, possibly infectious.", "recommendation": "Seek medical attention for treatment."},
+    "Cancer": {"description": "Suspicious masses suggest cancer; further imaging needed.", "recommendation": "Consult an oncologist."},
+    "TB": {"description": "Cavitary lesions indicate tuberculosis.", "recommendation": "Immediate medical evaluation required."},
+    "Other": {"description": "Unclear abnormality; further investigation needed.", "recommendation": "Consult a radiologist."},
+    "Coronary Artery Disease": {"description": "Narrowing of coronary arteries detected.", "recommendation": "Cardiology consultation required."},
+    "Aortic Aneurysm": {"description": "Abnormal enlargement of the aorta.", "recommendation": "Vascular surgery evaluation."},
+    "Stroke": {"description": "Signs of brain ischemia or hemorrhage.", "recommendation": "Urgent neurological evaluation."},
+    "Peripheral Artery Disease": {"description": "Reduced blood flow in peripheral arteries.", "recommendation": "Vascular specialist consultation."},
+    "Brain Tumor": {"description": "Abnormal mass in the brain detected.", "recommendation": "Consult a neurosurgeon."},
+    "Alzheimer's Disease": {"description": "Signs of neurodegenerative changes.", "recommendation": "Neurology consultation."},
+    "Multiple Sclerosis": {"description": "Demyelinating lesions in the CNS.", "recommendation": "Neurology consultation."},
+    "Epilepsy": {"description": "Signs of seizure activity.", "recommendation": "Neurology consultation."},
+    "COPD": {"description": "Lung damage from COPD observed.", "recommendation": "Pulmonary consultation."},
+    "Lung Cancer": {"description": "Malignant lung masses detected.", "recommendation": "Oncology consultation."},
+    "Pulmonary Embolism": {"description": "Blockage in pulmonary arteries.", "recommendation": "Urgent medical attention."},
+    "Fractures": {"description": "Bone break detected.", "recommendation": "Orthopedic evaluation."},
+    "Arthritis": {"description": "Joint inflammation detected.", "recommendation": "Rheumatology consultation."},
+    "Osteoporosis": {"description": "Reduced bone density observed.", "recommendation": "Bone health specialist consultation."},
+    "Appendicitis": {"description": "Inflammation of the appendix.", "recommendation": "Surgical evaluation."},
+    "Gallstones": {"description": "Stones in the gallbladder detected.", "recommendation": "Gastroenterology consultation."},
+    "Kidney Stones": {"description": "Stones in the kidneys detected.", "recommendation": "Urology consultation."},
+    "Infections": {"description": "Signs of infection observed.", "recommendation": "Infectious disease consultation."},
+    "Abdominal Aortic Aneurysm": {"description": "Enlargement of the abdominal aorta.", "recommendation": "Vascular surgery evaluation."},
+    "Diverticulitis": {"description": "Inflammation of diverticula in the colon.", "recommendation": "Gastroenterology consultation."}
+}
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = model.to(device)
+# Load model
+try:
+    model = models.densenet121(weights="IMAGENET1K_V1")
+except AttributeError:
+    model = models.densenet121(pretrained=True)
+
+model.classifier = torch.nn.Linear(model.classifier.in_features, len(conditions))
 model.eval()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-# Load CheXNet pre-trained weights
-model_path = "xray_model.pth"
+# Load model weights if available
+model_path = os.getenv("MODEL_PATH", "xray_model.pth")
 if os.path.exists(model_path):
     try:
-        checkpoint = torch.load(model_path, map_location=device)
-        if "state_dict" in checkpoint:
-            model.load_state_dict(checkpoint["state_dict"])
-        else:
-            model.load_state_dict(checkpoint)
-        logger.info("✅ Loaded CheXNet model weights.")
+        model.load_state_dict(torch.load(model_path, map_location=device))
+        logger.info("Model loaded from file.")
     except Exception as e:
-        logger.error(f"❌ Failed to load model weights: {e}")
+        logger.warning(f"Failed to load model weights: {e}")
 else:
-    logger.warning("⚠️ Model file 'xray_model.pth' not found!")
+    logger.info("No custom model weights found.")
 
-# Preprocessing for image input
+# Image preprocessing
 def preprocess_image(image):
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
@@ -48,7 +80,7 @@ def preprocess_image(image):
     ])
     return transform(image).unsqueeze(0).to(device)
 
-# X-ray prediction
+# X-ray prediction function with confidence threshold
 def predict_xray(image):
     try:
         if image is None:
@@ -58,34 +90,78 @@ def predict_xray(image):
         with torch.no_grad():
             output = model(img_tensor)
 
-        probs = torch.sigmoid(output)[0]  # sigmoid for multi-label prediction
-        result_lines = []
+        probs = torch.nn.functional.softmax(output, dim=1)[0]
+        results = {conditions[i]: probs[i].item() * 100 for i in range(len(conditions))}
 
-        for i, condition in enumerate(conditions):
-            confidence = probs[i].item() * 100
-            if confidence >= 10:  # only show confident predictions
-                result_lines.append(f"<b>{condition}:</b> {confidence:.2f}%")
+        top_condition = max(results, key=results.get)
+        confidence = results[top_condition]
 
-        if not result_lines:
-            return "<b>Uncertain:</b> No strong signs of any known condition detected."
+        if confidence < 50:
+            return f"""
+            <div style="font-family:Arial">
+                <h3>Prediction: <span style="color:#D62828;">Uncertain</span></h3>
+                <p><b>Confidence:</b> {confidence:.2f}%</p>
+                <p><b>Note:</b> The model is not confident enough to provide a clear diagnosis.</p>
+                <p><b>Recommendation:</b> Please consult a radiologist or upload a better-quality image.</p>
+            </div>
+            """
 
-        return "<br>".join(result_lines)
+        info = condition_details.get(top_condition, condition_details["Other"])
+        return f"""
+        <div style="font-family:Arial">
+            <h3>Prediction: <span style="color:#2A9D8F;">{top_condition}</span></h3>
+            <p><b>Confidence:</b> {confidence:.2f}%</p>
+            <p><b>Description:</b> {info['description']}</p>
+            <p><b>Recommendation:</b> {info['recommendation']}</p>
+        </div>
+        """
 
     except Exception as e:
-        logger.error(f"Prediction failed: {e}")
-        return f"<b>Error:</b> {str(e)}"
+        logger.error(f"Error in prediction: {e}")
+        return f"Error: {str(e)}"
+
+# Analyze PDF report
+def analyze_report(file):
+    if not file or not file.name.endswith(".pdf"):
+        return "Please upload a valid PDF file."
+    try:
+        doc = fitz.open(file.name)
+        text = "".join(page.get_text() for page in doc)
+        doc.close()
+
+        condition, disease, status = "Unclear", "Unknown", "Pending"
 
-# Gradio interface setup
+        if "stroke" in text.lower():
+            condition, disease, status = "Stroke", "Brain Disorder", "Urgent Care Needed"
+        elif "cancer" in text.lower():
+            condition, disease, status = "Cancer", "Malignant Growth", "Consult Oncologist"
+        elif "fracture" in text.lower():
+            condition, disease, status = "Fracture", "Bone Injury", "Orthopedic Attention Required"
+
+        preview = text[:300] + "..." if text else "No readable content."
+        return f"Condition: {condition}\nDisease: {disease}\nStatus: {status}\n\nPreview:\n{preview}"
+
+    except Exception as e:
+        return f"Failed to process PDF: {str(e)}"
+
+# Gradio interface
 def create_interface():
     with gr.Blocks() as demo:
-        gr.Markdown("<h1 style='text-align:center;'>🩻 RadiologyScan AI (CheXNet)</h1>")
-        with gr.Row():
-            image_input = gr.Image(label="Upload Chest X-ray", type="pil")
-            result_output = gr.HTML(label="Diagnosis Result")
-        gr.Button("Analyze X-ray").click(predict_xray, inputs=image_input, outputs=result_output)
+        gr.Markdown("<h1 style='text-align:center;'>🩻 RadiologyScan AI</h1><p style='text-align:center;'>AI-powered X-ray and Report Analysis</p>")
+
+        with gr.Tabs():
+            with gr.TabItem("X-ray Analysis"):
+                img_input = gr.Image(label="Upload Chest X-ray", type="pil")
+                img_output = gr.HTML()
+                gr.Button("Analyze X-ray").click(predict_xray, inputs=img_input, outputs=img_output)
+
+            with gr.TabItem("Report Analysis"):
+                pdf_input = gr.File(label="Upload PDF Report", file_types=[".pdf"])
+                pdf_output = gr.Textbox(label="Extracted Summary", lines=10)
+                gr.Button("Analyze Report").click(analyze_report, inputs=pdf_input, outputs=pdf_output)
+
     return demo
 
-# Run the app
 if __name__ == "__main__":
     demo = create_interface()
     demo.launch(server_port=7860, ssr_mode=False)