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,7 +8,7 @@ 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
+# Define the class labels in the correct order as used during training
 labels = ['Leggings', 'Jogger', 'Palazzo', 'Cargo', 'Dresspants', 'Chinos']
 logging.info(f"Labels: {labels}")
 
@@ -52,9 +52,6 @@ logging.info("Model set to evaluation mode")
 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")