update

app.py

@@ -1,4 +1,9 @@
-import shutil
+from zipfile import ZipFile
+import numpy as np
+from matplotlib import pyplot as plt
+
+import os, shutil
+
 import gradio as gr
 from PIL import Image
 import sys, os
@@ -9,34 +14,113 @@ def run_rembg(img):
     output = remove(img)
     return output
 
+def load_zip(inputs):
+    os.makedirs("results", exist_ok=True)
+
+    image_data_dict = {}
+    with ZipFile(inputs[0].name, "r") as zip_file:
+        image_names = zip_file.namelist()
+
+        prefix = ""
+        for name in image_names:
+            if prefix=="":
+                prefix = name.split("/")[0]
+            else:
+                break
+
+        image_files = []      
+          
+        for image_name in image_names:
+            if image_name[-3:] in "pngjpg":
+                try:
+                    with zip_file.open(image_name) as f:
+                        image = Image.open(f)
+                        image_files.append(image_name)
+                        
+                        image_array = np.array(image)
+                        image_name = image_name.split("/")[1]
+                        image_data_dict[image_name] = image_array
+
+                except Exception as e:
+                    print("Exception : ", e)
+    
+    
+    for image_name, image_data in image_data_dict.items():
+        img_array = remove(image_data)        
+
+        image_name = image_name.replace("jpg", "png")
+        plt.imsave(f"results/{image_name}", img_array)
+    
+    shutil.make_archive("results", "zip", "results")
+    with ZipFile("objects.zip", "w") as zip_file:
+        zip_file.write("results.zip")
+
+    shutil.rmtree("results")
+
+    return "objects.zip"
+
+
 
 if __name__=="__main__":
 
        
     with gr.Blocks() as demo:
 
-        gr.Markdown(
+        with gr.Tab("Zip"):
+            gr.Markdown(
             """
             <center><h1>Image Matting using U<sup>2</sup>-Net</h1></center>
             """
-        )        
-        with gr.Row():
-            with gr.Column():
-                gr.Markdown(
-                    """
-                    Input Image
-                    """
-                )        
-                image_input = gr.Image(type="numpy", label=" ")
-            with gr.Column():
-                gr.Markdown(
-                    """
-                    Output Image
-                    """
-                )        
-                image_output = gr.Image(type="numpy", label=" ")
-        
-        btn = gr.Button("Run!")
+            )        
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown(
+                        """
+                        Input Zip File
+                        """
+                    )        
+                    image_input = gr.File(file_count="multiple")
+                
+                with gr.Column():
+                    gr.Markdown(
+                        """
+                        Output Zip File
+                        """
+                    )        
+                    image_output = gr.File()
+            
+            btn = gr.Button("Run!")
+            
+            btn.click(
+                fn=load_zip,
+                inputs=image_input,
+                outputs=image_output
+                )
+
+
+        with gr.Tab("Image"):
+            gr.Markdown(
+                """
+                <center><h1>Image Matting using U<sup>2</sup>-Net</h1></center>
+                """
+            )        
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown(
+                        """
+                        Input Image
+                        """
+                    )        
+                    image_input = gr.Image(type="numpy")
+                with gr.Column():
+                    gr.Markdown(
+                        """
+                        Output Image
+                        """
+                    )        
+                    image_output = gr.Image(type="numpy")
+            
+            btn = gr.Button("Run!")
         
         btn.click(
             fn=run_rembg,