Update app.py

app.py

@@ -21,6 +21,7 @@ DB_PATH = os.path.join(OUTPUT_DIR, 'results.db')
 
 
 def init_db():
+    """Initialize SQLite database for storing results."""
     conn = sqlite3.connect(DB_PATH)
     cursor = conn.cursor()
     cursor.execute("""
@@ -38,24 +39,36 @@ def init_db():
 
 
 init_db()
+
+# ✅ Import your custom GLAM model
 from densenet_withglam import get_model_with_attention 
-# 1️⃣ Instantiate the model
+
+# ✅ Instantiate the model
 model = get_model_with_attention('densenet169', num_classes=3)  # Will have GLAM
 state_dict = torch.load('densenet169_seed40_best2.pt', map_location='cpu')
-model.load_state_dict(state_dict)  # Should now match
+model.load_state_dict(state_dict)  # Load your trained weights
 model.eval()
 
 # ✅ Class Names
 CLASS_NAMES = ["Normal", "Early Glaucoma", "Advanced Glaucoma"]
 
+# ✅ Transformation for input images
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),  # Adjust size based on training
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],  # Standard ImageNet stats
+                         std=[0.229, 0.224, 0.225])
+])
+
 @app.route('/')
 def home():
+    """Check that the API is working."""
     return "Glaucoma Detection Flask API (3-Class Model) is running!"
 
 @app.route("/test_file")
 def test_file():
     """Check if the .pt model file is present and readable."""
-    filepath = "densenet169_seed40_best.pt"
+    filepath = "densenet169_seed40_best2.pt"
     if os.path.exists(filepath):
         return f"✅ Model file found at: {filepath}"
     else: