first commit

.gitattributes

@@ -32,3 +32,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+vit_b_16_dout0.3_10epochs.pth filter=lfs diff=lfs merge=lfs -text
+examples/metal.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,6 +1,6 @@
 ---
 title: Garbage Sense
-emoji: 📈
+emoji: 👓
 colorFrom: purple
 colorTo: yellow
 sdk: gradio
@@ -10,4 +10,4 @@ pinned: false
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+A vision transformer trained to classify garbage into 6 categories, following DeIT’s training recipe.

app.py

@@ -0,0 +1,45 @@
+import torch
+import os
+import gradio as gr
+
+from model import create_vit
+from timeit import default_timer as timer
+from typing import Tuple, Dict
+
+class_names = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
+
+vit, vit_transform = create_vit(output_classes=len(class_names))
+
+vit.load_state_dict(torch.load(f="vit_b_16_dout0.3_10epochs.pth"))
+
+def predict(img) -> Tuple[Dict, float]:
+    start_time = timer()
+    img = vit_transform(img).unsqueeze(0)
+
+    vit.eval()
+    with torch.inference_mode():
+        pred_probs = torch.softmax(vit(img), dim=1)
+    
+    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
+    pred_time = round(timer() - start_time, 5)
+    return pred_labels_and_probs, pred_time
+
+title = "Garbage Sense"
+description = "A vision transformer trained to classify garbage into 6 categories on [trashnet](https://github.com/garythung/trashnet)."
+article = ""
+
+example_list = [["examples/" + example] for example in os.listdir("examples")]
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Label(num_top_classes=6, label="Predictions"),
+        gr.Number(label="Prediction time (s)"),
+    ],
+    examples=example_list,
+    title=title,
+    description=description
+)
+
+demo.launch()

model.py

@@ -0,0 +1,25 @@
+import torch
+from torchvision import transforms, models
+
+def create_vit(output_classes: int = 6, seed: int = 233):
+    
+    transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+    model = models.vit_b_16()
+
+    for param in model.parameters():
+        param.requires_grade = False
+    
+    torch.manual_seed(seed)
+
+    model.heads = torch.nn.Sequential(
+        torch.nn.Dropout(0.3),
+        torch.nn.Linear(in_features=768, out_features=output_classes)
+    )
+
+    return model, transform

requirements.txt

@@ -0,0 +1,2 @@
+torch>=1.12.0
+torchvision>=0.13.0

vit_b_16_dout0.3_10epochs.pth

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