import numpy as np import torch import torch.nn.functional as F import gradio as gr from ormbg import ORMBG from PIL import Image model_path = "ormbg.pth" net = ORMBG() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") net.to(device) if torch.cuda.is_available(): net.load_state_dict(torch.load(model_path)) net = net.cuda() else: net.load_state_dict(torch.load(model_path, map_location="cpu")) net.eval() def resize_image(image): image = image.convert("RGB") model_input_size = (1024, 1024) image = image.resize(model_input_size, Image.BILINEAR) return image def inference(image): # prepare input orig_image = Image.fromarray(image) w, h = orig_image.size image = resize_image(orig_image) im_np = np.array(image) im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1) im_tensor = torch.unsqueeze(im_tensor, 0) im_tensor = torch.divide(im_tensor, 255.0) if torch.cuda.is_available(): im_tensor = im_tensor.cuda() # inference result = net(im_tensor) # post process result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode="bilinear"), 0) ma = torch.max(result) mi = torch.min(result) result = (result - mi) / (ma - mi) # image to pil im_array = (result * 255).cpu().data.numpy().astype(np.uint8) pil_im = Image.fromarray(np.squeeze(im_array)) # paste the mask on the original image new_im = Image.new("RGBA", pil_im.size, (0, 0, 0, 0)) new_im.paste(orig_image, mask=pil_im) return new_im gr.Markdown("## Open Remove Background Model (ormbg)") gr.HTML( """

This is a demo for Open Remove Background Model (ormbg) that using Open Remove Background Model (ormbg) model as backbone.

""" ) title = "Transparent Background (Fast)" examples = ["./example1.png", "./example2.png", "./example3.png"] demo = gr.Interface( fn=inference, inputs="image", outputs="image", examples=examples, title=title, theme="Nymbo/Nymbo_Theme" ) if __name__ == "__main__": demo.launch(share=False)