├── .gitignore ├── LICENSE ├── README.md ├── classify_nsfw.py ├── data └── open_nsfw-weights.npy ├── eval ├── batch_classify.py └── eval.py ├── image_utils.py ├── model.py └── tools ├── create_predict_request.py ├── export_graph.py ├── export_savedmodel.py └── export_tflite.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Custom 2 | logs/ 3 | *.csv 4 | .DS_Store 5 | datasets/ 6 | *.tar.gz 7 | models/ 8 | # Byte-compiled / optimized / DLL files 9 | __pycache__/ 10 | *.py[cod] 11 | *$py.class 12 | eval/dataset 13 | 14 | # C extensions 15 | *.so 16 | 17 | # Distribution / packaging 18 | .Python 19 | env/ 20 | build/ 21 | develop-eggs/ 22 | dist/ 23 | downloads/ 24 | eggs/ 25 | .eggs/ 26 | lib/ 27 | lib64/ 28 | parts/ 29 | sdist/ 30 | var/ 31 | *.egg-info/ 32 | .installed.cfg 33 | *.egg 34 | 35 | # PyInstaller 36 | # Usually these files are written by a python script from a template 37 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 38 | *.manifest 39 | *.spec 40 | 41 | # Installer logs 42 | pip-log.txt 43 | pip-delete-this-directory.txt 44 | 45 | # Unit test / coverage reports 46 | htmlcov/ 47 | .tox/ 48 | .coverage 49 | .coverage.* 50 | .cache 51 | nosetests.xml 52 | coverage.xml 53 | *,cover 54 | .hypothesis/ 55 | 56 | # Translations 57 | *.mo 58 | *.pot 59 | 60 | # Django stuff: 61 | *.log 62 | local_settings.py 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # IPython Notebook 78 | .ipynb_checkpoints 79 | 80 | # pyenv 81 | .python-version 82 | 83 | # celery beat schedule file 84 | celerybeat-schedule 85 | 86 | # dotenv 87 | .env 88 | 89 | # virtualenv 90 | venv/ 91 | ENV/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | 96 | # Rope project settings 97 | .ropeproject 98 | 99 | # We dont want images 100 | *.jpg 101 | *.png 102 | *.jpeg 103 | *.zip 104 | *.gz 105 | *.xz 106 | *.tar 107 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2016, Yahoo Inc. 2 | 3 | Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 4 | 5 | 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 6 | 7 | 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. 8 | 9 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 10 | 11 | BSD 2-Clause License 12 | 13 | Copyright (c) 2017, Marc Dietrichstein 14 | All rights reserved. 15 | 16 | Redistribution and use in source and binary forms, with or without 17 | modification, are permitted provided that the following conditions are met: 18 | 19 | * Redistributions of source code must retain the above copyright notice, this 20 | list of conditions and the following disclaimer. 21 | 22 | * Redistributions in binary form must reproduce the above copyright notice, 23 | this list of conditions and the following disclaimer in the documentation 24 | and/or other materials provided with the distribution. 25 | 26 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 27 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 28 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 29 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 30 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 31 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 32 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 33 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 34 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 35 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Tensorflow Implementation of Yahoo's Open NSFW Model 2 | 3 | This repository contains an implementation of [Yahoo's Open NSFW Classifier](https://github.com/yahoo/open_nsfw) rewritten in tensorflow. 4 | 5 | The original caffe weights have been extracted using [Caffe to TensorFlow](https://github.com/ethereon/caffe-tensorflow). You can find them at `data/open_nsfw-weights.npy`. 6 | 7 | ## Prerequisites 8 | 9 | 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`. 10 | 11 | ### Usage 12 | 13 | ``` 14 | > python classify_nsfw.py -m data/open_nsfw-weights.npy test.jpg 15 | 16 | Results for 'test.jpg' 17 | SFW score: 0.9355766177177429 18 | NSFW score: 0.06442338228225708 19 | ``` 20 | 21 | __Note:__ Currently only jpeg images are supported. 22 | 23 | `classify_nsfw.py` accepts some optional parameters you may want to play around with: 24 | 25 | ``` 26 | usage: classify_nsfw.py [-h] -m MODEL_WEIGHTS [-l {yahoo,tensorflow}] 27 | [-t {tensor,base64_jpeg}] 28 | input_jpeg_file 29 | 30 | positional arguments: 31 | input_file Path to the input image. Only jpeg images are 32 | supported. 33 | 34 | optional arguments: 35 | -h, --help show this help message and exit 36 | -m MODEL_WEIGHTS, --model_weights MODEL_WEIGHTS 37 | Path to trained model weights file 38 | -l {yahoo,tensorflow}, --image_loader {yahoo,tensorflow} 39 | image loading mechanism 40 | -i {tensor,base64_jpeg}, --input_type {tensor,base64_jpeg} 41 | input type 42 | ``` 43 | 44 | __-l/--image-loader__ 45 | 46 | The classification tool supports two different image loading mechanisms. 47 | 48 | * `yahoo` (default) replicates yahoo's original image loading and preprocessing. Use this option if you want the same results as with the original implementation 49 | * `tensorflow` is an image loader which uses tensorflow exclusively (no dependencies on `PIL`, `skimage`, etc.). Tries to replicate the image loading mechanism used by the original caffe implementation, differs a bit though due to different jpeg and resizing implementations. See [this issue](https://github.com/mdietrichstein/tensorflow-open_nsfw/issues/2#issuecomment-346125345) for details. 50 | 51 | __Note:__ Classification results may vary depending on the selected image loader! 52 | 53 | __-i/--input_type__ 54 | 55 | Determines if the model internally uses a float tensor (`tensor` - `[None, 224, 224, 3]` - default) or a base64 encoded string tensor (`base64_jpeg` - `[None, ]`) as input. If `base64_jpeg` is used, then the `tensorflow` image loader will be used, regardless of the _-l/--image-loader_ argument. 56 | 57 | 58 | ### Tools 59 | 60 | The `tools` folder contains some utility scripts to test the model. 61 | 62 | __create_predict_request.py__ 63 | 64 | Takes an input image and generates a json file suitable for prediction requests to a Open NSFW Model deployed with [Google Cloud ML Engine](https://cloud.google.com/ml-engine/docs/concepts/prediction-overview) (`gcloud ml-engine predict`) or [tensorflow-serving](https://www.tensorflow.org/serving/). 65 | 66 | 67 | __export_savedmodel.py__ 68 | 69 | Exports the model using the tensorflow serving export api (`SavedModel`). The export can be used to deploy the model on [Google Cloud ML Engine](https://cloud.google.com/ml-engine/docs/concepts/prediction-overview), [Tensorflow Serving]() or on mobile (haven't tried that one yet). 70 | 71 | __export_tflite.py__ 72 | 73 | Exports the model in [TFLite format](https://www.tensorflow.org/lite/). Use this one if you want to run inference on mobile or IoT devices. Please note that the `base64_jpeg` input type does not work with TFLite since the standard runtime lacks a number of required tensorflow operations. 74 | 75 | __export_graph.py__ 76 | 77 | Exports the tensorflow graph and checkpoint. Freezes and optimizes the graph per default for improved inference and deployment usage (e.g. Android, iOS, etc.). Import the graph with `tf.import_graph_def`. -------------------------------------------------------------------------------- /classify_nsfw.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import sys 3 | import argparse 4 | import tensorflow as tf 5 | 6 | from model import OpenNsfwModel, InputType 7 | from image_utils import create_tensorflow_image_loader 8 | from image_utils import create_yahoo_image_loader 9 | 10 | import numpy as np 11 | 12 | 13 | IMAGE_LOADER_TENSORFLOW = "tensorflow" 14 | IMAGE_LOADER_YAHOO = "yahoo" 15 | 16 | 17 | def main(argv): 18 | parser = argparse.ArgumentParser() 19 | 20 | parser.add_argument("input_file", help="Path to the input image.\ 21 | Only jpeg images are supported.") 22 | 23 | parser.add_argument("-m", "--model_weights", required=True, 24 | help="Path to trained model weights file") 25 | 26 | parser.add_argument("-l", "--image_loader", 27 | default=IMAGE_LOADER_YAHOO, 28 | help="image loading mechanism", 29 | choices=[IMAGE_LOADER_YAHOO, IMAGE_LOADER_TENSORFLOW]) 30 | 31 | parser.add_argument("-i", "--input_type", 32 | default=InputType.TENSOR.name.lower(), 33 | help="input type", 34 | choices=[InputType.TENSOR.name.lower(), 35 | InputType.BASE64_JPEG.name.lower()]) 36 | 37 | args = parser.parse_args() 38 | 39 | model = OpenNsfwModel() 40 | 41 | with tf.Session() as sess: 42 | 43 | input_type = InputType[args.input_type.upper()] 44 | model.build(weights_path=args.model_weights, input_type=input_type) 45 | 46 | fn_load_image = None 47 | 48 | if input_type == InputType.TENSOR: 49 | if args.image_loader == IMAGE_LOADER_TENSORFLOW: 50 | fn_load_image = create_tensorflow_image_loader(tf.Session(graph=tf.Graph())) 51 | else: 52 | fn_load_image = create_yahoo_image_loader() 53 | elif input_type == InputType.BASE64_JPEG: 54 | import base64 55 | fn_load_image = lambda filename: np.array([base64.urlsafe_b64encode(open(filename, "rb").read())]) 56 | 57 | sess.run(tf.global_variables_initializer()) 58 | 59 | image = fn_load_image(args.input_file) 60 | 61 | predictions = \ 62 | sess.run(model.predictions, 63 | feed_dict={model.input: image}) 64 | 65 | print("Results for '{}'".format(args.input_file)) 66 | print("\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0])) 67 | 68 | if __name__ == "__main__": 69 | main(sys.argv) 70 | -------------------------------------------------------------------------------- /data/open_nsfw-weights.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdietrichstein/tensorflow-open_nsfw/ead9f4d1748e8bc80ab14bf0a36f696a5fe4109d/data/open_nsfw-weights.npy -------------------------------------------------------------------------------- /eval/batch_classify.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | import sys 4 | 5 | sys.path.append((os.path.normpath( 6 | os.path.join(os.path.dirname(os.path.realpath(__file__)), 7 | '..')))) 8 | 9 | import argparse 10 | import glob 11 | import tensorflow as tf 12 | from tqdm import tqdm 13 | 14 | from model import OpenNsfwModel, InputType 15 | from image_utils import create_tensorflow_image_loader 16 | from image_utils import create_yahoo_image_loader 17 | 18 | 19 | IMAGE_LOADER_TENSORFLOW = "tensorflow" 20 | IMAGE_LOADER_YAHOO = "yahoo" 21 | 22 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 23 | tf.logging.set_verbosity(tf.logging.ERROR) 24 | 25 | 26 | def create_batch_iterator(filenames, batch_size, fn_load_image): 27 | for i in range(0, len(filenames), batch_size): 28 | yield list(map(fn_load_image, filenames[i:i+batch_size])) 29 | 30 | 31 | def create_tf_batch_iterator(filenames, batch_size): 32 | for i in range(0, len(filenames), batch_size): 33 | with tf.Session(graph=tf.Graph()) as session: 34 | fn_load_image = create_tensorflow_image_loader(session, 35 | expand_dims=False) 36 | 37 | yield list(map(fn_load_image, filenames[i:i+batch_size])) 38 | 39 | 40 | def main(argv): 41 | parser = argparse.ArgumentParser() 42 | 43 | parser.add_argument("-s", "--source", required=True, 44 | help="Folder containing the images to classify") 45 | 46 | parser.add_argument("-o", "--output_file", required=True, 47 | help="Output file path") 48 | 49 | parser.add_argument("-m", "--model_weights", required=True, 50 | help="Path to trained model weights file") 51 | 52 | parser.add_argument("-b", "--batch_size", help="Number of images to \ 53 | classify simultaneously.", type=int, default=64) 54 | 55 | parser.add_argument("-l", "--image_loader", 56 | default=IMAGE_LOADER_YAHOO, 57 | help="image loading mechanism", 58 | choices=[IMAGE_LOADER_YAHOO, IMAGE_LOADER_TENSORFLOW]) 59 | 60 | args = parser.parse_args() 61 | batch_size = args.batch_size 62 | output_file = args.output_file 63 | 64 | input_type = InputType.TENSOR 65 | model = OpenNsfwModel() 66 | 67 | filenames = glob.glob(args.source + "/*.jpg") 68 | num_files = len(filenames) 69 | 70 | num_batches = int(num_files / batch_size) 71 | 72 | print("Found", num_files, " files") 73 | print("Split into", num_batches, " batches") 74 | 75 | config = tf.ConfigProto() 76 | config.gpu_options.allow_growth = True 77 | 78 | batch_iterator = None 79 | 80 | if args.image_loader == IMAGE_LOADER_TENSORFLOW: 81 | batch_iterator = create_tf_batch_iterator(filenames, batch_size) 82 | else: 83 | fn_load_image = create_yahoo_image_loader(expand_dims=False) 84 | batch_iterator = create_batch_iterator(filenames, batch_size, 85 | fn_load_image) 86 | 87 | with tf.Session(graph=tf.Graph(), config=config) as session: 88 | model.build(weights_path=args.model_weights, 89 | input_type=input_type) 90 | 91 | session.run(tf.global_variables_initializer()) 92 | 93 | with tqdm(total=num_files) as progress_bar: 94 | with open(output_file, 'w') as o: 95 | o.write('File\tSFW Score\tNSFW Score\n') 96 | 97 | for batch_num, images in enumerate(batch_iterator): 98 | predictions = \ 99 | session.run(model.predictions, 100 | feed_dict={model.input: images}) 101 | 102 | fi = (batch_num * batch_size) 103 | for i, prediction in enumerate(predictions): 104 | filename = os.path.basename(filenames[fi + i]) 105 | o.write('{}\t{}\t{}\n'.format(filename, 106 | prediction[0], 107 | prediction[1])) 108 | 109 | progress_bar.update(len(images)) 110 | 111 | if __name__ == "__main__": 112 | main(sys.argv) 113 | -------------------------------------------------------------------------------- /eval/eval.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import operator 3 | import argparse 4 | import numpy as np 5 | from scipy import stats 6 | 7 | 8 | def load_classifications(filename): 9 | is_first = True 10 | 11 | results = {} 12 | 13 | with open(filename, 'r') as f: 14 | for line in f: 15 | if is_first: 16 | is_first = False 17 | continue 18 | 19 | parts = line.split('\t') 20 | 21 | filename = parts[0] 22 | sfw_score = float(parts[1]) 23 | nsfw_score = float(parts[2]) 24 | 25 | results[filename] = (sfw_score, nsfw_score) 26 | 27 | return results 28 | 29 | 30 | def classification_matrix(classifications): 31 | results = np.zeros(shape=(len(classifications), 2)) 32 | 33 | for i, classification in enumerate(classifications): 34 | results[i] = np.array(classification[1]) 35 | 36 | return results 37 | 38 | 39 | def test(first, second): 40 | delta = np.abs(first - second) 41 | 42 | result = { 43 | 'min': np.amin(delta), 44 | 'max': np.amax(delta), 45 | 'median': np.median(delta), 46 | 'mean': np.mean(delta), 47 | 'std': np.std(delta), 48 | 'var': np.var(delta), 49 | 't-prob': stats.ttest_ind(first, second, equal_var=True)[1] 50 | } 51 | 52 | return result 53 | 54 | 55 | def main(argv): 56 | parser = argparse.ArgumentParser() 57 | 58 | parser.add_argument("original", 59 | help="File containing base classifications") 60 | 61 | parser.add_argument("other", 62 | help="File containing classifications to compare to\ 63 | base results") 64 | 65 | args = parser.parse_args() 66 | filename_original = args.original 67 | filename_other = args.other 68 | 69 | original = load_classifications(filename_original) 70 | other = load_classifications(filename_other) 71 | 72 | len(original) == len(other) 73 | 74 | original = sorted(original.items(), key=operator.itemgetter(0)) 75 | other = sorted(other.items(), key=operator.itemgetter(0)) 76 | 77 | print("Found", len(original), "entries") 78 | 79 | original_classifications = classification_matrix(original) 80 | other_classifications = classification_matrix(other) 81 | 82 | print('SFW:') 83 | print(test(original_classifications[:, 0], other_classifications[:, 0])) 84 | 85 | print() 86 | print('NSFW:') 87 | print(test(original_classifications[:, 1], other_classifications[:, 1])) 88 | 89 | if __name__ == "__main__": 90 | main(sys.argv) 91 | -------------------------------------------------------------------------------- /image_utils.py: -------------------------------------------------------------------------------- 1 | VGG_MEAN = [104, 117, 123] 2 | 3 | 4 | def create_yahoo_image_loader(expand_dims=True): 5 | """Yahoo open_nsfw image loading mechanism 6 | 7 | Approximation of the image loading mechanism defined in 8 | https://github.com/yahoo/open_nsfw/blob/79f77bcd45076b000df71742a59d726aa4a36ad1/classify_nsfw.py#L40 9 | """ 10 | import numpy as np 11 | import skimage 12 | import skimage.io 13 | from PIL import Image 14 | from io import BytesIO 15 | 16 | def load_image(image_path): 17 | pimg = open(image_path, 'rb').read() 18 | 19 | img_data = pimg 20 | im = Image.open(BytesIO(img_data)) 21 | 22 | if im.mode != "RGB": 23 | im = im.convert('RGB') 24 | 25 | imr = im.resize((256, 256), resample=Image.BILINEAR) 26 | 27 | fh_im = BytesIO() 28 | imr.save(fh_im, format='JPEG') 29 | fh_im.seek(0) 30 | 31 | image = (skimage.img_as_float(skimage.io.imread(fh_im, as_grey=False)) 32 | .astype(np.float32)) 33 | 34 | H, W, _ = image.shape 35 | h, w = (224, 224) 36 | 37 | h_off = max((H - h) // 2, 0) 38 | w_off = max((W - w) // 2, 0) 39 | image = image[h_off:h_off + h, w_off:w_off + w, :] 40 | 41 | # RGB to BGR 42 | image = image[:, :, :: -1] 43 | 44 | image = image.astype(np.float32, copy=False) 45 | image = image * 255.0 46 | image -= np.array(VGG_MEAN, dtype=np.float32) 47 | 48 | if expand_dims: 49 | image = np.expand_dims(image, axis=0) 50 | 51 | return image 52 | 53 | return load_image 54 | 55 | 56 | def create_tensorflow_image_loader(session, expand_dims=True, 57 | options=None, 58 | run_metadata=None): 59 | """Tensorflow image loader 60 | 61 | Results seem to deviate quite a bit from yahoo image loader due to 62 | different jpeg encoders/decoders and different image resize 63 | implementations between PIL, skimage and tensorflow 64 | 65 | Only supports jpeg images. 66 | 67 | Relevant tensorflow issues: 68 | * https://github.com/tensorflow/tensorflow/issues/6720 69 | * https://github.com/tensorflow/tensorflow/issues/12753 70 | """ 71 | import tensorflow as tf 72 | 73 | def load_image(image_path): 74 | image = tf.read_file(image_path) 75 | image = __tf_jpeg_process(image) 76 | 77 | if expand_dims: 78 | image_batch = tf.expand_dims(image, axis=0) 79 | return session.run(image_batch, 80 | options=options, 81 | run_metadata=run_metadata) 82 | 83 | return session.run(image, 84 | options=options, 85 | run_metadata=run_metadata) 86 | 87 | return load_image 88 | 89 | 90 | def load_base64_tensor(_input): 91 | import tensorflow as tf 92 | 93 | def decode_and_process(base64): 94 | _bytes = tf.decode_base64(base64) 95 | _image = __tf_jpeg_process(_bytes) 96 | 97 | return _image 98 | 99 | # we have to do some preprocessing with map_fn, since functions like 100 | # decode_*, resize_images and crop_to_bounding_box do not support 101 | # processing of batches 102 | image = tf.map_fn(decode_and_process, _input, 103 | back_prop=False, dtype=tf.float32) 104 | 105 | return image 106 | 107 | 108 | def __tf_jpeg_process(data): 109 | import tensorflow as tf 110 | 111 | # The whole jpeg encode/decode dance is neccessary to generate a result 112 | # that matches the original model's (caffe) preprocessing 113 | # (as good as possible) 114 | image = tf.image.decode_jpeg(data, channels=3, 115 | fancy_upscaling=True, 116 | dct_method="INTEGER_FAST") 117 | 118 | image = tf.image.convert_image_dtype(image, tf.float32, saturate=True) 119 | image = tf.image.resize_images(image, (256, 256), 120 | method=tf.image.ResizeMethod.BILINEAR, 121 | align_corners=True) 122 | 123 | image = tf.image.convert_image_dtype(image, tf.uint8, saturate=True) 124 | 125 | image = tf.image.encode_jpeg(image, format='', quality=75, 126 | progressive=False, optimize_size=False, 127 | chroma_downsampling=True, 128 | density_unit=None, 129 | x_density=None, y_density=None, 130 | xmp_metadata=None) 131 | 132 | image = tf.image.decode_jpeg(image, channels=3, 133 | fancy_upscaling=False, 134 | dct_method="INTEGER_ACCURATE") 135 | 136 | image = tf.cast(image, dtype=tf.float32) 137 | 138 | image = tf.image.crop_to_bounding_box(image, 16, 16, 224, 224) 139 | 140 | image = tf.reverse(image, axis=[2]) 141 | image -= VGG_MEAN 142 | 143 | return image 144 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import math 2 | import numpy as np 3 | import tensorflow as tf 4 | from enum import Enum, unique 5 | 6 | 7 | @unique 8 | class InputType(Enum): 9 | TENSOR = 1 10 | BASE64_JPEG = 2 11 | 12 | 13 | class OpenNsfwModel: 14 | """Tensorflow implementation of Yahoo's Open NSFW Model 15 | 16 | Original implementation: 17 | https://github.com/yahoo/open_nsfw 18 | 19 | Weights have been converted using caffe-tensorflow: 20 | https://github.com/ethereon/caffe-tensorflow 21 | """ 22 | 23 | def __init__(self): 24 | self.weights = {} 25 | self.bn_epsilon = 1e-5 # Default used by Caffe 26 | 27 | def build(self, weights_path="open_nsfw-weights.npy", 28 | input_type=InputType.TENSOR): 29 | 30 | self.weights = np.load(weights_path, encoding="latin1").item() 31 | self.input_tensor = None 32 | 33 | if input_type == InputType.TENSOR: 34 | self.input = tf.placeholder(tf.float32, 35 | shape=[None, 224, 224, 3], 36 | name="input") 37 | self.input_tensor = self.input 38 | elif input_type == InputType.BASE64_JPEG: 39 | from image_utils import load_base64_tensor 40 | 41 | self.input = tf.placeholder(tf.string, shape=(None,), name="input") 42 | self.input_tensor = load_base64_tensor(self.input) 43 | else: 44 | raise ValueError("invalid input type '{}'".format(input_type)) 45 | 46 | x = self.input_tensor 47 | 48 | x = tf.pad(x, [[0, 0], [3, 3], [3, 3], [0, 0]], 'CONSTANT') 49 | x = self.__conv2d("conv_1", x, filter_depth=64, 50 | kernel_size=7, stride=2, padding='valid') 51 | 52 | x = self.__batch_norm("bn_1", x) 53 | x = tf.nn.relu(x) 54 | 55 | x = tf.layers.max_pooling2d(x, pool_size=3, strides=2, padding='same') 56 | 57 | x = self.__conv_block(stage=0, block=0, inputs=x, 58 | filter_depths=[32, 32, 128], 59 | kernel_size=3, stride=1) 60 | 61 | x = self.__identity_block(stage=0, block=1, inputs=x, 62 | filter_depths=[32, 32, 128], kernel_size=3) 63 | x = self.__identity_block(stage=0, block=2, inputs=x, 64 | filter_depths=[32, 32, 128], kernel_size=3) 65 | 66 | x = self.__conv_block(stage=1, block=0, inputs=x, 67 | filter_depths=[64, 64, 256], 68 | kernel_size=3, stride=2) 69 | x = self.__identity_block(stage=1, block=1, inputs=x, 70 | filter_depths=[64, 64, 256], kernel_size=3) 71 | x = self.__identity_block(stage=1, block=2, inputs=x, 72 | filter_depths=[64, 64, 256], kernel_size=3) 73 | x = self.__identity_block(stage=1, block=3, inputs=x, 74 | filter_depths=[64, 64, 256], kernel_size=3) 75 | 76 | x = self.__conv_block(stage=2, block=0, inputs=x, 77 | filter_depths=[128, 128, 512], 78 | kernel_size=3, stride=2) 79 | x = self.__identity_block(stage=2, block=1, inputs=x, 80 | filter_depths=[128, 128, 512], kernel_size=3) 81 | x = self.__identity_block(stage=2, block=2, inputs=x, 82 | filter_depths=[128, 128, 512], kernel_size=3) 83 | x = self.__identity_block(stage=2, block=3, inputs=x, 84 | filter_depths=[128, 128, 512], kernel_size=3) 85 | x = self.__identity_block(stage=2, block=4, inputs=x, 86 | filter_depths=[128, 128, 512], kernel_size=3) 87 | x = self.__identity_block(stage=2, block=5, inputs=x, 88 | filter_depths=[128, 128, 512], kernel_size=3) 89 | 90 | x = self.__conv_block(stage=3, block=0, inputs=x, 91 | filter_depths=[256, 256, 1024], kernel_size=3, 92 | stride=2) 93 | x = self.__identity_block(stage=3, block=1, inputs=x, 94 | filter_depths=[256, 256, 1024], 95 | kernel_size=3) 96 | x = self.__identity_block(stage=3, block=2, inputs=x, 97 | filter_depths=[256, 256, 1024], 98 | kernel_size=3) 99 | 100 | x = tf.layers.average_pooling2d(x, pool_size=7, strides=1, 101 | padding="valid", name="pool") 102 | 103 | x = tf.reshape(x, shape=(-1, 1024)) 104 | 105 | self.logits = self.__fully_connected(name="fc_nsfw", 106 | inputs=x, num_outputs=2) 107 | self.predictions = tf.nn.softmax(self.logits, name="predictions") 108 | 109 | """Get weights for layer with given name 110 | """ 111 | def __get_weights(self, layer_name, field_name): 112 | if not layer_name in self.weights: 113 | raise ValueError("No weights for layer named '{}' found" 114 | .format(layer_name)) 115 | 116 | w = self.weights[layer_name] 117 | if not field_name in w: 118 | raise (ValueError("No entry for field '{}' in layer named '{}'" 119 | .format(field_name, layer_name))) 120 | 121 | return w[field_name] 122 | 123 | """Layer creation and weight initialization 124 | """ 125 | def __fully_connected(self, name, inputs, num_outputs): 126 | return tf.layers.dense( 127 | inputs=inputs, units=num_outputs, name=name, 128 | kernel_initializer=tf.constant_initializer( 129 | self.__get_weights(name, "weights"), dtype=tf.float32), 130 | bias_initializer=tf.constant_initializer( 131 | self.__get_weights(name, "biases"), dtype=tf.float32)) 132 | 133 | def __conv2d(self, name, inputs, filter_depth, kernel_size, stride=1, 134 | padding="same", trainable=False): 135 | 136 | if padding.lower() == 'same' and kernel_size > 1: 137 | if kernel_size > 1: 138 | oh = inputs.get_shape().as_list()[1] 139 | h = inputs.get_shape().as_list()[1] 140 | 141 | p = int(math.floor(((oh - 1) * stride + kernel_size - h)//2)) 142 | 143 | inputs = tf.pad(inputs, 144 | [[0, 0], [p, p], [p, p], [0, 0]], 145 | 'CONSTANT') 146 | else: 147 | raise Exception('unsupported kernel size for padding: "{}"' 148 | .format(kernel_size)) 149 | 150 | return tf.layers.conv2d( 151 | inputs, filter_depth, 152 | kernel_size=(kernel_size, kernel_size), 153 | strides=(stride, stride), padding='valid', 154 | activation=None, trainable=trainable, name=name, 155 | kernel_initializer=tf.constant_initializer( 156 | self.__get_weights(name, "weights"), dtype=tf.float32), 157 | bias_initializer=tf.constant_initializer( 158 | self.__get_weights(name, "biases"), dtype=tf.float32)) 159 | 160 | def __batch_norm(self, name, inputs, training=False): 161 | return tf.layers.batch_normalization( 162 | inputs, training=training, epsilon=self.bn_epsilon, 163 | gamma_initializer=tf.constant_initializer( 164 | self.__get_weights(name, "scale"), dtype=tf.float32), 165 | beta_initializer=tf.constant_initializer( 166 | self.__get_weights(name, "offset"), dtype=tf.float32), 167 | moving_mean_initializer=tf.constant_initializer( 168 | self.__get_weights(name, "mean"), dtype=tf.float32), 169 | moving_variance_initializer=tf.constant_initializer( 170 | self.__get_weights(name, "variance"), dtype=tf.float32), 171 | name=name) 172 | 173 | """ResNet blocks 174 | """ 175 | def __conv_block(self, stage, block, inputs, filter_depths, 176 | kernel_size=3, stride=2): 177 | filter_depth1, filter_depth2, filter_depth3 = filter_depths 178 | 179 | conv_name_base = "conv_stage{}_block{}_branch".format(stage, block) 180 | bn_name_base = "bn_stage{}_block{}_branch".format(stage, block) 181 | shortcut_name_post = "_stage{}_block{}_proj_shortcut" \ 182 | .format(stage, block) 183 | 184 | shortcut = self.__conv2d( 185 | name="conv{}".format(shortcut_name_post), stride=stride, 186 | inputs=inputs, filter_depth=filter_depth3, kernel_size=1, 187 | padding="same" 188 | ) 189 | 190 | shortcut = self.__batch_norm("bn{}".format(shortcut_name_post), 191 | shortcut) 192 | 193 | x = self.__conv2d( 194 | name="{}2a".format(conv_name_base), 195 | inputs=inputs, filter_depth=filter_depth1, kernel_size=1, 196 | stride=stride, padding="same", 197 | ) 198 | x = self.__batch_norm("{}2a".format(bn_name_base), x) 199 | x = tf.nn.relu(x) 200 | 201 | x = self.__conv2d( 202 | name="{}2b".format(conv_name_base), 203 | inputs=x, filter_depth=filter_depth2, kernel_size=kernel_size, 204 | padding="same", stride=1 205 | ) 206 | x = self.__batch_norm("{}2b".format(bn_name_base), x) 207 | x = tf.nn.relu(x) 208 | 209 | x = self.__conv2d( 210 | name="{}2c".format(conv_name_base), 211 | inputs=x, filter_depth=filter_depth3, kernel_size=1, 212 | padding="same", stride=1 213 | ) 214 | x = self.__batch_norm("{}2c".format(bn_name_base), x) 215 | 216 | x = tf.add(x, shortcut) 217 | 218 | return tf.nn.relu(x) 219 | 220 | def __identity_block(self, stage, block, inputs, 221 | filter_depths, kernel_size): 222 | filter_depth1, filter_depth2, filter_depth3 = filter_depths 223 | conv_name_base = "conv_stage{}_block{}_branch".format(stage, block) 224 | bn_name_base = "bn_stage{}_block{}_branch".format(stage, block) 225 | 226 | x = self.__conv2d( 227 | name="{}2a".format(conv_name_base), 228 | inputs=inputs, filter_depth=filter_depth1, kernel_size=1, 229 | stride=1, padding="same", 230 | ) 231 | 232 | x = self.__batch_norm("{}2a".format(bn_name_base), x) 233 | x = tf.nn.relu(x) 234 | 235 | x = self.__conv2d( 236 | name="{}2b".format(conv_name_base), 237 | inputs=x, filter_depth=filter_depth2, kernel_size=kernel_size, 238 | padding="same", stride=1 239 | ) 240 | x = self.__batch_norm("{}2b".format(bn_name_base), x) 241 | x = tf.nn.relu(x) 242 | 243 | x = self.__conv2d( 244 | name="{}2c".format(conv_name_base), 245 | inputs=x, filter_depth=filter_depth3, kernel_size=1, 246 | padding="same", stride=1 247 | ) 248 | x = self.__batch_norm("{}2c".format(bn_name_base), x) 249 | 250 | x = tf.add(x, inputs) 251 | 252 | return tf.nn.relu(x) 253 | -------------------------------------------------------------------------------- /tools/create_predict_request.py: -------------------------------------------------------------------------------- 1 | import base64 2 | import json 3 | import argparse 4 | 5 | import numpy as np 6 | import tensorflow as tf 7 | from tensorflow.python.saved_model.signature_constants import PREDICT_INPUTS 8 | 9 | import os 10 | import sys 11 | 12 | sys.path.append((os.path.normpath( 13 | os.path.join(os.path.dirname(os.path.realpath(__file__)), 14 | '..')))) 15 | 16 | from image_utils import create_tensorflow_image_loader 17 | from image_utils import create_yahoo_image_loader 18 | from model import InputType 19 | 20 | IMAGE_LOADER_TENSORFLOW = "tensorflow" 21 | IMAGE_LOADER_YAHOO = "yahoo" 22 | 23 | # Thanks to https://stackoverflow.com/a/47626762 24 | class NumpyEncoder(json.JSONEncoder): 25 | def default(self, obj): 26 | if isinstance(obj, np.ndarray): 27 | return obj.tolist() 28 | return json.JSONEncoder.default(self, obj) 29 | 30 | """Generates a json prediction request suitable for consumption by a model 31 | generated with 'export-model.py' and deployed on either ml-engine or tensorflow-serving 32 | """ 33 | if __name__ == "__main__": 34 | parser = argparse.ArgumentParser() 35 | parser.add_argument("input_file", help="Path to the input image file") 36 | 37 | parser.add_argument("-i", "--input_type", required=True, 38 | default=InputType.TENSOR.name.lower(), 39 | help="Input type", 40 | choices=[InputType.TENSOR.name.lower(), 41 | InputType.BASE64_JPEG.name.lower()]) 42 | 43 | parser.add_argument("-l", "--image_loader", required=False, 44 | default=IMAGE_LOADER_YAHOO, 45 | help="Image loading mechanism. Only relevant when using input_type 'tensor'", 46 | choices=[IMAGE_LOADER_YAHOO, IMAGE_LOADER_TENSORFLOW]) 47 | 48 | parser.add_argument("-t", "--target", required=True, 49 | choices=['ml-engine', 'tf-serving'], 50 | help="Create json request for ml-engine or tensorflow-serving") 51 | 52 | args = parser.parse_args() 53 | target = args.target 54 | 55 | input_type = InputType[args.input_type.upper()] 56 | 57 | image_data = None 58 | 59 | if input_type == InputType.TENSOR: 60 | fn_load_image = None 61 | 62 | if args.image_loader == IMAGE_LOADER_TENSORFLOW: 63 | with tf.Session() as sess: 64 | fn_load_image = create_tensorflow_image_loader(sess) 65 | sess.run(tf.global_variables_initializer()) 66 | image_data = fn_load_image(args.input_file)[0] 67 | else: 68 | image_data = create_yahoo_image_loader(tf.Session(graph=tf.Graph()))(args.input_file)[0] 69 | elif input_type == InputType.BASE64_JPEG: 70 | import base64 71 | image_data = base64.urlsafe_b64encode(open(args.input_file, "rb").read()).decode("ascii") 72 | 73 | if target == "ml-engine": 74 | print(json.dumps({PREDICT_INPUTS: image_data}, cls=NumpyEncoder)) 75 | elif target == "tf-serving": 76 | print(json.dumps({"instances": [image_data]}, cls=NumpyEncoder)) 77 | -------------------------------------------------------------------------------- /tools/export_graph.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import argparse 4 | 5 | import tensorflow as tf 6 | from tensorflow.python.tools import freeze_graph 7 | from tensorflow.python.tools import optimize_for_inference_lib 8 | 9 | sys.path.append((os.path.normpath( 10 | os.path.join(os.path.dirname(os.path.realpath(__file__)), 11 | '..')))) 12 | 13 | from model import OpenNsfwModel, InputType 14 | 15 | """Exports the graph so it can be imported via import_graph_def 16 | 17 | The exported model takes an base64 encoded string tensor as input 18 | """ 19 | if __name__ == "__main__": 20 | parser = argparse.ArgumentParser() 21 | 22 | parser.add_argument("target", help="output directory") 23 | 24 | parser.add_argument("-m", "--model_weights", required=True, 25 | help="Path to trained model weights file") 26 | 27 | parser.add_argument("-i", "--input_type", required=True, 28 | default=InputType.TENSOR.name.lower(), 29 | help="Input type", 30 | choices=[InputType.TENSOR.name.lower(), 31 | InputType.BASE64_JPEG.name.lower()]) 32 | 33 | parser.add_argument("-o", "--optimize", action='store_true', 34 | default=False, 35 | help="Optimize graph for inference") 36 | 37 | parser.add_argument("-f", "--freeze", action='store_true', 38 | required=False, default=False, 39 | help="Freeze graph: convert variables to ops") 40 | 41 | parser.add_argument("-t", "--text", action='store_true', 42 | required=False, default=False, 43 | help="Write graph as binary (.pb) or text (pbtext)") 44 | 45 | args = parser.parse_args() 46 | 47 | model = OpenNsfwModel() 48 | 49 | export_base_path = args.target 50 | do_freeze = args.freeze 51 | do_optimize = args.optimize 52 | as_binary = not args.text 53 | input_type = InputType[args.input_type.upper()] 54 | 55 | input_node_name = 'input' 56 | output_node_name = 'predictions' 57 | 58 | base_name = 'open_nsfw' 59 | 60 | checkpoint_path = os.path.join(export_base_path, base_name + '.ckpt') 61 | 62 | if as_binary: 63 | graph_name = base_name + '.pb' 64 | else: 65 | graph_name = base_name + '.pbtxt' 66 | 67 | graph_path = os.path.join(export_base_path, graph_name) 68 | frozen_graph_path = os.path.join(export_base_path, 69 | 'frozen_' + graph_name) 70 | optimized_graph_path = os.path.join(export_base_path, 71 | 'optimized_' + graph_name) 72 | 73 | with tf.Session() as sess: 74 | model.build(weights_path=args.model_weights, 75 | input_type=input_type) 76 | 77 | sess.run(tf.global_variables_initializer()) 78 | 79 | saver = tf.train.Saver() 80 | saver.save(sess, save_path=checkpoint_path) 81 | 82 | print('Checkpoint exported to {}'.format(checkpoint_path)) 83 | 84 | tf.train.write_graph(sess.graph_def, export_base_path, graph_name, 85 | as_text=not as_binary) 86 | 87 | print('Graph exported to {}'.format(graph_path)) 88 | 89 | if do_freeze: 90 | print('Freezing graph...') 91 | freeze_graph.freeze_graph( 92 | input_graph=graph_path, input_saver='', 93 | input_binary=as_binary, input_checkpoint=checkpoint_path, 94 | output_node_names=output_node_name, 95 | restore_op_name='save/restore_all', 96 | filename_tensor_name='save/Const:0', 97 | output_graph=frozen_graph_path, clear_devices=True, 98 | initializer_nodes='') 99 | 100 | print('Frozen graph exported to {}'.format(frozen_graph_path)) 101 | 102 | graph_path = frozen_graph_path 103 | 104 | if do_optimize: 105 | print('Optimizing graph...') 106 | input_graph_def = tf.GraphDef() 107 | 108 | with tf.gfile.Open(graph_path, 'rb') as f: 109 | data = f.read() 110 | input_graph_def.ParseFromString(data) 111 | 112 | output_graph_def =\ 113 | optimize_for_inference_lib.optimize_for_inference( 114 | input_graph_def, 115 | [input_node_name], 116 | [output_node_name], 117 | tf.float32.as_datatype_enum) 118 | 119 | f = tf.gfile.FastGFile(optimized_graph_path, 'wb') 120 | f.write(output_graph_def.SerializeToString()) 121 | 122 | print('Optimized graph exported to {}' 123 | .format(optimized_graph_path)) 124 | -------------------------------------------------------------------------------- /tools/export_savedmodel.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import argparse 4 | 5 | import tensorflow as tf 6 | from tensorflow.python.saved_model import builder as saved_model_builder 7 | from tensorflow.python.saved_model.signature_def_utils\ 8 | import predict_signature_def 9 | 10 | from tensorflow.python.saved_model.tag_constants import SERVING 11 | from tensorflow.python.saved_model.signature_constants\ 12 | import DEFAULT_SERVING_SIGNATURE_DEF_KEY 13 | 14 | from tensorflow.python.saved_model.signature_constants import PREDICT_INPUTS 15 | from tensorflow.python.saved_model.signature_constants import PREDICT_OUTPUTS 16 | 17 | sys.path.append((os.path.normpath( 18 | os.path.join(os.path.dirname(os.path.realpath(__file__)), 19 | '..')))) 20 | 21 | from model import OpenNsfwModel, InputType 22 | 23 | """Builds a SavedModel which can be used for deployment with 24 | gcloud ml-engine, tensorflow-serving, ... 25 | """ 26 | if __name__ == "__main__": 27 | parser = argparse.ArgumentParser() 28 | 29 | parser.add_argument("target", help="output directory") 30 | 31 | parser.add_argument("-i", "--input_type", required=True, 32 | default=InputType.TENSOR.name.lower(), 33 | help="Input type", 34 | choices=[InputType.TENSOR.name.lower(), 35 | InputType.BASE64_JPEG.name.lower()]) 36 | 37 | parser.add_argument("-v", "--export_version", 38 | help="export model version", 39 | default="1") 40 | 41 | parser.add_argument("-m", "--model_weights", required=True, 42 | help="Path to trained model weights file") 43 | 44 | args = parser.parse_args() 45 | 46 | model = OpenNsfwModel() 47 | 48 | export_base_path = args.target 49 | export_version = args.export_version 50 | input_type = InputType[args.input_type.upper()] 51 | 52 | export_path = os.path.join(export_base_path, export_version) 53 | 54 | with tf.Session() as sess: 55 | model.build(weights_path=args.model_weights, 56 | input_type=input_type) 57 | 58 | sess.run(tf.global_variables_initializer()) 59 | 60 | builder = saved_model_builder.SavedModelBuilder(export_path) 61 | 62 | builder.add_meta_graph_and_variables( 63 | sess, [SERVING], 64 | signature_def_map={ 65 | DEFAULT_SERVING_SIGNATURE_DEF_KEY: predict_signature_def( 66 | inputs={PREDICT_INPUTS: model.input}, 67 | outputs={PREDICT_OUTPUTS: model.predictions} 68 | ) 69 | } 70 | ) 71 | 72 | builder.save() 73 | -------------------------------------------------------------------------------- /tools/export_tflite.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import argparse 4 | 5 | import tensorflow as tf 6 | 7 | sys.path.append((os.path.normpath( 8 | os.path.join(os.path.dirname(os.path.realpath(__file__)), 9 | '..')))) 10 | 11 | from model import OpenNsfwModel, InputType 12 | 13 | """Exports a tflite version of tensorflow-open_nsfw 14 | 15 | Note: The standard TFLite runtime does not support all required ops when using the base64_jpeg input type. 16 | You will have to implement the missing ones by yourself. 17 | """ 18 | if __name__ == "__main__": 19 | parser = argparse.ArgumentParser() 20 | 21 | parser.add_argument("target", help="output filename, e.g. 'open_nsfw.tflite'") 22 | 23 | parser.add_argument("-i", "--input_type", required=True, 24 | default=InputType.TENSOR.name.lower(), 25 | help="Input type. Warning: base64_jpeg does not work with the standard TFLite runtime since a lot of operations are not supported", 26 | choices=[InputType.TENSOR.name.lower(), 27 | InputType.BASE64_JPEG.name.lower()]) 28 | 29 | parser.add_argument("-m", "--model_weights", required=True, 30 | help="Path to trained model weights file") 31 | 32 | args = parser.parse_args() 33 | 34 | model = OpenNsfwModel() 35 | 36 | export_path = args.target 37 | input_type = InputType[args.input_type.upper()] 38 | 39 | with tf.Session() as sess: 40 | model.build(weights_path=args.model_weights, 41 | input_type=input_type) 42 | 43 | sess.run(tf.global_variables_initializer()) 44 | 45 | converter = tf.contrib.lite.TFLiteConverter.from_session(sess, [model.input], [model.predictions]) 46 | tflite_model = converter.convert() 47 | 48 | with open(export_path, "wb") as f: 49 | f.write(tflite_model) 50 | --------------------------------------------------------------------------------