|
|
import cv2 |
|
|
import gradio as gr |
|
|
import os |
|
|
from PIL import Image |
|
|
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") |
|
|
|
|
|
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 |
|
|
''' |
|
|
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])]) |
|
|
|
|
|
def load_image(im_path, hypar): |
|
|
im = im_reader(im_path) |
|
|
im, im_shp = im_preprocess(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): |
|
|
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(image).convert("RGB") |
|
|
im_rgba = im_rgb.copy() |
|
|
im_rgba.putalpha(pil_mask) |
|
|
return [im_rgba, pil_mask] |
|
|
|
|
|
|
|
|
interface = gr.Interface( |
|
|
fn=inference, |
|
|
inputs=gr.Image(type='filepath'), |
|
|
outputs=[gr.Image(type='filepath', format="png"), gr.Image(type='filepath', format="png", visible=False)], |
|
|
flagging_mode="never", |
|
|
cache_mode="lazy", |
|
|
).queue(api_open=False).launch(show_error=False, show_api=False, share=False) |
|
|
|
