Load only image

ayooshkathuria · Jun 12, 2018 · fbb4ef9 · fbb4ef9
1 parent e42aa2f
commit fbb4ef9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/detect.py b/detect.py
@@ -134,7 +134,7 @@ def arg_parse():
     read_dir = time.time()
     #Detection phase
     try:
-        imlist = [osp.join(osp.realpath('.'), images, img) for img in os.listdir(images)]
+        imlist = [osp.join(osp.realpath('.'), images, img) for img in os.listdir(images) if os.path.splitext(img)[1] == '.png' or os.path.splitext(img)[1] =='.jpeg' or os.path.splitext(img)[1] =='.jpg']
     except NotADirectoryError:
         imlist = []
         imlist.append(osp.join(osp.realpath('.'), images))