Input IMage Filesを追加

app.py

@@ -7,6 +7,7 @@ import gradio as gr
 from PIL import Image
 import sys, os
 from rembg import remove
+from utils import functions
 
 
 def complete():
@@ -20,6 +21,15 @@ def clean(text_output):
         return ""
     else:
         pass
+
+def clean_by_name(text_output):
+    text_output, dir_name = text_output.split("+")
+    if text_output=="complete":
+        print("clean up")
+        os.remove(f"{dir_name}.zip")
+        return ""
+    else:
+        pass
     
     
 
@@ -32,7 +42,8 @@ def run_rembg(img):
 
     return cropped_image
 
-def load_zip(inputs):
+
+def from_zip(inputs):
     os.makedirs("objects", exist_ok=True)
 
     image_data_dict = {}
@@ -64,6 +75,7 @@ def load_zip(inputs):
     
     
     for image_name, image_data in image_data_dict.items():
+        print(type(image_data))
         output = remove(image_data)        
         output_pil = Image.fromarray(output)
         # Remove margins 
@@ -78,6 +90,33 @@ def load_zip(inputs):
     return "objects.zip", complete()
 
 
+def from_image_files(images):
+    dir_name = functions.get_random_name()
+    os.makedirs(dir_name, exist_ok=True)
+
+    for image in images:
+        image_name = image.name
+
+    
+        # 読み込み
+        image_data = np.array(Image.open(image_name))
+
+        output = remove(image_data)        
+        output_pil = Image.fromarray(output)
+        # Remove margins 
+        cropped_image = output_pil.crop(output_pil.getbbox())
+
+        image_name = image_name.split("/")[-1]
+        image_name = image_name.replace("jpg", "png")
+        cropped_image.save(f"{dir_name}/{image_name}")
+    
+    shutil.make_archive(f"{dir_name}", "zip", f"{dir_name}")
+    shutil.rmtree(f"{dir_name}")
+    
+    return f"{dir_name}.zip", complete()+"+"+dir_name
+
+
+
 
 if __name__=="__main__":
 
@@ -112,7 +151,7 @@ if __name__=="__main__":
             btn = gr.Button("Run!")
             
             btn.click(
-                fn=load_zip,
+                fn=from_zip,
                 inputs=image_input,
                 outputs=[image_output, text_output]
                 )
@@ -149,14 +188,59 @@ if __name__=="__main__":
                 image_output = gr.Image(type="pil")
             
             btn = gr.Button("Run!")
+
+
+            btn.click(
+                fn=run_rembg,
+                inputs=image_input,
+                outputs=image_output,
+                api_name="imageMatting"
+                )
+        
+
         
-        btn.click(
-            fn=run_rembg,
-            inputs=image_input,
-            outputs=image_output,
-            api_name="imageMatting"
+        with gr.Tab("Images"):
+            gr.Markdown(
+            """
+            <center><h1>Image Matting using U<sup>2</sup>-Net</h1></center>
+            """
+            )        
+            with gr.Row():
+                gr.Markdown(
+                    """
+                    ### Input Image Files
+                    <img src='file/assets/input_zip.png' width="85%" height="85%">
+                    """
+                )        
+                gr.Markdown(
+                    """
+                    ### Output Zip File
+                    <img src='file/assets/output_zip.png' width="85%" height="85%">
+                    """
+                )        
+            
+            with gr.Row():
+                image_input = gr.File(file_count="multiple")
+                image_output = gr.File()
+                text_output = gr.Textbox(visible=False)
+            
+            btn = gr.Button("Run!")
+            
+            btn.click(
+                fn=from_image_files,
+                inputs=image_input,
+                outputs=[image_output, text_output]
+                )
+            text_output.change(
+                fn=clean_by_name,
+                inputs=text_output,
+                outputs=text_output
+
             )
 
+
+
+
         gr.Markdown(
             """
             ---

utils/functions.py

@@ -0,0 +1,30 @@
+import random
+
+def get_random_name():
+    famous_painters = [
+    "Leonardo", "DaVinci",
+    "Michelangelo",
+    "Pablo", "Picasso",
+    "Vincent", "VanGogh",
+    "Rembrandt", "VanRijn",
+    "Claude", "Monet",
+    "Salvador", "Dali",
+    "Jackson", "Pollock",
+    "Andy", "Warhol",
+    "Henri", "Matisse",
+    "Georgia", "Keeffe",
+    "Edvard", "Munch",
+    "Wassily", "Kandinsky",
+    "Gustav", "Klimt",
+    "Rene", "Magritte",
+    "Frida", "Kahlo",
+    "Edgar", "Degas",
+    "Johannes", "Vermeer",
+    "Paul", "Cezanne",
+    "Marc", "Chagall",
+    ]
+
+    random_painter = random.choice(famous_painters)
+
+
+    return random_painter