Upload 2 files

requirements.txt

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

waste_sorting.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+import gradio as gr
+
+# Load model and image processor
+model_name = "watersplash/waste-classification"  # Change to a valid model
+model = AutoModelForImageClassification.from_pretrained(model_name)
+image_processor = AutoImageProcessor.from_pretrained(model_name)
+
+# Define preprocessing function
+def preprocess_image(image):
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
+    ])
+    return transform(image).unsqueeze(0)  # Add batch dimension
+
+# Define multi-label prediction function
+def predict_waste(image):
+    image = Image.fromarray(image)  # Convert NumPy array to PIL image
+    input_tensor = preprocess_image(image)
+
+    # Get model predictions
+    with torch.no_grad():
+        outputs = model(input_tensor)
+
+    # Apply sigmoid activation for multi-label classification
+    probabilities = torch.sigmoid(outputs.logits)[0]  # Convert logits to probabilities
+
+    # Set a threshold to select labels (e.g., >= 50%)
+    threshold = 0.5
+    predicted_labels = [label for idx, label in model.config.id2label.items() if probabilities[idx] >= threshold]
+    confidence_scores = [f"{probabilities[idx] * 100:.2f}%" for idx in range(len(probabilities)) if probabilities[idx] >= threshold]
+
+    if predicted_labels:
+        result = "\n".join([f"{label}: {score}" for label, score in zip(predicted_labels, confidence_scores)])
+    else:
+        result = "No clear classification (confidence below threshold)"
+
+    return result
+
+# Create Gradio interface
+interface = gr.Interface(
+    fn=predict_waste,
+    inputs=gr.Image(type="numpy"),
+    outputs="text",
+    title="Multi-Label Waste Sorting App",
+    description="Upload an image of waste. The model will classify it into multiple waste categories with confidence scores."
+)
+
+# Launch the app
+interface.launch()