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)