batch process and file name

.gitignore

@@ -160,4 +160,5 @@ cython_debug/
 #.idea/
 
 *.jpeg
-*.png
+*.png
+.DS_Store 

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import cv2
 import gradio as gr
 import os
+from pathlib import Path
 from PIL import Image
 import numpy as np
 import torch
@@ -10,6 +11,8 @@ from torchvision import transforms
 import torch.nn.functional as F
 import matplotlib.pyplot as plt
 import warnings
+import tempfile
+
 warnings.filterwarnings("ignore")
 
 
@@ -115,29 +118,74 @@ hypar["model"] = ISNetDIS()
 net = build_model(hypar, device)
 
 
-def inference(image: Image):
-  image_path = image
+def inference(image_path):
   
   image_tensor, orig_size = load_image(image_path, hypar) 
   mask = predict(net, image_tensor, orig_size, hypar, device)
   
   pil_mask = Image.fromarray(mask).convert('L')
-  im_rgb = Image.open(image).convert("RGB")
+  im_rgb = Image.open(image_path).convert("RGB")
   
   im_rgba = im_rgb.copy()
   im_rgba.putalpha(pil_mask)
-
-  return im_rgba
-
-def bw(image_file:Image):
+  file_name = Path(image_path).stem+"_nobg.png"
+  file_path = Path(Path(image_path).parent,file_name)
+  im_rgba.save(file_path)
+  return str(file_path.resolve())
+
+def bw(image_files):
+    print(image_files)
+    output = []
+    for idx, file in enumerate(image_files):
+        print(file.name)
+        img = Image.open(file.name)
+        img = img.convert("L")
+        output.append(img)
+    print(output)
+    return output
+
+def bw_single(image_file):
     img = Image.open(image_file)
     img = img.convert("L")
     return img
 
-iface = gr.Interface(fn=inference, 
-                    inputs=gr.Image(type='filepath'), 
-                    outputs=["image"],
-                    title="Remove Background",
-                    description="Uses <a href='https://github.com/xuebinqin/DIS'>DIS</a> to remove background"
-                    )
+def batch(image_files):
+    output = []
+    for idx, file in enumerate(image_files):
+        output.append(inference(file.name))
+    return output
+
+with gr.Blocks() as iface:
+    gr.Markdown("# Remove Background")
+    gr.HTML("Uses <a href='https://github.com/xuebinqin/DIS'>DIS</a> to remove background")
+    with gr.Tab("Single Image"):
+        with gr.Row():
+            with gr.Column():
+                image = gr.Image(type='filepath')
+            with gr.Column():
+                image_output = gr.Image(interactive=False)
+        with gr.Row():
+            with gr.Column():
+                single_removebg = gr.Button("Remove Bg")
+            with gr.Column():
+                single_clear = gr.Button("Clear")
+
+    
+    with gr.Tab("Batch"):
+        with gr.Row():
+            with gr.Column():
+                images = gr.File(file_count="multiple", file_types=["image"])
+            with gr.Column():
+                gallery = gr.Gallery()
+        with gr.Row():
+            with gr.Column():
+                batch_removebg = gr.Button("Batch Process")
+            with gr.Column():
+                batch_clear = gr.Button("Clear")
+    #Events
+    single_removebg.click(inference, inputs=image, outputs=image_output)
+    batch_removebg.click(batch, inputs=images, outputs=gallery)
+    single_clear.click(lambda: None, None, image, queue=False)
+    batch_clear.click(lambda: None, None, images, queue=False)
+
 iface.launch()