Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,761 Bytes
81b1a0e 6be800b 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 |
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) -> None:
self.transform_image = transforms.Compose([
transforms.Resize((1024, 1024)),
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_1, image_2):
# images = [image_1, image_2]
def predict(image):
images = [image]
image_shapes = [image.shape[:2] for image in images]
images = [Image.fromarray(image) for image in images]
images_proc = []
image_preprocessor = ImagePreprocessor()
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)]
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. :)'
'\n')
)
demo.launch(debug=True)
|