import torch
import torchvision.transforms as transforms
import gadio as gr
from huggingface_hub import hf_hub_download
from PIL import Image

Fruits = ['Acerola', 'Apple', 'Apricot', 'Avocado', 'Banana', 'Black Berry', 'Blue Berry', 'Cantaloupe', 'Cherry',
          'Coconut', 'Fig', 'Grapefruit', 'Grape', 'Guava', 'Kiwi Fruit', 'Lemon', 'Lime', 'Mango', 'Olive', 'Orange',
          'Passion Fruit', 'Peach', 'Pear', 'Pineapple', 'Plum', 'Pomegranate', 'Raspberry', 'Strawberry', 'Tomato',
          'Watermelon']

device = 'cuda' if torch.cuda.is_availabe() else 'cpu'

repo_name  = "VinayHajare/fruits30-resnet18"
file_name = "fruit_resnet18(99.40%).pt"
model_path = hf_hub_download(repo_id = repo_name, file_name = file_name)
model = torch.load(model_path).to(device)
model.eval()

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def predict_image(image):
    image_tensor = transform(Image.fromarray(image)).unsqueeze(0).to(device)
    with torch.no_grad():
        output = model(image_tensor)
        predicted = torch.argmax(output).item()
    return Fruits[predicted]

interface = gr.Interface(
    fn = predict_image,
    inputs = "image",
    outputs = "text",
    allow_flagging = "never"
    theme = gr.themes.soft()
) 

interface.launch(debug = True)