w3ndson/api-image-manipulation

1from fastapi import FastAPI, File, UploadFile2from fastapi.middleware.cors import CORSMiddleware3from fastapi.responses import StreamingResponse4from rembg import new_session, remove5from PIL import Image6import io7import base648import os9from dotenv import load_dotenv10import requests11 12load_dotenv()13 14app = FastAPI()15 16IMGBB_API_KEY = os.getenv("IMGBB_API_KEY")17 18if not IMGBB_API_KEY:19    raise ValueError("IMGBB_API_KEY não configurada")20 21app.add_middleware(22    CORSMiddleware,23    allow_origins=["*"],24    allow_credentials=True,25    allow_methods=["*"],26    allow_headers=["*"],27)28 29# =========================================30# 🔹 Criar sessão de modelo uma única vez31# =========================================32# "isnet-general-use" é leve e com boa qualidade33session = new_session("isnet-general-use")34 35 36# =========================================37# 🔹 ROTA 1 — REMOVER FUNDO38# =========================================39@app.post("/remove-background")40async def remove_background(file: UploadFile = File(...)):41    try:42        contents = await file.read()43 44        input_image = Image.open(io.BytesIO(contents)).convert("RGBA")45 46        # Remoção com sessão otimizada47        output_image = remove(48            input_image,49            session=session,50            alpha_matting=True,51            alpha_matting_foreground_threshold=240,52            alpha_matting_background_threshold=10,53            alpha_matting_erode_size=1054        )55 56        # 🔹 Pós-processamento leve para suavizar bordas57        output_image = output_image.convert("RGBA")58 59        buffer = io.BytesIO()60        output_image.save(buffer, format="PNG", optimize=True)61        buffer.seek(0)62 63        return StreamingResponse(64            buffer,65            media_type="image/png",66            headers={67                "Content-Disposition": "attachment; filename=removed.png"68            }69        )70 71    except Exception as e:72        return {"error": str(e)}73 74 75 76# =========================================77# 🔹 ROTA 2 — UPLOAD PARA IMGBB78# =========================================79@app.post("/upload-imgbb")80async def upload_imgbb(file: UploadFile = File(...)):81    try:82        contents = await file.read()83        img_base64 = base64.b64encode(contents)84 85        response = requests.post(86            "https://api.imgbb.com/1/upload",87            data={88                "key": IMGBB_API_KEY,89                "image": img_base6490            }91        )92 93        result = response.json()94        if "data" not in result:95            return {"error": result}96 97        return {"message": "Imagem enviada com sucesso", "imgbb_url": result["data"]["url"]}98 99    except Exception as e:100        return {"error": str(e)}101