Classifier model change

app.py

@@ -14,9 +14,52 @@ clip_processor = CLIPProcessor.from_pretrained(clip_model_name)
 
 clip_model.eval()
 
+CLIP_LABELS = {
+    "paper/cardboard": [
+        "a photo of cardboard packaging waste",
+        "a photo of paper waste or newspaper"
+    ],
+    "plastic": [
+        "a photo of plastic bottle or plastic packaging waste"
+    ],
+    "glass": [
+        "a photo of glass bottle or glass container waste"
+    ],
+    "metal": [
+        "a photo of metal can or metal packaging waste"
+    ],
+    "organic waste": [
+        "a photo of food waste or leftovers",
+        "a photo of organic waste such as fruit peels or leaves"
+    ],
+    "hazardous waste": [
+        "a photo of a used battery waste item",
+        "a photo of hazardous household waste"
+    ],
+    "electronic waste": [
+        "a photo of electronic waste such as cables, chargers, or devices"
+    ],
+    "textile": [
+        "a photo of textile waste or old clothing"
+    ],
+    "general trash": [
+        "a photo of mixed general trash or landfill waste"
+    ]
+}
+
+
 def classify_image(image):
+    all_prompts = []
+    prompt_to_category = {}
+
+    # Build prompt list
+    for category, prompts in CLIP_LABELS.items():
+        for p in prompts:
+            all_prompts.append(p)
+            prompt_to_category[p] = category
+
     inputs = clip_processor(
-        text=WASTE_LABELS,
+        text=all_prompts,
         images=image,
         return_tensors="pt",
         padding=True
@@ -28,15 +71,23 @@ def classify_image(image):
     logits = outputs.logits_per_image[0]
     probs = logits.softmax(dim=0)
 
-    # return top 3
-    topk = torch.topk(probs, k=3)
+    # merge probabilities by category
+    category_scores = {}
+    for i, p in enumerate(all_prompts):
+        cat = prompt_to_category[p]
+        category_scores[cat] = category_scores.get(cat, 0) + float(probs[i])
+    total = sum(category_scores.values())
+    for k in category_scores:
+        category_scores[k] /= total
 
-    results = {
-        WASTE_LABELS[i]: float(probs[i])
-        for i in topk.indices.tolist()
-    }
+    # Return top 3
+    top3 = sorted(
+        category_scores.items(),
+        key=lambda x: x[1],
+        reverse=True
+    )[:3]
 
-    return results
+    return dict(top3)
 
 
 # ------------------ LOAD CHAT MODEL ------------------
@@ -104,7 +155,6 @@ def explain_recycling(class_label):
 
 
 # ------------------ PIPELINE ------------------
-top_label = None
 waste_analyzation = {
     "cardboard": {
         "Recycling type": "Paper/Cardboard recycling",
@@ -221,19 +271,8 @@ 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"
-]
+top_label = None
+
 
 def classify_pipeline(image):
     global top_label