[feat] isolate uid by process

app.py

@@ -1,12 +1,11 @@
-import gradio as gr
-import shutil
 import os
+import shutil
 import subprocess
 import uuid
 from PIL import Image
-from huggingface_hub import hf_hub_download
-from shutil import copyfile
+import gradio as gr
 
+UPLOAD_DIR = "./sessions"
 RESULTS_DIR = "./results"
 CHECKPOINTS_DIR = "./checkpoints/SingleImageReflectionRemoval"
 SAMPLE_DIR = "./sample_images"
@@ -15,6 +14,9 @@ os.makedirs(RESULTS_DIR, exist_ok=True)
 os.makedirs(CHECKPOINTS_DIR, exist_ok=True)
 os.makedirs(SAMPLE_DIR, exist_ok=True)
 
+from huggingface_hub import hf_hub_download
+from shutil import copyfile
+
 REPO_ID = "hasnafk/SingleImageReflectionRemoval"
 MODEL_FILE = "310_net_G.pth"
 model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILE, cache_dir=CHECKPOINTS_DIR)
@@ -23,20 +25,33 @@ expected_model_path = os.path.join(CHECKPOINTS_DIR, MODEL_FILE)
 if not os.path.exists(expected_model_path):
     copyfile(model_path, expected_model_path)
 
-def reflection_removal(input_image, preprocess_type="resize_and_crop", session_id=None):
-    if not session_id:
-        return "Session ID missing. Please try again."
+def generate_session_id():
+    return str(uuid.uuid4())
+
+def randomize_file_name(original_name):
+    extension = os.path.splitext(original_name)[1]
+    new_name = f"{uuid.uuid4().hex}{extension}"
+    return new_name
 
-    upload_dir = os.path.join("./sessions", session_id, "uploads")
+def clear_session_files(session_id):
+    session_dir = os.path.join(UPLOAD_DIR, session_id)
+    if os.path.exists(session_dir):
+        shutil.rmtree(session_dir)
+
+def reflection_removal(input_image, preprocess_type="resize_and_crop"):
+    session_id = generate_session_id()
+    session_dir = os.path.join(UPLOAD_DIR, session_id)
+    upload_dir = os.path.join(session_dir, "uploads")
     os.makedirs(upload_dir, exist_ok=True)
 
     if not input_image or not os.path.exists(input_image):
         return "No image was provided or file was cleared. Please upload a valid image."
 
-    file_path = os.path.join(upload_dir, os.path.basename(input_image))
+    randomized_name = randomize_file_name(os.path.basename(input_image))
+    file_path = os.path.join(upload_dir, randomized_name)
     shutil.copy(input_image, file_path)
 
-    input_filename = os.path.splitext(os.path.basename(file_path))[0]
+    input_filename = os.path.splitext(randomized_name)[0]
 
     cmd = [
         "python", "test.py",
@@ -51,39 +66,28 @@ def reflection_removal(input_image, preprocess_type="resize_and_crop", session_i
     ]
     subprocess.run(cmd, check=True)
 
-    result_files = []
+    output_image = None
     for root, _, files in os.walk(RESULTS_DIR):
         for file in files:
             if file.startswith(input_filename) and file.endswith("_fake.png"):
                 result_path = os.path.join(root, file)
                 output_image = Image.open(result_path)
-                result_files.append(result_path)
-                
                 os.remove(result_path)
             elif file.startswith(input_filename) and file.endswith("_real.png"):
                 real_path = os.path.join(root, file)
-                result_files.append(real_path)
-                
                 os.remove(real_path)
 
     clear_session_files(session_id)
 
-    if result_files:
+    if output_image:
         return output_image
     return "No results found."
 
-
-def clear_session_files(session_id):
-    session_dir = os.path.join("./sessions", session_id)
-    if os.path.exists(session_dir):
-        shutil.rmtree(session_dir)
-    print(f"Session {session_id} files cleared.")
-
-def clear_action(session_id=None):
-    if session_id:
-        clear_session_files(session_id)
-        return "Upload cleared!"
-    return "No session found to clear."
+def use_sample_image(sample_image_name):
+    sample_image_path = os.path.join(SAMPLE_DIR, sample_image_name)
+    if not os.path.exists(sample_image_path):
+        return "Sample image not found."
+    return sample_image_path
 
 sample_images = [
     file for file in os.listdir(SAMPLE_DIR)
@@ -94,27 +98,20 @@ preprocess_options = [
     "resize_and_crop", "crop", "scale_width", "scale_width_and_crop", "none"
 ]
 
-def session_interface():
-    session_id = str(uuid.uuid4())
-
-    return gr.Interface(
-        fn=lambda img, prep: reflection_removal(img, prep, session_id),
-        inputs=[
-            gr.Image(type="filepath", label="Upload Image (JPG/PNG)", interactive=True),
-            gr.Dropdown(choices=preprocess_options, label="Preprocessing Type", value="resize_and_crop")
-        ],
-        outputs=gr.Image(label="Results after Reflection Removal"),
-        examples=[
-            [os.path.join("sample_images", img), "resize_and_crop"]
-            for img in os.listdir("sample_images") if img.endswith((".jpg", ".jpeg", ".png"))
-        ],
-        title="Reflection Remover with Pix2Pix",
-        description="Upload images to remove reflections using a Pix2Pix model. You can also try the sample images below.",
-        allow_flagging="never",
-        live=False
-    )
-
-os.makedirs("./sessions", exist_ok=True)
+iface = gr.Interface(
+    fn=reflection_removal,
+    inputs=[
+        gr.Image(type="filepath", label="Upload Image (JPG/PNG)"),
+        gr.Dropdown(choices=preprocess_options, label="Preprocessing Type", value="resize_and_crop")
+    ],
+    outputs=gr.Image(label="Result after Reflection Removal"),
+    examples=[
+        [os.path.join(SAMPLE_DIR, img), "resize_and_crop"]
+        for img in sample_images
+    ],
+    title="Reflection Remover with Pix2Pix",
+    description="Upload images to remove reflections using a Pix2Pix model. You can also try the sample images below."
+)
 
 if __name__ == "__main__":
-    session_interface().launch()
+    iface.launch()