Skip to content

Instantly share code, notes, and snippets.

@kaias1jp
Created December 4, 2019 08:02
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save kaias1jp/12ffb01ed28a9fdba06bf91ed0b9f508 to your computer and use it in GitHub Desktop.
Save kaias1jp/12ffb01ed28a9fdba06bf91ed0b9f508 to your computer and use it in GitHub Desktop.
app.py改造版
#!/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.
"""
import numpy as np
import os
import sys
import argparse
import glob
import time
import urllib2
import json
from bottle import get, run, template, request, HTTPResponse
from PIL import Image
from StringIO import StringIO
import caffe
def resize_image(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:
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))
if im.mode != "RGB":
im = im.convert('RGB')
imr = im.resize(sz, resample=Image.BILINEAR)
fh_im = StringIO()
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):
"""
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.
: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
are returned.
:return:
Returns the requested outputs from the Caffe net.
"""
if caffe_net is not None:
# Grab the default output names if none were requested specifically.
if output_layers is None:
output_layers = caffe_net.outputs
img_data_rs = resize_image(pimg, sz=(256, 256))
image = caffe.io.load_image(StringIO(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)
crop = image[h_off:h_off + h, w_off:w_off + w, :]
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})
outputs = all_outputs[output_layers[0]][0].astype(float)
return outputs
else:
return []
def main(argv):
pycaffe_dir = os.path.dirname(__file__)
parser = argparse.ArgumentParser()
# Required arguments: input file.
parser.add_argument(
"input_file",
help="Path to the input image file"
)
# Optional arguments.
parser.add_argument(
"--model_def",
help="Model definition file."
)
parser.add_argument(
"--pretrained_model",
help="Trained model weights file."
)
args = parser.parse_args()
image_data = open(args.input_file).read()
#response = urllib2.urlopen('https://image.kusokora.jp/v1/b7daff98c2b748dfa6cb8240c8ac7c33/kusokora/image-PI3KQm-1728.png')
#image_data = response.read()
# Pre-load caffe model.
nsfw_net = caffe.Net(args.model_def, # pylint: disable=invalid-name
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
# Classify.
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]
pycaffe_dir = os.path.dirname(__file__)
model_def = 'open_nsfw/nsfw_model/deploy.prototxt'
pretrained_model = 'open_nsfw/nsfw_model/resnet_50_1by2_nsfw.caffemodel'
# Pre-load caffe model.
nsfw_net = caffe.Net(model_def, # pylint: disable=invalid-name
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
@get('/nsfw_score')
def nsfw_score():
image_url = request.query.image_url
headers = {
"User-Agent": "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:47.0) Gecko/20100101 Firefox/47.0"
}
req = urllib2.Request(image_url, headers=headers)
response = urllib2.urlopen(req)
image_data = response.read()
scores = caffe_preprocess_and_compute(image_data, caffe_transformer=caffe_transformer, caffe_net=nsfw_net, output_layers=['prob'])
score = { 'score' :scores[1] }
body = json.dumps(score)
response = HTTPResponse(status=200, body=body)
response.set_header('Content-Type', 'application/json')
return response
run(host='0.0.0.0', port=80)
164,28 Bot
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment