Spaces:
Sleeping
Sleeping
from skimage import io | |
import torch, os | |
from PIL import Image | |
from briarmbg import BriaRMBG | |
from utilities import preprocess_image, postprocess_image | |
from huggingface_hub import hf_hub_download | |
import io as IO | |
import base64 | |
def example_inference(im_path, transprent_bg=False, color=(255, 255, 255, 255)): | |
net = BriaRMBG() | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
net = BriaRMBG.from_pretrained("briaai/RMBG-1.4") | |
net.to(device) | |
net.eval() | |
# prepare input | |
model_input_size = [1024,1024] | |
orig_im = io.imread(im_path, plugin='imageio') | |
orig_im_size = orig_im.shape[0:2] | |
image = preprocess_image(orig_im, model_input_size).to(device) | |
# inference | |
result=net(image) | |
# post process | |
result_image = postprocess_image(result[0][0], orig_im_size) | |
bgColor = (0,0,0, 0) if transprent_bg else color | |
# save result | |
pil_im = Image.fromarray(result_image) | |
no_bg_image = Image.new("RGBA", pil_im.size, bgColor) | |
orig_image = Image.open(IO.BytesIO(im_path)) | |
no_bg_image.paste(orig_image, mask=pil_im) | |
# Convert image to bytes and then to base64 | |
buffered = IO.BytesIO() | |
no_bg_image.save(buffered, format="PNG") | |
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") | |
return img_str |