feat: create ui using gradio

.gitignore

@@ -0,0 +1,3 @@
+.venv
+.gradio
+uv.lock

.python-version

@@ -0,0 +1 @@
+3.12

main.py

@@ -0,0 +1,109 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin
+from torchvision.models import mobilenet_v3_large
+from torchvision.transforms import v2
+from PIL import Image
+import os
+
+
+class TrashMobileNet(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, num_classes=6):
+        super(TrashMobileNet, self).__init__()
+        self.model = mobilenet_v3_large(weights="DEFAULT")
+        for param in self.model.parameters():
+            param.requires_grad = False
+        num_features = self.model.classifier[-1].in_features
+        self.model.classifier[-1] = nn.Linear(num_features, num_classes)
+        for param in self.model.classifier[-1].parameters():
+            param.requires_grad = True
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+
+# Load the model from Hugging Face Hub
+model_name = "pradanaadn/trash-clasification"
+model = TrashMobileNet.from_pretrained(model_name)
+model.eval()
+
+# Define the image transformations
+transform = v2.Compose([
+    v2.Resize((224, 224)),
+    v2.ToImage(),
+    v2.ToDtype(torch.float32, scale=True),
+])
+
+
+def predict(image):
+    """
+    Prediction function that takes a Gradio image input and returns class probabilities
+    """
+    labels = ["cardboard", "glass", "metal", "paper", "plastic", "trash"]
+    
+    # Convert Gradio image to PIL Image if it's not already
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    
+    # Apply transformations and add batch dimension
+    image_tensor = transform(image)
+    image_tensor = image_tensor.unsqueeze(0)
+    
+    # Get model predictions
+    with torch.no_grad():
+        outputs = model(image_tensor)
+        probabilities = torch.nn.functional.softmax(outputs, dim=1)
+        probabilities = probabilities[0].tolist()
+    
+    # Create dictionary of label-probability pairs
+    return {label: float(prob) for label, prob in zip(labels, probabilities)}
+
+
+
+
+# Create example images if they don't exist (you would need to provide these images)
+examples = [
+    ["examples/cardbox.jpeg", "A cardboard box"],
+    ["examples/glass.jpeg", "A glass bottle"],
+    ["examples/plastic.png", "Mixed trash"]
+]
+
+
+with gr.Blocks() as iface:
+   
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                label="Upload Image",
+                type="pil",
+                elem_id="image_upload"
+            )
+            submit_btn = gr.Button("Classify", variant="primary")
+        
+        with gr.Column():
+            output_label = gr.Label(
+                label="Classification Results",
+                num_top_classes=6
+            )
+    
+    gr.Markdown("### Example Images")
+    gr.Examples(
+        examples=examples,
+        inputs=input_image,
+        outputs=output_label,
+        fn=predict,
+        cache_examples=True
+    )
+    
+    submit_btn.click(
+        fn=predict,
+        inputs=input_image,
+        outputs=output_label
+    )
+    
+   
+# Launch the interface
+iface.launch(share=True)

pyproject.toml

@@ -0,0 +1,13 @@
+[project]
+name = "trash-classification"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "gradio>=5.3.0",
+    "huggingface-hub>=0.27.0",
+    "torch>=2.5.1",
+    "torchvision>=0.20.1",
+    "transformers>=4.47.1",
+]

requirements.txt

@@ -0,0 +1,12 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile pyproject.toml --no-deps -o requirements.txt
+gradio==5.9.1
+    # via trash-classification (pyproject.toml)
+huggingface-hub==0.27.0
+    # via trash-classification (pyproject.toml)
+torch==2.5.1
+    # via trash-classification (pyproject.toml)
+torchvision==0.20.1
+    # via trash-classification (pyproject.toml)
+transformers==4.47.1
+    # via trash-classification (pyproject.toml)