Update app.py

app.py

@@ -28,25 +28,29 @@ def download_model():
 # Load the model from Hugging Face
 
 def load_model(model_path):
-    try:
-        model = models.resnet50(pretrained=False)
-        num_features = model.fc.in_features
-        model.fc = nn.Sequential(
-            nn.Dropout(0.5),
-            nn.Linear(num_features, num_classes)
-        )
-        checkpoint = torch.load(model_path, map_location=torch.device("cpu"))
-
-        # Ensure compatibility by handling key mismatches
-        model_state_dict = checkpoint['model_state_dict']
-        for key in list(model_state_dict.keys()):
-            if key.startswith('fc.1'):
-                model_state_dict[key.replace('fc.1', 'fc')] = model_state_dict.pop(key)
-
-        model.load_state_dict(model_state_dict)
-        model.eval()
-        logging.info("Model loaded successfully.")
-        return model
+    model = models.resnet50(pretrained=False)
+    num_features = model.fc.in_features
+    model.fc = nn.Sequential(
+        nn.Dropout(0.5),
+        nn.Linear(num_features, 3)  # 3 classes
+    )
+    
+    # Load the checkpoint
+    checkpoint = torch.load(model_path, map_location=torch.device("cpu"))
+    
+    # Adjust for state dict mismatch by renaming keys
+    state_dict = checkpoint['model_state_dict']
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        if k == "fc.weight" or k == "fc.bias":
+            new_state_dict[f"fc.1.{k.split('.')[-1]}"] = v
+        else:
+            new_state_dict[k] = v
+
+    model.load_state_dict(new_state_dict, strict=False)
+    model.eval()
+    return model
+
     except Exception as e:
         logging.error(f"Failed to load model: {e}")
         raise