Update app.py

app.py

@@ -2,88 +2,116 @@ import gradio as gr
 from PIL import Image
 import torch
 from torchvision import models, transforms
-import PyPDF2  # For reading patient reports (PDFs)
+import PyPDF2
+import logging
 
-# Load the pre-trained model (ResNet18 with new weight parameter)
-model = models.resnet18(weights="IMAGENET1K_V1")  # Updated to the new weight parameter in torchvision 0.13+
+# 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
+num_features = model.classifier.in_features
+model.classifier = torch.nn.Linear(num_features, 5)  # Output layer for 5 conditions
 model.eval()
 
-# Define image preprocessing function
+# Define device (CPU or GPU)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+# Define image preprocessing function with normalization
 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
     ])
-    return transform(image).unsqueeze(0)
+    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
 def predict_xray(image):
-    image_tensor = preprocess_image(image)
-    with torch.no_grad():
-        outputs = model(image_tensor)
-    probs = torch.nn.functional.softmax(outputs[0], dim=0)
-    
-    # Define the conditions (replace these with the actual conditions your model predicts)
-    conditions = ["Normal", "Pneumonia", "Cancer", "TB", "Other"]
-    results = {conditions[i]: float(probs[i]) for i in range(len(conditions))}
-    
-    # Determine the most likely condition and provide detailed feedback
-    most_likely_condition = max(results, key=results.get)
-    confidence = results[most_likely_condition] * 100  # Convert to percentage
-    
-    # Provide a more detailed summary of the 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>."
-    
-    # Additional detailed descriptions and recommendations for each condition
-    condition_details = {
-        "Normal": {
-            "description": "The X-ray shows no abnormal signs, and the lungs appear healthy.",
-            "recommendation": "No further tests are required. Continue with regular health check-ups."
-        },
-        "Pneumonia": {
-            "description": "Pneumonia is an infection that causes inflammation in the lungs. It can be caused by bacteria, viruses, or fungi.",
-            "recommendation": "Consult a healthcare provider immediately for possible antibiotics and further treatment."
-        },
-        "Cancer": {
-            "description": "Lung cancer can manifest in the form of abnormal growths in the lungs. It may require more advanced diagnostic tools like biopsies.",
-            "recommendation": "Consult an oncologist for further tests and treatment options."
-        },
-        "TB": {
-            "description": "Tuberculosis (TB) is a serious bacterial infection that mainly affects the lungs. It is treatable but requires proper care.",
-            "recommendation": "Seek immediate medical attention for a proper treatment plan, which may include antibiotics."
-        },
-        "Other": {
-            "description": "The X-ray may show signs of other possible conditions that need further investigation.",
-            "recommendation": "Consult with a doctor to rule out any other serious conditions."
+    try:
+        if image is None:
+            return "Error: No image uploaded.", "", ""
+
+        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
+
+        # 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
+        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."
+            }
         }
-    }
-    
-    # Displaying the results in a structured way (bullet points)
-    detailed_results = "<ul>"
-    for condition, prob in results.items():
-        detailed_results += f"<li><b>{condition}:</b> {prob*100:.2f}%</li>"
-    detailed_results += "</ul>"
-    
-    # Additional advice based on the most likely condition
-    additional_feedback = condition_details.get(most_likely_condition, "Please consult with a doctor for further details.")
-    
-    return summary, detailed_results, additional_feedback
+
+        # 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.")
+
+        logger.info(f"Prediction: {most_likely_condition} with confidence {confidence:.2f}%")
+        return summary, detailed_results, additional_feedback
+
+    except Exception as e:
+        logger.error(f"Error in predict_xray: {str(e)}")
+        return f"Error: {str(e)}", "", ""
 
 # Define a function to read and analyze patient reports (PDFs)
 def analyze_report(file):
     text = ""
-    if file.name.endswith(".pdf"):
-        pdf_reader = PyPDF2.PdfReader(file)
-        for page in pdf_reader.pages:
-            text += page.extract_text()
-    # For simplicity, we are just summarizing the first 300 characters
-    report_summary = f"Patient Report (Preview): {text[:300]}..."
+    if file and file.name.endswith(".pdf"):
+        try:
+            pdf_reader = PyPDF2.PdfReader(file)
+            for page in pdf_reader.pages:
+                text += page.extract_text() or ""
+            report_summary = f"Patient Report (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
 def create_interface():
     with gr.Blocks() as demo:
-        # Custom CSS for UI
+        # Custom CSS for UI enhancement
         custom_css = """
         .gradio-container {
             background-color: #f4f6f9;
@@ -96,12 +124,18 @@ def create_interface():
             font-size: 30px;
             text-align: center;
             color: #4C6A92;
+            margin-bottom: 20px;
         }
         .gradio-button {
             background-color: #3B82F6;
             color: white;
             border-radius: 10px;
-            padding: 15px;
+            padding: 15px 30px;
+            font-size: 16px;
+            transition: background-color 0.3s;
+        }
+        .gradio-button:hover {
+            background-color: #2563EB;
         }
         .result-box {
             background-color: #ffffff;
@@ -109,43 +143,63 @@ def create_interface():
             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 section
+        # 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>")
         
-        # Upload X-ray image section
+        # Upload section with layout
         with gr.Row():
-            xray_input = gr.Image(label="Upload Chest X-ray", type="pil")
-            report_input = gr.File(label="Upload Patient Report (PDF)", file_count="single")
+            with gr.Column(scale=1):
+                xray_input = gr.Image(label="Upload Chest 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 for the X-ray image
-        xray_output = gr.HTML(label="X-ray Diagnosis Summary", elem_classes="result-box")  # Removed interactive=False
-        xray_result = gr.HTML(label="Detailed X-ray Results", elem_classes="result-box")  # Removed interactive=False
-        additional_feedback = gr.Textbox(label="Additional Feedback", interactive=False, elem_classes="result-box")
-        
-        # Results section for the patient report
-        report_output = gr.Textbox(label="Report Summary", interactive=False, elem_classes="result-box")
-        
-        # Event handlers for buttons
-        predict_button.click(predict_xray, inputs=xray_input, outputs=[xray_output, xray_result, additional_feedback])
-        report_button.click(analyze_report, inputs=report_input, outputs=report_output)
+        # 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,
+            outputs=[xray_output, xray_result, additional_feedback]
+        )
+        report_button.click(
+            fn=analyze_report,
+            inputs=report_input,
+            outputs=report_output
+        )
     
     return demo
 
 # Launch the Gradio interface
 demo = create_interface()
-demo.launch(share=True)
+demo.launch(share=True)