Update app.py

app.py

@@ -1,7 +1,8 @@
 # app.py
-from fastapi import FastAPI, Request, Response
-from PIL import Image
-import io, os, torch
+from fastapi import FastAPI, Request, Response, UploadFile, File
+from fastapi.middleware.cors import CORSMiddleware
+from PIL import Image, UnidentifiedImageError
+import io, os, torch, base64, traceback
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 
 # ========= CONFIG =========
@@ -21,55 +22,111 @@ 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)
+
 # ========= CARREGAMENTO =========
-processor = AutoImageProcessor.from_pretrained(MODEL_ID)
+# Removemos o aviso de processador "lento"
+processor = AutoImageProcessor.from_pretrained(MODEL_ID, use_fast=True)
 model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
 model.eval()
 
 app = FastAPI()
 
-def predict_image_bytes(img_bytes: bytes) -> str:
-    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
-    # Reduz um pouco para acelerar sem perder muito
-    img = img.resize((256, 256))
-
-    inputs = processor(images=img, return_tensors="pt")
-    logits = model(**inputs).logits
-    idx = int(logits.softmax(-1).argmax(-1))
-    label_en = model.config.id2label[idx].lower()
+# CORS (opcional; útil para testes via browser/front-end)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"], allow_credentials=True,
+    allow_methods=["*"], allow_headers=["*"],
+)
 
-    # Converte apenas se for uma das 4; senão marca como não identificado
+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")
 
+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
+
+    with torch.inference_mode():
+        inputs = processor(images=img, return_tensors="pt")
+        logits = model(**inputs).logits
+        idx = int(logits.softmax(-1).argmax(-1))
+        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")
 def health():
     return {"ok": True, "model": MODEL_ID}
 
 @app.post("/predict")
-async def predict(request: Request):
+async def predict(request: Request, file: UploadFile | None = File(default=None)):
     """
-    Espera: bytes JPEG (application/octet-stream)
+    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()
 
-        if "application/octet-stream" in ctype or "image/jpeg" in ctype:
+        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)
             img_bytes = await request.body()
+
         else:
-            # fallback opcional para JSON base64 (testes manuais)
+            # fallback para JSON base64 (útil em testes manuais via Postman)
             data = await request.json()
-            import base64
-            b64 = (data.get("image_b64") or "").split(",")[-1]
+            b64 = (data.get("image_b64") or "")
+            if "," in b64:
+                # aceita dataURL: "data:image/jpeg;base64,...."
+                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")
 
-        label = predict_image_bytes(img_bytes)
+        label = _predict_image_bytes(img_bytes)
         if label not in ALLOWED:
             label = "nao_identificado"
 
         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():
+            inputs = processor(images=dummy, return_tensors="pt")
+            _ = model(**inputs).logits
+        print("[startup] warm-up ok")
+    except Exception:
+        traceback.print_exc()
+        print("[startup] warm-up falhou (seguindo sem)")