Update u2net_utils.py

u2net_utils.py

@@ -1,7 +1,6 @@
 import torch
 import torchvision.transforms as transforms
 from PIL import Image
-from io import BytesIO
 import numpy as np
 import os
 import requests
@@ -46,29 +45,19 @@ def postprocess(mask, original_size):
     mask = Image.fromarray((mask * 255).astype(np.uint8)).resize(original_size, Image.BILINEAR)
     return mask
 
-def remove_background(image: Image.Image) -> Image.Image:
-    # Сохраняем оригинальный размер
-    original_size = image.size
-    # Преобразуем изображение в формат, подходящий для модели
-    transform = transforms.Compose([
-        transforms.Resize((320, 320)),  # только для подачи в модель
-        transforms.ToTensor(),
-    ])
-    image_tensor = transform(image).unsqueeze(0)
+def remove_background(image):
+    input_tensor = preprocess(image)
     with torch.no_grad():
-        output = model(image_tensor)[0][0]  # предположим, что это U2NET
-        mask = output.squeeze().cpu().numpy()
-        mask = (mask - mask.min()) / (mask.max() - mask.min())  # нормализация
-        # Маску надо вернуть к оригинальному размеру
-    mask = Image.fromarray((mask * 255).astype('uint8')).resize(original_size, Image.LANCZOS)
+        d1, *_ = model(input_tensor)
+    mask = postprocess(d1, image.size)
 
-    # Удаляем фон, сохраняя оригинальное изображение
     image = image.convert("RGBA")
-    mask = mask.convert("L")
-    new_image = Image.new("RGBA", original_size)
-    new_image.paste(image, (0, 0), mask=mask)
+    datas = image.getdata()
+    masks = mask.getdata()
+
+    new_data = []
+    for item, m in zip(datas, masks):
+        new_data.append((item[0], item[1], item[2], m))
 
-    buffer = BytesIO()
-new_image.save(buffer, format="PNG", optimize=False, compress_level=1)
-buffer.seek(0)
-return buffer
+    image.putdata(new_data)
+    return image