|
import numpy as np |
|
import torch |
|
import torch.nn.functional as F |
|
from torchvision.transforms.functional import normalize |
|
from briarmbg import BriaRMBG |
|
from huggingface_hub import hf_hub_download |
|
import gradio as gr |
|
from PIL import Image |
|
|
|
def resize_image(image): |
|
image = image.convert('RGB') |
|
model_input_size = (1024, 1024) |
|
image = image.resize(model_input_size, Image.BILINEAR) |
|
return image |
|
|
|
def process(image): |
|
|
|
if isinstance(image, np.ndarray): |
|
orig_image = Image.fromarray(image) |
|
else: |
|
|
|
orig_image = image |
|
|
|
w, h = orig_im_size = orig_image.size |
|
image = resize_image(orig_image) |
|
im_np = np.array(image) |
|
im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1) |
|
im_tensor = torch.unsqueeze(im_tensor, 0) |
|
im_tensor = torch.divide(im_tensor, 255.0) |
|
im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0]) |
|
if torch.cuda.is_available(): |
|
im_tensor = im_tensor.cuda() |
|
|
|
|
|
net = BriaRMBG() |
|
model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth') |
|
if torch.cuda.is_available(): |
|
net.load_state_dict(torch.load(model_path)) |
|
net = net.cuda() |
|
else: |
|
net.load_state_dict(torch.load(model_path, map_location="cpu")) |
|
net.eval() |
|
|
|
result = net(im_tensor) |
|
|
|
result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0) |
|
ma = torch.max(result) |
|
mi = torch.min(result) |
|
result = (result - mi) / (ma - mi) |
|
|
|
im_array = (result * 255).cpu().data.numpy().astype(np.uint8) |
|
pil_im = Image.fromarray(np.squeeze(im_array)) |
|
|
|
new_im = Image.new("RGBA", pil_im.size, (0, 0, 0, 0)) |
|
new_im.paste(orig_image, mask=pil_im) |
|
|
|
return new_im |
|
|
|
|
|
css = """ |
|
footer { |
|
visibility: hidden; |
|
} |
|
""" |
|
|
|
with gr.Blocks(css=css) as demo: |
|
with gr.Column(): |
|
input_image = gr.Image(type="pil") |
|
output_image = gr.Image() |
|
process_button = gr.Button("Remove Background Image") |
|
process_button.click(fn=process, inputs=input_image, outputs=output_image) |
|
|
|
demo.launch() |
|
|