Update app.py

app.py

@@ -2,15 +2,15 @@ import gradio as gr
 from PIL import Image
 import torch
 from torchvision import models, transforms
-import fitz  # PyMuPDF for better PDF parsing
+import fitz  # PyMuPDF
 import logging
 import os
 
-# Set up logging
-logging.basicConfig(level=logging.DEBUG)  # Set to DEBUG for detailed output
+# Logging setup
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Define conditions
+# List of possible conditions
 conditions = [
     "Normal", "Pneumonia", "Cancer", "TB", "Other",
     "Coronary Artery Disease", "Aortic Aneurysm", "Stroke", "Peripheral Artery Disease",
@@ -20,7 +20,7 @@ conditions = [
     "Appendicitis", "Gallstones", "Kidney Stones", "Infections", "Abdominal Aortic Aneurysm", "Diverticulitis"
 ]
 
-# Define condition details for all conditions
+# Details for each 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."},
@@ -52,155 +52,105 @@ condition_details = {
 # 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
+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 = model.to(device)
-logger.info(f"Using device: {device}")
+model.to(device)
 
-# Load model state if available
+# Load trained weights 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}")
+        logger.info("Loaded custom model weights.")
     except Exception as e:
-        logger.warning(f"Failed to load model from {model_path}: {str(e)}. Using random weights.")
+        logger.warning(f"Failed to load model weights: {e}")
 else:
-    logger.info("No pre-trained model found. Initializing with random weights. Training required.")
+    logger.info("No model weights found. Using random weights.")
 
-# Define image preprocessing function
+# Preprocess uploaded image
 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)),
         transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        transforms.Normalize([0.485, 0.456, 0.406], [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
+    return transform(image).unsqueeze(0).to(device)
 
-# Define prediction function
+# Predict X-ray condition
 def predict_xray(image):
     try:
         if image is None:
-            return "Error: No image uploaded.", "", ""
+            return "Please upload an image."
 
-        image_tensor = preprocess_image(image)
+        img_tensor = preprocess_image(image)
         with torch.no_grad():
-            outputs = model(image_tensor)
-        probs = torch.nn.functional.softmax(outputs, dim=1)[0]
-        results = {conditions[i]: float(probs[i].cpu().numpy()) * 100 for i in range(len(conditions))}
-
-        most_likely_condition = max(results, key=results.get)
-        confidence = results[most_likely_condition]
-
-        if most_likely_condition not in condition_details:
-            most_likely_condition = "Other"
-            confidence = 0.0
-
-        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>."
-        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}%")
-        return summary, detailed_results, additional_feedback
-
+            output = model(img_tensor)
+        probs = torch.nn.functional.softmax(output, dim=1)[0]
+        results = {conditions[i]: probs[i].item() * 100 for i in range(len(conditions))}
+
+        top_condition = max(results, key=results.get)
+        confidence = results[top_condition]
+
+        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"Error in predict_xray: {str(e)}")
-        return f"Error: {str(e)}", "", ""
+        return f"Prediction failed: {str(e)}"
 
-# Enhanced function to analyze patient reports (PDFs)
+# Analyze PDF medical report
 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"
+        return "Please upload a valid PDF report."
 
     try:
-        pdf_reader = fitz.open(file.name)
-        for page in pdf_reader:
-            text += page.get_text("text")
-        pdf_reader.close()
+        doc = fitz.open(file.name)
+        text = "".join(page.get_text() for page in doc)
+        doc.close()
+
+        condition, disease, status = "Unclear", "Unknown", "Pending evaluation"
 
         if "stroke" in text.lower():
-            patient_condition = "Stroke"
-            disease = "Brain Disorder"
-            status = "Urgent Care Needed"
+            condition, disease, status = "Stroke", "Brain Disorder", "Urgent Care Needed"
         elif "cancer" in text.lower():
-            patient_condition = "Cancer"
-            disease = "Malignant Growth"
-            status = "Consult Oncologist"
+            condition, disease, status = "Cancer", "Malignant Growth", "Consult Oncologist"
         elif "fracture" in text.lower():
-            patient_condition = "Fracture"
-            disease = "Bone Injury"
-            status = "Orthopedic Attention Required"
+            condition, disease, status = "Fracture", "Bone Injury", "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
+        return f"Condition: {condition}\nDisease: {disease}\nStatus: {status}\n\nPreview:\n{text[:300]}..."
 
     except Exception as e:
-        logger.error(f"Error reading PDF: {str(e)}")
-        return f"Error processing PDF: {str(e)}"
+        return f"Failed to analyze PDF: {str(e)}"
 
-# Gradio Interface with Tabs
+# Gradio interface
 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.Blocks() as demo:
+        gr.Markdown("<h1 style='text-align:center;'>🩻 RadiologyScan AI</h1><p style='text-align:center;'>AI-powered X-ray and PDF report analysis</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]
-                )
+                xray_input = gr.Image(label="Upload Chest X-ray", type="pil")
+                xray_output = gr.HTML()
+                gr.Button("Analyze X-ray").click(predict_xray, inputs=xray_input, outputs=xray_output)
 
             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")
+                pdf_input = gr.File(label="Upload Medical Report (PDF)", file_types=[".pdf"])
+                pdf_output = gr.Textbox(label="Report Summary", lines=10)
+                gr.Button("Analyze Report").click(analyze_report, inputs=pdf_input, outputs=pdf_output)
+
     return demo
 
+# Launch app
 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)}")
+    demo = create_interface()
+    demo.launch(server_port=7860, ssr_mode=False)