From 63ef10da69d8a48e73f3e720c65d71596b880add Mon Sep 17 00:00:00 2001 From: Colin Dean Date: Thu, 23 Feb 2023 23:17:40 -0500 Subject: [PATCH] Enables predictor to load a list of files as a list of strings Fixes #81 --- nsfw_detector/predict.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/nsfw_detector/predict.py b/nsfw_detector/predict.py index 4a63809..a1d41d6 100644 --- a/nsfw_detector/predict.py +++ b/nsfw_detector/predict.py @@ -17,7 +17,7 @@ def load_images(image_paths, image_size, verbose=True): ''' Function for loading images into numpy arrays for passing to model.predict inputs: - image_paths: list of image paths to load + image_paths: image paths to load as one file or one directory as a string, or a list of files as strings image_size: size into which images should be resized verbose: show all of the image path and sizes loaded @@ -29,11 +29,12 @@ def load_images(image_paths, image_size, verbose=True): loaded_images = [] loaded_image_paths = [] - if isdir(image_paths): - parent = abspath(image_paths) - image_paths = [join(parent, f) for f in listdir(image_paths) if isfile(join(parent, f))] - elif isfile(image_paths): - image_paths = [image_paths] + if type(image_paths) is str: + if isdir(image_paths): + parent = abspath(image_paths) + image_paths = [join(parent, f) for f in listdir(image_paths) if isfile(join(parent, f))] + elif isfile(image_paths): + image_paths = [image_paths] for img_path in image_paths: try: