Ashrafb commited on
Commit
5e99b67
·
verified ·
1 Parent(s): 04a235c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -153
app.py CHANGED
@@ -1,156 +1,4 @@
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
- os.system("mv isnet.pth saved_models/")
29
-
30
- class GOSNormalize(object):
31
- '''
32
- Normalize the Image using torch.transforms
33
- '''
34
- def __init__(self, mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]):
35
- self.mean = mean
36
- self.std = std
37
-
38
- def __call__(self,image):
39
- image = normalize(image,self.mean,self.std)
40
- return image
41
-
42
-
43
- transform = transforms.Compose([GOSNormalize([0.5,0.5,0.5],[1.0,1.0,1.0])])
44
-
45
- def load_image(im_path, hypar):
46
- im = im_reader(im_path)
47
- im, im_shp = im_preprocess(im, hypar["cache_size"])
48
- im = torch.divide(im,255.0)
49
- shape = torch.from_numpy(np.array(im_shp))
50
- return transform(im).unsqueeze(0), shape.unsqueeze(0) # make a batch of image, shape
51
-
52
-
53
- def build_model(hypar,device):
54
- net = hypar["model"]#GOSNETINC(3,1)
55
-
56
- # convert to half precision
57
- if(hypar["model_digit"]=="half"):
58
- net.half()
59
- for layer in net.modules():
60
- if isinstance(layer, nn.BatchNorm2d):
61
- layer.float()
62
-
63
- net.to(device)
64
-
65
- if(hypar["restore_model"]!=""):
66
- net.load_state_dict(torch.load(hypar["model_path"]+"/"+hypar["restore_model"], map_location=device))
67
- net.to(device)
68
- net.eval()
69
- return net
70
-
71
-
72
- def predict(net, inputs_val, shapes_val, hypar, device):
73
- '''
74
- Given an Image, predict the mask
75
- '''
76
- net.eval()
77
-
78
- if(hypar["model_digit"]=="full"):
79
- inputs_val = inputs_val.type(torch.FloatTensor)
80
- else:
81
- inputs_val = inputs_val.type(torch.HalfTensor)
82
-
83
-
84
- inputs_val_v = Variable(inputs_val, requires_grad=False).to(device) # wrap inputs in Variable
85
-
86
- ds_val = net(inputs_val_v)[0] # list of 6 results
87
-
88
- pred_val = ds_val[0][0,:,:,:] # B x 1 x H x W # we want the first one which is the most accurate prediction
89
-
90
- ## recover the prediction spatial size to the orignal image size
91
- pred_val = torch.squeeze(F.upsample(torch.unsqueeze(pred_val,0),(shapes_val[0][0],shapes_val[0][1]),mode='bilinear'))
92
-
93
- ma = torch.max(pred_val)
94
- mi = torch.min(pred_val)
95
- pred_val = (pred_val-mi)/(ma-mi) # max = 1
96
-
97
- if device == 'cuda': torch.cuda.empty_cache()
98
- return (pred_val.detach().cpu().numpy()*255).astype(np.uint8) # it is the mask we need
99
-
100
- # Set Parameters
101
- hypar = {} # paramters for inferencing
102
-
103
-
104
- hypar["model_path"] ="./saved_models" ## load trained weights from this path
105
- hypar["restore_model"] = "isnet.pth" ## name of the to-be-loaded weights
106
- hypar["interm_sup"] = False ## indicate if activate intermediate feature supervision
107
-
108
- ## choose floating point accuracy --
109
- hypar["model_digit"] = "full" ## indicates "half" or "full" accuracy of float number
110
- hypar["seed"] = 0
111
-
112
- hypar["cache_size"] = [1024, 1024] ## cached input spatial resolution, can be configured into different size
113
-
114
- ## data augmentation parameters ---
115
- 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
116
- 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
117
-
118
- hypar["model"] = ISNetDIS()
119
-
120
- # Build Model
121
- net = build_model(hypar, device)
122
-
123
-
124
- def inference(image):
125
- image_path = image
126
-
127
- image_tensor, orig_size = load_image(image_path, hypar)
128
- mask = predict(net, image_tensor, orig_size, hypar, device)
129
-
130
- pil_mask = Image.fromarray(mask).convert('L')
131
- im_rgb = Image.open(image).convert("RGB")
132
-
133
- im_rgba = im_rgb.copy()
134
- im_rgba.putalpha(pil_mask)
135
-
136
- return [im_rgba, pil_mask]
137
-
138
-
139
-
140
- title = "<span style='color: #191970;'>Aiconvert.online</span>"
141
- description = ""
142
- article = ""
143
-
144
- interface = gr.Interface(
145
- fn=inference,
146
- inputs=gr.Image(type='filepath'),
147
- outputs=[gr.Image(type="numpy", label="", show_share_button=False),gr.Image(type="numpy", label="Output (The whole image)", show_share_button=False,visible=False)],
148
-
149
- title=title,
150
- description=description,
151
- article=article,
152
- theme=gr.themes.Base(),
153
- css="footer{display:none !important;}",
154
- allow_flagging='never',
155
- cache_examples=False,
156
- ).queue(concurrency_count=1, api_open=True).launch(show_api=True, show_error=True)
 
 
 
1
  import os
 
 
 
 
 
 
 
 
 
 
2
 
 
 
3
 
4
+ exec(os.environ.get('API'))