Update app.py

app.py

@@ -2,22 +2,34 @@ import torch
 from transformers import ViTForImageClassification, ViTFeatureExtractor, ViTConfig
 import gradio as gr
 from PIL import Image
+import os
 
-
-print("p1")
 # Define the class labels as used during training
 labels = ['Leggings', 'Jogger', 'Palazzo', 'Cargo', 'Dresspants', 'Chinos']
 
-# Load the configuration
-config = ViTConfig.from_pretrained("DumbledoreWiz/PantsShape")
+# Define the path to the uploaded model file
+model_path = "final_fine_tuned_vit_Leggings_Jogger_Palazzo_Cargo_Dresspants_Chinos_2024-08-14.pth"
+
+# Check if config.json exists, otherwise use default config
+if os.path.exists("config.json"):
+    config = ViTConfig.from_pretrained(".")
+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)}
 
-# Load the ViT model with the configuration
-model = ViTForImageClassification.from_pretrained("DumbledoreWiz/PantsShape", config=config)
-feature_extractor = ViTFeatureExtractor.from_pretrained("DumbledoreWiz/PantsShape")
-print("Model loaded")
-# Set the model to evaluation mode
+# Load the model
+model = ViTForImageClassification(config)
+model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
 model.eval()
 
+# 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")
+
 # Define the prediction function
 def predict(image):
     # Preprocess the image