From 56d930a8548e7a604f8dd327fb6c672095fb0085 Mon Sep 17 00:00:00 2001
From: zaqxs123456 <18200842@life.hkbu.edu.hk>
Date: Fri, 25 Oct 2024 11:43:28 +0800
Subject: [PATCH] fix edge cases

---
 app.py            | 23 +++++++++++++++++++----
 postprocessing.py |  1 +
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/app.py b/app.py
index 5bf8607..218c078 100644
--- a/app.py
+++ b/app.py
@@ -25,8 +25,6 @@ info = json.load(open('info.json'))
 comfyui_address = info['comfyui_address']
 expo_raw_sd_dir = info['expo_raw_sd_dir']
 expo_openpose_dir = info['expo_openpose_dir']
-expo_postprocessed_dir = info['expo_postprocessed_dir']
-expo_postprocess_temp_dir = info['expo_postprocess_temp_dir']
 
 on_postprocessing = False
 on_testing = False
@@ -142,6 +140,14 @@ async def expo_postprocess():
     if not await wait_for_files_to_match(expo_openpose_dir, expo_raw_sd_dir):
         print("Timeout reached, proceeding with postprocessing")
 
+    # Check if directories exist and are not empty
+    if not os.path.exists(expo_openpose_dir) or not os.listdir(expo_openpose_dir):
+        on_postprocessing = False
+        return jsonify({"status": "error", "message": "No images to process in expo_openpose_dir"}), 404
+    if not os.path.exists(expo_raw_sd_dir) or not os.listdir(expo_raw_sd_dir):
+        on_postprocessing = False
+        return jsonify({"status": "error", "message": "No images to process in expo_raw_sd_dir"}), 404
+    
     await asyncio.to_thread(expo_postprocess_main)
     await asyncio.to_thread(expo_clear_images)
     on_postprocessing = False
@@ -152,14 +158,23 @@ async def expo_postprocess():
 async def wait_for_files_to_match(dir1: str, dir2: str, timeout: int = 60, interval: int = 1) -> bool:
     start_time = asyncio.get_event_loop().time()
     while asyncio.get_event_loop().time() - start_time < timeout:
-        files1 = set(os.listdir(dir1))
-        files2 = set(os.listdir(dir2))
+        files1 = get_all_files(dir1)
+        files2 = get_all_files(dir2)
         if files1 == files2:
             return True
         await asyncio.sleep(interval)
     return False
 
 
+def get_all_files(directory):
+    all_files = set()
+    for root, _, files in os.walk(directory):
+        for file in files:
+            # Store the relative path of the file
+            relative_path = os.path.relpath(os.path.join(root, file), directory)
+            all_files.add(relative_path)
+    return all_files
+
 @app.route('/gen_image', methods=['POST'])
 def gen_image():
     if request.is_json:
diff --git a/postprocessing.py b/postprocessing.py
index 135fb96..0605a13 100644
--- a/postprocessing.py
+++ b/postprocessing.py
@@ -191,6 +191,7 @@ def current_session():
 
 def expo_postprocess_main():
     print("Postprocessing")
+    os.makedirs(expo_postprocessed_dir, exist_ok=True)
     os.makedirs(expo_postprocess_temp_dir, exist_ok=True)
     
     left_fencer_raw_image_dir = os.path.join(expo_raw_sd_dir, 'left_fencer')