Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import ViTForImageClassification, ViTFeatureExtractor, ViTConfig
+from transformers import ViTForImageClassification, ViTFeatureExtractor
 import gradio as gr
 from PIL import Image
 import os
@@ -8,49 +8,34 @@ import logging
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Define the class labels in the correct order as used during training
+# Define the labels in the correct order 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 = "best_fine_tuned_vit_Leggings_Jogger_Palazzo_Cargo_Dresspants_Chinos_93.90243902439025_2024-08-26.pth"
-logging.info(f"Looking for model file: {model_path}")
+# Define paths to the model files (all in the same directory as app.py)
+model_path = "model.safetensors"
+config_path = "config.json"
+preprocessor_path = "preprocessor_config.json"
 
-if os.path.exists(model_path):
-    logging.info(f"Model file found: {model_path}")
-else:
-    logging.error(f"Model file not found: {model_path}")
-    raise FileNotFoundError(f"Model file not found: {model_path}")
+# Check if all required files exist
+for path in [model_path, config_path, preprocessor_path]:
+    if not os.path.exists(path):
+        logging.error(f"File not found: {path}")
+        raise FileNotFoundError(f"Required file not found: {path}")
+    else:
+        logging.info(f"Found file: {path}")
 
-# 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)}
-
-# Create a configuration for the model
-config = ViTConfig.from_pretrained("google/vit-base-patch16-224-in21k")
-config.num_labels = len(labels)
-config.id2label = id2label
-config.label2id = label2id
+# Load the configuration and feature extractor
+config = ViTForImageClassification.from_pretrained(".", config=config_path)
+feature_extractor = ViTFeatureExtractor.from_pretrained(".")
 
 # Initialize the model with the configuration
-model = ViTForImageClassification(config)
-
-try:
-    # Load the state dict of the fine-tuned model
-    state_dict = torch.load(model_path, map_location=torch.device('cpu'))
-    model.load_state_dict(state_dict)
-    logging.info("Fine-tuned model loaded successfully")
-except Exception as e:
-    logging.error(f"Error loading model: {str(e)}")
-    raise
+model = ViTForImageClassification.from_pretrained(".", config=config)
 
+# Ensure the model is in evaluation mode
 model.eval()
 logging.info("Model set to evaluation mode")
 
-# Load feature extractor
-feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
-logging.info("Feature extractor loaded")
-
 # Define the prediction function
 def predict(image):
     logging.info("Starting prediction")
@@ -88,4 +73,4 @@ gradio_app = gr.Interface(
 # Launch the app
 if __name__ == "__main__":
     logging.info("Launching the app")
-    gradio_app.launch()
+    gradio_app.launch()