Update app.py

app.py

@@ -1,12 +1,11 @@
 import gradio as gr
-from PIL import Image
 import torch
 from torchvision import models, transforms
-import fitz  # PyMuPDF for better PDF parsing
-import logging
 import os
+import time
 
-# Set up logging
+# Set up logging (optional for debugging)
+import logging
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
@@ -20,23 +19,23 @@ conditions = [
     "Appendicitis", "Gallstones", "Kidney Stones", "Infections", "Abdominal Aortic Aneurysm", "Diverticulitis"
 ]
 
-# Load and configure the model
-model = models.densenet121(weights="IMAGENET1K_V1")  # DenseNet pre-trained on ImageNet
-num_features = model.classifier.in_features
-model.classifier = torch.nn.Linear(num_features, len(conditions))  # Output for all 24 conditions
-model.eval()
-
-# Define device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = model.to(device)
+# Function to load the model efficiently
+def load_model():
+    model_path = "/mnt/data/densenet121-a639ec97.pth"  # Set the model path
+    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
+        model.eval()
+        logger.info("Downloaded and cached the model.")
+    return model
 
-# Load model state if available, otherwise initialize
-model_path = "xray_model.pth"
-if os.path.exists(model_path):
-    model.load_state_dict(torch.load(model_path))
-    logger.info(f"Loaded model from {model_path}")
-else:
-    logger.info("No pre-trained model found. Initializing with random weights. Training required.")
+# Load the model at the beginning (this will take time but only happens once)
+model = load_model()
 
 # Define image preprocessing function
 def preprocess_image(image):
@@ -161,6 +160,6 @@ def create_interface():
         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))
 
-        gr.TabbedInterface([xray_tab, report_tab], tab_names=["X-ray Analysis", "Report Analysis"]).launch(share=True)
+        gr.TabbedInterface([xray_tab, report_tab], tab_names=["X-ray Analysis", "Report Analysis"]).launch(share=False)
 
 demo = create_interface()