Upload 3 files

app.py

@@ -0,0 +1,74 @@
+import numpy as np
+import gradio as gr
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image
+
+
+class LargeNet(nn.Module):
+    def __init__(self):
+        super(LargeNet, self).__init__()
+        self.name = "large"
+        self.conv1 = nn.Conv2d(3, 5, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(5, 10, 5)
+        self.fc1 = nn.Linear(10 * 29 * 29, 32)
+        self.fc2 = nn.Linear(32, 7)
+
+    def forward(self, x):
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = x.view(-1, 10 * 29 * 29)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        x = x.squeeze(1) # Flatten to [batch_size]
+        return x
+
+ORIGINAL_LABELS = ['battery', 'biological', 'cardboard', 'glass', 'metal', 'paper', 'plastic']
+NEW_LABELS = ["paper", "cardboard", "other"]
+
+# Load best model
+net = LargeNet()
+net.load_state_dict(torch.load("net.pth"))
+net.eval()
+
+# Transform uploaded images
+transform = transforms.Compose([
+    transforms.Resize((128, 128)),  # Resize to 128x128
+    transforms.ToTensor(),         # Convert to Tensor
+    transforms.Normalize((0.5,), (0.5,))  # Normalize to [-1, 1]
+])
+
+# Make predictions on uploaded images
+def predict(img):
+    img_t = transform(img).unsqueeze(0)
+
+    with torch.no_grad():
+        output = net(img_t)
+        probs = F.softmax(output[0], dim=0)
+
+    new_probs = [0.0, 0.0, 0.0]
+    for original_label, prob in zip(ORIGINAL_LABELS, probs):
+        if original_label == "paper":
+          new_probs[0] = prob
+        elif original_label == "cardboard":
+          new_probs[1] = prob
+        else:
+          # Accumulate probabilities if the model do not predict "paper" or "cardboard"
+          new_probs[2] += prob
+
+    return {label: float(prob) for label, prob in zip(NEW_LABELS, new_probs)}
+
+# Create interface for uploading image and showing predictions
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=3),
+    title="Garbage Classifier",
+    description="Classify images into Paper, Cardboard, or Other."
+)
+
+if __name__ == "__main__":
+    iface.launch(share=True)

net.pth

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

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+torch
+torchvision
+numpy
+pillow