Spaces:
Sleeping
Sleeping
File size: 2,628 Bytes
869b2b3 5c80b12 869b2b3 5c80b12 869b2b3 5c80b12 869b2b3 5c80b12 869b2b3 5c80b12 869b2b3 5c80b12 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
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
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
model_name = "pradanaadn/trash-clasification"
model = TrashMobileNet.from_pretrained(model_name)
model.eval()
transform = v2.Compose([
v2.Resize((224, 224)),
v2.ToImage(),
v2.ToDtype(torch.float32, scale=True),
])
def predict(image):
labels = ["cardboard", "glass", "metal", "paper", "plastic", "trash"]
if not isinstance(image, Image.Image):
image = Image.fromarray(image)
image_tensor = transform(image)
image_tensor = image_tensor.unsqueeze(0)
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)}
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() |