Update app.py

app.py

@@ -1,40 +1,82 @@
+# app.py
 from fastapi import FastAPI, Request, Response
 from PIL import Image
-import io, torch
+import io, os, torch
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 
-MODEL_ID = "huggingface-projects/resnet-50"
-PT_MAP = {
-    "plastic":"plastico", "paper":"papel", "glass":"vidro",
-    "metal":"metal", "cardboard":"papel", "trash":"nao_identificado"
+# ========= CONFIG =========
+MODEL_ID = "AmadFR/ecovision_mobilenetv3"
+
+# Mapeamento para português (apenas as 4 classes desejadas)
+MAP_PT = {
+    "glass": "vidro",
+    "metal": "metal",
+    "paper": "papel",
+    "plastic": "plastico"
 }
+ALLOWED = set(MAP_PT.values())
+
+# ========= OTIMIZAÇÕES =========
+torch.set_grad_enabled(False)
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
 
+# ========= CARREGA MODELO =========
 processor = AutoImageProcessor.from_pretrained(MODEL_ID)
 model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
 model.eval()
 
 app = FastAPI()
 
-def predict_bytes(img_bytes: bytes) -> str:
+def predict_image_bytes(img_bytes: bytes) -> str:
+    """Recebe bytes JPEG e devolve um dos rótulos simplificados."""
     img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+    img = img.resize((224, 224))  # acelera a inferência
+
     inputs = processor(images=img, return_tensors="pt")
-    with torch.no_grad():
-        logits = model(**inputs).logits
-    idx = int(logits.softmax(-1).argmax(-1))
-    label_en = model.config.id2label[idx].lower()
-    return PT_MAP.get(label_en, "nao_identificado")
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class_idx = int(logits.argmax(-1))
+    label_en = model.config.id2label[predicted_class_idx].lower()
+
+    # Converte para português
+    label_pt = MAP_PT.get(label_en, "nao_identificado")
+    return label_pt
 
 @app.get("/health")
 def health():
-    return {"ok": True}
+    return {"ok": True, "model": MODEL_ID}
 
 @app.post("/predict")
 async def predict(request: Request):
+    """
+    Espera: imagem JPEG (application/octet-stream)
+    Retorna: texto puro - 'vidro', 'papel', 'plastico', 'metal' ou 'nao_identificado'
+    """
     try:
-        img_bytes = await request.body()
+        content_type = (request.headers.get("content-type") or "").lower()
+
+        # ESP32 envia como application/octet-stream
+        if "application/octet-stream" in content_type or "image/jpeg" in content_type:
+            img_bytes = await request.body()
+        else:
+            # fallback para JSON base64 (para testes)
+            data = await request.json()
+            import base64
+            b64 = data.get("image_b64", "").split(",")[-1]
+            img_bytes = base64.b64decode(b64) if b64 else b""
+
         if not img_bytes:
             return Response("nao_identificado", media_type="text/plain")
-        label = predict_bytes(img_bytes)
+
+        label = predict_image_bytes(img_bytes)
+
+        # filtro final — só 4 materiais
+        if label not in ALLOWED:
+            label = "nao_identificado"
+
         return Response(label, media_type="text/plain")
-    except Exception:
+
+    except Exception as e:
+        print("Erro:", e)
         return Response("nao_identificado", media_type="text/plain")