Update app.py

app.py

@@ -3,8 +3,7 @@ os.environ["NUMBA_CACHE_DIR"] = "/tmp/numba_cache"
 os.environ["U2NET_HOME"] = "/tmp/u2net"
 
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse
-import requests
+from pydantic import BaseModel
 from io import BytesIO
 from PIL import Image, ImageFilter
 import rembg
@@ -12,6 +11,7 @@ import onnxruntime as ort
 from concurrent.futures import ThreadPoolExecutor
 import asyncio
 import gc
+import base64
 
 app = FastAPI()
 
@@ -32,13 +32,9 @@ def resize_image(image, max_size=512):
         image = image.resize(new_size, Image.Resampling.LANCZOS)
     return image
 
-def process_image(image_url):
-    # Baixa a imagem da URL fornecida
-    response = requests.get(image_url)
-    response.raise_for_status()
-    
-    # Abre a imagem usando Pillow
-    image = Image.open(BytesIO(response.content))
+def process_image(image_data):
+    # Decodifica a imagem base64
+    image = Image.open(BytesIO(base64.b64decode(image_data)))
     
     # Pré-processamento: apenas redimensiona para 512px
     image = resize_image(image, max_size=512)
@@ -54,17 +50,21 @@ def process_image(image_url):
     output.save(img_byte_arr, format='PNG')
     img_byte_arr.seek(0)
     
-    return img_byte_arr
+    # Codifica a imagem processada em base64
+    return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
+
+class ImageBase64(BaseModel):
+    image: str  # A imagem em base64
 
-@app.get("/remove-background")
-async def remove_background(image_url: str):
+@app.post("/remove-background")
+async def remove_background(image_data: ImageBase64):
     try:
         # Executa o processamento da imagem em um thread separado
         loop = asyncio.get_event_loop()
-        img_byte_arr = await loop.run_in_executor(executor, process_image, image_url)
+        processed_image_base64 = await loop.run_in_executor(executor, process_image, image_data.image)
         
-        # Retorna a imagem processada diretamente no navegador
-        return StreamingResponse(img_byte_arr, media_type="image/png")
+        # Retorna a imagem processada em base64
+        return {"processed_image": processed_image_base64}
     
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))