itsyogesh commited on
Commit
4857a1c
·
verified ·
1 Parent(s): 2a3b300

Fixed autocorrect issues

Browse files
Files changed (1) hide show
  1. app.py +41 -57
app.py CHANGED
@@ -1,5 +1,3 @@
1
- from data_loader_cache import normalize, im_reader, im_preprocess
2
- from models import *
3
  import cv2
4
  from skimage.restoration import denoise_nl_means, estimate_sigma
5
  import gradio as gr
@@ -19,48 +17,45 @@ os.system("git clone https://github.com/xuebinqin/DIS")
19
  os.system("mv DIS/IS-Net/* .")
20
 
21
  # project imports
 
 
22
 
23
- # Helpers
24
  device = 'cuda' if torch.cuda.is_available() else 'cpu'
25
 
26
  # Download official weights
27
  if not os.path.exists("saved_models"):
28
  os.mkdir("saved_models")
29
  os.system("mv isnet.pth saved_models/")
30
-
31
-
32
  class GOSNormalize(object):
33
  '''
34
  Normalize the Image using torch.transforms
35
  '''
36
-
37
- def __init__(self, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
38
  self.mean = mean
39
  self.std = std
40
 
41
- def __call__(self, image):
42
- image = normalize(image, self.mean, self.std)
43
  return image
44
 
45
 
46
- transform = transforms.Compose(
47
- [GOSNormalize([0.5, 0.5, 0.5], [1.0, 1.0, 1.0])])
48
-
49
 
50
  def load_image(im_path, hypar):
51
  im = im_reader(im_path)
52
  im, im_shp = im_preprocess(im, hypar["cache_size"])
53
- im = torch.divide(im, 255.0)
54
  shape = torch.from_numpy(np.array(im_shp))
55
- # make a batch of image, shape
56
- return transform(im).unsqueeze(0), shape.unsqueeze(0)
57
 
58
 
59
- def build_model(hypar, device):
60
- net = hypar["model"] # GOSNETINC(3,1)
61
 
62
  # convert to half precision
63
- if (hypar["model_digit"] == "half"):
64
  net.half()
65
  for layer in net.modules():
66
  if isinstance(layer, nn.BatchNorm2d):
@@ -68,11 +63,10 @@ def build_model(hypar, device):
68
 
69
  net.to(device)
70
 
71
- if (hypar["restore_model"] != ""):
72
- net.load_state_dict(torch.load(
73
- hypar["model_path"]+"/"+hypar["restore_model"], map_location=device))
74
  net.to(device)
75
- net.eval()
76
  return net
77
 
78
 
@@ -121,65 +115,56 @@ def smooth_and_denoise(mask):
121
  return denoised_mask
122
 
123
 
 
124
  def predict(net, inputs_val, shapes_val, hypar, device):
125
  '''
126
  Given an Image, predict the mask
127
  '''
128
  net.eval()
129
 
130
- if (hypar["model_digit"] == "full"):
131
  inputs_val = inputs_val.type(torch.FloatTensor)
132
  else:
133
  inputs_val = inputs_val.type(torch.HalfTensor)
134
 
135
- inputs_val_v = Variable(inputs_val, requires_grad=False).to(
136
- device) # wrap inputs in Variable
137
-
138
- ds_val = net(inputs_val_v)[0] # list of 6 results
139
 
140
- # B x 1 x H x W # we want the first one which is the most accurate prediction
141
- pred_val = ds_val[0][0, :, :, :]
142
 
143
- # recover the prediction spatial size to the orignal image size
144
- pred_val = torch.squeeze(F.upsample(torch.unsqueeze(
145
- pred_val, 0), (shapes_val[0][0], shapes_val[0][1]), mode='bilinear'))
146
 
147
  ma = torch.max(pred_val)
148
  mi = torch.min(pred_val)
149
- pred_val = (pred_val-mi)/(ma-mi) # max = 1
150
-
151
- if device == 'cuda':
152
- torch.cuda.empty_cache()
153
- # it is the mask we need
154
- return (pred_val.detach().cpu().numpy()*255).astype(np.uint8)
155
-
156
 
 
 
 
157
  # Set Parameters
158
- hypar = {} # paramters for inferencing
159
 
160
 
161
- hypar["model_path"] = "./saved_models" # load trained weights from this path
162
- hypar["restore_model"] = "isnet.pth" # name of the to-be-loaded weights
163
- # indicate if activate intermediate feature supervision
164
- hypar["interm_sup"] = False
165
 
166
- # choose floating point accuracy --
167
- # indicates "half" or "full" accuracy of float number
168
- hypar["model_digit"] = "full"
169
  hypar["seed"] = 0
170
 
171
- # cached input spatial resolution, can be configured into different size
172
- hypar["cache_size"] = [1024, 1024]
173
 
174
- # data augmentation parameters ---
175
- # mdoel input spatial size, usually use the same value hypar["cache_size"], which means we don't further resize the images
176
- hypar["input_size"] = [1024, 1024]
177
- # random crop size from the input, it is usually set as smaller than hypar["cache_size"], e.g., [920,920] for data augmentation
178
- hypar["crop_size"] = [1024, 1024]
179
 
180
  hypar["model"] = ISNetDIS()
181
 
182
- # Build Model
183
  net = build_model(hypar, device)
184
 
185
 
@@ -208,7 +193,6 @@ def inference(image):
208
 
209
  return [im_rgba, pil_mask, im_dark]
210
 
211
-
212
  title = "Mysign.id - Signature Background removal based on DIS"
213
  description = "ML Model based on ECCV2022/dis-background-removal specifically made for removing background from signatures."
214
 
@@ -221,4 +205,4 @@ interface = gr.Interface(
221
  description=description,
222
  allow_flagging='never',
223
  cache_examples=False,
224
- ).queue(api_open=True).launch(show_api=True, show_error=True)
 
 
 
1
  import cv2
2
  from skimage.restoration import denoise_nl_means, estimate_sigma
3
  import gradio as gr
 
17
  os.system("mv DIS/IS-Net/* .")
18
 
19
  # project imports
20
+ from data_loader_cache import normalize, im_reader, im_preprocess
21
+ from models import *
22
 
23
+ #Helpers
24
  device = 'cuda' if torch.cuda.is_available() else 'cpu'
25
 
26
  # Download official weights
27
  if not os.path.exists("saved_models"):
28
  os.mkdir("saved_models")
29
  os.system("mv isnet.pth saved_models/")
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):
 
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
 
 
115
  return denoised_mask
116
 
117
 
118
+
119
  def predict(net, inputs_val, shapes_val, hypar, device):
120
  '''
121
  Given an Image, predict the mask
122
  '''
123
  net.eval()
124
 
125
+ if(hypar["model_digit"]=="full"):
126
  inputs_val = inputs_val.type(torch.FloatTensor)
127
  else:
128
  inputs_val = inputs_val.type(torch.HalfTensor)
129
 
130
+
131
+ inputs_val_v = Variable(inputs_val, requires_grad=False).to(device) # wrap inputs in Variable
132
+
133
+ ds_val = net(inputs_val_v)[0] # list of 6 results
134
 
135
+ pred_val = ds_val[0][0,:,:,:] # B x 1 x H x W # we want the first one which is the most accurate prediction
 
136
 
137
+ ## recover the prediction spatial size to the orignal image size
138
+ pred_val = torch.squeeze(F.upsample(torch.unsqueeze(pred_val,0),(shapes_val[0][0],shapes_val[0][1]),mode='bilinear'))
 
139
 
140
  ma = torch.max(pred_val)
141
  mi = torch.min(pred_val)
142
+ pred_val = (pred_val-mi)/(ma-mi) # max = 1
 
 
 
 
 
 
143
 
144
+ if device == 'cuda': torch.cuda.empty_cache()
145
+ return (pred_val.detach().cpu().numpy()*255).astype(np.uint8) # it is the mask we need
146
+
147
  # Set Parameters
148
+ hypar = {} # paramters for inferencing
149
 
150
 
151
+ hypar["model_path"] ="./saved_models" ## load trained weights from this path
152
+ hypar["restore_model"] = "isnet.pth" ## name of the to-be-loaded weights
153
+ hypar["interm_sup"] = False ## indicate if activate intermediate feature supervision
 
154
 
155
+ ## choose floating point accuracy --
156
+ hypar["model_digit"] = "full" ## indicates "half" or "full" accuracy of float number
 
157
  hypar["seed"] = 0
158
 
159
+ hypar["cache_size"] = [1024, 1024] ## cached input spatial resolution, can be configured into different size
 
160
 
161
+ ## data augmentation parameters ---
162
+ 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
163
+ 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
 
 
164
 
165
  hypar["model"] = ISNetDIS()
166
 
167
+ # Build Model
168
  net = build_model(hypar, device)
169
 
170
 
 
193
 
194
  return [im_rgba, pil_mask, im_dark]
195
 
 
196
  title = "Mysign.id - Signature Background removal based on DIS"
197
  description = "ML Model based on ECCV2022/dis-background-removal specifically made for removing background from signatures."
198
 
 
205
  description=description,
206
  allow_flagging='never',
207
  cache_examples=False,
208
+ ).queue(api_open=True).launch(show_api=True, show_error=True)