Update app.py

app.py

@@ -3,9 +3,10 @@ import torch
 from torchvision import models, transforms
 import os
 import time
+import logging
+import fitz  # PyMuPDF for better PDF parsing
 
 # Set up logging (optional for debugging)
-import logging
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
@@ -19,9 +20,16 @@ conditions = [
     "Appendicitis", "Gallstones", "Kidney Stones", "Infections", "Abdominal Aortic Aneurysm", "Diverticulitis"
 ]
 
+# Define path to store the model manually
+model_path = "/home/user/.cache/torch/hub/checkpoints/densenet121-a639ec97.pth"  # Adjusted to a valid path
+
+# Ensure the parent directory exists
+parent_dir = os.path.dirname(model_path)
+if not os.path.exists(parent_dir):
+    os.makedirs(parent_dir)  # Create the parent directory if it doesn't exist
+
 # 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
@@ -30,13 +38,16 @@ def load_model():
     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 the model at the beginning (this will take time but only happens once)
 model = load_model()
 
+# Define device for model inference
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
 # Define image preprocessing function
 def preprocess_image(image):
     transform = transforms.Compose([