|
import os |
|
from glob import glob |
|
import cv2 |
|
import numpy as np |
|
from PIL import Image |
|
import torch |
|
from torchvision import transforms |
|
import gradio as gr |
|
|
|
from models.baseline import BiRefNet |
|
from config import Config |
|
|
|
|
|
config = Config() |
|
device = config.device |
|
|
|
|
|
class ImagePreprocessor(): |
|
def __init__(self, resolution=(1024, 1024)) -> None: |
|
self.transform_image = transforms.Compose([ |
|
transforms.Resize(resolution), |
|
transforms.ToTensor(), |
|
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), |
|
]) |
|
|
|
def proc(self, image): |
|
image = self.transform_image(image) |
|
return image |
|
|
|
|
|
model = BiRefNet().to(device) |
|
state_dict = './birefnet_dis.pth' |
|
if os.path.exists(state_dict): |
|
birefnet_dict = torch.load(state_dict, map_location=device) |
|
unwanted_prefix = '_orig_mod.' |
|
for k, v in list(birefnet_dict.items()): |
|
if k.startswith(unwanted_prefix): |
|
birefnet_dict[k[len(unwanted_prefix):]] = birefnet_dict.pop(k) |
|
model.load_state_dict(birefnet_dict) |
|
model.eval() |
|
|
|
|
|
|
|
|
|
def predict(image, resolution='1024x1024'): |
|
images = [image] |
|
image_shapes = [image.shape[:2] for image in images] |
|
images = [Image.fromarray(image) for image in images] |
|
|
|
resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')] |
|
image_preprocessor = ImagePreprocessor(resolution=resolution) |
|
images_proc = [] |
|
for image in images: |
|
images_proc.append(image_preprocessor.proc(image)) |
|
images_proc = torch.cat([image_proc.unsqueeze(0) for image_proc in images_proc]) |
|
|
|
with torch.no_grad(): |
|
scaled_preds_tensor = model(images_proc.to(device))[-1].sigmoid() |
|
preds = [] |
|
for image_shape, pred_tensor in zip(image_shapes, scaled_preds_tensor): |
|
if device == 'cuda': |
|
pred_tensor = pred_tensor.cpu() |
|
preds.append(torch.nn.functional.interpolate(pred_tensor.unsqueeze(0), size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy()) |
|
image_preds = [] |
|
for image, pred in zip(images, preds): |
|
image_preds.append( |
|
cv2.cvtColor((pred*255).astype(np.uint8), cv2.COLOR_GRAY2RGB) |
|
) |
|
return image_preds[:] if len(images) > 1 else image_preds[0] |
|
|
|
|
|
examples = [[_] for _ in glob('materials/examples/*')][:] |
|
|
|
N = 1 |
|
ipt = [gr.Image() for _ in range(N)] |
|
opt = [gr.Image() for _ in range(N)] |
|
ipt += [gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `512x512`. Higher resolutions can be much slower for inference.", label="Resolution")] |
|
demo = gr.Interface( |
|
fn=predict, |
|
inputs=ipt, |
|
outputs=opt, |
|
examples=examples, |
|
title='Online demo for `Bilateral Reference for High-Resolution Dichotomous Image Segmentation`', |
|
description=('Upload a picture, our model will give you the binary maps of the highly accurate segmentation of the salient objects in it. :)' |
|
'\nThe resolution used in our training was `1024x1024`, which is too much burden for the huggingface free spaces like this one (cost ~500s). Please set resolution as more than `768x768` for images with many texture details to obtain good results!') |
|
) |
|
demo.launch(debug=True) |
|
|
