Update app.py

app.py

@@ -22,33 +22,24 @@ else:
     logging.error(f"Model file not found: {model_path}")
     raise FileNotFoundError(f"Model file not found: {model_path}")
 
-# Create a custom configuration
-config = ViTConfig.from_pretrained("google/vit-base-patch16-224-in21k")
-config.num_labels = len(labels)
-config.id2label = {str(i): label for i, label in enumerate(labels)}
-config.label2id = {label: str(i) for i, label in enumerate(labels)}
-logging.info(f"Custom config created with {len(labels)} labels")
+# Create label mappings consistent with training
+id2label = {str(i): label for i, label in enumerate(labels)}
+label2id = {label: str(i) for i, label in enumerate(labels)}
 
-# Load the model with the custom configuration
-logging.info("Loading the model with custom configuration")
-model = ViTForImageClassification(config)
+# Load the model with custom label mapping
+logging.info("Loading the model with custom label mapping")
+model = ViTForImageClassification.from_pretrained(
+    "google/vit-base-patch16-224-in21k",
+    num_labels=len(labels),
+    id2label=id2label,
+    label2id=label2id,
+    ignore_mismatched_sizes=True
+)
 
 try:
     # Load the state dict
     state_dict = torch.load(model_path, map_location=torch.device('cpu'))
-    
-    # Check if the state dict keys match the model's keys
-    model_keys = set(model.state_dict().keys())
-    loaded_keys = set(state_dict.keys())
-    
-    if model_keys != loaded_keys:
-        logging.warning("Mismatch in state dict keys. Attempting to adjust...")
-        # Adjust keys if necessary (e.g., remove 'module.' prefix if it exists)
-        new_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
-        model.load_state_dict(new_state_dict)
-    else:
-        model.load_state_dict(state_dict)
-    
+    model.load_state_dict(state_dict)
     logging.info("Model loaded successfully")
 except Exception as e:
     logging.error(f"Error loading model: {str(e)}")
@@ -57,10 +48,11 @@ except Exception as e:
 model.eval()
 logging.info("Model set to evaluation mode")
 
-# Load or create feature extractor
+# Load feature extractor
 feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
 logging.info("Feature extractor loaded")
 
+
 logging.info("Model and feature extractor loaded successfully")
 
 # Define the prediction function