Update app.py

app.py

@@ -3,48 +3,93 @@ from transformers import ViTForImageClassification, ViTFeatureExtractor, ViTConf
 import gradio as gr
 from PIL import Image
 import os
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Define the class labels as used during training
 labels = ['Leggings', 'Jogger', 'Palazzo', 'Cargo', 'Dresspants', 'Chinos']
+logging.info(f"Labels: {labels}")
 
 # Define the path to the uploaded model file
 model_path = "final_fine_tuned_vit_Leggings_Jogger_Palazzo_Cargo_Dresspants_Chinos_2024-08-14.pth"
+logging.info(f"Looking for model file: {model_path}")
 
-# Check if config.json exists, otherwise use default config
-if os.path.exists("config.json"):
-    config = ViTConfig.from_pretrained(".")
+if os.path.exists(model_path):
+    logging.info(f"Model file found: {model_path}")
 else:
-    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.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")
 
-# Load the model
+# Load the model with the custom configuration
+logging.info("Loading the model with custom configuration")
 model = ViTForImageClassification(config)
-model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+
+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)
+    
+    logging.info("Model loaded successfully")
+except Exception as e:
+    logging.error(f"Error loading model: {str(e)}")
+    raise
+
 model.eval()
+logging.info("Model set to evaluation mode")
 
 # Load or create feature extractor
-if os.path.exists("preprocessor_config.json"):
-    feature_extractor = ViTFeatureExtractor.from_pretrained(".")
-else:
-    feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
+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
 def predict(image):
+    logging.info("Starting prediction")
+    logging.info(f"Input image shape: {image.size}")
+    
     # Preprocess the image
+    logging.info("Preprocessing image")
     inputs = feature_extractor(images=image, return_tensors="pt")
+    logging.info(f"Preprocessed input shape: {inputs['pixel_values'].shape}")
     
+    logging.info("Running inference")
     with torch.no_grad():
         outputs = model(**inputs)
         logits = outputs.logits
         probabilities = torch.nn.functional.softmax(logits[0], dim=0)
     
+    logging.info(f"Raw logits: {logits}")
+    logging.info(f"Probabilities: {probabilities}")
+    
     # Prepare the output dictionary
     result = {labels[i]: float(probabilities[i]) for i in range(len(labels))}
+    logging.info(f"Prediction result: {result}")
+    
     return result
 
 # Set up the Gradio Interface
+logging.info("Setting up Gradio interface")
 gradio_app = gr.Interface(
     fn=predict, 
     inputs=gr.Image(type="pil"), 
@@ -54,4 +99,5 @@ gradio_app = gr.Interface(
 
 # Launch the app
 if __name__ == "__main__":
+    logging.info("Launching the app")
     gradio_app.launch()