Spaces:
Sleeping
Sleeping
File size: 2,676 Bytes
add1294 afac93b add1294 4b2ecc5 d77e4f1 1aab113 d77e4f1 96eebff 3581736 4b2ecc5 1aab113 96eebff 1aab113 3581736 96eebff 3581736 96eebff 3581736 96eebff 3581736 d77e4f1 4b2ecc5 d77e4f1 4b2ecc5 d77e4f1 1aab113 d77e4f1 96eebff 3581736 1aab113 96eebff 1aab113 d77e4f1 4b2ecc5 d77e4f1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import os
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 io import BytesIO
from PIL import Image, ImageFilter, ImageEnhance
import rembg
import onnxruntime as ort
import cv2
import numpy as np
app = FastAPI()
# Limita o número de threads do onnxruntime
options = ort.SessionOptions()
options.intra_op_num_threads = 2
def resize_image(image, max_size=1024):
width, height = image.size
if width > max_size or height > max_size:
ratio = min(max_size / width, max_size / height)
new_size = (int(width * ratio), int(height * ratio))
image = image.resize(new_size, Image.Resampling.LANCZOS)
return image
def adjust_brightness_contrast(image, brightness=1.2, contrast=1.2):
enhancer = ImageEnhance.Brightness(image)
image = enhancer.enhance(brightness)
enhancer = ImageEnhance.Contrast(image)
image = enhancer.enhance(contrast)
return image
def remove_shadows(image):
# Converte a imagem para escala de cinza
gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGBA2GRAY)
# Aplica limiarização para remover áreas escuras (sombras)
_, mask = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY)
# Aplica a máscara na imagem original
result = cv2.bitwise_and(np.array(image), np.array(image), mask=mask)
return Image.fromarray(result)
@app.get("/remove-background")
async def remove_background(image_url: str):
try:
# 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))
# Pré-processamento: redimensiona e ajusta brilho/contraste
image = resize_image(image, max_size=1024)
image = adjust_brightness_contrast(image)
# Remove o fundo da imagem usando rembg
output = rembg.remove(image, model="u2net_human_seg", session_options=options)
# Remove sombras
output = remove_shadows(output)
# Suaviza a imagem
output = output.filter(ImageFilter.SMOOTH_MORE)
# Converte a imagem de volta para bytes
img_byte_arr = BytesIO()
output.save(img_byte_arr, format='PNG')
img_byte_arr.seek(0)
# Retorna a imagem processada diretamente no navegador
return StreamingResponse(img_byte_arr, media_type="image/png")
except Exception as e:
raise HTTPException(status_code=400, detail=str(e)) |