Update app.py

app.py

@@ -5,18 +5,25 @@ import gradio as gr
 # Load model and processor from Hugging Face
 model_name = "ombhojane/healthyPlantsModel"
 
-# Load model (ensure it's a classification model)
+# Load model and processor
 model = AutoModelForImageClassification.from_pretrained(model_name)
 processor = AutoProcessor.from_pretrained(model_name)
 
+# Get class labels from model config
+id2label = model.config.id2label  # Mapping of indices to class names
+
 # Define a function to run inference
 def classify_image(image):
     inputs = processor(images=image, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
     logits = outputs.logits
-    predicted_class = logits.argmax(-1).item()
-    return f"Predicted Class: {predicted_class}"
+    predicted_class_idx = logits.argmax(-1).item()
+    
+    # Get human-readable class label
+    predicted_class_name = id2label.get(predicted_class_idx, "Unknown")
+    
+    return f"Predicted Class: {predicted_class_name}"
 
 # Create a Gradio interface
 demo = gr.Interface(fn=classify_image, inputs="image", outputs="text")