Update app.py

app.py

@@ -1,23 +1,22 @@
 # app.py
 from fastapi import FastAPI, Request, Response
-from PIL import Image, ImageOps
+from PIL import Image
 import io, os, torch
-import torch.nn.functional as F
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 
 # ========= CONFIG =========
 MODEL_ID = "prithivMLmods/Trash-Net"
 
-# PT-BR map (somente 4 classes principais)
+# Mantemos só estas 4 classes em PT-BR; o resto vira "nao_identificado"
 MAP_PT = {
     "glass": "vidro",
     "metal": "metal",
     "paper": "papel",
     "plastic": "plastico",
 }
-TARGETS_EN = list(MAP_PT.keys())
+ALLOWED = set(MAP_PT.values())
 
-# ========= OTIMIZAÇÕES =========
+# ========= OTIMIZAÇÕES (CPU do Space) =========
 torch.set_grad_enabled(False)
 torch.set_num_threads(1)
 torch.set_num_interop_threads(1)
@@ -27,113 +26,50 @@ processor = AutoImageProcessor.from_pretrained(MODEL_ID)
 model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
 model.eval()
 
-# --- Trata o caso label2id invertido ---
-id2label_raw = model.config.id2label
-label2id_raw = model.config.label2id
-
-id2label = {}
-label2id = {}
-
-for k, v in id2label_raw.items():
-    # Normaliza chaves e valores para int→str
-    try:
-        id2label[int(k)] = str(v)
-    except Exception:
-        id2label[int(v)] = str(k)
-
-for k, v in label2id_raw.items():
-    # Normaliza para str→int
-    try:
-        label2id[str(k).lower()] = int(v)
-    except Exception:
-        label2id[str(v).lower()] = int(k)
-
-# Descobre índices das 4 classes-alvo
-target_indices = []
-target_indices_en = []
-for en in TARGETS_EN:
-    if en in label2id:
-        target_indices.append(label2id[en])
-        target_indices_en.append(en)
-
-if len(target_indices) < 4:
-    for en in TARGETS_EN:
-        if en in target_indices_en:
-            continue
-        found = None
-        en_low = en.lower()
-        for i, lab in id2label.items():
-            if en_low in lab.lower():
-                found = i
-                break
-        if found is not None and found not in target_indices:
-            target_indices.append(found)
-            target_indices_en.append(en)
-
-# ========= APP =========
 app = FastAPI()
 
-# ========= FUNÇÕES =========
-def _prepare_image(img_bytes: bytes) -> Image.Image:
+def predict_image_bytes(img_bytes: bytes) -> str:
     img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
-    img = ImageOps.exif_transpose(img)
-    w, h = img.size
-    side = min(w, h)
-    left = (w - side) // 2
-    top = (h - side) // 2
-    img = img.crop((left, top, left + side, top + side))
-    return img
+    # Reduz um pouco para acelerar sem perder muito
+    img = img.resize((256, 256))
 
-def predict_image_bytes(img_bytes: bytes):
-    img = _prepare_image(img_bytes)
     inputs = processor(images=img, return_tensors="pt")
     logits = model(**inputs).logits
-    probs = F.softmax(logits, dim=-1)[0]
+    idx = int(logits.softmax(-1).argmax(-1))
+    label_en = model.config.id2label[idx].lower()
 
-    if target_indices:
-        subset = probs[target_indices]
-        j = int(torch.argmax(subset).item())
-        best_idx_global = target_indices[j]
-        best_en = id2label[best_idx_global].lower()
-        conf = float(subset[j].item())
-        label_pt = MAP_PT.get(best_en, MAP_PT[target_indices_en[j]])
-        return label_pt, conf
-    else:
-        i = int(torch.argmax(probs).item())
-        best_en = id2label[i].lower()
-        conf = float(probs[i].item())
-        if "glass" in best_en:
-            label_pt = "vidro"
-        elif "metal" in best_en or "steel" in best_en or "aluminum" in best_en:
-            label_pt = "metal"
-        elif "paper" in best_en or "cardboard" in best_en:
-            label_pt = "papel"
-        else:
-            label_pt = "plastico"
-        return label_pt, conf
+    # Converte apenas se for uma das 4; senão marca como não identificado
+    return MAP_PT.get(label_en, "nao_identificado")
 
-# ========= ROTAS =========
 @app.get("/health")
 def health():
-    return {"ok": True, "model": MODEL_ID, "targets": list(MAP_PT.values())}
+    return {"ok": True, "model": MODEL_ID}
 
 @app.post("/predict")
 async def predict(request: Request):
+    """
+    Espera: bytes JPEG (application/octet-stream)
+    Retorna: texto puro — 'vidro' | 'papel' | 'plastico' | 'metal' | 'nao_identificado'
+    """
     try:
         ctype = (request.headers.get("content-type") or "").lower()
-        if "application/octet-stream" in ctype or "image/jpeg" in ctype or "image/png" in ctype:
+
+        if "application/octet-stream" in ctype or "image/jpeg" in ctype:
             img_bytes = await request.body()
         else:
+            # fallback opcional para JSON base64 (testes manuais)
             data = await request.json()
             import base64
             b64 = (data.get("image_b64") or "").split(",")[-1]
             img_bytes = base64.b64decode(b64) if b64 else b""
 
         if not img_bytes:
-            return Response("plastico", media_type="text/plain", headers={"X-Confidence": "0.00"})
+            return Response("nao_identificado", media_type="text/plain")
+
+        label = predict_image_bytes(img_bytes)
+        if label not in ALLOWED:
+            label = "nao_identificado"
 
-        label_pt, conf = predict_image_bytes(img_bytes)
-        return Response(label_pt, media_type="text/plain", headers={"X-Confidence": f"{conf:.4f}"})
-    except Exception as e:
-        print("predict error:", e)
-        return Response("plastico", media_type="text/plain", headers={"X-Confidence": "0.00"})
+        return Response(label, media_type="text/plain")
+    except Exception:
+        return Response("nao_identificado", media_type="text/plain")