Merge pull request #59 from ezavesky/feature/nd_interface

nd_array compute, simplify prediction render, fix CLI

README.md

@@ -45,15 +45,18 @@ predict.classify(model, '/Users/bedapudi/Desktop/')
 
 ```
 
-If you've installed the package, this should work, too...
+If you've installed the package or use the command-line this should work, too...
 
 ```sh
 # a single image
-nsfw-predict --saved_model_path mobilenet_v2_140_224 --image_source images
+nsfw-predict --saved_model_path mobilenet_v2_140_224 --image_source test.jpg
 
 # an image directory
 nsfw-predict --saved_model_path mobilenet_v2_140_224 --image_source images
 
+# a single image (from code/CLI)
+python3 nsfw_detector/predict.py --saved_model_path mobilenet_v2_140_224 --image_source test.jpg
+
 ```
 
 
@@ -125,10 +128,18 @@ This project follows the [all-contributors](https://github.com/kentcdodds/all-co
 
 # Changes
 
+## 1.1.1
+
+- break out numpy (nd array) function
+- remove classic app run modes for argparse
+- one more example in README for running
+- turn down verbosity in image load via file
+- fix requirements for clean system (needs PIL)
+
 ## 1.1.0
 
 - update to tensorflow 2.1.0 and updated mobilenet-based model
 
 ## 1.0.0
 
-- initial creation
+- initial creation

nsfw_detector/predict.py

@@ -1,6 +1,5 @@
 #! python
 
-from absl import app
 import argparse
 import json
 from os import listdir
@@ -12,12 +11,13 @@
 
 IMAGE_DIM = 224   # required/default image dimensionality
 
-def load_images(image_paths, image_size):
+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_size: size into which images should be resized
+        verbose: show all of the image path and sizes loaded
     
     outputs:
         loaded_images: loaded images on which keras model can run predictions
@@ -35,7 +35,8 @@ def load_images(image_paths, image_size):
 
     for img_path in image_paths:
         try:
-            print(img_path, "size:", image_size)
+            if verbose:
+                print(img_path, "size:", image_size)
             image = keras.preprocessing.image.load_img(img_path, target_size=image_size)
             image = keras.preprocessing.image.img_to_array(image)
             image /= 255
@@ -58,30 +59,25 @@ def load_model(model_path):
 def classify(model, input_paths, image_dim=IMAGE_DIM):
     """ Classify given a model, input paths (could be single string), and image dimensionality...."""
     images, image_paths = load_images(input_paths, (image_dim, image_dim))
+    probs = classify_nd(model, images)
+    return dict(zip(image_paths, probs))
 
-    model_preds = model.predict(images)
-
-    preds = np.argsort(model_preds, axis = 1).tolist()
-
-    probs = []
 
+def classify_nd(model, nd_images):
+    """ Classify given a model, image array (numpy)...."""
+
+    model_preds = model.predict(nd_images)
+    # preds = np.argsort(model_preds, axis = 1).tolist()
+
     categories = ['drawings', 'hentai', 'neutral', 'porn', 'sexy']
 
-    for i, single_preds in enumerate(preds):
-        single_probs = []
+    probs = []
+    for i, single_preds in enumerate(model_preds):
+        single_probs = {}
         for j, pred in enumerate(single_preds):
-            single_probs.append(model_preds[i][pred])
-            preds[i][j] = categories[pred]
-
+            single_probs[categories[j]] = float(pred)
         probs.append(single_probs)
-
-    images_preds = {}
-
-    for i, loaded_image_path in enumerate(image_paths):
-        images_preds[loaded_image_path] = {}
-        for _ in range(len(preds[i])):
-            images_preds[loaded_image_path][preds[i][_]] = str(probs[i][_])
-    return images_preds
+    return probs
 
 
 def main(args=None):
@@ -109,11 +105,8 @@ def main(args=None):
 
     model = load_model(config['saved_model_path'])    
     image_preds = classify(model, config['image_source'], config['image_dim'])
-    print(json.dumps(image_preds, sort_keys=True, indent=2), '\n')
-
+    print(json.dumps(image_preds, indent=2), '\n')
 
-def run_main():
-	app.run(main)
 
 if __name__ == "__main__":
-	run_main()
+	main()

requirements.txt

@@ -1,2 +1,3 @@
 tensorflow>=2.1.0
 tensorflow-hub==0.7.0
+pillow

setup.py

@@ -18,7 +18,7 @@
 EMAIL = 'gantman@gmail.com'
 AUTHOR = 'Gant Laborde'
 REQUIRES_PYTHON = '>=3.5.0'
-VERSION = '1.1.0'
+VERSION = '1.1.1'
 
 # What packages are optional?
 EXTRAS = {