import torch
from transformers import ViTForImageClassification, ViTFeatureExtractor
import gradio as gr
from PIL import Image

# Check if GPU is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load pre-trained ViT model from Hugging Face
model = ViTForImageClassification.from_pretrained('Dhahlan2000/banana_ripeness_level_detection', num_labels=20, ignore_mismatched_sizes=True)
model.to(device)
model.eval()

# Load ViT feature extractor
feature_extractor = ViTFeatureExtractor.from_pretrained('Dhahlan2000/banana_ripeness_level_detection')

# Class labels
predicted_classes = ['Overripe', 'ripe', 'rotten', 'unripe', 'new']

# Function for inference
def classify_fruit(image):
    inputs = feature_extractor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()
    return predicted_classes[predicted_class_idx]

# Gradio UI
demo = gr.Interface(
    fn=classify_fruit,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(),
    title="Fruit Ripeness Detection",
    description="Upload an image of a fruit to determine whether it's fresh or rotten."
)

demo.launch()