File size: 2,094 Bytes
f2b8e83 be6a051 fb6147c be6a051 fb6147c be6a051 fe98e45 be6a051 f2b8e83 be6a051 |
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 |
from typing import Dict, Any
import os
import requests
from io import BytesIO
from PIL import Image
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
# Configuración
torch.set_float32_matmul_precision(["high", "highest"][0])
device = "cuda" if torch.cuda.is_available() else "cpu"
class EndpointHandler():
def __init__(self, path=''):
# Cargamos el modelo BiRefNet original (Efectivo y rápido)
self.model = AutoModelForImageSegmentation.from_pretrained(
'zhengpeng7/BiRefNet',
trust_remote_code=True
)
self.model.to(device)
self.model.eval()
self.model.half()
def __call__(self, data: Dict[str, Any]):
# 1. RECIBIR IMAGEN (Entrada Blindada)
image_src = data["inputs"]
image = None
if isinstance(image_src, Image.Image):
image = image_src
elif isinstance(image_src, str):
if image_src.startswith('http'):
image = Image.open(BytesIO(requests.get(image_src).content))
else:
image = Image.open(image_src)
else:
image = Image.open(BytesIO(image_src))
# 2. LIMPIEZA: Aseguramos RGB (Color Real)
image = image.convert("RGB")
orig_size = image.size
# 3. PROCESAMIENTO IA
transform = transforms.Compose([
transforms.Resize((1024, 1024)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
input_tensor = transform(image).unsqueeze(0).to(device).half()
with torch.no_grad():
preds = self.model(input_tensor)[-1].sigmoid().cpu()
# 4. MÁSCARA
pred = preds[0].squeeze()
mask_pil = transforms.ToPILImage()(pred)
mask_pil = mask_pil.resize(orig_size, resample=Image.Resampling.LANCZOS)
# 5. APLICACIÓN FINAL (Sin tocar colores)
image.putalpha(mask_pil)
return image |