doevent commited on
Commit
2e2f07b
β€’
1 Parent(s): 6b216cc

Add new file

Browse files
Files changed (4) hide show
  1. app.py +154 -0
  2. requirements.txt +8 -0
  3. robot.png +0 -0
  4. ship.png +0 -0
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