Upload 6 files

app.py

@@ -0,0 +1,42 @@
+import torch
+from torchvision import transforms
+import gradio as gr
+import torch
+from efficientnet_pytorch import EfficientNet
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+model = EfficientNet.from_name('efficientnet-b0')
+
+in_features = model._fc.in_features
+model._fc = torch.nn.Linear(in_features, 2) 
+model.load_state_dict(torch.load('model_transfer.pt'))
+
+model.to(device)
+
+model.eval()
+
+
+labels = ["Organic Waste","Recyclable Waste"]
+
+def predict(inp):
+    inp = transforms.ToTensor()(inp).unsqueeze(0)
+    
+    inp = inp.to(device)
+    
+    with torch.no_grad():
+        prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
+        
+        confidences = {labels[i]: float(prediction[i]) for i in range(len(prediction))}
+    
+    return confidences
+
+
+gr.Interface(
+    fn=predict, 
+    inputs=gr.components.Image(type="pil"),
+    outputs=gr.components.Label(num_top_classes=2),
+    examples=["tissue.jpg", "carrots.jpg"],
+    theme="default",
+    css=".footer{display:none !important}"
+).launch()

model_transfer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:669c0f182354391b318a80f17014f07373dadb7d4f955e23c6632739202433cb
+size 16335318

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+efficientnet_pytorch
+gradio
+Pillow