Update app.py

app.py

@@ -42,18 +42,22 @@ init_db()
 model = models.densenet169(pretrained=False)
 model.classifier = nn.Linear(model.classifier.in_features, 3)
 
-# 2️⃣ Load checkpoint
-checkpoint = torch.load('densenet169_seed40_best.pt', map_location='cpu')
+# 2️⃣ Load your checkpoint
+state_dict = torch.load('densenet169_seed40_best.pt', map_location='cpu')
 
-# ✅ If checkpoint contains state_dict directly
-# model.load_state_dict(checkpoint)
+# 3️⃣ Strip the "features.0." prefix
+new_state_dict = {}
+for k, v in state_dict.items():
+    new_key = k.replace("features.0.", "")
+    new_state_dict[new_key] = v
 
-# ✅ If checkpoint contains a dict
-if 'model_state_dict' in checkpoint:
-    model.load_state_dict(checkpoint['model_state_dict'])
-else:
-    model.load_state_dict(checkpoint)
+# 4️⃣ Load the stripped state dict
+model.load_state_dict(new_state_dict)
 
+# 5️⃣ Set eval mode
+model.eval()
+
+print("✅ Model successfully loaded and ready for inference.")
 # ✅ Preprocess Image
 transform = transforms.Compose([
     transforms.Resize((224, 224)),  # Adjust if needed