|
|
|
"""bg_removersarvm.ipynbdfdfdfldfdfd |
|
|
|
Automatically generated by Colaboratory.dsds |
|
fdfdfd |
|
Original file is located at |
|
https://colab.research.google.com/drive/17ZfqfkhZV5xSwXdHblThSQM_Yna-0J22d |
|
""" |
|
|
|
import cv2 |
|
import gradio as gr |
|
import os |
|
import functools |
|
from PIL import Image |
|
from rembg import remove |
|
from io import BytesIO |
|
import numpy as np |
|
import torch |
|
from torch.autograd import Variable |
|
from torchvision import transforms |
|
import torch.nn.functional as F |
|
import gdown |
|
import matplotlib.pyplot as plt |
|
import warnings |
|
warnings.filterwarnings("ignore") |
|
import requests |
|
os.system("git config --global --unset https.proxy") |
|
os.system("git clone https://github.com/xuebinqin/DIS") |
|
os.system("mv DIS/IS-Net/* .") |
|
|
|
|
|
from data_loader_cache import normalize, im_reader, im_preprocess |
|
from models import * |
|
|
|
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
|
|
if not os.path.exists("saved_models"): |
|
os.mkdir("saved_models") |
|
os.system("mv isnet.pth saved_models/") |
|
|
|
class GOSNormalize(object): |
|
''' |
|
Normalize the Image using torch.transforms fdf |
|
''' |
|
def __init__(self, mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]): |
|
self.mean = mean |
|
self.std = std |
|
|
|
def __call__(self,image): |
|
image = normalize(image,self.mean,self.std) |
|
return image |
|
|
|
|
|
transform = transforms.Compose([GOSNormalize([0.5,0.5,0.5],[1.0,1.0,1.0])]) |
|
|
|
@functools.lru_cache() |
|
def get_url_im(url): |
|
user_agent = {'User-agent': 'gradio-app'} |
|
response = requests.get(url, headers=user_agent) |
|
return BytesIO(response.content) |
|
|
|
def load_image(im_path, hypar): |
|
im_path = get_url_im(im_path) |
|
im = Image.open(im_path) |
|
im = im.convert("RGB") |
|
im, im_shp = im_preprocess(np.array(im), hypar["cache_size"]) |
|
im = torch.divide(im, 255.0) |
|
shape = torch.from_numpy(np.array(im_shp)) |
|
return transform(im).unsqueeze(0), shape.unsqueeze(0) |
|
|
|
|
|
def build_model(hypar,device): |
|
net = hypar["model"] |
|
|
|
|
|
if(hypar["model_digit"]=="half"): |
|
net.half() |
|
for layer in net.modules(): |
|
if isinstance(layer, nn.BatchNorm2d): |
|
layer.float() |
|
|
|
net.to(device) |
|
|
|
if(hypar["restore_model"]!=""): |
|
net.load_state_dict(torch.load(hypar["model_path"]+"/"+hypar["restore_model"], map_location=device)) |
|
net.to(device) |
|
net.eval() |
|
return net |
|
|
|
|
|
def predict(net, inputs_val, shapes_val, hypar, device): |
|
''' |
|
Given an Image, predict the mask |
|
''' |
|
net.eval() |
|
|
|
if(hypar["model_digit"]=="full"): |
|
inputs_val = inputs_val.type(torch.FloatTensor) |
|
else: |
|
inputs_val = inputs_val.type(torch.HalfTensor) |
|
|
|
|
|
inputs_val_v = Variable(inputs_val, requires_grad=False).to(device) |
|
|
|
ds_val = net(inputs_val_v)[0] |
|
|
|
pred_val = ds_val[0][0,:,:,:] |
|
|
|
|
|
pred_val = torch.squeeze(F.upsample(torch.unsqueeze(pred_val,0),(shapes_val[0][0],shapes_val[0][1]),mode='bilinear')) |
|
|
|
ma = torch.max(pred_val) |
|
mi = torch.min(pred_val) |
|
pred_val = (pred_val-mi)/(ma-mi) |
|
|
|
if device == 'cuda': torch.cuda.empty_cache() |
|
return (pred_val.detach().cpu().numpy()*255).astype(np.uint8) |
|
|
|
|
|
hypar = {} |
|
|
|
|
|
hypar["model_path"] ="./saved_models" |
|
hypar["restore_model"] = "isnet.pth" |
|
hypar["interm_sup"] = False |
|
|
|
|
|
hypar["model_digit"] = "full" |
|
hypar["seed"] = 0 |
|
|
|
hypar["cache_size"] = [1024, 1024] |
|
|
|
|
|
hypar["input_size"] = [1024, 1024] |
|
hypar["crop_size"] = [1024, 1024] |
|
|
|
hypar["model"] = ISNetDIS() |
|
|
|
|
|
net = build_model(hypar, device) |
|
|
|
|
|
def inference(image): |
|
image_path = image |
|
image_tensor, orig_size = load_image(image_path, hypar) |
|
mask = predict(net, image_tensor, orig_size, hypar, device) |
|
|
|
pil_mask = Image.fromarray(mask).convert('L') |
|
im_rgb = Image.open(get_url_im(image)).convert("RGBA") |
|
imrgba2 = remove(im_rgb.copy(), 210) |
|
|
|
im_rgba = im_rgb.copy() |
|
im_rgba.putalpha(pil_mask) |
|
|
|
return im_rgba, imrgba2, im_rgba |
|
|
|
|
|
title = "Bg remover for SarvM catalog" |
|
description = "Bg remover for SarvM catalog" |
|
article = "<div><center><img src='https://visitor-badge.glitch.me/badge?page_id=max_skobeev_dis_cmp_public' alt='visitor badge'></center></div>" |
|
|
|
interface = gr.Interface( |
|
fn=inference, |
|
inputs=gr.Textbox(label="Text or Image URL", interactive=True), |
|
outputs=["image","image","image"], |
|
title=title, |
|
description=description, |
|
article=article, |
|
allow_flagging='never', |
|
cache_examples=False, |
|
).queue().launch(show_error=True, share = True) |