Update app.py

app.py

@@ -1,16 +1,16 @@
 import gradio as gr
+from PIL import Image
 import torch
 from torchvision import models, transforms
-import os
-import time
-import logging
 import fitz  # PyMuPDF for better PDF parsing
+import logging
+import os
 
-# Set up logging (optional for debugging)
-logging.basicConfig(level=logging.DEBUG)
+# Set up logging
+logging.basicConfig(level=logging.DEBUG)  # Set to DEBUG for detailed output
 logger = logging.getLogger(__name__)
 
-# Define conditions based on provided radiology information
+# Define conditions
 conditions = [
     "Normal", "Pneumonia", "Cancer", "TB", "Other",
     "Coronary Artery Disease", "Aortic Aneurysm", "Stroke", "Peripheral Artery Disease",
@@ -20,38 +20,71 @@ conditions = [
     "Appendicitis", "Gallstones", "Kidney Stones", "Infections", "Abdominal Aortic Aneurysm", "Diverticulitis"
 ]
 
-# Define path to store the model manually
-model_path = "/home/user/.cache/torch/hub/checkpoints/densenet121-a639ec97.pth"  # Adjusted to a valid path
-
-# Ensure the parent directory exists
-parent_dir = os.path.dirname(model_path)
-if not os.path.exists(parent_dir):
-    os.makedirs(parent_dir)  # Create the parent directory if it doesn't exist
-
-# Function to load the model efficiently
-def load_model():
-    if os.path.exists(model_path):
-        model = models.densenet121()
-        model.load_state_dict(torch.load(model_path))  # Load from cached path
-        model.eval()  # Set to evaluation mode
-        logger.info("Loaded model from cache.")
-    else:
-        model = models.densenet121(weights="IMAGENET1K_V1")  # If not cached, download model
-        torch.save(model.state_dict(), model_path)  # Cache the model locally
-        logger.info("Downloaded and cached the model.")
-    return model
-
-# Load the model at the beginning (this will take time but only happens once)
-model = load_model()
-
-# Define device for model inference
+# Define condition details for all conditions
+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."}
+}
+
+# Load and configure the model
+try:
+    model = models.densenet121(weights="IMAGENET1K_V1")
+    num_features = model.classifier.in_features
+    model.classifier = torch.nn.Linear(num_features, len(conditions))
+    model.eval()
+except AttributeError:
+    model = models.densenet121(pretrained=True)
+    num_features = model.classifier.in_features
+    model.classifier = torch.nn.Linear(num_features, len(conditions))
+    model.eval()
+
+# Define device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
+logger.info(f"Using device: {device}")
+
+# Load model state if available
+model_path = os.getenv("MODEL_PATH", "xray_model.pth")
+if os.path.exists(model_path):
+    try:
+        model.load_state_dict(torch.load(model_path, map_location=device))
+        logger.info(f"Loaded model from {model_path}")
+    except Exception as e:
+        logger.warning(f"Failed to load model from {model_path}: {str(e)}. Using random weights.")
+else:
+    logger.info("No pre-trained model found. Initializing with random weights. Training required.")
 
 # Define image preprocessing function
 def preprocess_image(image):
+    if not isinstance(image, Image.Image):
+        logger.error("Invalid image format. Expected PIL Image.")
+        raise ValueError("Uploaded file is not a valid image.")
+    
     transform = transforms.Compose([
-        transforms.Resize((224, 224)),  # Resize to fit the model input size
+        transforms.Resize((224, 224)),
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
     ])
@@ -59,7 +92,7 @@ def preprocess_image(image):
     logger.debug(f"Preprocessed image tensor shape: {image_tensor.shape}")
     return image_tensor
 
-# Define prediction function with detailed output and error handling
+# Define prediction function
 def predict_xray(image):
     try:
         if image is None:
@@ -68,40 +101,21 @@ 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 all conditions
+        probs = torch.nn.functional.softmax(outputs, dim=1)[0]
         results = {conditions[i]: float(probs[i].cpu().numpy()) * 100 for i in range(len(conditions))}
 
-        # Handle case where predicted class might not be in our list of conditions
         most_likely_condition = max(results, key=results.get)
-        
-        # Ensure the predicted condition is in the valid conditions list
-        if most_likely_condition not in conditions:
+        confidence = results[most_likely_condition]
+
+        if most_likely_condition not in condition_details:
             most_likely_condition = "Other"
             confidence = 0.0
-        else:
-            confidence = results[most_likely_condition]
 
-        # Create a detailed summary of results
         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>."
-
-        # Enhanced condition details for each disease/condition
-        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."},
-            "Fractures": {"description": "Bone break detected.", "recommendation": "Orthopedic evaluation."},
-            "COPD": {"description": "Lung damage from COPD observed.", "recommendation": "Pulmonary consultation."},
-            # Add the rest of the conditions as needed...
-        }
-
-        # Display results in a clear 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 = 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}%")
@@ -111,66 +125,82 @@ def predict_xray(image):
         logger.error(f"Error in predict_xray: {str(e)}")
         return f"Error: {str(e)}", "", ""
 
-# Enhanced function to analyze patient reports (PDFs) using PyMuPDF (fitz)
+# Enhanced function to analyze patient reports (PDFs)
 def analyze_report(file):
+    if not file or not file.name.endswith(".pdf"):
+        return "Please upload a valid PDF file."
+
     text = ""
     patient_condition = "Unclear"
     disease = "Unknown"
     status = "Pending further tests"
 
-    if file and file.name.endswith(".pdf"):
-        try:
-            # Open the PDF using PyMuPDF
-            pdf_reader = fitz.open(file.name)
-            for page in pdf_reader:
-                text += page.get_text("text")  # Extract text from each page
-
-            # Example: Let's search for conditions in the text
-            if "stroke" in text.lower():
-                patient_condition = "Stroke"
-                disease = "Brain Disorder"
-                status = "Urgent Care Needed"
-            elif "cancer" in text.lower():
-                patient_condition = "Cancer"
-                disease = "Malignant Growth"
-                status = "Consult Oncologist"
-            elif "fracture" in text.lower():
-                patient_condition = "Fracture"
-                disease = "Bone Injury"
-                status = "Orthopedic Attention Required"
-            # You can add more conditions here based on keyword matching
-            
-            report_summary = f"Patient's Condition: {patient_condition}\nDisease: {disease}\nCondition Status: {status}\n\nReport Preview: {text[:300]}..." if text else "No readable text found in the PDF."
-        
-        except Exception as e:
-            logger.error(f"Error reading PDF: {str(e)}")
-            report_summary = f"Error processing PDF: {str(e)}"
-    else:
-        report_summary = "Please upload a valid PDF file."
-
-    return report_summary
-
-# Gradio Interface with enhanced UI using Tabs
-def create_interface():
-    with gr.Blocks() as demo:
-        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; }
-        """
-        
-        gr.Markdown("<h1 class='title'>RadiologyScan AI</h1>")
-        gr.Markdown("<p style='text-align: center; color: #666;'>AI-powered analysis for X-rays and patient reports</p>")
-        
-        # Correctly provide interface_list to TabbedInterface
-        xray_tab = gr.Interface(fn=predict_xray, inputs=gr.Image(label="Upload X-ray", type="pil"), outputs=[gr.HTML(), gr.HTML(), gr.HTML()])
-        report_tab = gr.Interface(fn=analyze_report, inputs=gr.File(label="Upload Patient Report (PDF)", file_count="single"), outputs=gr.Textbox(label="Report Summary", interactive=False))
+    try:
+        pdf_reader = fitz.open(file.name)
+        for page in pdf_reader:
+            text += page.get_text("text")
+        pdf_reader.close()
+
+        if "stroke" in text.lower():
+            patient_condition = "Stroke"
+            disease = "Brain Disorder"
+            status = "Urgent Care Needed"
+        elif "cancer" in text.lower():
+            patient_condition = "Cancer"
+            disease = "Malignant Growth"
+            status = "Consult Oncologist"
+        elif "fracture" in text.lower():
+            patient_condition = "Fracture"
+            disease = "Bone Injury"
+            status = "Orthopedic Attention Required"
+
+        report_summary = f"Patient's Condition: {patient_condition}\nDisease: {disease}\nCondition Status: {status}\n\nReport Preview: {text[:300]}..." if text else "No readable text found in the PDF."
+        return report_summary
 
-        gr.TabbedInterface([xray_tab, report_tab], tab_names=["X-ray Analysis", "Report Analysis"]).launch(share=False)
+    except Exception as e:
+        logger.error(f"Error reading PDF: {str(e)}")
+        return f"Error processing PDF: {str(e)}"
 
-demo = create_interface()
+# Gradio Interface with Tabs
+def create_interface():
+    logger.debug("Initializing Gradio interface")
+    # Minimal CSS to avoid styling issues
+    custom_css = """
+    .title { font-size: 30px; text-align: center; color: #4C6A92; }
+    .subtitle { text-align: center; color: #666; }
+    """
+    with gr.Blocks(css=custom_css) as demo:
+        gr.Markdown("<h1 class='title'>RadiologyScan AI</h1>")
+        gr.Markdown("<p class='subtitle'>AI-powered analysis for X-rays and patient reports</p>")
+
+        with gr.Tabs():
+            with gr.TabItem("X-ray Analysis"):
+                image_input = gr.Image(label="Upload X-ray", type="pil")
+                output_summary = gr.HTML(label="Summary")
+                output_details = gr.HTML(label="Detailed Results")
+                output_feedback = gr.HTML(label="Additional Feedback")
+                gr.Button("Analyze X-ray").click(
+                    fn=predict_xray,
+                    inputs=image_input,
+                    outputs=[output_summary, output_details, output_feedback]
+                )
+
+            with gr.TabItem("Report Analysis"):
+                file_input = gr.File(label="Upload Patient Report (PDF)", file_count="single")
+                output_report = gr.Textbox(label="Report Summary", interactive=False)
+                gr.Button("Analyze Report").click(
+                    fn=analyze_report,
+                    inputs=file_input,
+                    outputs=output_report
+                )
+        logger.debug("Gradio interface initialized")
+    return demo
+
+if __name__ == "__main__":
+    logger.debug("Starting Gradio application")
+    try:
+        demo = create_interface()
+        demo.launch(server_port=7860, ssr_mode=False)  # Explicit port, disable SSR
+        logger.debug("Gradio application launched")
+    except Exception as e:
+        logger.error(f"Failed to launch Gradio application: {str(e)}")