from PIL import Image
import torch
import gradio as gr
from torchvision.transforms import Compose, Normalize, ToTensor, Resize, CenterCrop


inference = torch.load('fine_tune_resnet.pth')
inference.eval()

def classifier(image):
    if torch.cuda.is_available():
        inference.cpu()

    with torch.no_grad():
        output = inference(test_transform(Image.open('/content/1000-ml-plastic-water-bottle-500x500.webp')).unsqueeze(0))
        _, prediction = torch.max(output,1)
        confidence = round(torch.softmax(output,1).max().item(),4)*100

    return f'{label_dict[prediction.item()]} (Confidence: {confidence}%)'

iface = gr.Interface(fc=classifier,
                    inputs=gr.Image(type="pil"),
                    outputs='text')
iface.launch(share=True)