doevent
commited on
Commit
β’
2e2f07b
1
Parent(s):
6b216cc
Add new file
Browse files
app.py
ADDED
@@ -0,0 +1,154 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import cv2
|
2 |
+
import gradio as gr
|
3 |
+
import os
|
4 |
+
from PIL import Image
|
5 |
+
import numpy as np
|
6 |
+
import torch
|
7 |
+
from torch.autograd import Variable
|
8 |
+
from torchvision import transforms
|
9 |
+
import torch.nn.functional as F
|
10 |
+
import gdown
|
11 |
+
import matplotlib.pyplot as plt
|
12 |
+
import warnings
|
13 |
+
warnings.filterwarnings("ignore")
|
14 |
+
|
15 |
+
os.system("git clone https://github.com/xuebinqin/DIS")
|
16 |
+
os.system("mv DIS/IS-Net/* .")
|
17 |
+
|
18 |
+
# project imports
|
19 |
+
from data_loader_cache import normalize, im_reader, im_preprocess
|
20 |
+
from models import *
|
21 |
+
|
22 |
+
#Helpers
|
23 |
+
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
24 |
+
|
25 |
+
# Download official weights
|
26 |
+
if not os.path.exists("saved_models"):
|
27 |
+
os.mkdir("saved_models")
|
28 |
+
MODEL_PATH_URL = "https://drive.google.com/uc?id=1KyMpRjewZdyYfxHPYcd-ZbanIXtin0Sn"
|
29 |
+
gdown.download(MODEL_PATH_URL, "saved_models/isnet.pth", use_cookies=False)
|
30 |
+
|
31 |
+
class GOSNormalize(object):
|
32 |
+
'''
|
33 |
+
Normalize the Image using torch.transforms
|
34 |
+
'''
|
35 |
+
def __init__(self, mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]):
|
36 |
+
self.mean = mean
|
37 |
+
self.std = std
|
38 |
+
|
39 |
+
def __call__(self,image):
|
40 |
+
image = normalize(image,self.mean,self.std)
|
41 |
+
return image
|
42 |
+
|
43 |
+
|
44 |
+
transform = transforms.Compose([GOSNormalize([0.5,0.5,0.5],[1.0,1.0,1.0])])
|
45 |
+
|
46 |
+
def load_image(im_path, hypar):
|
47 |
+
im = im_reader(im_path)
|
48 |
+
im, im_shp = im_preprocess(im, hypar["cache_size"])
|
49 |
+
im = torch.divide(im,255.0)
|
50 |
+
shape = torch.from_numpy(np.array(im_shp))
|
51 |
+
return transform(im).unsqueeze(0), shape.unsqueeze(0) # make a batch of image, shape
|
52 |
+
|
53 |
+
|
54 |
+
def build_model(hypar,device):
|
55 |
+
net = hypar["model"]#GOSNETINC(3,1)
|
56 |
+
|
57 |
+
# convert to half precision
|
58 |
+
if(hypar["model_digit"]=="half"):
|
59 |
+
net.half()
|
60 |
+
for layer in net.modules():
|
61 |
+
if isinstance(layer, nn.BatchNorm2d):
|
62 |
+
layer.float()
|
63 |
+
|
64 |
+
net.to(device)
|
65 |
+
|
66 |
+
if(hypar["restore_model"]!=""):
|
67 |
+
net.load_state_dict(torch.load(hypar["model_path"]+"/"+hypar["restore_model"], map_location=device))
|
68 |
+
net.to(device)
|
69 |
+
net.eval()
|
70 |
+
return net
|
71 |
+
|
72 |
+
|
73 |
+
def predict(net, inputs_val, shapes_val, hypar, device):
|
74 |
+
'''
|
75 |
+
Given an Image, predict the mask
|
76 |
+
'''
|
77 |
+
net.eval()
|
78 |
+
|
79 |
+
if(hypar["model_digit"]=="full"):
|
80 |
+
inputs_val = inputs_val.type(torch.FloatTensor)
|
81 |
+
else:
|
82 |
+
inputs_val = inputs_val.type(torch.HalfTensor)
|
83 |
+
|
84 |
+
|
85 |
+
inputs_val_v = Variable(inputs_val, requires_grad=False).to(device) # wrap inputs in Variable
|
86 |
+
|
87 |
+
ds_val = net(inputs_val_v)[0] # list of 6 results
|
88 |
+
|
89 |
+
pred_val = ds_val[0][0,:,:,:] # B x 1 x H x W # we want the first one which is the most accurate prediction
|
90 |
+
|
91 |
+
## recover the prediction spatial size to the orignal image size
|
92 |
+
pred_val = torch.squeeze(F.upsample(torch.unsqueeze(pred_val,0),(shapes_val[0][0],shapes_val[0][1]),mode='bilinear'))
|
93 |
+
|
94 |
+
ma = torch.max(pred_val)
|
95 |
+
mi = torch.min(pred_val)
|
96 |
+
pred_val = (pred_val-mi)/(ma-mi) # max = 1
|
97 |
+
|
98 |
+
if device == 'cuda': torch.cuda.empty_cache()
|
99 |
+
return (pred_val.detach().cpu().numpy()*255).astype(np.uint8) # it is the mask we need
|
100 |
+
|
101 |
+
# Set Parameters
|
102 |
+
hypar = {} # paramters for inferencing
|
103 |
+
|
104 |
+
|
105 |
+
hypar["model_path"] ="./saved_models" ## load trained weights from this path
|
106 |
+
hypar["restore_model"] = "isnet.pth" ## name of the to-be-loaded weights
|
107 |
+
hypar["interm_sup"] = False ## indicate if activate intermediate feature supervision
|
108 |
+
|
109 |
+
## choose floating point accuracy --
|
110 |
+
hypar["model_digit"] = "full" ## indicates "half" or "full" accuracy of float number
|
111 |
+
hypar["seed"] = 0
|
112 |
+
|
113 |
+
hypar["cache_size"] = [1024, 1024] ## cached input spatial resolution, can be configured into different size
|
114 |
+
|
115 |
+
## data augmentation parameters ---
|
116 |
+
hypar["input_size"] = [1024, 1024] ## mdoel input spatial size, usually use the same value hypar["cache_size"], which means we don't further resize the images
|
117 |
+
hypar["crop_size"] = [1024, 1024] ## random crop size from the input, it is usually set as smaller than hypar["cache_size"], e.g., [920,920] for data augmentation
|
118 |
+
|
119 |
+
hypar["model"] = ISNetDIS()
|
120 |
+
|
121 |
+
# Build Model
|
122 |
+
net = build_model(hypar, device)
|
123 |
+
|
124 |
+
|
125 |
+
def inference(image: Image):
|
126 |
+
image_path = image
|
127 |
+
|
128 |
+
image_tensor, orig_size = load_image(image_path, hypar)
|
129 |
+
mask = predict(net, image_tensor, orig_size, hypar, device)
|
130 |
+
|
131 |
+
pil_mask = Image.fromarray(mask).convert('L')
|
132 |
+
im_rgb = Image.open(image).convert("RGB")
|
133 |
+
|
134 |
+
im_rgba = im_rgb.copy()
|
135 |
+
im_rgba.putalpha(pil_mask)
|
136 |
+
|
137 |
+
return [im_rgba, pil_mask]
|
138 |
+
|
139 |
+
|
140 |
+
title = "Highly Accurate Dichotomous Image Segmentation"
|
141 |
+
description = "This is an unofficial demo for DIS, a model that can remove the background from a given image. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below.<br>GitHub: https://github.com/xuebinqin/DIS<br>[![](https://img.shields.io/twitter/follow/DoEvent?label=@DoEvent&style=social)](https://twitter.com/DoEvent)"
|
142 |
+
article = "<div><center><img src='https://visitor-badge.glitch.me/badge?page_id=max_skobeev_dis_public' alt='visitor badge'></center></div>"
|
143 |
+
|
144 |
+
interface = gr.Interface(
|
145 |
+
fn=inference,
|
146 |
+
inputs=gr.Image(type='filepath'),
|
147 |
+
outputs=["image", "image"],
|
148 |
+
examples=[['robot.png'], ['ship.png']],
|
149 |
+
title=title,
|
150 |
+
description=description,
|
151 |
+
article=article,
|
152 |
+
allow_flagging='never',
|
153 |
+
theme="default",
|
154 |
+
).launch(enable_queue=True, debug=True)
|
requirements.txt
ADDED
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
torch
|
2 |
+
torchvision
|
3 |
+
requests
|
4 |
+
gdown
|
5 |
+
matplotlib
|
6 |
+
opencv-python
|
7 |
+
Pillow==8.0.0
|
8 |
+
scikit-image==0.15.0
|
robot.png
ADDED
![]() |
ship.png
ADDED
![]() |