Update app.py

app.py

@@ -6,11 +6,11 @@ import fitz  # PyMuPDF
 import logging
 import os
 
-# Logging setup
+# Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# List of possible conditions
+# Condition list
 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"
 ]
 
-# Details for each condition
+# 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."},
@@ -49,7 +49,7 @@ condition_details = {
     "Diverticulitis": {"description": "Inflammation of diverticula in the colon.", "recommendation": "Gastroenterology consultation."}
 }
 
-# Load and configure the model
+# Load model
 try:
     model = models.densenet121(weights="IMAGENET1K_V1")
 except AttributeError:
@@ -57,22 +57,21 @@ except AttributeError:
 
 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 trained weights if available
+# Load model 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("Loaded custom model weights.")
+        logger.info("Model loaded from file.")
     except Exception as e:
         logger.warning(f"Failed to load model weights: {e}")
 else:
-    logger.info("No model weights found. Using random weights.")
+    logger.info("No custom model weights found.")
 
-# Preprocess uploaded image
+# Image preprocessing
 def preprocess_image(image):
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
@@ -81,7 +80,7 @@ def preprocess_image(image):
     ])
     return transform(image).unsqueeze(0).to(device)
 
-# Predict X-ray condition
+# X-ray prediction function with confidence threshold
 def predict_xray(image):
     try:
         if image is None:
@@ -90,12 +89,23 @@ def predict_xray(image):
         img_tensor = preprocess_image(image)
         with torch.no_grad():
             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]
 
+        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>
+            """
+
         info = condition_details.get(top_condition, condition_details["Other"])
         return f"""
         <div style="font-family:Arial">
@@ -105,20 +115,21 @@ def predict_xray(image):
             <p><b>Recommendation:</b> {info['recommendation']}</p>
         </div>
         """
+
     except Exception as e:
-        return f"Prediction failed: {str(e)}"
+        logger.error(f"Error in prediction: {e}")
+        return f"Error: {str(e)}"
 
-# Analyze PDF medical report
+# Analyze PDF report
 def analyze_report(file):
     if not file or not file.name.endswith(".pdf"):
-        return "Please upload a valid PDF report."
-
+        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 evaluation"
+        condition, disease, status = "Unclear", "Unknown", "Pending"
 
         if "stroke" in text.lower():
             condition, disease, status = "Stroke", "Brain Disorder", "Urgent Care Needed"
@@ -127,30 +138,30 @@ def analyze_report(file):
         elif "fracture" in text.lower():
             condition, disease, status = "Fracture", "Bone Injury", "Orthopedic Attention Required"
 
-        return f"Condition: {condition}\nDisease: {disease}\nStatus: {status}\n\nPreview:\n{text[:300]}..."
+        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 analyze PDF: {str(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</h1><p style='text-align:center;'>AI-powered X-ray and PDF report analysis</p>")
+        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"):
-                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)
+                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 Medical Report (PDF)", file_types=[".pdf"])
-                pdf_output = gr.Textbox(label="Report Summary", lines=10)
+                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
 
-# Launch app
 if __name__ == "__main__":
     demo = create_interface()
     demo.launch(server_port=7860, ssr_mode=False)