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('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)