aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorRahiel Kasim <rahielkasim@gmail.com>2017-06-13 16:08:23 +0200
committerRahiel Kasim <rahielkasim@gmail.com>2017-06-13 16:15:36 +0200
commit53f7043cc6c4e85d5b83484ab136c29bcd33958f (patch)
tree0564ff3ab3b52e4d77e28f9f57d8c911afe2bbe9
parent49100b4a02691cdebc1a94c6d60b104c19edeb9b (diff)
implement a simple HTTP API
-rw-r--r--.gitignore1
-rw-r--r--LICENSE-yahoo.txt23
-rw-r--r--LICENSE.txt41
-rw-r--r--api.py41
-rw-r--r--classify_nsfw.py54
-rw-r--r--requirements.txt4
6 files changed, 123 insertions, 41 deletions
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..bee8a64
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+__pycache__
diff --git a/LICENSE-yahoo.txt b/LICENSE-yahoo.txt
new file mode 100644
index 0000000..635fbec
--- /dev/null
+++ b/LICENSE-yahoo.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/LICENSE.txt b/LICENSE.txt
index 635fbec..6293d9b 100644
--- a/LICENSE.txt
+++ b/LICENSE.txt
@@ -1,23 +1,28 @@
-Copyright 2016, Yahoo Inc.
+Copyright (c) 2017, Rahiel Kasim
+All rights reserved.
-Redistribution and use in source and binary forms, with or without modification,
-are permitted provided that the following conditions are met:
+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.
+* 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.
+* 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.
+* Neither the name of the copyright holder nor the names of its
+ contributors may be used to endorse or promote products derived from
+ this software without specific prior written permission.
+
+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/api.py b/api.py
new file mode 100644
index 0000000..0221752
--- /dev/null
+++ b/api.py
@@ -0,0 +1,41 @@
+import asyncio
+
+import aiohttp
+import async_timeout
+import numpy as np
+import uvloop
+from aiohttp import web
+
+from classify_nsfw import caffe_preprocess_and_compute, load_model
+
+
+nsfw_net, caffe_transformer = load_model()
+
+
+def classify(image: bytes) -> np.float64:
+ scores = caffe_preprocess_and_compute(image, caffe_transformer=caffe_transformer, caffe_net=nsfw_net, output_layers=["prob"])
+ return scores[1]
+
+async def fetch(session, url):
+ with async_timeout.timeout(10):
+ async with session.get(url) as response:
+ return await response.read()
+
+class API(web.View):
+ async def post(self):
+ request = self.request
+ data = await request.post()
+ try:
+ image = await fetch(session, data["url"])
+ nsfw_prob = classify(image)
+ text = nsfw_prob.astype(str)
+ return web.Response(text=text)
+ except:
+ raise web.HTTPBadRequest()
+
+
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+session = aiohttp.ClientSession()
+app = web.Application()
+app.router.add_route("*", "/api", API)
+web.run_app(app)
diff --git a/classify_nsfw.py b/classify_nsfw.py
index b585361..437c420 100644
--- a/classify_nsfw.py
+++ b/classify_nsfw.py
@@ -1,7 +1,7 @@
"""
Copyright 2016 Yahoo Inc.
Licensed under the terms of the 2 clause BSD license.
-Please see LICENSE.txt in the project root for the terms.
+Please see LICENSE-yahoo.txt in the project root for the terms.
"""
import argparse
@@ -13,13 +13,13 @@ import numpy as np
from PIL import Image
-def resize_image(img_data, sz=(256, 256)):
+def resize_image(img_data, size=(256, 256)):
"""
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:
+ :param size tuple:
The resized image dimensions
:returns bytearray:
A byte array with the resized image
@@ -27,7 +27,7 @@ def resize_image(img_data, sz=(256, 256)):
im = Image.open(BytesIO(img_data))
if im.mode != "RGB":
im = im.convert("RGB")
- imr = im.resize(sz, resample=Image.BILINEAR)
+ imr = im.resize(size, resample=Image.BILINEAR)
fh_im = BytesIO()
imr.save(fh_im, format="JPEG")
fh_im.seek(0)
@@ -54,7 +54,7 @@ def caffe_preprocess_and_compute(pimg, caffe_transformer=None, caffe_net=None, o
if output_layers is None:
output_layers = caffe_net.outputs
- img_data_rs = resize_image(pimg, sz=(256, 256))
+ img_data_rs = resize_image(pimg, size=(256, 256))
image = caffe.io.load_image(BytesIO(img_data_rs))
H, W, _ = image.shape
@@ -74,6 +74,24 @@ def caffe_preprocess_and_compute(pimg, caffe_transformer=None, caffe_net=None, o
else:
return []
+def load_model(model_def=None, pretrained_model=None):
+ if model_def is None:
+ model_def = "nsfw_model/deploy.prototxt"
+ if pretrained_model is None:
+ pretrained_model = "nsfw_model/resnet_50_1by2_nsfw.caffemodel"
+
+ # Pre-load caffe model.
+ nsfw_net = caffe.Net(model_def, 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
+
+ return nsfw_net, caffe_transformer
def main(argv):
parser = argparse.ArgumentParser()
@@ -86,35 +104,25 @@ def main(argv):
# Optional arguments.
parser.add_argument(
"--model_def",
- help="Model definition file.",
- default="nsfw_model/deploy.prototxt"
+ help="Model definition file."
)
parser.add_argument(
"--pretrained_model",
- help="Trained model weights file.",
- default="nsfw_model/resnet_50_1by2_nsfw.caffemodel"
+ help="Trained model weights file."
)
args = parser.parse_args()
- image_data = open(args.input_file, "rb").read()
- # Pre-load caffe model.
- 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
+ nsfw_net, caffe_transformer = load_model(args.model_def, args.pretrained_model)
+ image = open(args.input_file, "rb").read()
# 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, 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: {}".format(scores[1]))
+
+ nsfw_prob = scores[1]
+ print("NSFW score: {}".format(nsfw_prob.astype(str)))
if __name__ == "__main__":
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..c564afb
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,4 @@
+aiodns
+aiohttp
+cchardet
+uvloop