Spaces:
Sleeping
Sleeping
File size: 1,157 Bytes
98fe69e f5de5fe 1ca7c37 f5de5fe b1ae1a1 d6ef684 98fe69e 1ca7c37 7f16b2f 1ca7c37 98fe69e d6ef684 b0f2d06 b1ae1a1 98fe69e b1ae1a1 98fe69e 1ca7c37 98fe69e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
from PIL import Image
import torch
import gradio as gr
from torchvision.transforms import Compose, Normalize, ToTensor, Resize, CenterCrop
from pathlib import Path
labels = {0: 'glass',
1: 'metal',
2: 'organic-waste',
3: 'organice-waste',
4: 'paper',
5: 'plastic',
6: 'textiles'}
inference = torch.load('fine_tune_resnet.pth', map_location=torch.device('cpu'))
inference.eval()
example = [str(i) for i in Path('examples').glob('*')]
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])
])
with torch.no_grad():
output = inference(test_transform(image).unsqueeze(0))
out = torch.softmax(output,1)
values,indices = torch.topk(out[0],k=7)
return {labels[i.item()]: v.item() for i, v in zip(indices,values)}
iface = gr.Interface(fn=classifier,
inputs=gr.Image(type="pil"),
outputs='label',
examples = example,
title = 'Garbage Image Classification')
iface.launch(share=True) |