Update app.py

app.py

@@ -1,8 +1,8 @@
 # app.py
-from fastapi import FastAPI, Request, Response, UploadFile, File
+from fastapi import FastAPI, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image, UnidentifiedImageError
-import io, os, torch, base64, traceback
+import io, torch, base64, traceback
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 
 # ========= CONFIG =========
@@ -22,18 +22,19 @@ torch.set_grad_enabled(False)
 torch.set_num_threads(1)
 torch.set_num_interop_threads(1)
 
-# Evita ataques com imagens gigantes (e economiza RAM)
-Image.MAX_IMAGE_PIXELS = 25_000_000  # ~25MP de teto (bem alto p/ segurança)
+# Evita imagens gigantes
+Image.MAX_IMAGE_PIXELS = 25_000_000
 
 # ========= CARREGAMENTO =========
-# Removemos o aviso de processador "lento"
+# 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()
 
 app = FastAPI()
 
-# CORS (opcional; útil para testes via browser/front-end)
+# CORS (opcional)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"], allow_credentials=True,
@@ -41,27 +42,23 @@ app.add_middleware(
 )
 
 def _to_label_pt(label_en: str) -> str:
-    # Normaliza e converte somente se estiver mapeado
-    label_en = (label_en or "").strip().lower()
-    return MAP_PT.get(label_en, "nao_identificado")
+    return MAP_PT.get((label_en or "").strip().lower(), "nao_identificado")
 
 def _predict_image_bytes(img_bytes: bytes) -> str:
-    # Lê, converte e reduz um pouco p/ acelerar, mantendo boa acurácia
     with Image.open(io.BytesIO(img_bytes)) as img:
-        img = img.convert("RGB")  # garante 3 canais
-        img = img.resize((256, 256))  # tradeoff bom p/ CPU do Space
+        img = img.convert("RGB")
+        img = img.resize((256, 256))  # tradeoff bom para CPU do Space
 
     with torch.inference_mode():
         inputs = processor(images=img, return_tensors="pt")
         logits = model(**inputs).logits
-        idx = int(logits.softmax(-1).argmax(-1))
+        idx = int(logits.argmax(-1))  # softmax não precisa para argmax
         label_en = model.config.id2label[idx]
         return _to_label_pt(label_en)
 
 # ========= ROTAS =========
 @app.get("/")
 def root():
-    # Retornar 200 aqui ajuda o “wake up” do Space pelo firmware
     return {"ok": True, "message": "TrashNet classifier up", "model": MODEL_ID}
 
 @app.get("/health")
@@ -69,34 +66,26 @@ def health():
     return {"ok": True, "model": MODEL_ID}
 
 @app.post("/predict")
-async def predict(request: Request, file: UploadFile | None = File(default=None)):
+async def predict(request: Request):
     """
     Aceita:
     - application/octet-stream (raw JPEG no corpo)
     - image/jpeg (raw JPEG no corpo)
-    - multipart/form-data (campo: file)
     - application/json {"image_b64": "..."} (dataURL ou base64 puro)
 
     Retorna: texto puro — 'vidro' | 'papel' | 'plastico' | 'metal' | 'nao_identificado'
     """
     try:
-        img_bytes: bytes = b""
         ctype = (request.headers.get("content-type") or "").lower()
+        img_bytes: bytes = b""
 
-        if file is not None:
-            # multipart/form-data
-            img_bytes = await file.read()
-
-        elif "application/octet-stream" in ctype or "image/jpeg" in ctype:
-            # raw bytes (ESP32 manda deste jeito)
+        if "application/octet-stream" in ctype or "image/jpeg" in ctype:
             img_bytes = await request.body()
-
         else:
-            # fallback para JSON base64 (útil em testes manuais via Postman)
+            # fallback: JSON base64
             data = await request.json()
             b64 = (data.get("image_b64") or "")
-            if "," in b64:
-                # aceita dataURL: "data:image/jpeg;base64,...."
+            if "," in b64:  # dataURL
                 b64 = b64.split(",", 1)[1]
             img_bytes = base64.b64decode(b64) if b64 else b""
 
@@ -110,17 +99,14 @@ async def predict(request: Request, file: UploadFile | None = File(default=None)
         return Response(label, media_type="text/plain")
 
     except UnidentifiedImageError:
-        # bytes não eram uma imagem válida
         return Response("nao_identificado", media_type="text/plain")
     except Exception:
-        # loga stack trace no console do Space para debug
         traceback.print_exc()
         return Response("nao_identificado", media_type="text/plain")
 
 # ========= WARM-UP =========
 @app.on_event("startup")
 def _warmup():
-    # Faz uma inferência boba para “carregar” tudo e reduzir a latência da 1ª chamada real
     try:
         dummy = Image.new("RGB", (256, 256), (127, 127, 127))
         with torch.inference_mode():
@@ -129,4 +115,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)")