Fix: Model loading error

app.py

@@ -20,41 +20,45 @@ def load_model(model_path: str):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     print(f"Using device: {device}")
     
-    # Load the model with default weights first
-    model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
+    # Initialize a fresh model without pretrained weights
+    model = models.resnet50(weights=None)
     model = model.to(device)
 
     # Load custom weights
     state_dict = torch.load(model_path, map_location=device)
     
-    # Debug: Print state dict info
-    print("\nState dict keys:", list(state_dict['model_state_dict'].keys())[:5])
-    print("Model state dict keys:", list(model.state_dict().keys())[:5])
-    
-    # Check if the final layer weights match
-    fc_weight_shape = state_dict['model_state_dict']['fc.weight'].shape
-    print(f"\nFC layer weight shape: {fc_weight_shape}")
+    # Debug: Print original state dict keys
+    print("\nOriginal state dict keys:", list(state_dict['model_state_dict'].keys())[:5])
     
-    filtered_state_dict = {k: v for k, v in state_dict['model_state_dict'].items() if k in model.state_dict()}
-    print(f"Filtered state dict size: {len(filtered_state_dict)} / {len(state_dict['model_state_dict'])}")
+    # Remove the 'model.' prefix from state dict keys
+    new_state_dict = {}
+    for key, value in state_dict['model_state_dict'].items():
+        new_key = key.replace('model.', '')
+        new_state_dict[new_key] = value
     
-    model.load_state_dict(filtered_state_dict, strict=False)
-    model.eval()
+    # Debug: Print modified state dict keys
+    print("Modified state dict keys:", list(new_state_dict.keys())[:5])
+    print("Model state dict keys:", list(model.state_dict().keys())[:5])
     
-    # Verify model
-    print("\nModel architecture:")
-    print(model)
+    # Load the modified state dict
+    try:
+        model.load_state_dict(new_state_dict)
+        print("Successfully loaded model weights")
+    except Exception as e:
+        print(f"Error loading state dict: {str(e)}")
+        raise e
     
+    model.eval()
     return model
 
 
 def load_classes():
     """
-    Load the classes.
+    Load the ImageNet classes
     """
-    # Load classes from the same weights version as the model was trained with
-    weights = models.ResNet50_Weights.IMAGENET1K_V1  # Try V1 instead of V2
+    weights = models.ResNet50_Weights.IMAGENET1K_V1
     classes = weights.meta["categories"]
+    print(f"Loaded {len(classes)} classes")
     return classes