Update app.py

app.py

@@ -9,10 +9,6 @@ from safetensors.torch import load_file  # Import safetensors loading function
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# 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 directory containing the model files
 model_dir = "."  # Use current directory
 
@@ -32,6 +28,10 @@ for path in [model_path, config_path, preprocessor_path]:
 # Load the configuration
 config = AutoConfig.from_pretrained(config_path)
 
+# Ensure the labels are consistent with the model's config
+labels = list(config.id2label.values())
+logging.info(f"Labels: {labels}")
+
 # Load the feature extractor
 feature_extractor = ViTFeatureExtractor.from_pretrained(preprocessor_path)