1Classification Report:
2 precision recall f1-score support
3
4 cardboard 0.9912 0.9739 0.9825 806
5 glass 0.9564 0.9641 0.9602 1002
6 metal 0.9523 0.9744 0.9632 820
7 paper 0.9520 0.9848 0.9681 1188
8 plastic 0.9835 0.9274 0.9546 964
9 trash 0.9127 0.9161 0.9144 274
10
11 accuracy 0.9626 5054
12 macro avg 0.9580 0.9568 0.9572 5054
13weighted avg 0.9631 0.9626 0.9626 5054
1import gradio as gr
2from transformers import AutoImageProcessor
3from transformers import SiglipForImageClassification
4from transformers.image_utils import load_image
5from PIL import Image
6import torch
7
8# Load model and processor
9model_name = "prithivMLmods/Trash-Net"
10model = SiglipForImageClassification.from_pretrained(model_name)
11processor = AutoImageProcessor.from_pretrained(model_name)
12
13def trash_classification(image):
14 """Predicts the category of waste material in the image."""
15 image = Image.fromarray(image).convert("RGB")
16 inputs = processor(images=image, return_tensors="pt")
17
18 with torch.no_grad():
19 outputs = model(**inputs)
20 logits = outputs.logits
21 probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
22
23 labels = {
24 "0": "cardboard",
25 "1": "glass",
26 "2": "metal",
27 "3": "paper",
28 "4": "plastic",
29 "5": "trash"
30 }
31 predictions = {labels[str(i)]: round(probs[i], 3) for i in range(len(probs))}
32
33 return predictions
34
35# Create Gradio interface
36iface = gr.Interface(
37 fn=trash_classification,
38 inputs=gr.Image(type="numpy"),
39 outputs=gr.Label(label="Prediction Scores"),
40 title="Trash Classification",
41 description="Upload an image to classify the type of waste material."
42)
43
44# Launch the app
45if __name__ == "__main__":
46 iface.launch()