[update] only allowed 1 image at a time and add sample image

app.py

@@ -9,52 +9,68 @@ from shutil import copyfile
 UPLOAD_DIR = "./uploaded_images"
 RESULTS_DIR = "./results"
 CHECKPOINTS_DIR = "./checkpoints/SingleImageReflectionRemoval"
+SAMPLE_DIR = "./sample_images"
 
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(RESULTS_DIR, exist_ok=True)
 os.makedirs(CHECKPOINTS_DIR, exist_ok=True)
+os.makedirs(SAMPLE_DIR, exist_ok=True)
 
 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)
 
 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_images):
-    for input_image in input_images:
-        file_path = os.path.join(UPLOAD_DIR, os.path.basename(input_image.name))
-        shutil.copy(input_image.name, file_path)
-            
+def reflection_removal(input_image):
+    if not input_image.name.endswith((".jpg", ".jpeg", ".png")):
+        return "FIle is not supported (only .jpg, .jpeg, .png)."
+    
+    file_path = os.path.join(UPLOAD_DIR, os.path.basename(input_image.name))
+    shutil.copy(input_image.name, file_path)
     
     cmd = [
-        "python", "test.py", 
+        "python", "test.py",
         "--dataroot", UPLOAD_DIR,
         "--name", "SingleImageReflectionRemoval",
         "--model", "test", "--netG", "unet_256",
         "--direction", "AtoB", "--dataset_mode", "single",
         "--norm", "batch", "--epoch", "310",
-        "--num_test", str(len(input_images)),
+        "--num_test", "1",
         "--gpu_ids", "-1"
     ]
     subprocess.run(cmd, check=True)
     
-    output_images = []
     for root, _, files in os.walk(RESULTS_DIR):
         for file in files:
             if file.endswith("_fake.png"):
-                output_images.append(Image.open(os.path.join(root, file)))
+                result_path = os.path.join(root, file)
+                return Image.open(result_path)
+
+    return "No results found."
+
+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
 
-    return output_images
+sample_images = [
+    file for file in os.listdir(SAMPLE_DIR)
+    if file.endswith((".jpg", ".jpeg", ".png"))
+]
 
 iface = gr.Interface(
     fn=reflection_removal,
-    inputs=gr.File(file_types=["image"], file_count="multiple"),
-    outputs=gr.Gallery(label="Results after Reflection Removal"),
+    inputs=gr.File(file_types=["image"], label="Upload Image (JPG/PNG)", live=True),
+    outputs=gr.Image(type="pil", label="Processed Image"),
     title="Reflection Remover with Pix2Pix",
-    description="Upload images to remove reflections using a Pix2Pix model."
+    description="Upload a single image to remove reflections. No GIFs allowed.",
+    examples=[
+        os.path.join(SAMPLE_DIR, img) for img in sample_images
+    ],
 )
 
 if __name__ == "__main__":