Spaces:

---

chuuhtetnaing

/

background-remover

like 10

Running

Hi raphael, I'm not the original author and my solution may not be optimal. However, you can do it in this way.

images = [im1, im2]

image_sizes = [image.size for image in images]
input_images = [transform_image(img) for img in images]
input_batch = torch.stack(input_images).to("cuda")

# Prediction
with torch.no_grad():
    preds = birefnet(input_batch)[-1].sigmoid().cpu()

preds = preds.squeeze()

pred_pils = [transforms.ToPILImage()(pred) for pred in preds]

masks = [pred_pil.resize(image_size) for pred_pil, image_size in zip(pred_pils, image_sizes)]
white_backgrounds = [Image.new("RGBA", image_size, (255, 255, 255, 255)) for image_size in image_sizes]

for image, mask in zip(images, masks):
    image.putalpha(mask)

combineds = [Image.alpha_composite(white_background, image) for white_background, image in zip(white_backgrounds, images)]

If you need complete code, can check here. I tested it on kaggle.