Update app.py

app.py

@@ -3,7 +3,13 @@ import torch
 import gradio as gr
 from torchvision.transforms import Compose, Normalize, ToTensor, Resize, CenterCrop
 
-
+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()
 
@@ -14,18 +20,17 @@ def classifier(image):
             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
+        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)}
 
-    return f'{label_dict[prediction.item()]} (Confidence: {confidence}%)'
 
-iface = gr.Interface(fc=classifier,
+iface = gr.Interface(fn=classifier,
                     inputs=gr.Image(type="pil"),
-                    outputs='text')
+                    outputs='label')
 iface.launch(share=True)