File size: 3,948 Bytes
81b1a0e 327742a 81b1a0e f406b44 81b1a0e 327742a 81b1a0e 7dd89ba 81b1a0e f406b44 327742a 81b1a0e 327742a 81b1a0e f406b44 81b1a0e 6be800b 81b1a0e de5ed42 f406b44 de5ed42 81b1a0e 4cf55f4 81b1a0e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
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
def array_to_pil_image(image, size=(1024, 1024)):
image = cv2.resize(image, size, interpolation=cv2.INTER_LINEAR)
image = Image.fromarray(image).convert('RGB')
return image
class ImagePreprocessor():
def __init__(self, resolution=(1024, 1024)) -> None:
self.transform_image = transforms.Compose([
# transforms.Resize(resolution), # 1. keep consistent with the cv2.resize used in training 2. redundant with that in path_to_image()
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_ep580.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_1, image_2):
# images = [image_1, image_2]
def predict(image, resolution='1024x1024'):
# Image is a RGB numpy array.
resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
images = [image]
image_shapes = [image.shape[:2] for image in images]
images = [array_to_pil_image(image, resolution) for image in images]
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() # BiRefNet needs an sigmoid activation outside the forward.
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)]
# Add the option of resolution in a text box.
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")]
for idx_example, example in enumerate(examples):
examples[idx_example].append('1024x1024')
examples.append(examples[-1].copy())
examples[-1][1] = '512x512'
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!\n Ours codes can be found at https://github.com/ZhengPeng7/BiRefNet.')
)
demo.launch(debug=True)
|