[debug]

.gitignore

@@ -5,7 +5,6 @@ checkpoints/
 results/
 build/
 dist/
-*.png
 torch.egg-info/
 */**/__pycache__
 torch/version.py

app.py

@@ -23,6 +23,9 @@ model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILE, cache_dir=CHE
 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 count_files(directory):
+    return sum([len(files) for _, _, files in os.walk(directory)])
 
 def reflection_removal(input_image):
     if not input_image.lower().endswith((".jpg", ".jpeg", ".png")):
@@ -41,7 +44,7 @@ def reflection_removal(input_image):
         "--model", "test", "--netG", "unet_256",
         "--direction", "AtoB", "--dataset_mode", "single",
         "--norm", "batch", "--epoch", "310",
-        "--num_test", "1",
+        "--num_test", count_files(UPLOAD_DIR),
         "--gpu_ids", "-1"
     ]
     subprocess.run(cmd, check=True)