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

print("Stage 0: Completed")

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:0' if torch.cuda.is_available() else 'cpu'

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

print("Stage 1: Completed ")

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)
    image_tensor = image_tensor.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",
    examples = [
        os.path.join(os.path.dirname(__file__), "0.jpg"),
        os.path.join(os.path.dirname(__file__), "1.jpg"),
        os.path.join(os.path.dirname(__file__), "3.jpg"),
        os.path.join(os.path.dirname(__file__), "8.jpg"),
        os.path.join(os.path.dirname(__file__), "10.jpg"),
        os.path.join(os.path.dirname(__file__), "11.jpg"),
        os.path.join(os.path.dirname(__file__), "12.jpg"),
        os.path.join(os.path.dirname(__file__), "14.jpg"),
        os.path.join(os.path.dirname(__file__), "21.jpg"),
        os.path.join(os.path.dirname(__file__), "27.jpg"),
        os.path.join(os.path.dirname(__file__), "37.jpg"),        
    ],
    theme = gr.themes.Soft()
) 

print("Stage 2: Completed")

interface.launch(debug = True)