Update app.py

app.py

@@ -1,9 +1,3 @@
-# app.py
-# Purpose: FastAPI API for RecycloMate
-# - POST /predict with (file, category=regular|ewaste)
-# - regular  -> HuggingFace Transformers model (Trash-Net)
-# - ewaste   -> Roboflow serverless API (no inference-sdk, no cv2)
-
 import io
 import os
 from typing import Dict, Any
@@ -22,13 +16,49 @@ HF_MODEL_ID = "prithivMLmods/Trash-Net"
 
 ROBOFLOW_API_KEY = os.getenv("ROBOFLOW_API_KEY", "")
 ROBOFLOW_MODEL_ID = os.getenv("ROBOFLOW_MODEL_ID", "e-waste-2ecoq/2")
-ROBOFLOW_URL = f"https://detect.roboflow.com/{ROBOFLOW_MODEL_ID}?api_key={ROBOFLOW_API_KEY}"
 
 CONFIDENCE_THRESHOLD = 0.70
 MARGIN_THRESHOLD = 0.15
 
 RECYCLABLE = {"cardboard", "glass", "metal", "paper", "plastic"}
 
+# ----------------------------
+# ✅ E-waste label control (ADD THIS)
+# ----------------------------
+EWASTE_ALLOWED = {
+    "keyboards",
+    "mobile",
+    "mouses",
+    "tv",
+    "camera",
+    "laptop",
+    "microwave",
+}
+
+# Map common label variations -> your allowed labels
+EWASTE_ALIAS = {
+    "keyboard": "keyboards",
+    "keyboards": "keyboards",
+    "mouse": "mouses",
+    "mice": "mouses",
+    "mouses": "mouses",
+    "phone": "mobile",
+    "phones": "mobile",
+    "cellphone": "mobile",
+    "cell phone": "mobile",
+    "smartphone": "mobile",
+    "mobile": "mobile",
+    "television": "tv",
+    "tvs": "tv",
+    "tv": "tv",
+    "camera": "camera",
+    "cameras": "camera",
+    "laptop": "laptop",
+    "laptops": "laptop",
+    "microwave": "microwave",
+    "microwaves": "microwave",
+}
+
 # ----------------------------
 # Load HF model once
 # ----------------------------
@@ -44,7 +74,6 @@ CLASS_NAMES = [id2label[i] for i in range(len(id2label))]
 # ----------------------------
 app = FastAPI()
 
-# If you call this Space from Firebase Hosting, keep CORS open or restrict later
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],  # later you can put your firebase domain only
@@ -114,7 +143,6 @@ def roboflow_predict(image_bytes: bytes) -> Dict[str, Any]:
     if not preds:
         return {"class": "unknown", "recyclable": False, "confidence": 0.0}
 
-    # pick best prediction
     best = max(preds, key=lambda p: p.get("confidence", 0) or 0)
     raw_label = str(best.get("class", "unknown")).strip().lower()
     conf = float(best.get("confidence", 0.0) or 0.0)
@@ -132,13 +160,11 @@ def roboflow_predict(image_bytes: bytes) -> Dict[str, Any]:
         }
 
     return {
-        "class": label,
+        "class": label,  # return the normalized label
         "recyclable": False,
         "confidence": round(conf, 4),
     }
 
-
-
 @app.post("/predict")
 async def predict(
     file: UploadFile = File(...),
@@ -151,7 +177,6 @@ async def predict(
         if category == "ewaste":
             return roboflow_predict(image_bytes)
 
-        # default: regular
         return hf_predict(image)
 
     except Exception as e: