from transformers import AutoModelForImageClassification, AutoProcessor
import torch
import gradio as gr

# Load model and processor from Hugging Face
model_name = "ombhojane/healthyPlantsModel"

# 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_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")
demo.launch()