habulaj commited on
Commit
96eebff
·
verified ·
1 Parent(s): 3581736

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -21
app.py CHANGED
@@ -8,9 +8,9 @@ import requests
8
  from io import BytesIO
9
  from PIL import Image, ImageFilter, ImageEnhance
10
  import rembg
 
11
  import cv2
12
  import numpy as np
13
- import onnxruntime as ort
14
 
15
  app = FastAPI()
16
 
@@ -18,16 +18,14 @@ app = FastAPI()
18
  options = ort.SessionOptions()
19
  options.intra_op_num_threads = 2
20
 
21
- # Função para redimensionar a imagem
22
  def resize_image(image, max_size=1024):
23
  width, height = image.size
24
  if width > max_size or height > max_size:
25
  ratio = min(max_size / width, max_size / height)
26
  new_size = (int(width * ratio), int(height * ratio))
27
- image = image.resize(new_size, Image.Resampling.LANCZOS) # Usando LANCZOS para redimensionamento
28
  return image
29
 
30
- # Função para ajustar brilho e contraste
31
  def adjust_brightness_contrast(image, brightness=1.2, contrast=1.2):
32
  enhancer = ImageEnhance.Brightness(image)
33
  image = enhancer.enhance(brightness)
@@ -35,29 +33,17 @@ def adjust_brightness_contrast(image, brightness=1.2, contrast=1.2):
35
  image = enhancer.enhance(contrast)
36
  return image
37
 
38
- # Função para remover sombras
39
  def remove_shadows(image):
40
- # Converte a imagem para um array numpy (compatível com OpenCV)
41
- image_np = np.array(image)
42
-
43
  # Converte a imagem para escala de cinza
44
- gray = cv2.cvtColor(image_np, cv2.COLOR_RGBA2GRAY)
45
 
46
- # Aplica limiarização para remover sombras
47
  _, mask = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY)
48
 
49
- # Aplica operações morfológicas para melhorar a máscara
50
- kernel = np.ones((3, 3), np.uint8)
51
- mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) # Remove ruídos
52
- mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) # Fecha buracos
53
-
54
  # Aplica a máscara na imagem original
55
- result = cv2.bitwise_and(image_np, image_np, mask=mask)
56
-
57
- # Converte o resultado de volta para uma imagem PIL
58
  return Image.fromarray(result)
59
 
60
- # Endpoint da API
61
  @app.get("/remove-background")
62
  async def remove_background(image_url: str):
63
  try:
@@ -73,12 +59,12 @@ async def remove_background(image_url: str):
73
  image = adjust_brightness_contrast(image)
74
 
75
  # Remove o fundo da imagem usando rembg
76
- output = rembg.remove(image, model="u2net", session_options=options)
77
 
78
  # Remove sombras
79
  output = remove_shadows(output)
80
 
81
- # Pós-processamento: suaviza as bordas
82
  output = output.filter(ImageFilter.SMOOTH_MORE)
83
 
84
  # Converte a imagem de volta para bytes
 
8
  from io import BytesIO
9
  from PIL import Image, ImageFilter, ImageEnhance
10
  import rembg
11
+ import onnxruntime as ort
12
  import cv2
13
  import numpy as np
 
14
 
15
  app = FastAPI()
16
 
 
18
  options = ort.SessionOptions()
19
  options.intra_op_num_threads = 2
20
 
 
21
  def resize_image(image, max_size=1024):
22
  width, height = image.size
23
  if width > max_size or height > max_size:
24
  ratio = min(max_size / width, max_size / height)
25
  new_size = (int(width * ratio), int(height * ratio))
26
+ image = image.resize(new_size, Image.Resampling.LANCZOS)
27
  return image
28
 
 
29
  def adjust_brightness_contrast(image, brightness=1.2, contrast=1.2):
30
  enhancer = ImageEnhance.Brightness(image)
31
  image = enhancer.enhance(brightness)
 
33
  image = enhancer.enhance(contrast)
34
  return image
35
 
 
36
  def remove_shadows(image):
 
 
 
37
  # Converte a imagem para escala de cinza
38
+ gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGBA2GRAY)
39
 
40
+ # Aplica limiarização para remover áreas escuras (sombras)
41
  _, mask = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY)
42
 
 
 
 
 
 
43
  # Aplica a máscara na imagem original
44
+ result = cv2.bitwise_and(np.array(image), np.array(image), mask=mask)
 
 
45
  return Image.fromarray(result)
46
 
 
47
  @app.get("/remove-background")
48
  async def remove_background(image_url: str):
49
  try:
 
59
  image = adjust_brightness_contrast(image)
60
 
61
  # Remove o fundo da imagem usando rembg
62
+ output = rembg.remove(image, model="u2net_human_seg", session_options=options)
63
 
64
  # Remove sombras
65
  output = remove_shadows(output)
66
 
67
+ # Suaviza a imagem
68
  output = output.filter(ImageFilter.SMOOTH_MORE)
69
 
70
  # Converte a imagem de volta para bytes