Update app.py

app.py

@@ -49,28 +49,13 @@ condition_details = {
     "Diverticulitis": {"description": "Inflammation of diverticula in the colon.", "recommendation": "Gastroenterology consultation."}
 }
 
-# Load model
-try:
-    model = models.densenet121(weights="IMAGENET1K_V1")
-except AttributeError:
-    model = models.densenet121(pretrained=True)
-
-model.classifier = torch.nn.Linear(model.classifier.in_features, len(conditions))
+# Load model (using a smaller model like MobileNetV2 for faster inference)
+model = models.mobilenet_v2(pretrained=True)
+model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, len(conditions))  # Adjust the classifier for our condition count
 model.eval()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# 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("Model loaded from file.")
-    except Exception as e:
-        logger.warning(f"Failed to load model weights: {e}")
-else:
-    logger.info("No custom model weights found.")
-
 # Image preprocessing
 def preprocess_image(image):
     transform = transforms.Compose([
@@ -80,7 +65,7 @@ def preprocess_image(image):
     ])
     return transform(image).unsqueeze(0).to(device)
 
-# X-ray prediction function with confidence threshold
+# X-ray prediction function with summary output
 def predict_xray(image):
     try:
         if image is None:
@@ -96,23 +81,24 @@ def predict_xray(image):
         top_condition = max(results, key=results.get)
         confidence = results[top_condition]
 
+        # Construct a summary based on prediction
         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>
+                <h3>Summary</h3>
+                <p><b>Disease Identified:</b> Uncertain</p>
+                <p><b>Cause/Status:</b> The model is not confident enough to provide a clear diagnosis.</p>
+                <p><b>Treatment/Recommendation:</b> Please consult a radiologist or upload a better-quality image for better accuracy.</p>
             </div>
             """
 
         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>
+            <h3>Summary</h3>
+            <p><b>Disease Identified:</b> {top_condition}</p>
+            <p><b>Cause/Status:</b> {info['description']}</p>
+            <p><b>Treatment/Recommendation:</b> {info['recommendation']}</p>
         </div>
         """
 
@@ -153,12 +139,12 @@ def create_interface():
             with gr.TabItem("X-ray Analysis"):
                 img_input = gr.Image(label="Upload Chest X-ray", type="pil")
                 summary_output = gr.HTML(label="Summary Result")
-                gr.Button("Analyze X-ray", elem_id="analyze_button", scale=0.5).click(predict_xray, inputs=img_input, outputs=summary_output)
+                gr.Button("Analyze X-ray", elem_id="analyze_button", scale=0.3).click(predict_xray, inputs=img_input, outputs=summary_output)
 
             with gr.TabItem("Report Analysis"):
                 pdf_input = gr.File(label="Upload PDF Report", file_types=[".pdf"])
                 summary_output_report = gr.Textbox(label="Summary Result", lines=5)
-                gr.Button("Analyze Report", elem_id="analyze_button", scale=0.5).click(analyze_report, inputs=pdf_input, outputs=summary_output_report)
+                gr.Button("Analyze Report", elem_id="analyze_button", scale=0.3).click(analyze_report, inputs=pdf_input, outputs=summary_output_report)
 
     return demo