minor chnages in demo.py so only images in the specified directory (n…

…ot deeper) gets ran
akanazawa · May 2, 2019 · ed9cd3c · ed9cd3c
1 parent 984ab41
commit ed9cd3c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/demo.py b/demo.py
@@ -7,6 +7,7 @@
 import torch.utils.data
 import numpy as np
 from opt import opt
+from glob import glob
 
 from dataloader import ImageLoader, DetectionLoader, DetectionProcessor, DataWriter, Mscoco
 from yolo.util import write_results, dynamic_write_results
@@ -36,8 +37,10 @@
     if len(inputlist):
         im_names = open(inputlist, 'r').readlines()
     elif len(inputpath) and inputpath != '/':
-        for root, dirs, files in os.walk(inputpath):
-            im_names = files
+        # @kanazawa: only run on images in the directory.
+        im_names = sorted(glob(inputpath + '/*.png'))
+        # for root, dirs, files in os.walk(inputpath):
+        #     im_names = files
     else:
         raise IOError('Error: must contain either --indir/--list')
 

diff --git a/requirements.txt b/requirements.txt
@@ -6,4 +6,4 @@ pandas
 tqdm
 matplotlib
 opencv-python
-ntpath
+#ntpath