Update app.py

app.py

@@ -1,43 +1,49 @@
-import gradio as gr
-from transformers import AutoImageProcessor, SiglipForImageClassification
-from PIL import Image
-import torch
-
-# Load processor and model
-model_name = "prithivMLmods/Recycling-Net-11"
-processor = AutoImageProcessor.from_pretrained(model_name)
-model = SiglipForImageClassification.from_pretrained(model_name)
-id2label = model.config.id2label
-
-# Label groups
-recyclable_labels = ["cardboard", "glass", "metal", "paper", "plastic", "can", "carton"]
-non_recyclable_labels = ["food waste", "trash", "garbage", "organic"]
-
-# Inference function
-def classify_image(image):
-    img = image.convert("RGB") if not isinstance(image, Image.Image) else image
-    inputs = processor(images=img, return_tensors="pt")
-
-    with torch.no_grad():
-        logits = model(**inputs).logits
-        probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
-
-    pred_idx = max(range(len(probs)), key=lambda i: probs[i])
-    pred_label = id2label[pred_idx].lower()
-
-    category = "Recyclable" if any(label in pred_label for label in recyclable_labels) else "Non-Recyclable"
-    confidence = probs[pred_idx] * 100
-
-    return f"{category} ({pred_label}) with {confidence:.2f}% confidence"
-
-# Gradio interface using Camera
-interface = gr.Interface(
-    fn=classify_image,
-    inputs=gr.Camera(label="Take a photo of the waste"),
-    outputs="text",
-    live=False,
-    title="♻️ Live Waste Classifier",
-    description="Capture waste via webcam to classify as recyclable or non-recyclable using a HuggingFace model."
-)
-
-interface.launch()
+from transformers import AutoImageProcessor, SiglipForImageClassification
+from PIL import Image
+import torch
+import gradio as gr
+
+# Load model and processor from HuggingFace
+model_name = "prithivMLmods/Recycling-Net-11"
+processor = AutoImageProcessor.from_pretrained(model_name)
+model = SiglipForImageClassification.from_pretrained(model_name)
+
+# Define recyclable and non-recyclable categories
+recyclable_labels = [
+    "cardboard", "glass", "metal", "paper", "plastic", "can", "carton"
+]
+non_recyclable_labels = [
+    "food waste", "trash", "garbage", "organic"
+]
+
+# Get model class label mapping
+id2label = model.config.id2label
+
+def classify_frame(frame):
+    if frame is None:
+        return "No frame detected"
+    
+    img = Image.fromarray(frame)
+    inputs = processor(images=img, return_tensors="pt")
+
+    with torch.no_grad():
+        logits = model(**inputs).logits
+        probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
+
+    pred_idx = max(range(len(probs)), key=lambda i: probs[i])
+    pred_label = id2label[pred_idx].lower()
+
+    if any(word in pred_label for word in recyclable_labels):
+        return f"♻️ Recyclable ({probs[pred_idx]*100:.1f}%)"
+    else:
+        return f"🗑️ Non-Recyclable ({probs[pred_idx]*100:.1f}%)"
+
+# Gradio Interface
+gr.Interface(
+    fn=classify_frame,
+    inputs=gr.Image(source="webcam", streaming=True, label="Live Waste Feed"),
+    outputs=gr.Text(label="Prediction"),
+    live=True,
+    title="Live Waste Classification",
+    description="Classifies live webcam input into Recyclable or Non-Recyclable using 11-class model."
+).launch()