Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -12,7 +12,6 @@ import onnxruntime as ort
|
|
12 |
from concurrent.futures import ThreadPoolExecutor
|
13 |
import asyncio
|
14 |
import gc
|
15 |
-
from asyncio import Semaphore
|
16 |
|
17 |
app = FastAPI()
|
18 |
|
@@ -21,15 +20,9 @@ options = ort.SessionOptions()
|
|
21 |
options.intra_op_num_threads = 2 # Limita o número de threads para evitar sobrecarga
|
22 |
options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL # Execução sequencial para melhor desempenho em CPU
|
23 |
|
24 |
-
# Cria uma sessão global do rembg para reutilização
|
25 |
-
session = rembg.new_session(model_name="u2net", session_options=options)
|
26 |
-
|
27 |
# Pool de threads para executar tarefas bloqueantes
|
28 |
executor = ThreadPoolExecutor(max_workers=4)
|
29 |
|
30 |
-
# Semáforo para limitar requisições simultâneas
|
31 |
-
semaphore = Semaphore(4) # Ajuste conforme a capacidade da sua CPU/RAM
|
32 |
-
|
33 |
def resize_image(image, max_size=512):
|
34 |
"""Redimensiona a imagem para uma largura máxima de 512px, mantendo a proporção."""
|
35 |
width, height = image.size
|
@@ -40,46 +33,42 @@ def resize_image(image, max_size=512):
|
|
40 |
return image
|
41 |
|
42 |
def process_image(image_url):
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
|
56 |
-
|
57 |
-
|
58 |
-
|
59 |
-
|
60 |
-
|
61 |
-
|
62 |
-
|
63 |
-
|
64 |
-
|
65 |
-
# Libera memória explicitamente
|
66 |
-
del image, output
|
67 |
-
gc.collect() # Força a liberação de memória
|
68 |
-
|
69 |
-
return img_byte_arr
|
70 |
-
except Exception as e:
|
71 |
-
raise e
|
72 |
|
73 |
@app.get("/remove-background")
|
74 |
async def remove_background(image_url: str):
|
75 |
-
|
76 |
-
|
77 |
-
|
78 |
-
|
79 |
-
img_byte_arr = await loop.run_in_executor(executor, process_image, image_url)
|
80 |
-
|
81 |
-
# Retorna a imagem processada diretamente no navegador
|
82 |
-
return StreamingResponse(img_byte_arr, media_type="image/png")
|
83 |
|
84 |
-
|
85 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
12 |
from concurrent.futures import ThreadPoolExecutor
|
13 |
import asyncio
|
14 |
import gc
|
|
|
15 |
|
16 |
app = FastAPI()
|
17 |
|
|
|
20 |
options.intra_op_num_threads = 2 # Limita o número de threads para evitar sobrecarga
|
21 |
options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL # Execução sequencial para melhor desempenho em CPU
|
22 |
|
|
|
|
|
|
|
23 |
# Pool de threads para executar tarefas bloqueantes
|
24 |
executor = ThreadPoolExecutor(max_workers=4)
|
25 |
|
|
|
|
|
|
|
26 |
def resize_image(image, max_size=512):
|
27 |
"""Redimensiona a imagem para uma largura máxima de 512px, mantendo a proporção."""
|
28 |
width, height = image.size
|
|
|
33 |
return image
|
34 |
|
35 |
def process_image(image_url):
|
36 |
+
# Baixa a imagem da URL fornecida
|
37 |
+
response = requests.get(image_url)
|
38 |
+
response.raise_for_status()
|
39 |
+
|
40 |
+
# Abre a imagem usando Pillow
|
41 |
+
image = Image.open(BytesIO(response.content))
|
42 |
+
|
43 |
+
# Pré-processamento: apenas redimensiona para 512px
|
44 |
+
image = resize_image(image, max_size=512)
|
45 |
+
|
46 |
+
# Remove o fundo da imagem usando rembg
|
47 |
+
output = rembg.remove(image, session_options=options)
|
48 |
+
|
49 |
+
# Pós-processamento: suaviza as bordas
|
50 |
+
output = output.filter(ImageFilter.SMOOTH_MORE)
|
51 |
+
|
52 |
+
# Converte a imagem de volta para bytes
|
53 |
+
img_byte_arr = BytesIO()
|
54 |
+
output.save(img_byte_arr, format='PNG')
|
55 |
+
img_byte_arr.seek(0)
|
56 |
+
|
57 |
+
return img_byte_arr
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
58 |
|
59 |
@app.get("/remove-background")
|
60 |
async def remove_background(image_url: str):
|
61 |
+
try:
|
62 |
+
# Executa o processamento da imagem em um thread separado
|
63 |
+
loop = asyncio.get_event_loop()
|
64 |
+
img_byte_arr = await loop.run_in_executor(executor, process_image, image_url)
|
|
|
|
|
|
|
|
|
65 |
|
66 |
+
# Retorna a imagem processada diretamente no navegador
|
67 |
+
return StreamingResponse(img_byte_arr, media_type="image/png")
|
68 |
+
|
69 |
+
except Exception as e:
|
70 |
+
raise HTTPException(status_code=400, detail=str(e))
|
71 |
+
|
72 |
+
finally:
|
73 |
+
# Força a coleta de lixo após cada requisição
|
74 |
+
gc.collect()
|