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/ripeness_detection', num_labels=20)
model.to(device)
model.eval()

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

# Class labels
predicted_classes = [
    'FreshApple', 'FreshBanana', 'FreshBellpepper', 'FreshCarrot', 'FreshCucumber', 'FreshMango', 'FreshOrange', 
    'FreshPotato', 'FreshStrawberry', 'FreshTomato', 'RottenApple', 'RottenBanana', 'RottenBellpepper', 'RottenCarrot', 
    'RottenCucumber', 'RottenMango', 'RottenOrange', 'RottenPotato', 'RottenStrawberry', 'RottenTomato']

# 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()