Ignore non-files in the input directory

aftiela · Oct 19, 2020 · 9686071 · 9686071
1 parent 43e5cb7
commit 9686071
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/Global/detection.py b/Global/detection.py
@@ -106,7 +106,11 @@ def main(config):
         print("processing", image_name)
 
         results = []
-        scratch_image = Image.open(os.path.join(config.test_path, image_name)).convert("RGB")
+        scratch_file = os.path.join(config.test_path, image_name)
+        if not os.path.isfile(scratch_file):
+            print("Skipping non-file %s" % image_name)
+            continue
+        scratch_image = Image.open(scratch_file).convert("RGB")
 
         w, h = scratch_image.size
 

diff --git a/Global/test.py b/Global/test.py
@@ -108,7 +108,7 @@ def parameter_set(opt):
     dataset_size = 0
 
     input_loader = os.listdir(opt.test_input)
-    dataset_size = len(os.listdir(opt.test_input))
+    dataset_size = len(input_loader)
     input_loader.sort()
 
     if opt.test_mask != "":
@@ -124,7 +124,11 @@ def parameter_set(opt):
     for i in range(dataset_size):
 
         input_name = input_loader[i]
-        input = Image.open(os.path.join(opt.test_input, input_name)).convert("RGB")
+        input_file = os.path.join(opt.test_input, input_name)
+        if not os.path.isfile(input_file):
+            print("Skipping non-file %s" % input_name)
+            continue
+        input = Image.open(input_file).convert("RGB")
 
         print("Now you are processing %s" % (input_name))