Spaces:
Sleeping
Sleeping
File size: 2,596 Bytes
c3a6888 587862b c3a6888 587862b c3a6888 587862b c3a6888 587862b c3a6888 587862b c3a6888 587862b c3a6888 587862b c3a6888 587862b c3a6888 587862b c3a6888 587862b c3a6888 587862b c3a6888 587862b c3a6888 587862b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
# app.py
from fastapi import FastAPI, Request, Response
from PIL import Image
import io, os, torch
from transformers import AutoImageProcessor, AutoModelForImageClassification
# ========= 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_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")
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, "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:
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_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)
return Response("nao_identificado", media_type="text/plain")
|