diff --git a/examples/test_python_folder_classify.py b/examples/test_python_folder_classify.py index a5e1d66..be5ae8a 100644 --- a/examples/test_python_folder_classify.py +++ b/examples/test_python_folder_classify.py @@ -1,6 +1,6 @@ from tensorpy import image_base classifications = image_base.classify_folder_images('./images') -print("*** Displaying Image Classification Results: ***") +print("*** Displaying Image Classification Results as list: ***") for classification in classifications: print classification diff --git a/examples/test_python_folder_classify_dict.py b/examples/test_python_folder_classify_dict.py new file mode 100644 index 0000000..58683f5 --- /dev/null +++ b/examples/test_python_folder_classify_dict.py @@ -0,0 +1,5 @@ +from tensorpy import image_base + +dictionary = image_base.classify_folder_images('./images', return_dict=True) +print("*** Displaying Image Classification Results as dictionary: ***") +print dictionary diff --git a/setup.py b/setup.py index 977b511..0bcbb42 100644 --- a/setup.py +++ b/setup.py @@ -8,7 +8,7 @@ setup( name='tensorpy', - version='1.0.15', + version='1.0.16', url='http://tensorpy.com', author='Michael Mintz', author_email='@mintzworld', diff --git a/tensorpy/image_base.py b/tensorpy/image_base.py index e65e1e2..1c39e16 100644 --- a/tensorpy/image_base.py +++ b/tensorpy/image_base.py @@ -118,8 +118,9 @@ def classify_local_image(file_path): return best_guess -def classify_folder_images(folder_path): - classified_images = [] +def classify_folder_images(folder_path, return_dict=False): + classified_images_list = [] + classified_images_dict = {} files = [f for f in listdir(folder_path) if isfile(join(folder_path, f))] images = [f for f in files if (f.endswith('.jpg') or f.endswith('.png'))] total = len(images) @@ -128,10 +129,13 @@ def classify_folder_images(folder_path): counter += 1 sys.stdout.write("\rClassifying Image %d of %s..." % (counter, total)) sys.stdout.flush() - classified_images.append( - classify_local_image(os.path.join(folder_path, image))) + result = classify_local_image(os.path.join(folder_path, image)) + classified_images_list.append(result) + classified_images_dict[image] = result sys.stdout.write("\rAll classifications have been completed!\n") - return classified_images + if return_dict: + return classified_images_dict + return classified_images_list def classify(image_url):