<chores>: add pipfile

Pipfile

@@ -0,0 +1,14 @@
+[[source]]
+url = "https://pypi.org/simple"
+verify_ssl = true
+name = "pypi"
+
+[packages]
+yapf = "*"
+tensorflow = "*"
+scikit-image = "*"
+
+[dev-packages]
+
+[requires]
+python_version = "3.6"

README.md

@@ -8,6 +8,9 @@ The original caffe weights have been extracted using [Caffe to TensorFlow](https
 
 All code should be compatible with `Python 3.6` and `Tensorflow 1.x` (tested with 1.12). The model implementation can be found in `model.py`.
 
+1. `pip3 install -U pipenv`
+2. `pipenv install`
+
 ### Usage
 
 ```

classify_nsfw.py

@@ -1,14 +1,13 @@
 #!/usr/bin/env python
-import sys
 import argparse
-import tensorflow as tf
-
-from model import OpenNsfwModel, InputType
-from image_utils import create_tensorflow_image_loader
-from image_utils import create_yahoo_image_loader
+import sys
 
 import numpy as np
 
+import tensorflow as tf
+from image_utils import (create_tensorflow_image_loader,
+                         create_yahoo_image_loader)
+from model import InputType, OpenNsfwModel
 
 IMAGE_LOADER_TENSORFLOW = "tensorflow"
 IMAGE_LOADER_YAHOO = "yahoo"
@@ -17,22 +16,33 @@
 def main(argv):
     parser = argparse.ArgumentParser()
 
-    parser.add_argument("input_file", help="Path to the input image.\
+    parser.add_argument(
+        "input_file",
+        help="Path to the input image.\
                         Only jpeg images are supported.")
 
-    parser.add_argument("-m", "--model_weights", required=True,
-                        help="Path to trained model weights file")
-
-    parser.add_argument("-l", "--image_loader",
-                        default=IMAGE_LOADER_YAHOO,
-                        help="image loading mechanism",
-                        choices=[IMAGE_LOADER_YAHOO, IMAGE_LOADER_TENSORFLOW])
-
-    parser.add_argument("-i", "--input_type",
-                        default=InputType.TENSOR.name.lower(),
-                        help="input type",
-                        choices=[InputType.TENSOR.name.lower(),
-                                 InputType.BASE64_JPEG.name.lower()])
+    parser.add_argument(
+        "-m",
+        "--model_weights",
+        required=True,
+        help="Path to trained model weights file")
+
+    parser.add_argument(
+        "-l",
+        "--image_loader",
+        default=IMAGE_LOADER_YAHOO,
+        help="image loading mechanism",
+        choices=[IMAGE_LOADER_YAHOO, IMAGE_LOADER_TENSORFLOW])
+
+    parser.add_argument(
+        "-i",
+        "--input_type",
+        default=InputType.TENSOR.name.lower(),
+        help="input type",
+        choices=[
+            InputType.TENSOR.name.lower(),
+            InputType.BASE64_JPEG.name.lower()
+        ])
 
     args = parser.parse_args()
 
@@ -47,7 +57,8 @@ def main(argv):
 
         if input_type == InputType.TENSOR:
             if args.image_loader == IMAGE_LOADER_TENSORFLOW:
-                fn_load_image = create_tensorflow_image_loader(tf.Session(graph=tf.Graph()))
+                fn_load_image = create_tensorflow_image_loader(
+                    tf.Session(graph=tf.Graph()))
             else:
                 fn_load_image = create_yahoo_image_loader()
         elif input_type == InputType.BASE64_JPEG:
@@ -65,5 +76,6 @@ def main(argv):
         print("Results for '{}'".format(args.input_file))
         print("\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0]))
 
+
 if __name__ == "__main__":
     main(sys.argv)

eval/batch_classify.py

@@ -1,20 +1,17 @@
-
+import argparse
+import glob
 import os
 import sys
 
-sys.path.append((os.path.normpath(
-                 os.path.join(os.path.dirname(os.path.realpath(__file__)),
-                              '..'))))
-
-import argparse
-import glob
-import tensorflow as tf
 from tqdm import tqdm
 
-from model import OpenNsfwModel, InputType
-from image_utils import create_tensorflow_image_loader
-from image_utils import create_yahoo_image_loader
+import tensorflow as tf
+from image_utils import (create_tensorflow_image_loader,
+                         create_yahoo_image_loader)
+from model import InputType, OpenNsfwModel
 
+sys.path.append((os.path.normpath(
+    os.path.join(os.path.dirname(os.path.realpath(__file__)), '..'))))
 
 IMAGE_LOADER_TENSORFLOW = "tensorflow"
 IMAGE_LOADER_YAHOO = "yahoo"
@@ -25,37 +22,50 @@
 
 def create_batch_iterator(filenames, batch_size, fn_load_image):
     for i in range(0, len(filenames), batch_size):
-        yield list(map(fn_load_image, filenames[i:i+batch_size]))
+        yield list(map(fn_load_image, filenames[i:i + batch_size]))
 
 
 def create_tf_batch_iterator(filenames, batch_size):
     for i in range(0, len(filenames), batch_size):
         with tf.Session(graph=tf.Graph()) as session:
-            fn_load_image = create_tensorflow_image_loader(session,
-                                                           expand_dims=False)
+            fn_load_image = create_tensorflow_image_loader(
+                session, expand_dims=False)
 
-            yield list(map(fn_load_image, filenames[i:i+batch_size]))
+            yield list(map(fn_load_image, filenames[i:i + batch_size]))
 
 
 def main(argv):
     parser = argparse.ArgumentParser()
 
-    parser.add_argument("-s", "--source", required=True,
-                        help="Folder containing the images to classify")
-
-    parser.add_argument("-o", "--output_file", required=True,
-                        help="Output file path")
-
-    parser.add_argument("-m", "--model_weights", required=True,
-                        help="Path to trained model weights file")
-
-    parser.add_argument("-b", "--batch_size", help="Number of images to \
-                        classify simultaneously.", type=int, default=64)
-
-    parser.add_argument("-l", "--image_loader",
-                        default=IMAGE_LOADER_YAHOO,
-                        help="image loading mechanism",
-                        choices=[IMAGE_LOADER_YAHOO, IMAGE_LOADER_TENSORFLOW])
+    parser.add_argument(
+        "-s",
+        "--source",
+        required=True,
+        help="Folder containing the images to classify")
+
+    parser.add_argument(
+        "-o", "--output_file", required=True, help="Output file path")
+
+    parser.add_argument(
+        "-m",
+        "--model_weights",
+        required=True,
+        help="Path to trained model weights file")
+
+    parser.add_argument(
+        "-b",
+        "--batch_size",
+        help="Number of images to \
+                        classify simultaneously.",
+        type=int,
+        default=64)
+
+    parser.add_argument(
+        "-l",
+        "--image_loader",
+        default=IMAGE_LOADER_YAHOO,
+        help="image loading mechanism",
+        choices=[IMAGE_LOADER_YAHOO, IMAGE_LOADER_TENSORFLOW])
 
     args = parser.parse_args()
     batch_size = args.batch_size
@@ -85,8 +95,7 @@ def main(argv):
                                                fn_load_image)
 
     with tf.Session(graph=tf.Graph(), config=config) as session:
-        model.build(weights_path=args.model_weights,
-                    input_type=input_type)
+        model.build(weights_path=args.model_weights, input_type=input_type)
 
         session.run(tf.global_variables_initializer())
 
@@ -102,11 +111,11 @@ def main(argv):
                     fi = (batch_num * batch_size)
                     for i, prediction in enumerate(predictions):
                         filename = os.path.basename(filenames[fi + i])
-                        o.write('{}\t{}\t{}\n'.format(filename,
-                                                      prediction[0],
+                        o.write('{}\t{}\t{}\n'.format(filename, prediction[0],
                                                       prediction[1]))
 
                     progress_bar.update(len(images))
 
+
 if __name__ == "__main__":
     main(sys.argv)

eval/eval.py

@@ -1,7 +1,9 @@
-import sys
-import operator
 import argparse
+import operator
+import sys
+
 import numpy as np
+
 from scipy import stats
 
 
@@ -55,11 +57,12 @@ def test(first, second):
 def main(argv):
     parser = argparse.ArgumentParser()
 
-    parser.add_argument("original",
-                        help="File containing base classifications")
+    parser.add_argument(
+        "original", help="File containing base classifications")
 
-    parser.add_argument("other",
-                        help="File containing classifications to compare to\
+    parser.add_argument(
+        "other",
+        help="File containing classifications to compare to\
                         base results")
 
     args = parser.parse_args()
@@ -86,5 +89,6 @@ def main(argv):
     print('NSFW:')
     print(test(original_classifications[:, 1], other_classifications[:, 1]))
 
+
 if __name__ == "__main__":
     main(sys.argv)

image_utils.py

@@ -28,8 +28,8 @@ def load_image(image_path):
         imr.save(fh_im, format='JPEG')
         fh_im.seek(0)
 
-        image = (skimage.img_as_float(skimage.io.imread(fh_im, as_grey=False))
-                        .astype(np.float32))
+        image = (skimage.img_as_float(skimage.io.imread(
+            fh_im, as_grey=False)).astype(np.float32))
 
         H, W, _ = image.shape
         h, w = (224, 224)
@@ -39,7 +39,7 @@ def load_image(image_path):
         image = image[h_off:h_off + h, w_off:w_off + w, :]
 
         # RGB to BGR
-        image = image[:, :, :: -1]
+        image = image[:, :, ::-1]
 
         image = image.astype(np.float32, copy=False)
         image = image * 255.0
@@ -53,7 +53,8 @@ def load_image(image_path):
     return load_image
 
 
-def create_tensorflow_image_loader(session, expand_dims=True,
+def create_tensorflow_image_loader(session,
+                                   expand_dims=True,
                                    options=None,
                                    run_metadata=None):
     """Tensorflow image loader
@@ -76,13 +77,10 @@ def load_image(image_path):
 
         if expand_dims:
             image_batch = tf.expand_dims(image, axis=0)
-            return session.run(image_batch,
-                               options=options,
-                               run_metadata=run_metadata)
+            return session.run(
+                image_batch, options=options, run_metadata=run_metadata)
 
-        return session.run(image,
-                           options=options,
-                           run_metadata=run_metadata)
+        return session.run(image, options=options, run_metadata=run_metadata)
 
     return load_image
 
@@ -99,8 +97,8 @@ def decode_and_process(base64):
     # we have to do some preprocessing with map_fn, since functions like
     # decode_*, resize_images and crop_to_bounding_box do not support
     # processing of batches
-    image = tf.map_fn(decode_and_process, _input,
-                      back_prop=False, dtype=tf.float32)
+    image = tf.map_fn(
+        decode_and_process, _input, back_prop=False, dtype=tf.float32)
 
     return image
 
@@ -111,27 +109,34 @@ def __tf_jpeg_process(data):
     # The whole jpeg encode/decode dance is neccessary to generate a result
     # that matches the original model's (caffe) preprocessing
     # (as good as possible)
-    image = tf.image.decode_jpeg(data, channels=3,
-                                 fancy_upscaling=True,
-                                 dct_method="INTEGER_FAST")
+    image = tf.image.decode_jpeg(
+        data, channels=3, fancy_upscaling=True, dct_method="INTEGER_FAST")
 
     image = tf.image.convert_image_dtype(image, tf.float32, saturate=True)
-    image = tf.image.resize_images(image, (256, 256),
-                                   method=tf.image.ResizeMethod.BILINEAR,
-                                   align_corners=True)
+    image = tf.image.resize_images(
+        image, (256, 256),
+        method=tf.image.ResizeMethod.BILINEAR,
+        align_corners=True)
 
     image = tf.image.convert_image_dtype(image, tf.uint8, saturate=True)
 
-    image = tf.image.encode_jpeg(image, format='', quality=75,
-                                 progressive=False, optimize_size=False,
-                                 chroma_downsampling=True,
-                                 density_unit=None,
-                                 x_density=None, y_density=None,
-                                 xmp_metadata=None)
-
-    image = tf.image.decode_jpeg(image, channels=3,
-                                 fancy_upscaling=False,
-                                 dct_method="INTEGER_ACCURATE")
+    image = tf.image.encode_jpeg(
+        image,
+        format='',
+        quality=75,
+        progressive=False,
+        optimize_size=False,
+        chroma_downsampling=True,
+        density_unit=None,
+        x_density=None,
+        y_density=None,
+        xmp_metadata=None)
+
+    image = tf.image.decode_jpeg(
+        image,
+        channels=3,
+        fancy_upscaling=False,
+        dct_method="INTEGER_ACCURATE")
 
     image = tf.cast(image, dtype=tf.float32)