Update app.py

app.py

@@ -2,20 +2,20 @@
 from fastapi import FastAPI, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image, UnidentifiedImageError
-import io, torch, base64, traceback
+import io, torch, base64, traceback, random
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 
 # ========= CONFIG =========
 MODEL_ID = "prithivMLmods/Trash-Net"
 
-# Mantemos só estas 4 classes em PT-BR; o resto vira "nao_identificado"
+# Mapa EN -> PT (apenas 4 classes desejadas)
 MAP_PT = {
     "glass": "vidro",
     "metal": "metal",
     "paper": "papel",
     "plastic": "plastico",
 }
-ALLOWED = set(MAP_PT.values())
+ALLOWED = ["plastico", "papel", "vidro", "metal"]  # ordem fixa p/ random
 
 # ========= OTIMIZAÇÕES (CPU do Space) =========
 torch.set_grad_enabled(False)
@@ -26,8 +26,6 @@ torch.set_num_interop_threads(1)
 Image.MAX_IMAGE_PIXELS = 25_000_000
 
 # ========= CARREGAMENTO =========
-# Se aparecer aviso "use_fast=True mas torchvision não disponível",
-# é só um warning; pode trocar para use_fast=False se quiser ocultar.
 processor = AutoImageProcessor.from_pretrained(MODEL_ID, use_fast=True)
 model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
 model.eval()
@@ -41,8 +39,14 @@ app.add_middleware(
     allow_methods=["*"], allow_headers=["*"],
 )
 
-def _to_label_pt(label_en: str) -> str:
-    return MAP_PT.get((label_en or "").strip().lower(), "nao_identificado")
+def _force_allowed(label_en: str | None) -> str:
+    """Converte label EN para PT se mapeado; caso contrário, escolhe aleatoriamente uma das 4."""
+    if label_en:
+        pt = MAP_PT.get(label_en.strip().lower())
+        if pt in ALLOWED:
+            return pt
+    # fallback forçado
+    return random.choice(ALLOWED)
 
 def _predict_image_bytes(img_bytes: bytes) -> str:
     with Image.open(io.BytesIO(img_bytes)) as img:
@@ -52,9 +56,9 @@ def _predict_image_bytes(img_bytes: bytes) -> str:
     with torch.inference_mode():
         inputs = processor(images=img, return_tensors="pt")
         logits = model(**inputs).logits
-        idx = int(logits.argmax(-1))  # softmax não precisa para argmax
+        idx = int(logits.argmax(-1))
         label_en = model.config.id2label[idx]
-        return _to_label_pt(label_en)
+        return _force_allowed(label_en)
 
 # ========= ROTAS =========
 @app.get("/")
@@ -73,36 +77,40 @@ async def predict(request: Request):
     - image/jpeg (raw JPEG no corpo)
     - application/json {"image_b64": "..."} (dataURL ou base64 puro)
 
-    Retorna: texto puro — 'vidro' | 'papel' | 'plastico' | 'metal' | 'nao_identificado'
+    Retorna SEMPRE: texto puro — 'vidro' | 'papel' | 'plastico' | 'metal'
+    (nunca 'nao_identificado')
     """
     try:
-        ctype = (request.headers.get("content-type") or "").lower()
-        img_bytes: bytes = b""
+      ctype = (request.headers.get("content-type") or "").lower()
+      img_bytes: bytes = b""
 
-        if "application/octet-stream" in ctype or "image/jpeg" in ctype:
-            img_bytes = await request.body()
-        else:
-            # fallback: JSON base64
-            data = await request.json()
-            b64 = (data.get("image_b64") or "")
-            if "," in b64:  # dataURL
-                b64 = b64.split(",", 1)[1]
-            img_bytes = base64.b64decode(b64) if b64 else b""
+      if "application/octet-stream" in ctype or "image/jpeg" in ctype:
+          img_bytes = await request.body()
+      else:
+          # fallback: JSON base64
+          data = await request.json()
+          b64 = (data.get("image_b64") or "")
+          if "," in b64:  # dataURL
+              b64 = b64.split(",", 1)[1]
+          img_bytes = base64.b64decode(b64) if b64 else b""
 
-        if not img_bytes:
-            return Response("nao_identificado", media_type="text/plain")
+      # Se veio vazio, ainda assim devolve um dos 4
+      if not img_bytes:
+          return Response(random.choice(ALLOWED), media_type="text/plain")
 
-        label = _predict_image_bytes(img_bytes)
-        if label not in ALLOWED:
-            label = "nao_identificado"
+      label = _predict_image_bytes(img_bytes)
 
-        return Response(label, media_type="text/plain")
+      # Por garantia, força para uma das 4
+      if label not in ALLOWED:
+          label = random.choice(ALLOWED)
+
+      return Response(label, media_type="text/plain")
 
     except UnidentifiedImageError:
-        return Response("nao_identificado", media_type="text/plain")
+      return Response(random.choice(ALLOWED), media_type="text/plain")
     except Exception:
-        traceback.print_exc()
-        return Response("nao_identificado", media_type="text/plain")
+      traceback.print_exc()
+      return Response(random.choice(ALLOWED), media_type="text/plain")
 
 # ========= WARM-UP =========
 @app.on_event("startup")
@@ -115,4 +123,4 @@ def _warmup():
         print("[startup] warm-up ok")
     except Exception:
         traceback.print_exc()
-        print("[startup] warm-up falhou (seguindo sem)")
+        print("[startup] warm-up falhou (seguindo sem)")