Update app.py

app.py

@@ -1,10 +1,46 @@
 import gradio as gr
-from transformers import pipeline
-
-pipe = pipeline(task="image-classification", 
-                model="KhadijaAsehnoune12/OrangeLeafDiseaseDetector")
-gr.Interface.from_pipeline(pipe, 
-                           title="Orange Disease Image Classification",
-                           description="Detect diseases in orange leaves and fruits.",
-                           examples = ['MoucheB.jpg', 'verdissement.jpg',],
-                           ).launch()
+import torch
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from PIL import Image
+
+# Define the model and feature extractor
+model_name = "KhadijaAsehnoune12/OrangeLeafDiseaseDetector"
+model = ViTForImageClassification.from_pretrained(model_name)
+feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
+
+# Define the label mapping
+id2label = {
+    "0": "Aleurocanthus spiniferus",
+    "1": "Chancre citrique",
+    "2": "Cochenille blanche",
+    "3": "Dépérissement des agrumes",
+    "4": "Feuille saine",
+    "5": "Jaunissement des feuilles",
+    "6": "Maladie de l'oïdium",
+    "7": "Maladie du dragon jaune",
+    "8": "Mineuse des agrumes",
+    "9": "Trou de balle"
+}
+
+def predict(image):
+    # Preprocess the image
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    
+    # Forward pass through the model
+    outputs = model(**inputs)
+    
+    # Get the predicted label
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
+    
+    # Get the label name
+    predicted_label = id2label[str(predicted_class_idx)]
+    
+    return predicted_label
+
+# Create the Gradio interface
+image = gr.Image(type="pil")
+label = gr.Label(num_top_classes=3)
+
+gr.Interface(fn=predict, inputs=image, outputs=label, title="Citrus Disease Classification",
+             description="Upload an image of a citrus leaf or fruit to classify its disease.").launch()