Update app.py

app.py

@@ -1,22 +1,27 @@
 import gradio as gr
 from transformers import ViTForImageClassification, ViTFeatureExtractor
 from PIL import Image
-import torch
 
-# Load model and feature extractor from Hugging Face Hub
+# Load model and feature extractor
 model = ViTForImageClassification.from_pretrained("iamomtiwari/VITPEST")
 feature_extractor = ViTFeatureExtractor.from_pretrained("iamomtiwari/VITPEST")
 
-# Define inference function
-def classify_image(img):
-    inputs = feature_extractor(images=img, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-    predicted_class_idx = torch.argmax(logits, dim=-1).item()
-    class_name = model.config.id2label[predicted_class_idx]
-    return class_name
+def predict(image):
+    # Preprocess image
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
+    class_labels = model.config.id2label  # Assuming you have a class label mapping
+    return class_labels[predicted_class_idx]
+
+# Gradio interface
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="Crop Disease Classifier",
+    description="Upload an image of a crop leaf to classify the disease."
+)
 
-# Create Gradio interface
-interface = gr.Interface(fn=classify_image, inputs=gr.Image(type="pil"), outputs="text")
 interface.launch()