Classifier model change

app.py

@@ -1,24 +1,41 @@
 import gradio as gr
 import random
 import torch
-from transformers import (
-    AutoTokenizer, AutoModelForCausalLM,
-    T5Tokenizer,
-    T5ForConditionalGeneration,
-    pipeline
-)
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import CLIPProcessor, CLIPModel
+
+
+
 # ------------------ LOAD CLASSIFIER ------------------
-cls_model_name = "yangy50/garbage-classification"
-classifier = pipeline("image-classification", model=cls_model_name)
+clip_model_name = "openai/clip-vit-base-patch32"
+
+clip_model = CLIPModel.from_pretrained(clip_model_name)
+clip_processor = CLIPProcessor.from_pretrained(clip_model_name)
+
+clip_model.eval()
 
 def classify_image(image):
-    preds = classifier(image)
+    inputs = clip_processor(
+        text=WASTE_LABELS,
+        images=image,
+        return_tensors="pt",
+        padding=True
+    )
+
+    with torch.no_grad():
+        outputs = clip_model(**inputs)
+
+    logits = outputs.logits_per_image[0]
+    probs = logits.softmax(dim=0)
+
+    # return top 3
+    topk = torch.topk(probs, k=3)
+
+    results = {
+        WASTE_LABELS[i]: float(probs[i])
+        for i in topk.indices.tolist()
+    }
 
-    results = {}
-    for item in preds[:3]:
-        label = item["label"]
-        score = float(item["score"])
-        results[label] = score
     return results
 
 
@@ -204,6 +221,20 @@ waste_analyzation_v3 = {
 }
 analysis = [waste_analyzation_v3, waste_analyzation_v2, waste_analyzation]
 
+WASTE_LABELS = [
+    "cardboard",
+    "paper",
+    "plastic",
+    "glass",
+    "metal",
+    "food waste",
+    "organic waste",
+    "battery",
+    "electronic waste",
+    "textile",
+    "trash"
+]
+
 def classify_pipeline(image):
     global top_label
 
@@ -217,7 +248,7 @@ def analyze_pipeline():
     if top_label is None:
         return "Please classify an image first."
     
-    explanation = explain_recycling(top_label)
+    explanation = explain_recycling(top_label)  #EXCEEDS CPU, WILL GIVE INACCURATE ANSWER. DO NOT RETURN
     
     choice = random.randint(0, 2)
     info = analysis[choice][top_label]  

requirements.txt

@@ -4,5 +4,4 @@ Pillow
 torch
 transformers
 sentencepiece
-accelerate
-
+accelerate