Update app.py

app.py

@@ -1,21 +1,29 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+import torch
 from PIL import Image
 
-# Load a pre-trained model suitable for general plant disease classification
-# Replace with an appropriate model from Hugging Face's hub trained on PlantVillage or similar
-model = pipeline("image-classification", model="PlantVillage/plant-disease-model")  # Replace with actual general plant disease model name if available
+# Load the image processor and model
+processor = AutoImageProcessor.from_pretrained("ozair23/swin-tiny-patch4-window7-224-finetuned-plantdisease")
+model = AutoModelForImageClassification.from_pretrained("ozair23/swin-tiny-patch4-window7-224-finetuned-plantdisease")
 
+# Define the function to process the image and make predictions
 def classify_leaf_disease(image):
+    # Preprocess the image
+    inputs = processor(images=image, return_tensors="pt")
+    
     # Run the model on the image
-    results = model(image)
+    with torch.no_grad():
+        outputs = model(**inputs)
     
-    # Get the top prediction
-    disease_name = results[0]['label']
-    confidence_score = results[0]['score']
+    # Get the predicted label and confidence score
+    probs = torch.softmax(outputs.logits, dim=1)
+    predicted_class_idx = probs.argmax(dim=1).item()
+    predicted_label = model.config.id2label[predicted_class_idx]
+    confidence_score = probs[0][predicted_class_idx].item()
     
     # Format the output
-    return disease_name, f"{confidence_score:.2f}", image
+    return predicted_label, f"{confidence_score:.2f}", image
 
 # Create Gradio Interface
 interface = gr.Interface(