app.py

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', map_location=torch.device('cpu'))
inference.eval()

def classifier(image):
    test_transform = Compose([
            Resize(256),
            CenterCrop(224),
            ToTensor(),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    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)