Update u2net_utils.py

u2net_utils.py

@@ -45,19 +45,26 @@ def postprocess(mask, original_size):
     mask = Image.fromarray((mask * 255).astype(np.uint8)).resize(original_size, Image.BILINEAR)
     return mask
 
-def remove_background(image):
-    input_tensor = preprocess(image)
+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)
     with torch.no_grad():
-        d1, *_ = model(input_tensor)
-    mask = postprocess(d1, image.size)
+        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)
 
+    # Удаляем фон, сохраняя оригинальное изображение
     image = image.convert("RGBA")
-    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))
+    mask = mask.convert("L")
+    new_image = Image.new("RGBA", original_size)
+    new_image.paste(image, (0, 0), mask=mask)
 
-    image.putdata(new_data)
-    return image
+    return new_image