Update app.py

app.py

@@ -4,33 +4,52 @@ import torch
 from torchvision import models, transforms
 import PyPDF2
 import logging
+import os
 
 # Set up logging
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-# Load the pre-trained model (DenseNet121, suitable for medical imaging)
-model = models.densenet121(pretrained=True)  # Using ImageNet weights as a starting point
+# Define conditions based on provided radiology information
+conditions = [
+    "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 and configure the model
+model = models.densenet121(pretrained=False)  # Start without pre-trained weights for custom training
 num_features = model.classifier.in_features
-model.classifier = torch.nn.Linear(num_features, 5)  # Output layer for 5 conditions
+model.classifier = torch.nn.Linear(num_features, len(conditions))  # Output for all 16 conditions
 model.eval()
 
-# Define device (CPU or GPU)
+# Define device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
 
-# Define image preprocessing function with normalization
+# Load model state if available, otherwise initialize
+model_path = "xray_model.pth"
+if os.path.exists(model_path):
+    model.load_state_dict(torch.load(model_path))
+    logger.info(f"Loaded model from {model_path}")
+else:
+    logger.info("No pre-trained model found. Initializing with random weights. Training required.")
+
+# Define image preprocessing function
 def preprocess_image(image):
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # ImageNet normalization
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
     ])
     image_tensor = transform(image).unsqueeze(0).to(device)
     logger.debug(f"Preprocessed image tensor shape: {image_tensor.shape}")
     return image_tensor
 
-# Define a prediction function for X-ray images with detailed output
+# Define prediction function with detailed output
 def predict_xray(image):
     try:
         if image is None:
@@ -39,51 +58,49 @@ def predict_xray(image):
         image_tensor = preprocess_image(image)
         with torch.no_grad():
             outputs = model(image_tensor)
-        probs = torch.nn.functional.softmax(outputs, dim=1)[0]  # Softmax over the 5 classes
+        probs = torch.nn.functional.softmax(outputs, dim=1)[0]  # Softmax over all conditions
+        results = {conditions[i]: float(probs[i].cpu().numpy()) * 100 for i in range(len(conditions))}
 
-        # Define the conditions
-        conditions = ["Normal", "Pneumonia", "Cancer", "TB", "Other"]
-        results = {conditions[i]: float(probs[i].cpu().numpy()) * 100 for i in range(5)}
-        
-        # Determine the most likely condition and confidence
         most_likely_condition = max(results, key=results.get)
         confidence = results[most_likely_condition]
 
-        # Generate summary
         summary = f"**Summary**: Based on the X-ray analysis, the most likely diagnosis is: <b>{most_likely_condition}</b> with a confidence of <b>{confidence:.2f}%</b>."
 
-        # Condition details with enhanced descriptions
+        # Enhanced condition details
         condition_details = {
-            "Normal": {
-                "description": "The X-ray shows no abnormal signs, indicating healthy lung tissue with clear structures.",
-                "recommendation": "No immediate action required. Schedule routine check-ups to monitor lung health."
-            },
-            "Pneumonia": {
-                "description": "Pneumonia is detected, showing lung inflammation, possibly due to bacterial or viral infection, with visible opacities.",
-                "recommendation": "Seek medical attention promptly; treatment may include antibiotics or antiviral medication."
-            },
-            "Cancer": {
-                "description": "Suspicious masses or nodules suggest lung cancer, requiring advanced imaging (e.g., CT) for confirmation.",
-                "recommendation": "Urgently consult an oncologist for a biopsy and personalized treatment plan."
-            },
-            "TB": {
-                "description": "Tuberculosis is indicated by cavitary lesions or consolidation, a contagious bacterial infection.",
-                "recommendation": "Contact a healthcare provider immediately for a treatment regimen, likely involving multiple antibiotics."
-            },
-            "Other": {
-                "description": "Unclear abnormalities detected; could indicate conditions like fibrosis or heart-related issues.",
-                "recommendation": "Refer to a radiologist for specialized imaging and diagnosis."
-            }
+            "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": "Blockages in heart arteries detected.", "recommendation": "Cardiology consultation advised."},
+            "Aortic Aneurysm": {"description": "Aortic dilation observed.", "recommendation": "Monitor with imaging; surgical consult if large."},
+            "Stroke": {"description": "Brain damage from stroke detected.", "recommendation": "Urgent neurological care needed."},
+            "Peripheral Artery Disease": {"description": "Reduced limb blood flow observed.", "recommendation": "Vascular specialist consultation."},
+            "Brain Tumor": {"description": "Abnormal growth in brain detected.", "recommendation": "Neurological evaluation required."},
+            "Alzheimer's Disease": {"description": "Brain atrophy suggestive of Alzheimer's.", "recommendation": "Neurological assessment."},
+            "Multiple Sclerosis": {"description": "Lesions in brain/spinal cord detected.", "recommendation": "Consult neurologist."},
+            "Epilepsy": {"description": "Seizure source possibly identified.", "recommendation": "Neurological workup needed."},
+            "COPD": {"description": "Lung damage from COPD observed.", "recommendation": "Pulmonary consultation."},
+            "Lung Cancer": {"description": "Nodules suggest lung cancer.", "recommendation": "Oncologist referral."},
+            "Pulmonary Embolism": {"description": "Blood clot in lungs detected.", "recommendation": "Emergency care required."},
+            "Fractures": {"description": "Bone break detected.", "recommendation": "Orthopedic evaluation."},
+            "Arthritis": {"description": "Joint damage observed.", "recommendation": "Rheumatology consult."},
+            "Osteoporosis": {"description": "Reduced bone density detected.", "recommendation": "Bone health assessment."},
+            "Appendicitis": {"description": "Inflammation of appendix observed.", "recommendation": "Surgical evaluation."},
+            "Gallstones": {"description": "Stones in gallbladder detected.", "recommendation": "Gastroenterology consult."},
+            "Kidney Stones": {"description": "Stones in kidneys observed.", "recommendation": "Urology evaluation."},
+            "Infections": {"description": "Signs of infection detected.", "recommendation": "Infectious disease consult."},
+            "Abdominal Aortic Aneurysm": {"description": "Abdominal aortic dilation observed.", "recommendation": "Vascular surgery consult."},
+            "Diverticulitis": {"description": "Digestive tract inflammation detected.", "recommendation": "Gastroenterology evaluation."}
         }
 
-        # Detailed results in a structured format
         detailed_results = "<ul class='result-list'>"
         for condition, prob in results.items():
             detailed_results += f"<li><b>{condition}:</b> {prob:.2f}%</li>"
         detailed_results += "</ul>"
 
-        # Additional feedback based on the condition
-        additional_feedback = condition_details.get(most_likely_condition, "Please consult a medical professional for a detailed evaluation.")
+        additional_feedback = f"<div class='feedback-box'><b>Description:</b> {condition_details[most_likely_condition]['description']}<br><b>Recommendation:</b> {condition_details[most_likely_condition]['recommendation']}</div>"
 
         logger.info(f"Prediction: {most_likely_condition} with confidence {confidence:.2f}%")
         return summary, detailed_results, additional_feedback
@@ -92,7 +109,7 @@ def predict_xray(image):
         logger.error(f"Error in predict_xray: {str(e)}")
         return f"Error: {str(e)}", "", ""
 
-# Define a function to read and analyze patient reports (PDFs)
+# Define function to read and analyze patient reports (PDFs)
 def analyze_report(file):
     text = ""
     if file and file.name.endswith(".pdf"):
@@ -111,82 +128,36 @@ def analyze_report(file):
 # Gradio Interface with enhanced UI
 def create_interface():
     with gr.Blocks() as demo:
-        # Custom CSS for UI enhancement
         custom_css = """
-        .gradio-container {
-            background-color: #f4f6f9;
-            border-radius: 15px;
-            box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
-            padding: 30px;
-            font-family: 'Segoe UI', sans-serif;
-        }
-        .title {
-            font-size: 30px;
-            text-align: center;
-            color: #4C6A92;
-            margin-bottom: 20px;
-        }
-        .gradio-button {
-            background-color: #3B82F6;
-            color: white;
-            border-radius: 10px;
-            padding: 15px 30px;
-            font-size: 16px;
-            transition: background-color 0.3s;
-        }
-        .gradio-button:hover {
-            background-color: #2563EB;
-        }
-        .result-box {
-            background-color: #ffffff;
-            border-radius: 10px;
-            padding: 20px;
-            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
-            margin-top: 20px;
-            max-width: 100%;
-        }
-        .result-list {
-            padding-left: 20px;
-            margin: 10px 0;
-        }
-        .result-summary {
-            font-size: 18px;
-            color: #2F4F4F;
-            font-weight: 500;
-        }
-        .feedback-box {
-            background-color: #F0FFF4;
-            padding: 10px;
-            border-left: 4px solid #38A169;
-            border-radius: 5px;
-            margin-top: 10px;
-        }
+        .gradio-container { background-color: #f4f6f9; border-radius: 15px; box-shadow: 0 4px 15px rgba(0,0,0,0.1); padding: 30px; font-family: 'Segoe UI', sans-serif; }
+        .title { font-size: 30px; text-align: center; color: #4C6A92; margin-bottom: 20px; }
+        .gradio-button { background-color: #3B82F6; color: white; border-radius: 10px; padding: 15px 30px; font-size: 16px; transition: background-color 0.3s; }
+        .gradio-button:hover { background-color: #2563EB; }
+        .result-box { background-color: #ffffff; border-radius: 10px; padding: 20px; box-shadow: 0 4px 8px rgba(0,0,0,0.1); margin-top: 20px; max-width: 100%; }
+        .result-list { padding-left: 20px; margin: 10px 0; }
+        .result-summary { font-size: 18px; color: #2F4F4F; font-weight: 500; }
+        .feedback-box { background-color: #F0FFF4; padding: 10px; border-left: 4px solid #38A169; border-radius: 5px; margin-top: 10px; }
         """
         
-        # Title and description
         gr.Markdown("<h1 class='title'>RadiologyScan AI</h1>")
-        gr.Markdown("<p style='text-align: center; color: #666;'>Advanced X-ray and patient report analysis powered by AI</p>")
+        gr.Markdown("<p style='text-align: center; color: #666;'>AI-powered analysis for X-rays and patient reports</p>")
         
-        # Upload section with layout
         with gr.Row():
             with gr.Column(scale=1):
-                xray_input = gr.Image(label="Upload Chest X-ray", type="pil", elem_id="xray-input")
+                xray_input = gr.Image(label="Upload X-ray", type="pil", elem_id="xray-input")
             with gr.Column(scale=1):
                 report_input = gr.File(label="Upload Patient Report (PDF)", file_count="single", elem_id="report-input")
         
-        # Buttons for analysis
         with gr.Row():
             predict_button = gr.Button("Analyze X-ray", elem_classes="gradio-button")
             report_button = gr.Button("Analyze Report", elem_classes="gradio-button")
         
-        # Results section
         with gr.Column():
             xray_output = gr.HTML(label="X-ray Diagnosis Summary", elem_classes="result-box")
             xray_result = gr.HTML(label="Detailed X-ray Results", elem_classes="result-box")
             additional_feedback = gr.HTML(label="Additional Feedback", elem_classes="result-box feedback-box")
             report_output = gr.Textbox(label="Report Summary", interactive=False, elem_classes="result-box")
 
-        # Event handlers
         predict_button.click(
             fn=predict_xray,
             inputs=xray_input,
@@ -195,11 +166,10 @@ def create_interface():
         report_button.click(
             fn=analyze_report,
             inputs=report_input,
-            outputs=report_output
-        )
+            outputs=report_output)
     
     return demo
 
-# Launch the Gradio interface
+# Launch the interface and save model after training (to be implemented by user)
 demo = create_interface()
 demo.launch(share=True)