Update app.py

app.py

@@ -1,61 +1,105 @@
-from fastapi import FastAPI, UploadFile, File
+import os
+from fastapi import FastAPI, UploadFile, File, Form
 from fastapi.middleware.cors import CORSMiddleware
-from transformers import AutoImageProcessor, SiglipForImageClassification
 from PIL import Image
+import io
 import torch
+from transformers import AutoImageProcessor, SiglipForImageClassification
 
-MODEL_ID = "prithivMLmods/Trash-Net"
-
-PROCESSOR = AutoImageProcessor.from_pretrained(MODEL_ID)
-MODEL = SiglipForImageClassification.from_pretrained(MODEL_ID)
-MODEL.eval()
-
-id2label = MODEL.config.id2label
-CLASS_NAMES = [id2label[i] for i in range(len(id2label))]
-RECYCLABLE = {"cardboard", "glass", "metal", "paper", "plastic"}
-
-CONFIDENCE_THRESHOLD = 0.70
-MARGIN_THRESHOLD = 0.15
+# Roboflow client
+from inference_sdk import InferenceHTTPClient
 
 app = FastAPI()
+
+# Allow frontend (Firebase) to call this API
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
+    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
+# ===============================
+# 🔹 Everyday recyclables model (HF)
+# ===============================
+MODEL_NAME = "prithivMLmods/Trash-Net"
+
+processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+model = SiglipForImageClassification.from_pretrained(MODEL_NAME)
+model.eval()
+
+LABELS = model.config.id2label
+
+
+# ===============================
+# 🔹 Roboflow E-waste client
+# ===============================
+RF_CLIENT = InferenceHTTPClient(
+    api_url="https://serverless.roboflow.com",
+    api_key=os.getenv("ROBOFLOW_API_KEY")   # set in Space secrets
+)
+
+RF_MODEL_ID = os.getenv("ROBOFLOW_MODEL_ID", "e-waste-2ecoq/2")
+
+
+# ===============================
+# 🔹 Health check
+# ===============================
 @app.get("/health")
 def health():
-    return {"ok": True, "model": MODEL_ID}
+    return {"status": "ok"}
+
 
+# ===============================
+# 🔹 Unified predict endpoint
+# ===============================
 @app.post("/predict")
-async def predict(file: UploadFile = File(...)):
-    img = Image.open(file.file).convert("RGB")
-    inputs = PROCESSOR(images=img, return_tensors="pt")
+async def predict(
+    file: UploadFile = File(...),
+    category: str = Form("regular")   # "regular" or "ewaste"
+):
+    image_bytes = await file.read()
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-    with torch.no_grad():
-        probs = torch.softmax(MODEL(**inputs).logits, dim=1)[0]
-        top2 = torch.topk(probs, k=2)
+    # ---------------------------
+    # 🔸 E-WASTE → Roboflow API
+    # ---------------------------
+    if category == "ewaste":
+        result = RF_CLIENT.infer(image, model_id=RF_MODEL_ID)
+
+        if not result["predictions"]:
+            return {"error": "No object detected"}
 
-    top1_idx = int(top2.indices[0].item())
-    top2_idx = int(top2.indices[1].item())
-    top1_prob = float(top2.values[0].item())
-    top2_prob = float(top2.values[1].item())
-    margin = top1_prob - top2_prob
+        pred = result["predictions"][0]
 
-    pred_class = CLASS_NAMES[top1_idx]
+        label = pred["class"]
+        confidence = pred["confidence"]
 
-    if top1_prob < CONFIDENCE_THRESHOLD or margin < MARGIN_THRESHOLD:
         return {
-            "class": "unknown",
-            "recyclable": False,
-            "confidence": round(top1_prob, 4),
-            "message": "Not confident enough. Retake photo.",
+            "class": label.lower(),
+            "recyclable": True,
+            "confidence": confidence
         }
 
+    # ---------------------------
+    # 🔸 REGULAR → HF model
+    # ---------------------------
+    inputs = processor(images=image, return_tensors="pt")
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probs = torch.softmax(outputs.logits, dim=1)
+
+    score, idx = torch.max(probs, dim=1)
+
+    label = LABELS[idx.item()].lower()
+    confidence = float(score.item())
+
+    recyclable_classes = ["plastic", "paper", "metal", "glass", "cardboard"]
+
     return {
-        "class": pred_class,
-        "recyclable": pred_class in RECYCLABLE,
-        "confidence": round(top1_prob, 4),
+        "class": label,
+        "recyclable": label in recyclable_classes,
+        "confidence": confidence
     }