Upload 3 files

app.py

@@ -0,0 +1,55 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+from torchvision import transforms, models
+from PIL import Image
+
+# Define classes
+classes = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
+
+# Load model architecture
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+base_model = models.resnet50(weights=None)  # No pretrained weights here since we load ours
+in_features = base_model.fc.in_features
+base_model.fc = nn.Sequential(
+    nn.Linear(in_features, 512),
+    nn.ReLU(),
+    nn.Dropout(0.4),
+    nn.Linear(512, len(classes))
+)
+
+# Load trained weights
+base_model.load_state_dict(torch.load("best_model.pth", map_location=device))
+base_model = base_model.to(device)
+base_model.eval()
+
+# Define transforms (same as training but without augmentation)
+transform = transforms.Compose([
+    transforms.Lambda(lambda x: x.convert('RGB')),
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5, 0.5, 0.5],
+                         std=[0.5, 0.5, 0.5])
+])
+
+# Prediction function
+def predict(img):
+    img = transform(img).unsqueeze(0).to(device)
+    with torch.no_grad():
+        outputs = base_model(img)
+        probs = torch.softmax(outputs, dim=1).cpu().numpy()[0]
+    # Create a dictionary of class probabilities
+    return {classes[i]: float(probs[i]) for i in range(len(classes))}
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=6),
+    title="Garbage Classification",
+    description="Upload an image of garbage and the model will predict its type (cardboard, glass, metal, paper, plastic, trash)."
+)
+
+if __name__ == "__main__":
+    demo.launch(debug=True)

best_model.pth

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

requirements.txt

@@ -0,0 +1,4 @@
+torch
+torchvision
+gradio
+Pillow