briaai

/

RMBG-1.4

like 1.63k

Follow

BRIA AI 431

Yes, the results really don't look very good, if you have relevant data, you can train the model that will give better results for your case.
Our model was not trained on this type of data

Since your images are very simple, you can remove the background without deep learning.

Here is an example using the remove_background_bicolor function from the PyMatting library.

from pymatting import *
import numpy as np
import urllib.request

url = "https://cdn-uploads.huggingface.co/production/uploads/666bd373726e7ee7113c689a/2QsaGjluZS3BG4swfP2vW.png"

# Assuming that foreground color is white
# You may have to adjust this for other images
fg_color = np.array([1.0, 1.0, 1.0])

image_path = url.split("/")[-1]

# Download image
urllib.request.urlretrieve(url, image_path)

image = load_image(image_path, "RGB")

h, w = image.shape[:2]

# Assuming that background color is median color of top 5% rows of image
# You may have to adjust this for other images
bg_color = np.median(image[:round(0.05 * h)], axis=(0, 1))

# Remove background
image = remove_background_bicolor(image, fg_color, bg_color)

# Assume that foreground color is white everywhere
image[:, :, :3] = fg_color

# Bost alpha a bit to fix errors due to low quality input image
image[:, :, 3] *= 1.1

print("Saving cutout.png")
save_image("cutout.png", image)