habulaj commited on
Commit
3231e07
·
verified ·
1 Parent(s): 72a2ce8

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +35 -46
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
- try:
44
- # Baixa a imagem da URL fornecida
45
- response = requests.get(image_url)
46
- response.raise_for_status()
47
-
48
- # Abre a imagem usando Pillow
49
- image = Image.open(BytesIO(response.content))
50
-
51
- # Pré-processamento: apenas redimensiona para 512px
52
- image = resize_image(image, max_size=512)
53
-
54
- # Remove o fundo da imagem usando rembg (reutiliza a sessão global)
55
- output = rembg.remove(image, session=session)
56
-
57
- # Pós-processamento: suaviza as bordas
58
- output = output.filter(ImageFilter.SMOOTH_MORE)
59
-
60
- # Converte a imagem de volta para bytes
61
- img_byte_arr = BytesIO()
62
- output.save(img_byte_arr, format='PNG')
63
- img_byte_arr.seek(0)
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
- async with semaphore: # Limita o número de requisições simultâneas
76
- try:
77
- # Executa o processamento da imagem em um thread separado
78
- loop = asyncio.get_event_loop()
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
- except Exception as e:
85
- raise HTTPException(status_code=400, detail=str(e))
 
 
 
 
 
 
 
 
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()