Fixed autocorrect issues
Browse files
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 |
-
#
|
| 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,
|
| 42 |
-
image = normalize(image,
|
| 43 |
return image
|
| 44 |
|
| 45 |
|
| 46 |
-
transform =
|
| 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,
|
| 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,
|
| 60 |
-
net = hypar["model"]
|
| 61 |
|
| 62 |
# convert to half precision
|
| 63 |
-
if
|
| 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
|
| 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
|
| 131 |
inputs_val = inputs_val.type(torch.FloatTensor)
|
| 132 |
else:
|
| 133 |
inputs_val = inputs_val.type(torch.HalfTensor)
|
| 134 |
|
| 135 |
-
|
| 136 |
-
|
| 137 |
-
|
| 138 |
-
ds_val = net(inputs_val_v)[0]
|
| 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 |
-
|
| 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)
|
| 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 = {}
|
| 159 |
|
| 160 |
|
| 161 |
-
hypar["model_path"] =
|
| 162 |
-
hypar["restore_model"] = "isnet.pth"
|
| 163 |
-
|
| 164 |
-
hypar["interm_sup"] = False
|
| 165 |
|
| 166 |
-
|
| 167 |
-
|
| 168 |
-
hypar["model_digit"] = "full"
|
| 169 |
hypar["seed"] = 0
|
| 170 |
|
| 171 |
-
|
| 172 |
-
hypar["cache_size"] = [1024, 1024]
|
| 173 |
|
| 174 |
-
|
| 175 |
-
|
| 176 |
-
hypar["
|
| 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)
|