Merge pull request #12 from WillBrennan/bug/find-images

Fixes #11 - replaced ispath with isfile plus testing this case

FromFile.py

@@ -4,33 +4,13 @@ __author__ = 'Will Brennan'
 
 import argparse
 import logging
-import os
+
 import cv2
+
 import skin_detector
 
 logger = logging.getLogger('main')
 
-
-def find_images(path, recursive=False, ignore=True):
-    if os.path.exists(path):
-        yield path
-    elif os.path.isdir(path):
-        assert os.path.isdir(path), 'FileIO - get_images: Directory does not exist'
-        assert isinstance(recursive, bool), 'FileIO - get_images: recursive must be a boolean variable'
-        ext, result = ['png', 'jpg', 'jpeg'], []
-        for path_a in os.listdir(path):
-            path_a = path + '/' + path_a
-            if os.path.isdir(path_a) and recursive:
-                for path_b in find_images(path_a):
-                    yield path_b
-            check_a = path_a.split('.')[-1] in ext
-            check_b = ignore or ('-' not in path_a.split('/')[-1])
-            if check_a and check_b:
-                yield path_a
-    else:
-        raise ValueError('error! path is not a valid path or directory')
-
-
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument('image_paths', type=str, nargs='+', help="paths to one or more images or image directories")
@@ -48,7 +28,7 @@ if __name__ == '__main__':
     logger = logging.getLogger("main")
 
     for image_arg in args.image_paths:
-        for image_path in find_images(image_arg):
+        for image_path in skin_detector.find_images(image_arg):
             logging.info("loading image from {0}".format(image_path))
             img_col = cv2.imread(image_path, 1)
 

skin_detector/__init__.py

 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-__author__ = 'willbrennan'
 
 from .skin_detector import *
 from .scripts import display
-
-__all__ = [
+from .scripts import find_images
\ No newline at end of file
-    "process", "display", "get_hsv_mask", "get_rgb_mask", "get_ycrcb_mask", "grab_cut_mask", "grab_cut_mask", "closing"
-]

skin_detector/scripts.py

@@ -2,11 +2,32 @@
 # -*- coding: utf-8 -*-
 __author__ = 'Will Brennan'
 
+import os
 
 import cv2
 import numpy
 
 
+def find_images(path, recursive=False, ignore=True):
+    if os.path.isfile(path):
+        yield path
+    elif os.path.isdir(path):
+        assert os.path.isdir(path), 'FileIO - get_images: Directory does not exist'
+        assert isinstance(recursive, bool), 'FileIO - get_images: recursive must be a boolean variable'
+        ext, result = ['png', 'jpg', 'jpeg'], []
+        for path_a in os.listdir(path):
+            path_a = path + '/' + path_a
+            if os.path.isdir(path_a) and recursive:
+                for path_b in find_images(path_a):
+                    yield path_b
+            check_a = path_a.split('.')[-1] in ext
+            check_b = ignore or ('-' not in path_a.split('/')[-1])
+            if check_a and check_b:
+                yield path_a
+    else:
+        raise ValueError('error! path is not a valid path or directory')
+
+
 def display(title, img, max_size=200000):
     assert isinstance(img, numpy.ndarray), 'img must be a numpy array'
     assert isinstance(title, str), 'title must be a string'

skin_detector/skin_detector.py

@@ -45,7 +45,6 @@ def get_rgb_mask(img, debug=False):
     # msk_rgb = cv2.bitwise_and(mask_c, cv2.bitwise_and(mask_a, mask_b))
     mask_d = numpy.bitwise_and(numpy.uint64(mask_a), numpy.uint64(mask_b))
     msk_rgb = numpy.bitwise_and(numpy.uint64(mask_c), numpy.uint64(mask_d))
-    
 
     msk_rgb[msk_rgb < 128] = 0
     msk_rgb[msk_rgb >= 128] = 1

tests/skin_detector/test_unit.py

-import cv2
 import os
+import shutil
+
+import cv2
 import numpy
+
 import skin_detector
 
 
@@ -42,3 +45,51 @@ def test_process():
     img = cv2.imread(img_path)
     mask = skin_detector.process(img)
     assert img.shape[:2] == mask.shape
+
+
+def test_find_images():
+    image_dir = os.path.abspath("./tests/images/")
+    if os.path.isdir(image_dir):
+        shutil.rmtree(image_dir)
+    os.makedirs(image_dir)
+
+    image_paths = sorted(["hello.png", "world.jpg", "simon.png", "says.jpeg"])
+    image_paths = [os.path.join(image_dir, image_path) for image_path in image_paths]
+    image_paths = [os.path.abspath(image_path) for image_path in image_paths]
+
+    for image_path in image_paths:
+        img = numpy.random.randint(255, size=(1920, 1080, 3))
+        cv2.imwrite(image_path, img)
+
+    assert sorted(skin_detector.find_images(image_dir)) == image_paths
+    shutil.rmtree(image_dir)
+
+
+def test_find_images_recursive():
+    image_dir = os.path.abspath("./tests/images/")
+    if os.path.isdir(image_dir):
+        shutil.rmtree(image_dir)
+
+    os.makedirs(image_dir)
+
+    image_paths = sorted(["hello.png", "world.jpg", "simon.png", "says.jpeg"])
+    image_paths = [os.path.join(image_dir, image_path) for image_path in image_paths]
+
+    for image_path in image_paths:
+        img = numpy.random.randint(255, size=(1920, 1080, 3))
+        cv2.imwrite(image_path, img)
+
+    recursive_dir = os.path.abspath("./tests/images/recurse")
+    os.makedirs(recursive_dir)
+
+    recursive_images = sorted(["alpha.png", "beta.jpeg", "gamma.png"])
+    recursive_images = [os.path.join(recursive_dir, image_path) for image_path in recursive_images]
+
+    for image_path in recursive_images:
+        img = numpy.random.randint(255, size=(1920, 1080, 3))
+        cv2.imwrite(image_path, img)
+
+    all_images = sorted(recursive_images + image_paths)
+
+    assert sorted(skin_detector.find_images(image_dir, recursive=True)) == all_images
+    shutil.rmtree(image_dir)