Update app.py

app.py

@@ -38,30 +38,13 @@ def init_db():
 
 
 init_db()
-
+from densenet_withglam import get_model_with_attention 
 # 1️⃣ Instantiate the model
-model = models.densenet169(pretrained=False)
-model.classifier = nn.Linear(model.classifier.in_features, 3)
-
-# 2️⃣ Load checkpoint
-checkpoint = torch.load('densenet169_seed40_best2.pt', map_location='cpu')
-state_dict = checkpoint['state_dict']
-
-# 3️⃣ Fix state dict
-new_state_dict = {}
-for k, v in state_dict.items():
-    # Check if it's prefixed with 'features.0.'
-    if k.startswith('features.0.'):
-        new_key = 'features.' + k[len('features.0.'):]  # Remove the '0.' segment
-    else:
-        new_key = k
-    new_state_dict[new_key] = v
-
-# 4️⃣ Load into the model
-model.load_state_dict(new_state_dict)
-
-# Done!
+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.eval()
+
 # ✅ Class Names
 CLASS_NAMES = ["Normal", "Early Glaucoma", "Advanced Glaucoma"]