update working directory

app.py

@@ -8,29 +8,50 @@ from PIL import Image
 import sys, os
 from rembg import remove
 from utils import functions
+import os, sys
+import argparse
 
+from logging import getLogger, StreamHandler, DEBUG
+logger = getLogger(__name__)
+handler = StreamHandler(); handler.setLevel(DEBUG)
+logger.setLevel(DEBUG)
+logger.addHandler(handler)
+logger.propagate = False
 
-def complete():
-    return "complete"
+# logger.debug('variable')
+# logger.info('system')
+
+
+
+def complete(work_dir):
+    work_dir = work_dir
+    # logger.debug(f"complete :", work_dir)
+    return work_dir
 
 
 def clean(text_output):
-    if text_output=="complete":
-        print("clean up")
-        os.remove("objects.zip")
-        return ""
+    # logger.debug(f"text_output : {text_output}")
+
+    if text_output!="idle_state":
+        logger.info(f"clean up : {text_output}.zip")
+        os.remove(f"{text_output}.zip")
+        return "idle_state"
     else:
-        pass
+        logger.info(f"reset")
+        return "idle_state"
+
 
 
 def clean_by_name(text_output):
-    text_output, dir_name = text_output.split("+")
-    if text_output=="complete":
-        print("clean up")
+    # logger.debug(f"text_output : {text_output}")
+    if text_output!="idle_state":
+        text_output, dir_name = text_output.split("+")
+        logger.info(f"clean up : {dir_name}.zip")
         os.remove(f"{dir_name}.zip")
-        return ""
+        return "idle_state"
     else:
-        pass
+        logger.info(f"reset")
+        return "idle_state"
     
     
 
@@ -45,7 +66,8 @@ def run_rembg(img):
 
 
 def from_zip(inputs):
-    os.makedirs("objects", exist_ok=True)
+    work_dir = functions.get_random_name()
+    os.makedirs(work_dir, exist_ok=True)
 
     image_data_dict = {}
     with ZipFile(inputs[0].name, "r") as zip_file:
@@ -68,27 +90,30 @@ def from_zip(inputs):
                         image_files.append(image_name)
                         
                         image_array = np.array(image)
-                        image_name = image_name.split("/")[1]
+                        logger.debug(f"image name : {image_name}")
+                        category_dir = image_name.split("/")[0]
+                        # image_name = image_name.split("/")[1]
+                        os.makedirs(f"{work_dir}/{category_dir}", exist_ok=True)
                         image_data_dict[image_name] = image_array
 
                 except Exception as e:
-                    print("Exception : ", e)
+                    logger.info(f"Exception : {e}")
     
     
     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 
         cropped_image = output_pil.crop(output_pil.getbbox())
 
         image_name = image_name.replace("jpg", "png")
-        cropped_image.save(f"objects/{image_name}")
+        cropped_image.save(f"{work_dir}/{image_name}")
     
-    shutil.make_archive("objects", "zip", "objects")
-    shutil.rmtree("objects")
+    shutil.make_archive(work_dir, "zip", work_dir)
+    shutil.rmtree(work_dir)
     
-    return "objects.zip", complete()
+    return f"{work_dir}.zip", complete(work_dir)
 
 
 def from_image_files(images, text_class_name):
@@ -97,11 +122,12 @@ def from_image_files(images, text_class_name):
         dir_name = text_class_name
     else:
         dir_name = functions.get_random_name()
-        os.makedirs(dir_name, exist_ok=True)
+
+    os.makedirs(dir_name, exist_ok=True)
 
     for image in images:
         image_name = image.name
-
+        # logger.debug(f"image name : {image_name}")
     
         # 読み込み
         image_data = np.array(Image.open(image_name))
@@ -112,13 +138,14 @@ def from_image_files(images, text_class_name):
         cropped_image = output_pil.crop(output_pil.getbbox())
 
         image_name = image_name.split("/")[-1]
-        image_name = image_name.replace("jpg", "png")
+        image_name = image_name[:image_name.find("_", image_name.find("_") + 1)] + ".png"
+        # logger.debug(f"save image name : {image_name}")
         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
+    return f"{dir_name}.zip", complete("complete")+"+"+dir_name
 
 
 
@@ -172,7 +199,6 @@ if __name__=="__main__":
 
 
 
-
         with gr.Tab("Zip"):
             gr.Markdown(
             """
@@ -196,7 +222,7 @@ if __name__=="__main__":
             with gr.Row():
                 image_input = gr.File(file_count="multiple")
                 image_output = gr.File()
-                text_output = gr.Textbox(visible=False)
+                text_output = gr.Textbox(visible=False, value="idle_state")
             
             btn = gr.Button("Run!")