|
|
import gradio as gr |
|
|
from gradio_imageslider import ImageSlider |
|
|
from loadimg import load_img |
|
|
import spaces |
|
|
from transformers import AutoModelForImageSegmentation |
|
|
import torch |
|
|
from torchvision import transforms |
|
|
|
|
|
torch.set_float32_matmul_precision(["high", "highest"][0]) |
|
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
|
birefnet = AutoModelForImageSegmentation.from_pretrained( |
|
|
"ZhengPeng7/BiRefNet", trust_remote_code=True |
|
|
) |
|
|
birefnet.to(device) |
|
|
transform_image = transforms.Compose( |
|
|
[ |
|
|
transforms.Resize((1024, 1024)), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), |
|
|
] |
|
|
) |
|
|
|
|
|
|
|
|
@spaces.GPU |
|
|
def fn(image): |
|
|
im = load_img(image, output_type="pil") |
|
|
im = im.convert("RGB") |
|
|
image_size = im.size |
|
|
origin = im.copy() |
|
|
image = load_img(im) |
|
|
input_images = transform_image(image).unsqueeze(0).to(device) |
|
|
|
|
|
with torch.no_grad(): |
|
|
preds = birefnet(input_images)[-1].sigmoid().cpu() |
|
|
pred = preds[0].squeeze() |
|
|
pred_pil = transforms.ToPILImage()(pred) |
|
|
mask = pred_pil.resize(image_size) |
|
|
image.putalpha(mask) |
|
|
return image |
|
|
|
|
|
|
|
|
chameleon = load_img("chameleon.jpg", output_type="pil") |
|
|
|
|
|
url = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg" |
|
|
demo = gr.Interface( |
|
|
fn, |
|
|
inputs=gr.Image(label="Upload an image"), |
|
|
outputs=gr.Image(label="birefnet", format="png"), |
|
|
examples=[chameleon], |
|
|
api_name="image", |
|
|
flagging_mode="never", |
|
|
cache_mode="lazy", |
|
|
) |
|
|
|
|
|
demo.queue(default_concurrency_limit=1).launch(show_error=True) |
|
|
|
