Update app.py

app.py

@@ -4,8 +4,8 @@ from transformers import ViTForImageClassification, ViTFeatureExtractor
 from PIL import Image
 
 # Load model and feature extractor
-model = ViTForImageClassification.from_pretrained('Dhahlan2000/fruit_and_veg_chategorizer', num_labels=10, ignore_mismatched_sizes=True)
-feature_extractor = ViTFeatureExtractor.from_pretrained('Dhahlan2000/fruit_and_veg_chategorizer')
+model = ViTForImageClassification.from_pretrained('shahmi0519/banana_artificial', num_labels=2, ignore_mismatched_sizes=True)
+feature_extractor = ViTFeatureExtractor.from_pretrained('shahmi0519/banana_artificial')
 
 # Move to GPU if available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -14,16 +14,8 @@ model.eval()
 
 # Class labels (modify according to your model)
 class_labels = [
-    "Apple",
-    "Banana",
-    "Bellpepper",
-    "Carrot",
-    "Cucumber",
-    "Potato",
-    "Tomato",
-    "mango",
-    "oranges",
-    "strawberry"
+    "Artificial",
+    "Natural"
 ]
 
 def predict_freshness(image):