add prithivMLmods/Trash-Net

app.py

@@ -1,46 +1,46 @@
 import gradio as gr
-import tensorflow as tf
+from transformers import AutoImageProcessor
+from transformers import SiglipForImageClassification
+from transformers.image_utils import load_image
 from PIL import Image
-import numpy as np
-import requests
-import os
+import torch
 
-# Ensure model folder exists
-os.makedirs("model", exist_ok=True)
+# Load model and processor
+model_name = "prithivMLmods/Trash-Net"
+model = SiglipForImageClassification.from_pretrained(model_name)
+processor = AutoImageProcessor.from_pretrained(model_name)
 
-# Download the model from Hugging Face if not already present
-model_path = "model/mobnet_model.keras"
-if not os.path.exists(model_path):
-    url = "https://huggingface.co/ahmzakif/TrashNet-Classification/resolve/main/model/mobnet_model.keras"
-    r = requests.get(url)
-    with open(model_path, "wb") as f:
-        f.write(r.content)
-
-# Load Keras model
-model = tf.keras.models.load_model(model_path)
-
-# TrashNet classes
-classes = ["cardboard", "glass", "metal", "paper", "plastic", "trash"]
-
-# Image preprocessing
-def predict(image: Image.Image):
-    image = image.convert("RGB").resize((224, 224))
-    x = np.array(image, dtype=np.float32) / 255.0
-    x = np.expand_dims(x, axis=0)
+def trash_classification(image):
+    """Predicts the category of waste material in the image."""
+    image = Image.fromarray(image).convert("RGB")
+    inputs = processor(images=image, return_tensors="pt")
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+        probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
     
-    preds = model.predict(x)[0]
-    scores = {classes[i]: float(preds[i]) for i in range(len(classes))}
-    top_class = max(scores, key=scores.get)
+    labels = {
+        "0": "cardboard", 
+        "1": "glass", 
+        "2": "metal", 
+        "3": "paper", 
+        "4": "plastic", 
+        "5": "trash"
+    }
+    predictions = {labels[str(i)]: round(probs[i], 3) for i in range(len(probs))}
     
-    return {"prediction": top_class, "scores": scores}
+    return predictions
 
-# Gradio interface
+# Create Gradio interface
 iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs="json",
-    title="TrashNet Classification API",
-    description="Upload an image of trash to get its classification."
+    fn=trash_classification,
+    inputs=gr.Image(type="numpy"),
+    outputs=gr.Label(label="Prediction Scores"),
+    title="Trash Classification",
+    description="Upload an image to classify the type of waste material."
 )
 
-iface.launch()
+# Launch the app
+if __name__ == "__main__":
+    iface.launch()

requirements.txt

@@ -1,13 +1,4 @@
 gradio
-tqdm
-imutils
-numpy
-pandas
+transformers
+torch
 pillow
-matplotlib
-seaborn
-albumentations
-opencv-python
-tensorflow>=2.15.0
-scikit-learn
-wandb

runtime.txt

@@ -1 +1 @@
-python-3.11.5
+python-3.12