Update app.py

app.py

@@ -5,23 +5,23 @@ import io, os, torch
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 
 # ========= CONFIG =========
-MODEL_ID = "AmadFR/ecovision_mobilenetv3"
+MODEL_ID = "prithivMLmods/Trash-Net"
 
-# Mapeamento para português (apenas as 4 classes desejadas)
+# Mantemos só estas 4 classes em PT-BR; o resto vira "nao_identificado"
 MAP_PT = {
     "glass": "vidro",
     "metal": "metal",
     "paper": "papel",
-    "plastic": "plastico"
+    "plastic": "plastico",
 }
 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)
 
-# ========= CARREGA MODELO =========
+# ========= CARREGAMENTO =========
 processor = AutoImageProcessor.from_pretrained(MODEL_ID)
 model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
 model.eval()
@@ -29,19 +29,17 @@ model.eval()
 app = FastAPI()
 
 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
+    # Reduz um pouco para acelerar sem perder muito
+    img = img.resize((256, 256))
 
     inputs = processor(images=img, return_tensors="pt")
-    outputs = model(**inputs)
-    logits = outputs.logits
-    predicted_class_idx = int(logits.argmax(-1))
-    label_en = model.config.id2label[predicted_class_idx].lower()
+    logits = model(**inputs).logits
+    idx = int(logits.softmax(-1).argmax(-1))
+    label_en = model.config.id2label[idx].lower()
 
-    # Converte para português
-    label_pt = MAP_PT.get(label_en, "nao_identificado")
-    return label_pt
+    # Converte apenas se for uma das 4; senão marca como não identificado
+    return MAP_PT.get(label_en, "nao_identificado")
 
 @app.get("/health")
 def health():
@@ -50,33 +48,28 @@ def health():
 @app.post("/predict")
 async def predict(request: Request):
     """
-    Espera: imagem JPEG (application/octet-stream)
-    Retorna: texto puro - 'vidro', 'papel', 'plastico', 'metal' ou 'nao_identificado'
+    Espera: bytes JPEG (application/octet-stream)
+    Retorna: texto puro — 'vidro' | 'papel' | 'plastico' | 'metal' | 'nao_identificado'
     """
     try:
-        content_type = (request.headers.get("content-type") or "").lower()
+        ctype = (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:
+        if "application/octet-stream" in ctype or "image/jpeg" in ctype:
             img_bytes = await request.body()
         else:
-            # fallback para JSON base64 (para testes)
+            # fallback opcional para JSON base64 (testes manuais)
             data = await request.json()
             import base64
-            b64 = data.get("image_b64", "").split(",")[-1]
+            b64 = (data.get("image_b64") or "").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_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 as e:
-        print("Erro:", e)
+    except Exception:
         return Response("nao_identificado", media_type="text/plain")