gaviego commited on
Commit
62b27da
1 Parent(s): 25d8185

batch process and file name

Browse files
Files changed (2) hide show
  1. .gitignore +2 -1
  2. app.py +61 -13
.gitignore CHANGED
@@ -160,4 +160,5 @@ cython_debug/
160
  #.idea/
161
 
162
  *.jpeg
163
- *.png
 
 
160
  #.idea/
161
 
162
  *.jpeg
163
+ *.png
164
+ .DS_Store
app.py CHANGED
@@ -2,6 +2,7 @@ import gradio as gr
2
  import cv2
3
  import gradio as gr
4
  import os
 
5
  from PIL import Image
6
  import numpy as np
7
  import torch
@@ -10,6 +11,8 @@ from torchvision import transforms
10
  import torch.nn.functional as F
11
  import matplotlib.pyplot as plt
12
  import warnings
 
 
13
  warnings.filterwarnings("ignore")
14
 
15
 
@@ -115,29 +118,74 @@ hypar["model"] = ISNetDIS()
115
  net = build_model(hypar, device)
116
 
117
 
118
- def inference(image: Image):
119
- image_path = image
120
 
121
  image_tensor, orig_size = load_image(image_path, hypar)
122
  mask = predict(net, image_tensor, orig_size, hypar, device)
123
 
124
  pil_mask = Image.fromarray(mask).convert('L')
125
- im_rgb = Image.open(image).convert("RGB")
126
 
127
  im_rgba = im_rgb.copy()
128
  im_rgba.putalpha(pil_mask)
129
-
130
- return im_rgba
131
-
132
- def bw(image_file:Image):
 
 
 
 
 
 
 
 
 
 
 
 
 
133
  img = Image.open(image_file)
134
  img = img.convert("L")
135
  return img
136
 
137
- iface = gr.Interface(fn=inference,
138
- inputs=gr.Image(type='filepath'),
139
- outputs=["image"],
140
- title="Remove Background",
141
- description="Uses <a href='https://github.com/xuebinqin/DIS'>DIS</a> to remove background"
142
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
143
  iface.launch()
 
2
  import cv2
3
  import gradio as gr
4
  import os
5
+ from pathlib import Path
6
  from PIL import Image
7
  import numpy as np
8
  import torch
 
11
  import torch.nn.functional as F
12
  import matplotlib.pyplot as plt
13
  import warnings
14
+ import tempfile
15
+
16
  warnings.filterwarnings("ignore")
17
 
18
 
 
118
  net = build_model(hypar, device)
119
 
120
 
121
+ def inference(image_path):
 
122
 
123
  image_tensor, orig_size = load_image(image_path, hypar)
124
  mask = predict(net, image_tensor, orig_size, hypar, device)
125
 
126
  pil_mask = Image.fromarray(mask).convert('L')
127
+ im_rgb = Image.open(image_path).convert("RGB")
128
 
129
  im_rgba = im_rgb.copy()
130
  im_rgba.putalpha(pil_mask)
131
+ file_name = Path(image_path).stem+"_nobg.png"
132
+ file_path = Path(Path(image_path).parent,file_name)
133
+ im_rgba.save(file_path)
134
+ return str(file_path.resolve())
135
+
136
+ def bw(image_files):
137
+ print(image_files)
138
+ output = []
139
+ for idx, file in enumerate(image_files):
140
+ print(file.name)
141
+ img = Image.open(file.name)
142
+ img = img.convert("L")
143
+ output.append(img)
144
+ print(output)
145
+ return output
146
+
147
+ def bw_single(image_file):
148
  img = Image.open(image_file)
149
  img = img.convert("L")
150
  return img
151
 
152
+ def batch(image_files):
153
+ output = []
154
+ for idx, file in enumerate(image_files):
155
+ output.append(inference(file.name))
156
+ return output
157
+
158
+ with gr.Blocks() as iface:
159
+ gr.Markdown("# Remove Background")
160
+ gr.HTML("Uses <a href='https://github.com/xuebinqin/DIS'>DIS</a> to remove background")
161
+ with gr.Tab("Single Image"):
162
+ with gr.Row():
163
+ with gr.Column():
164
+ image = gr.Image(type='filepath')
165
+ with gr.Column():
166
+ image_output = gr.Image(interactive=False)
167
+ with gr.Row():
168
+ with gr.Column():
169
+ single_removebg = gr.Button("Remove Bg")
170
+ with gr.Column():
171
+ single_clear = gr.Button("Clear")
172
+
173
+
174
+ with gr.Tab("Batch"):
175
+ with gr.Row():
176
+ with gr.Column():
177
+ images = gr.File(file_count="multiple", file_types=["image"])
178
+ with gr.Column():
179
+ gallery = gr.Gallery()
180
+ with gr.Row():
181
+ with gr.Column():
182
+ batch_removebg = gr.Button("Batch Process")
183
+ with gr.Column():
184
+ batch_clear = gr.Button("Clear")
185
+ #Events
186
+ single_removebg.click(inference, inputs=image, outputs=image_output)
187
+ batch_removebg.click(batch, inputs=images, outputs=gallery)
188
+ single_clear.click(lambda: None, None, image, queue=False)
189
+ batch_clear.click(lambda: None, None, images, queue=False)
190
+
191
  iface.launch()