aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorRahiel Kasim <rahielkasim@gmail.com>2017-06-12 16:37:05 +0200
committerRahiel Kasim <rahielkasim@gmail.com>2017-06-12 16:43:50 +0200
commit49100b4a02691cdebc1a94c6d60b104c19edeb9b (patch)
treec8e12a6da779d300ec5180b0121702d869aa9759
parentd43a104503c6c186eb599524426453bab836f222 (diff)
cleanup, port to Python 3
-rw-r--r--LICENSE.md11
-rw-r--r--LICENSE.txt23
-rw-r--r--classify_nsfw.py85
3 files changed, 62 insertions, 57 deletions
diff --git a/LICENSE.md b/LICENSE.md
deleted file mode 100644
index d1124b0..0000000
--- a/LICENSE.md
+++ /dev/null
@@ -1,11 +0,0 @@
-
-Copyright 2016, Yahoo Inc.
-
-Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
-
-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.
-
-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.
-
diff --git a/LICENSE.txt b/LICENSE.txt
new file mode 100644
index 0000000..635fbec
--- /dev/null
+++ b/LICENSE.txt
@@ -0,0 +1,23 @@
+
+Copyright 2016, Yahoo Inc.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+list of conditions and the following disclaimer.
+
+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.
+
+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.
diff --git a/classify_nsfw.py b/classify_nsfw.py
index 155478c..b585361 100644
--- a/classify_nsfw.py
+++ b/classify_nsfw.py
@@ -1,54 +1,49 @@
-#!/usr/bin/env python
"""
Copyright 2016 Yahoo Inc.
-Licensed under the terms of the 2 clause BSD license.
-Please see LICENSE file in the project root for terms.
+Licensed under the terms of the 2 clause BSD license.
+Please see LICENSE.txt in the project root for the terms.
"""
-import numpy as np
-import os
-import sys
import argparse
-import glob
-import time
-from PIL import Image
-from StringIO import StringIO
+import sys
+from io import BytesIO
+
import caffe
+import numpy as np
+from PIL import Image
-def resize_image(data, sz=(256, 256)):
+def resize_image(img_data, sz=(256, 256)):
"""
- Resize image. Please use this resize logic for best results instead of the
- caffe, since it was used to generate training dataset
- :param str data:
+ Resize image. Please use this resize logic for best results, as it was used
+ to generate the training dataset.
+ :param bytes data:
The image data
:param sz tuple:
The resized image dimensions
:returns bytearray:
A byte array with the resized image
"""
- img_data = str(data)
- im = Image.open(StringIO(img_data))
+ im = Image.open(BytesIO(img_data))
if im.mode != "RGB":
- im = im.convert('RGB')
+ im = im.convert("RGB")
imr = im.resize(sz, resample=Image.BILINEAR)
- fh_im = StringIO()
- imr.save(fh_im, format='JPEG')
+ fh_im = BytesIO()
+ imr.save(fh_im, format="JPEG")
fh_im.seek(0)
return bytearray(fh_im.read())
-def caffe_preprocess_and_compute(pimg, caffe_transformer=None, caffe_net=None,
- output_layers=None):
+def caffe_preprocess_and_compute(pimg, caffe_transformer=None, caffe_net=None, output_layers=None):
"""
Run a Caffe network on an input image after preprocessing it to prepare
it for Caffe.
:param PIL.Image pimg:
PIL image to be input into Caffe.
:param caffe.Net caffe_net:
- A Caffe network with which to process pimg afrer preprocessing.
+ A Caffe network with which to process pimg after preprocessing.
:param list output_layers:
A list of the names of the layers from caffe_net whose outputs are to
- to be returned. If this is None, the default outputs for the network
+ to be returned. If this is None, the default outputs for the network
are returned.
:return:
Returns the requested outputs from the Caffe net.
@@ -60,19 +55,19 @@ def caffe_preprocess_and_compute(pimg, caffe_transformer=None, caffe_net=None,
output_layers = caffe_net.outputs
img_data_rs = resize_image(pimg, sz=(256, 256))
- image = caffe.io.load_image(StringIO(img_data_rs))
+ image = caffe.io.load_image(BytesIO(img_data_rs))
H, W, _ = image.shape
- _, _, h, w = caffe_net.blobs['data'].data.shape
- h_off = max((H - h) / 2, 0)
- w_off = max((W - w) / 2, 0)
+ _, _, h, w = caffe_net.blobs["data"].data.shape
+ h_off = max((H - h) // 2, 0)
+ w_off = max((W - w) // 2, 0)
crop = image[h_off:h_off + h, w_off:w_off + w, :]
- transformed_image = caffe_transformer.preprocess('data', crop)
+ transformed_image = caffe_transformer.preprocess("data", crop)
transformed_image.shape = (1,) + transformed_image.shape
input_name = caffe_net.inputs[0]
all_outputs = caffe_net.forward_all(blobs=output_layers,
- **{input_name: transformed_image})
+ **{input_name: transformed_image})
outputs = all_outputs[output_layers[0]][0].astype(float)
return outputs
@@ -81,8 +76,6 @@ def caffe_preprocess_and_compute(pimg, caffe_transformer=None, caffe_net=None,
def main(argv):
- pycaffe_dir = os.path.dirname(__file__)
-
parser = argparse.ArgumentParser()
# Required arguments: input file.
parser.add_argument(
@@ -93,36 +86,36 @@ def main(argv):
# Optional arguments.
parser.add_argument(
"--model_def",
- help="Model definition file."
+ help="Model definition file.",
+ default="nsfw_model/deploy.prototxt"
)
parser.add_argument(
"--pretrained_model",
- help="Trained model weights file."
+ help="Trained model weights file.",
+ default="nsfw_model/resnet_50_1by2_nsfw.caffemodel"
)
args = parser.parse_args()
- image_data = open(args.input_file).read()
+ image_data = open(args.input_file, "rb").read()
# Pre-load caffe model.
- nsfw_net = caffe.Net(args.model_def, # pylint: disable=invalid-name
- args.pretrained_model, caffe.TEST)
+ nsfw_net = caffe.Net(args.model_def, args.pretrained_model, caffe.TEST)
# Load transformer
# Note that the parameters are hard-coded for best results
- caffe_transformer = caffe.io.Transformer({'data': nsfw_net.blobs['data'].data.shape})
- caffe_transformer.set_transpose('data', (2, 0, 1)) # move image channels to outermost
- caffe_transformer.set_mean('data', np.array([104, 117, 123])) # subtract the dataset-mean value in each channel
- caffe_transformer.set_raw_scale('data', 255) # rescale from [0, 1] to [0, 255]
- caffe_transformer.set_channel_swap('data', (2, 1, 0)) # swap channels from RGB to BGR
+ caffe_transformer = caffe.io.Transformer({"data": nsfw_net.blobs["data"].data.shape})
+ caffe_transformer.set_transpose("data", (2, 0, 1)) # move image channels to outermost
+ caffe_transformer.set_mean("data", np.array([104, 117, 123])) # subtract the dataset-mean value in each channel
+ caffe_transformer.set_raw_scale("data", 255) # rescale from [0, 1] to [0, 255]
+ caffe_transformer.set_channel_swap("data", (2, 1, 0)) # swap channels from RGB to BGR
# Classify.
- scores = caffe_preprocess_and_compute(image_data, caffe_transformer=caffe_transformer, caffe_net=nsfw_net, output_layers=['prob'])
+ scores = caffe_preprocess_and_compute(image_data, caffe_transformer=caffe_transformer, caffe_net=nsfw_net, output_layers=["prob"])
# Scores is the array containing SFW / NSFW image probabilities
# scores[1] indicates the NSFW probability
- print "NSFW score: " , scores[1]
-
+ print("NSFW score: {}".format(scores[1]))
-if __name__ == '__main__':
- main(sys.argv) \ No newline at end of file
+if __name__ == "__main__":
+ main(sys.argv)