batch process and file name
Browse files- .gitignore +2 -1
- 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(
|
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(
|
126 |
|
127 |
im_rgba = im_rgb.copy()
|
128 |
im_rgba.putalpha(pil_mask)
|
129 |
-
|
130 |
-
|
131 |
-
|
132 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
133 |
img = Image.open(image_file)
|
134 |
img = img.convert("L")
|
135 |
return img
|
136 |
|
137 |
-
|
138 |
-
|
139 |
-
|
140 |
-
|
141 |
-
|
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()
|