def main(argv):
global sess
global model
parser = argparse.ArgumentParser()
parser.add_argument("-m",
"--model_weights",
required=True,
help="Path to trained model weights file")
parser.add_argument("-t",
"--input_type",
default=InputType.TENSOR.name.lower(),
help="input type",
choices=[
InputType.TENSOR.name.lower(),
InputType.BASE64_JPEG.name.lower()
])
args = parser.parse_args()
model = OpenNsfwModel()
sess = tf.Session()
if not (sess):
exit(1)
input_type = InputType[args.input_type.upper()]
model.build(weights_path=args.model_weights, input_type=input_type)
sess.run(tf.global_variables_initializer())
print("Session initialized. Running flask")
app.run(debug=True, host='0.0.0.0')
def classifier(img, url):
image_loader = 'yahoo'
input_file = img
input_type = 'tensor'
model_weights = 'data/open_nsfw-weights.npy'
model = OpenNsfwModel()
with tf.Session() as sess:
input_type = InputType[input_type.upper()]
model.build(weights_path=model_weights, input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if image_loader == IMAGE_LOADER_TENSORFLOW:
fn_load_image = create_tensorflow_image_loader(sess)
else:
fn_load_image = create_yahoo_image_loader()
elif input_type == InputType.BASE64_JPEG:
import base64
fn_load_image = lambda filename: np.array(
[base64.urlsafe_b64encode(open(filename, "rb").read())])
# fn_load_image = img
sess.run(tf.global_variables_initializer())
image = fn_load_image(input_file)
predictions = sess.run(model.predictions,
feed_dict={model.input: image})
sess.close()
return {
'url': url,
'sfw': str(predictions[0][0]),
'nsfw': str(predictions[0][1])
}
def main(argv):
parser = argparse.ArgumentParser()
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()
])
args = parser.parse_args()
model = OpenNsfwModel()
with tf.Session() as sess:
input_type = InputType[args.input_type.upper()]
model.build(weights_path=args.model_weights, input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if args.image_loader == IMAGE_LOADER_TENSORFLOW:
fn_load_image = create_tensorflow_image_loader(
tf.Session(graph=tf.get_default_graph()))
else:
fn_load_image = create_yahoo_image_loader()
elif input_type == InputType.BASE64_JPEG:
import base64
fn_load_image = lambda filename: np.array(
[base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.global_variables_initializer())
image = fn_load_image(args.input_file)
predictions = \
sess.run(model.predictions,
feed_dict={model.input: image})
print("Results for '{}'".format(args.input_file))
print("\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0]))
def main(argv):
parser = argparse.ArgumentParser()
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()])
args = parser.parse_args()
model = OpenNsfwModel()
with tf.Session() as sess:
input_type = InputType[args.input_type.upper()]
model.build(weights_path=args.model_weights, input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if args.image_loader == IMAGE_LOADER_TENSORFLOW:
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:
import base64
fn_load_image = lambda filename: np.array([base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.global_variables_initializer())
image = fn_load_image(args.input_file)
predictions = \
sess.run(model.predictions,
feed_dict={model.input: image})
print("Results for '{}'".format(args.input_file))
print("\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0]))
def predict_nsfw_faster(image_path):
print("predicting nsfw for the image: ", image_path)
model = OpenNsfwModel()
with tf.Session() as sess:
itype = InputType.TENSOR.name.lower()
image_loader = IMAGE_LOADER_TENSORFLOW
input_type = InputType[itype.upper()]
model.build(weights_path="open_nsfw-weights.npy", input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if image_loader == IMAGE_LOADER_TENSORFLOW:
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:
import base64
fn_load_image = lambda filename: np.array([base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.global_variables_initializer())
image = fn_load_image(image_path)
predictions = \
sess.run(model.predictions,
feed_dict={model.input: image})
sfw_score = predictions[0][0]
print("\tSFW score:\t{}".format(predictions[0][0]))
print("\tNSFW score:\t{}".format(predictions[0][1]))
if sfw_score > 0.94:
return "sfw"
else:
return "nsfw"
def get(self):
model = OpenNsfwModel()
fname = request.args.get('fn')
with tf.Session() as sess:
model.build(weights_path=model_weights, input_type=input_type)
import base64
fn_load_image = lambda filename: np.array([base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.global_variables_initializer())
image = fn_load_image(fname)
predictions = \
sess.run(model.predictions,
feed_dict={model.input: image})
sfw, nsw = "{}:{}".format(*predictions[0]).split(":")
scores = {
"fileName" : fname,
"sfw" : sfw,
"nsw" : nsw
}
return jsonify(scores)
class nsfw:
def __init__(self):
"""Cart bar detection class
Arguments
---------
frozen_path: str
Path to .pb model. Default="carts_model/frozen_inference_graph.pb"
gpu_memory_fraction: float
Percentage of how much gpu should the model use. If too little model can fail to load. Default=0.5
gpu_device: str
Required to set env variable CUDA_VISIBLE_DEVICES so Tensorflow only uses desired gpu. Default='0'
"""
"""Limit visible GPUs"""
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
self.model = OpenNsfwModel()
self.model.build()
self.sess = tf.Session()
self.sess.run(tf.initialize_all_variables())
def classify(self, image_str):
img = decode_img_base64(image_str)
cv2.imwrite("img.jpg", img)
fn_load_image = create_tensorflow_image_loader(self.sess)
image = fn_load_image("img.jpg")
predictions = self.sess.run(self.model.predictions,feed_dict={self.model.input: image})
# print("\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0]))
return predictions[0]
def main(argv):
global sess
global model
parser = argparse.ArgumentParser()
parser.add_argument("-m",
"--model_weights",
required=True,
help="Path to trained model weights file")
parser.add_argument("-t",
"--input_type",
default=InputType.TENSOR.name.lower(),
help="input type",
choices=[
InputType.TENSOR.name.lower(),
InputType.BASE64_JPEG.name.lower()
])
parser.add_argument("-p", "--port", default=6000, help="port number")
args = parser.parse_args()
model = OpenNsfwModel()
gpuConfig = tf.ConfigProto(
gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.3),
device_count={'GPU': 1})
sess = tf.Session(config=gpuConfig)
if not (sess):
exit(1)
input_type = InputType[args.input_type.upper()]
model.build(weights_path=args.model_weights, input_type=input_type)
sess.run(tf.global_variables_initializer())
print("Session initialized. Running flask")
port = detect_port(int(args.port))
app.run(debug=False, host='0.0.0.0', port=port)
def nsfw_main():
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
IMAGE_LOADER_TENSORFLOW = "tensorflow"
class args:
pass
args.input_file = "girl.jpg"
args.model_weights = "data/open_nsfw-weights.npy"
args.image_loader = IMAGE_LOADER_TENSORFLOW
args.input_type = InputType.TENSOR.name.lower()
model = OpenNsfwModel()
# This is important for reset graph
tf.reset_default_graph()
with tf.Session() as sess:
input_type = InputType[args.input_type.upper()]
model.build(weights_path=args.model_weights, input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if args.image_loader == IMAGE_LOADER_TENSORFLOW:
fn_load_image = create_tensorflow_image_loader(sess)
else:
fn_load_image = create_yahoo_image_loader()
elif input_type == InputType.BASE64_JPEG:
import base64
fn_load_image = lambda filename: np.array([base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.global_variables_initializer())
image = fn_load_image(args.input_file)
predictions = \
sess.run(model.predictions,
feed_dict={model.input: image})
print("Results for '{}'".format(args.input_file))
print("\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0]))
class nsfw_detector():
def __init__(self,
model_weights = '/home/citao/github/tensorflow-open_nsfw/data/open_nsfw-weights.npy',
image_loader = 'IMAGE_LOADER_YAHOO',
input_type = InputType.TENSOR.name.lower()
):
self._sess = tf.Session()
self._model = OpenNsfwModel()
input_type = InputType[input_type.upper()]
self._model.build(weights_path = model_weights,
input_type = input_type)
self.fn_load_image = None
if input_type == InputType.TENSOR:
if image_loader == IMAGE_LOADER_TENSORFLOW:
self.fn_load_image = create_tensorflow_image_loader(tf.Session(graph=tf.Graph()))
else:
self.fn_load_image = create_yahoo_image_loader()
elif input_type == InputType.BASE64_JPEG:
import base64
self.fn_load_image = lambda filename: np.array([base64.urlsafe_b64encode(open(filename, "rb").read())])
self._sess.run(tf.global_variables_initializer())
def predict(self, input_file):
image = self.fn_load_image(input_file)
predictions = self._sess.run(
self._model.predictions,
feed_dict = {self._model.input: image}
)
probs = predictions[0]
result = {
'SFW': probs[0],
'NSFW': probs[1]
}
return result
def main(argv):
parser = argparse.ArgumentParser()
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("-cb", "--callback", default='', help="Callback Url")
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()
model = OpenNsfwModel()
current_path = os.getcwd()
dir_path = '%s/image_temp' % (current_path)
if 'http' in args.input_file:
image_file_path = utils.download(args.input_file, dir_path)
logger.info("image download to: " + image_file_path)
else:
image_file_path = args.input_file
if '.jpg' not in image_file_path:
jpg_image_file_path = utils.convPNG2JPG(image_file_path)
if False == jpg_image_file_path:
logger.error('Conv Image Fail!' + image_file_path)
exit(1)
os.remove(image_file_path)
image_file_path = jpg_image_file_path
with tf.compat.v1.Session() as sess:
input_type = InputType[args.input_type.upper()]
model.build(weights_path=args.model_weights, input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if args.image_loader == IMAGE_LOADER_TENSORFLOW:
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:
import base64
def fn_load_image(filename):
return np.array(
[base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.compat.v1.global_variables_initializer())
image = fn_load_image(image_file_path)
predictions = \
sess.run(model.predictions,
feed_dict={model.input: image})
logger.info("Results for '{}'".format(args.input_file))
logger.info(
"\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0]))
if '' != args.callback:
param = {
'sfw': str(predictions[0][0]),
'nsfw': str(predictions[0][1])
}
ret = utils.get(args.callback, param)
logger.info(ret)
if 'http' in args.input_file:
os.remove(image_file_path)
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])
args = parser.parse_args()
batch_size = args.batch_size
output_file = args.output_file
input_type = InputType.TENSOR
model = OpenNsfwModel()
filenames = glob.glob(args.source + "/*.jpg")
num_files = len(filenames)
num_batches = int(num_files / batch_size)
print("Found", num_files, " files")
print("Split into", num_batches, " batches")
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
batch_iterator = None
if args.image_loader == IMAGE_LOADER_TENSORFLOW:
batch_iterator = create_tf_batch_iterator(filenames, batch_size)
else:
fn_load_image = create_yahoo_image_loader(expand_dims=False)
batch_iterator = create_batch_iterator(filenames, batch_size,
fn_load_image)
with tf.Session(graph=tf.Graph(), config=config) as session:
model.build(weights_path=args.model_weights, input_type=input_type)
session.run(tf.global_variables_initializer())
with tqdm(total=num_files) as progress_bar:
with open(output_file, 'w') as o:
o.write('File\tSFW Score\tNSFW Score\n')
for batch_num, images in enumerate(batch_iterator):
predictions = \
session.run(model.predictions,
feed_dict={model.input: images})
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],
prediction[1]))
progress_bar.update(len(images))
def run():
global BATCH_SIZE
global testImg
global LR
global logs_train_dir
prepare()
images = []
labels = []
for line in testImg:
images.append(line)
filename = os.path.splitext(os.path.split(line)[1])[0]
true_index = 0
if filename.split('P**n')[0] == 'v':
true_index = 1
labels.append(true_index)
images = tf.constant(images)
labels = tf.constant(labels)
images = tf.random_shuffle(images, seed=0)
labels = tf.random_shuffle(labels, seed=0)
data = tf.data.Dataset.from_tensor_slices((images, labels))
data = data.shuffle(len(testImg))
data = data.map(process, num_parallel_calls=16)
data = data.prefetch(buffer_size=BATCH_SIZE * 8)
batched_data = data.apply(
tf.contrib.data.batch_and_drop_remainder(BATCH_SIZE)).repeat(
40) #num of epoch
iterator = tf.data.Iterator.from_structure(batched_data.output_types,
batched_data.output_shapes)
init_op = iterator.make_initializer(batched_data)
Y = tf.placeholder(tf.float32, shape=[
BATCH_SIZE,
], name='Y')
model = OpenNsfwModel()
model.build(weights_path=None, batchsize=BATCH_SIZE, tag='')
Y_pred = model.logits
# sup1 = model.sup1
# sup2 = model.sup2
# sup3 = model.sup3
# loss = supervision_loss(Y, Y_pred, sup1, sup2, sup3)
loss = log_loss(Y, Y_pred)
accuracy = evaluation(Y, Y_pred)
train_op = trainning(loss, LR)
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.6
step = 0
epoch = 0
aver_acc = 0
with tf.Session(config=config) as sess:
saver = tf.train.Saver(max_to_keep=0)
sess.run(tf.global_variables_initializer())
sess.run(init_op)
images, filenames, labels = iterator.get_next()
ckpt = tf.train.get_checkpoint_state(logs_train_dir)
if ckpt and ckpt.model_checkpoint_path:
global_step = ckpt.model_checkpoint_path.split('/')[-1].split(
'-')[-1]
print('global_step', global_step)
if global_step != '0':
step = int(global_step) + 1
saver.restore(sess, ckpt.model_checkpoint_path)
record_loss = 0
record_acc = 0
while True:
try:
print('step', step)
name, label_Y, input = sess.run([filenames, labels, images])
get_loss, get_op, get_acc = sess.run(
[loss, train_op, accuracy],
feed_dict={
model.input: input,
Y: label_Y
})
print('loss', get_loss, 'accuracy', get_acc)
record_loss = record_loss + get_loss
record_acc = record_acc + get_acc
aver_acc = aver_acc + get_acc
if (step + 1) % (len(testImg) / BATCH_SIZE) == 0:
epoch = epoch + 1
###save record
record_loss = float(record_loss) / (len(testImg) /
BATCH_SIZE)
record_acc = float(record_acc) / (len(testImg) /
BATCH_SIZE)
update_csv(record_loss, record_acc, epoch)
###
print('epoch', epoch)
print('train_learning_rate', LR)
checkpoint_path = os.path.join(logs_train_dir,
'model.ckpt')
saver.save(sess, checkpoint_path, global_step=step)
step = step + 1
except tf.errors.OutOfRangeError:
sess.run(init_op)
break
print(float(aver_acc) / step)
def main(argv):
# parse input
parser = argparse.ArgumentParser()
parser.add_argument("input_file",
help="Path to the input video.\
Only mp4 is supported.")
args = parser.parse_args()
# initialize NSFW Model
model = OpenNsfwModel()
# initialize variables
frameTotal = 0
frameNsfw = 0
frame_skip = 0
# load video (argument)
videoFile = args.input_file
st = args.input_file
st = st[:-4]
outfilename = st + "_Clean.avi"
with tf.Session() as sess:
# set weights and initialize SFW model
model.build(weights_path="pretrained_models/open_nsfw-weights.npy")
fn_load_image = None
fn_load_image = create_yahoo_image_loader()
sess.run(tf.global_variables_initializer())
cap = cv2.VideoCapture(videoFile)
frameRate = cap.get(5) #frame rate
ret, frame = cap.read()
height, width, nchannels = frame.shape
fourcc = cv2.VideoWriter_fourcc(*'MJPG')
out = cv2.VideoWriter(outfilename, fourcc, math.floor(frameRate),
(width, height))
global flag
global frame_skip
while (True):
ret, frame = cap.read()
if (ret != True): # if there is no video frame detected then exit
break
else:
cv2.imwrite('./temp_files/temp.jpg', frame)
image = fn_load_image('./temp_files/temp.jpg')
frameTotal = frameTotal + 1
predictions = sess.run(model.predictions,
feed_dict={model.input: image})
if (predictions[0][1] <= 0.50):
out.write(frame)
else:
frameNsfw = frameNsfw + 1
# print summary after program runs
if (frameNsfw > 0):
print("Video contained NSFW content.")
else:
print("Video is SFW.")
print((frameNsfw / frameTotal) * 100)
cap.release()
out.release()
print("Done")
checkpoint_path = os.path.join(export_base_path, base_name + '.ckpt')
if as_binary:
graph_name = base_name + '.pb'
else:
graph_name = base_name + '.pbtxt'
graph_path = os.path.join(export_base_path, graph_name)
frozen_graph_path = os.path.join(export_base_path,
'frozen_' + graph_name)
optimized_graph_path = os.path.join(export_base_path,
'optimized_' + graph_name)
with tf.Session() as sess:
model.build(weights_path=args.model_weights,
input_type=input_type)
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.save(sess, save_path=checkpoint_path)
print('Checkpoint exported to {}'.format(checkpoint_path))
tf.train.write_graph(sess.graph_def, export_base_path, graph_name,
as_text=not as_binary)
print('Graph exported to {}'.format(graph_path))
if do_freeze:
print('Freezing graph...')
def main(argv):
parser = argparse.ArgumentParser()
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("-t", "--input_type",
default=InputType.TENSOR.name.lower(),
help="input type",
choices=[InputType.TENSOR.name.lower(),
InputType.BASE64_JPEG.name.lower()])
args = parser.parse_args()
model = OpenNsfwModel()
with tf.Session() as sess:
input_type = InputType[args.input_type.upper()]
model.build(weights_path=args.model_weights, input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if args.image_loader == IMAGE_LOADER_TENSORFLOW:
fn_load_image = create_tensorflow_image_loader(sess)
else:
fn_load_image = create_yahoo_image_loader()
elif input_type == InputType.BASE64_JPEG:
import base64
fn_load_image = lambda filename: np.array([base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.global_variables_initializer())
image = fn_load_image(args.input_file)
predictions = \
sess.run(model.predictions,
feed_dict={model.input: image})
print("Results for '{}'".format(args.input_file))
print(predictions[0][0])
print(predictions[0][1])
print("\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0]))
# SFW : Safe For Work , NSFW : Not Safe For Work
nude_json = {'SFW' : predictions[0][0] , 'NSFW' : predictions[0][1] }
result= json.dumps(nude_json, cls=MyEncoder)
loaded_json = json.loads(result)
#for x in loaded_json:
# print("%s: %f" % (x, loaded_json[x]))
print(loaded_json)
f = open('data.txt', 'r+')
f.truncate()
with open('data.txt', 'w') as outfile:
json.dump(loaded_json, outfile)
return loaded_json
import tensorflow as tf
from image_utils import yahoo_load_image
from model import OpenNsfwModel
import ai_integration
model = OpenNsfwModel()
# the default session behavior is to consume the entire GPU RAM during inference!
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.065
with tf.Session(config=config) as sess:
model.build()
sess.run(tf.global_variables_initializer())
while True:
with ai_integration.get_next_input(inputs_schema={
"image": {
"type": "image"
}
}) as inputs_dict:
image = yahoo_load_image(inputs_dict['image'])
predictions = sess.run(model.predictions, feed_dict={model.input: image})
result_data = {
IMAGE_LOADER_YAHOO = "yahoo"
MODEL = "open_nsfw-weights.npy"
INPUT_TYPE = "tensor"
image_loader = IMAGE_LOADER_YAHOO
print("[INFO] loading model ...")
now = time.time()
model = OpenNsfwModel()
start_time = time.time()
total_time = 0
sess = tf.Session()
input_type = InputType[INPUT_TYPE.upper()]
model.build(weights_path=MODEL, input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if image_loader == IMAGE_LOADER_TENSORFLOW:
fn_load_image = create_tensorflow_image_loader(sess)
else:
fn_load_image = create_yahoo_image_loader()
elif input_type == InputType.BASE64_JPEG:
import base64
fn_load_image = lambda filename: np.array(
[base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.global_variables_initializer())
print("[INFO] Time of model upload: {:.2} s".format(time.time() - start_time))
def main():
model = OpenNsfwModel()
with tf.Session() as sess:
# input_type = InputType[args.input_type.upper()]
input_type = InputType.TENSOR ###############
model.build(weights_path=model_weights, input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if image_loader == IMAGE_LOADER_TENSORFLOW:
fn_load_image = create_tensorflow_image_loader(sess)
else:
fn_load_image = create_yahoo_image_loader()
elif input_type == InputType.BASE64_JPEG:
import base64
fn_load_image = lambda filename: np.array([base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.global_variables_initializer())
##########
result_list = []
from os_utils import get_file_list
file_list = get_file_list("dataset")
remark = []
for file in file_list:
input_image_file = os.path.join(".",file)
if input_image_file[-4::] == ".png":
from skimage import io
from skimage.color import rgba2rgb
im = io.imread(file)
im_rgb = rgba2rgb(im)
io.imsave("current_temp.jpeg", im_rgb)
input_file = "current_temp.jpeg"
else:
input_file = file
image = fn_load_image(input_file)
predictions = \
sess.run(model.predictions,
feed_dict={model.input: image})
# print(predictions[0])
result_list.append(predictions[0][1])
print("Results for '{}'".format(input_file))
print("\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0]))
if predictions[0][1] >= 0.03:
print("\t低俗图片")
remark.append("低俗")
else:
print("\t不是低俗图片")
remark.append("正常")
print(predictions[0][1])
os.system("rm -rf current_temp.jpeg")
result_dict = {"image_list":file_list, "result_list":result_list, "remark":remark}
return result_dict
def main(argv):
parser = argparse.ArgumentParser()
parser.add_argument(
"path",
help="Path to the input images. Only jpeg images are supported.")
parser.add_argument("-b",
"--bucket",
required=True,
help="AWS S3 bucket name")
parser.add_argument("-m",
"--model_weights",
required=True,
help="Path to trained model weights file")
parser.add_argument("-o",
"--output",
required=True,
help="Path to output result 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("-t",
"--input_type",
default=InputType.TENSOR.name.lower(),
help="input type",
choices=[
InputType.TENSOR.name.lower(),
InputType.BASE64_JPEG.name.lower()
])
args = parser.parse_args()
s3 = boto3.client('s3')
bucket_name = ''
for bucket in s3.list_buckets().get('Buckets'):
if bucket.get('Name') == args.bucket:
bucket_name = bucket.get('Name')
if not bucket_name:
print("Bucket {} not available".format(args.bucket))
exit(-1)
images = []
next_token = ''
while True:
if next_token:
response = s3.list_objects_v2(Bucket=bucket_name,
Delimiter='|',
EncodingType='url',
MaxKeys=1000,
Prefix=args.path,
ContinuationToken=next_token,
FetchOwner=False)
else:
response = s3.list_objects_v2(Bucket=bucket_name,
Delimiter='|',
EncodingType='url',
MaxKeys=1000,
Prefix=args.path,
FetchOwner=False)
content = response.get('Contents')
next_token = response.get('NextContinuationToken')
for item in content:
images.append(item.get('Key'))
if not next_token:
break
print(next_token)
# if len(images) > 100:
# break
model = OpenNsfwModel()
with tf.Session() as sess:
input_type = InputType[args.input_type.upper()]
model.build(weights_path=args.model_weights, input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if args.image_loader == IMAGE_LOADER_TENSORFLOW:
fn_load_image = create_tensorflow_image_loader(sess)
else:
fn_load_image = create_yahoo_image_loader()
elif input_type == InputType.BASE64_JPEG:
import base64
fn_load_image = lambda filename: np.array(
[base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.global_variables_initializer())
output = open(args.output, "a")
for item in images:
temp_file = tempfile.mkstemp()
s3.download_file(bucket_name, item, temp_file[1])
try:
image = fn_load_image(temp_file[1])
except IOError:
print("Read Image Error")
pass
predictions = sess.run(model.predictions,
feed_dict={model.input: image})
output.write(
"https://www.themebeta.com/media/cache/400x225/files/{}, {}\r\n"
.format(item, predictions[0][0]))
os.remove(temp_file[1])
print(
"Results for https://www.themebeta.com/media/cache/400x225/files/{} : {}"
.format(item, predictions[0][0]))
output.close()
print("parameter error")
elif INPUT_TYPE == 'BASE64_JPEG':
if FN_LOAD_IMAGE == IMAGE_LOADER_YAHOO:
print("parameter mismatch")
elif FN_LOAD_IMAGE == IMAGE_LOADER_TENSORFLOW:
import base64
fn_load_image = lambda filename: np.array(
[base64.urlsafe_b64encode(open(filename, "rb").read())])
else:
print("parameter error")
else:
print("parameter error")
model = OpenNsfwModel()
input_type = InputType[INPUT_TYPE]
model.build(weights_path="data/open_nsfw-weights.npy", input_type=input_type)
sess = tf.Session()
def process_start(conn, addr):
logging.info("Connection address:" + str(addr))
rec_d = bytes([])
try:
while True:
data = conn.recv(BUFSIZ)
if not data or len(data) == 0:
break
else:
rec_d = rec_d + data
except Exception as e:
logging.error(str(e))
def main(argv):
parser = argparse.ArgumentParser()
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")
args = parser.parse_args()
model = OpenNsfwModel()
frameTotal = 0
frameNsfw = 0
with tf.compat.v1.Session() as sess:
input_type = InputType[args.input_type.upper()]
model.build(weights_path=args.model_weights, input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if args.image_loader == IMAGE_LOADER_TENSORFLOW:
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:
import base64
fn_load_image = lambda filename: np.array(
[base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.compat.v1.global_variables_initializer())
#image = fn_load_image(args.input_file)
videoFile = args.input_file
cap = cv2.VideoCapture(videoFile)
frameRate = cap.get(5) #frame rate
while (cap.isOpened()):
frameId = cap.get(1) #current frame number
ret, frame = cap.read()
if (ret != True):
break
if (frameId % math.floor(frameRate) == 0):
cv2.imwrite('./images/temp.jpg', frame)
image = fn_load_image('./images/temp.jpg')
frameTotal = frameTotal + 1
predictions = \
sess.run(model.predictions,
feed_dict={model.input: image})
if (predictions[0][1] >= 0.50):
frameNsfw = frameNsfw + 1
#print("\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0]))
cap.release()
if (frameNsfw > 0):
print("NSFW")
else:
print("SFW")
print(str((frameNsfw / frameTotal) * 100))
import argparse
import tensorflow as tf
from model import OpenNsfwModel, InputType
from image_utils import create_yahoo_image_loader
from flask import Flask, request, jsonify
app = Flask(__name__)
@app.route('/classify', methods=['GET'])
def classify():
filename = request.args["image_path"]
image = create_yahoo_image_loader()(filename)
predictions = sess.run(model.predictions, feed_dict={model.input: image})
# print("\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0]))
predictions = predictions[0].tolist()
return jsonify(dict(sfw=predictions[0], nsfw=predictions[1]))
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-p", "--port", default=8082, help="server http port")
args = parser.parse_args()
model = OpenNsfwModel()
with tf.compat.v1.Session() as sess:
model.build(weights_path="data/open_nsfw-weights.npy", input_type=InputType["TENSOR"])
sess.run(tf.compat.v1.global_variables_initializer())
app.run(port=args.port)
# ### Png to Jpg Image
# In[ ]:
pngs = glob('./*.png')
for j in pngs:
img = cv2.imread(j)
cv2.imwrite((j[:-3] + 'jpg'), img)
# ### Model Load and Build
# In[ ]:
model = OpenNsfwModel()
model.build(weights_path="data/open_nsfw-weights.npy",
input_type=InputType.BASE64_JPEG)
# In[ ]:
fn_load_image = None
fn_load_image = lambda filename: np.array(
[base64.urlsafe_b64encode(open(filename, "rb").read())])
# In[ ]:
with tf.Session() as sess:
jpgs = glob('./*.jpg')
for j in range(len(jpgs)):
if os.stat(jpgs[j]).st_size != 0:
input_file = jpgs[j]
sess.run(tf.global_variables_initializer())
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])
args = parser.parse_args()
batch_size = args.batch_size
output_file = args.output_file
input_type = InputType.TENSOR
model = OpenNsfwModel()
filenames = glob.glob(args.source + "/*.jpg")
num_files = len(filenames)
num_batches = int(num_files / batch_size)
print("Found", num_files, " files")
print("Split into", num_batches, " batches")
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
batch_iterator = None
if args.image_loader == IMAGE_LOADER_TENSORFLOW:
batch_iterator = create_tf_batch_iterator(filenames, batch_size)
else:
fn_load_image = create_yahoo_image_loader(expand_dims=False)
batch_iterator = create_batch_iterator(filenames, batch_size,
fn_load_image)
with tf.Session(graph=tf.Graph(), config=config) as session:
model.build(weights_path=args.model_weights,
input_type=input_type)
session.run(tf.global_variables_initializer())
with tqdm(total=num_files) as progress_bar:
with open(output_file, 'w') as o:
o.write('File\tSFW Score\tNSFW Score\n')
for batch_num, images in enumerate(batch_iterator):
predictions = \
session.run(model.predictions,
feed_dict={model.input: images})
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],
prediction[1]))
progress_bar.update(len(images))
class NSFW:
def __init__(self):
self.fn_load_image = None
self.sess = None
self.model = OpenNsfwModel()
self.data_path = None
def session_run(self, source_type):
with tf.compat.v1.Session() as self.sess:
input_type = InputType[InputType.TENSOR.name.lower().upper()]
self.model.build(
weights_path='data/open_nsfw-weights.npy', input_type=input_type)
if input_type == InputType.TENSOR:
if IMAGE_LOADER_YAHOO == IMAGE_LOADER_TENSORFLOW:
self.fn_load_image = create_tensorflow_image_loader(
tf.Session(graph=tf.Graph()))
else:
self.fn_load_image = create_yahoo_image_loader()
elif input_type == InputType.BASE64_JPEG:
import base64
self.fn_load_image = lambda filename: np.array(
[base64.urlsafe_b64encode(open(filename, "rb").read())])
self.sess.run(tf.compat.v1.global_variables_initializer())
# ======
data_txt = pd.read_table(
self.data_path, sep='\t', header=None).values.tolist()
data_result = []
for d in data_txt:
urls = get_urls(str(d[0]))
# print('网站 = ', str(d[0]), '图片数量:', len(urls))
if len(urls) > 0:
treu_sum = 0
for url in urls:
headers = {
'User-Agent': 'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36 QIHU 360SE'}
try:
res = requests.get(
url, headers=headers, timeout=20)
except:
continue
with open('img/text.jpg', 'wb') as f:
f.write(res.content)
f.close()
try:
image = self.fn_load_image('img/text.jpg')
except:
continue
predictions = self.sess.run(self.model.predictions, feed_dict={
self.model.input: image})
# print(float(predictions[0][1]), type(float(predictions[0][1])))
if float(predictions[0][1]) >= 0.8:
treu_sum += 1
if treu_sum / len(urls) >= 0.2:
data_result.append(str(d[0]))
sday = time.strftime(
'%Y-%m-%d', time.localtime(time.time()))
# print("INSERT INTO `pornographic_website_detection_zsf` (`sday`, `type`, `url`) VALUES ('" + sday + "', '"+source_type+"', '" + str(d[0]) + "' ,'" + str(d[1]) + "')")
# mysql_insert("INSERT INTO `pornographic_website_detection_zsf` (`sday`, `type`, `url`, `sum`) VALUES ('" +
sday + "', '"+source_type+"', '" + str(d[0]) + "' ,'" + str(d[1]) + "')")
print('网站 = ', str(d[0]), '图片数量:', len(urls), '色情图片情况', treu_sum, len(urls) - treu_sum, '占比:', treu_sum / len(urls))
return data_result
def main(argv):
# parse inputs
parser = argparse.ArgumentParser()
parser.add_argument("input_file", help="Path to the input video.")
parser.add_argument("id_folder",
type=str,
nargs="+",
help="Folder containing ID folders")
args = parser.parse_args()
# initialize NSFW Model
model = OpenNsfwModel()
with tf.Graph().as_default():
with tf.Session() as sess:
# set variable defaults
videoFile = args.input_file
cap = cv2.VideoCapture(videoFile)
frameRate = cap.get(5) # get the frame rate
totalFrameCount = cap.get(7) # get the total number of frames
img_size = 64
margin = 0.4
frameNsfw = 0
isMinor = False
minorDetected = False
# set weights and initialize SFW model IsSFW
with tf.variable_scope('IsSFW'):
model.build(
weights_path="pretrained_models/open_nsfw-weights.npy")
fn_load_image = None
fn_load_image = create_yahoo_image_loader()
sess.run(tf.global_variables_initializer())
# initialize dlib face detector model and set variables
detector = dlib.get_frontal_face_detector()
model2 = WideResNet(img_size, 16, 8)()
model2.load_weights("pretrained_models/weights.29-3.76_utk.hdf5")
# initialize face identification model
mtcnn = detect_and_align.create_mtcnn(sess, None)
load_model("model/20170512-110547.pb")
threshold = 1.0
images_placeholder = tf.get_default_graph().get_tensor_by_name(
"input:0")
embeddings = tf.get_default_graph().get_tensor_by_name(
"embeddings:0")
phase_train_placeholder = tf.get_default_graph(
).get_tensor_by_name("phase_train:0")
# Load anchor IDs for face identification model
id_data = IdData(args.id_folder[0], mtcnn, sess, embeddings,
images_placeholder, phase_train_placeholder,
threshold)
while (cap.isOpened()):
ret, frame = cap.read()
frameId = cap.get(1) # get the current frame number
if (ret !=
True): # if there is no video frame detected then exit
break
# write video frame to disk and load as an image
cv2.imwrite('./temp_files/temp.jpg', frame)
image = fn_load_image('./temp_files/temp.jpg')
# determine SFW status
predictions = sess.run(model.predictions,
feed_dict={model.input: image})
if (predictions[0][1] >= 0.50):
frameNsfw = frameNsfw + 1
display_lbl = "NSFW"
AlertColor = [0, 0, 255]
else:
display_lbl = "SFW"
AlertColor = [255, 0, 0]
# detect faces in dlib face detection model
image2 = frame
image2_h, image2_w, _ = np.shape(image2)
detected = detector(image2, 0)
faces = np.empty((len(detected), img_size, img_size, 3))
if len(detected
) > 0: # one or more faces were found in the frame
for i, d in enumerate(detected):
# extract the coordinates of the face
x1, y1, x2, y2, w, h = d.left(), d.top(), d.right(
) + 1, d.bottom() + 1, d.width(), d.height()
xw1 = max(int(x1 - margin * w), 0)
yw1 = max(int(y1 - margin * h), 0)
xw2 = min(int(x2 + margin * w), image2_w - 1)
yw2 = min(int(y2 + margin * h), image2_h - 1)
# draw a rectangle around the face
cv2.rectangle(image2, (x1, y1), (x2, y2), (255, 0, 0),
2)
faces[i, :, :, :] = cv2.resize(
image2[yw1:yw2 + 1, xw1:xw2 + 1, :],
(img_size, img_size))
# determine the height of the rectangle in case is near top of frame
rectangle_height = y2 - y1
# predict ages and genders of faces using dlib model
results = model2.predict(faces)
predicted_genders = results[0]
ages = np.arange(0, 101).reshape(101, 1)
predicted_ages = results[1].dot(ages).flatten()
# draw predictions by faces using dlib model
for i, d in enumerate(detected):
isMinor = False
if (int(predicted_ages[i] < 18)
): # detect if a minor is present in the video
isMinor = True
minorDetected = True
label = "{},{},{}".format(
int(predicted_ages[i]),
"M" if predicted_genders[i][0] < 0.5 else "F",
"-MINOR" if isMinor else "")
draw_label(image2, (d.left(), d.top()), label,
rectangle_height)
# Locate faces and landmarks in frame for identification
face_patches, padded_bounding_boxes, landmarks = detect_and_align.detect_faces(
frame, mtcnn)
if len(face_patches) > 0:
face_patches = np.stack(face_patches)
feed_dict = {
images_placeholder: face_patches,
phase_train_placeholder: False
}
embs = sess.run(embeddings, feed_dict=feed_dict)
matching_ids, matching_distances = id_data.find_matching_ids(
embs)
for bb, landmark, matching_id, dist in zip(
padded_bounding_boxes, landmarks, matching_ids,
matching_distances):
font = cv2.FONT_HERSHEY_COMPLEX_SMALL
cv2.putText(frame, matching_id,
(bb[0] + 30, bb[3] + 5), font, 1,
(255, 0, 255), 1, cv2.LINE_AA)
# display whether frame is SFW or not
percentageComplete = round((frameId) / (totalFrameCount) * 100)
display_lbl = display_lbl + " " + str(
percentageComplete) + "% fps= " + str(round(frameRate, 2))
size = cv2.getTextSize(display_lbl, cv2.FONT_HERSHEY_SIMPLEX,
0.4, 1)[0]
cv2.rectangle(image2, (1, 15 - size[1]), (1 + size[0], 20),
AlertColor, cv2.FILLED)
cv2.putText(image2,
display_lbl, (1, 19),
cv2.FONT_HERSHEY_SIMPLEX,
0.4, (255, 255, 255),
1,
lineType=cv2.LINE_AA)
# display the frame as processed as quickly as possible
cv2.imshow('frame2', image2)
cv2.waitKey(1)
# end of video
cap.release()
cv2.destroyAllWindows()
if os.path.isfile('temp_files/temp.jpg'):
os.remove("temp_files/temp.jpg")
# print summary
if totalFrameCount > 0:
if (frameNsfw > 0):
if (minorDetected):
print("This video contains minors, and " +
str(round((frameNsfw / totalFrameCount * 100), 1)) +
"% of the video contains NSFW elements.")
else:
print(
str(round((frameNsfw / totalFrameCount * 100), 1)) +
"% of the video contains NSFW elements.")
else:
print("Video is SFW.")
else:
print(
"No video frames were detected! Please check the file type or file name."
)
def main(argv):
parser = argparse.ArgumentParser()
parser.add_argument(
'-i',
"--input_dir",
default=config.test_data_set[data_source],
help="Path to the input image. Only jpeg images are supported.")
parser.add_argument("-m",
"--model_weights",
default=config.nsfw_model_weight_file,
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("-t",
"--input_type",
default=InputType.TENSOR.name.lower(),
help="input type",
choices=[
InputType.TENSOR.name.lower(),
InputType.BASE64_JPEG.name.lower()
])
args = parser.parse_args()
with utils.timer('Load image files'):
image_files, labels = data_utils.load_files(args.input_dir,
data_source, sample_rate)
print('image files %s' % len(image_files))
model = OpenNsfwModel()
predictions = []
with tf.Session() as sess:
input_type = InputType[args.input_type.upper()]
with utils.timer('Load model weight'):
model.build(weights_path=args.model_weights, input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if args.image_loader == IMAGE_LOADER_TENSORFLOW:
fn_load_image = create_tensorflow_image_loader(sess)
else:
fn_load_image = create_yahoo_image_loader()
elif input_type == InputType.BASE64_JPEG:
import base64
fn_load_image = lambda filename: np.array(
[base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.global_variables_initializer())
with utils.timer('Prediction'):
start = 0
end = start + config.batch_size
while (start < len(image_files)):
if (end > len(image_files)):
end = len(image_files)
with utils.timer('Batch[%s] prediction' % config.batch_size):
batch_images = [
fn_load_image(image_files[i])
for i in range(start, end)
]
predictions.extend(
sess.run(model.predictions,
feed_dict={model.input: batch_images})[:, 1])
print('Prediction %s done.' % end)
start = end
end = start + config.batch_size
# save
PredictOutputFile = '%s/%s.csv' % (config.TestOutputDir, data_source)
with utils.timer('Save predictions'):
data_utils.save_predictions(image_files, labels, predictions,
PredictOutputFile)
# visualization on threshold for f1/precision/recall
if (data_source == 'hisotry'):
output_image_file = '%s/%s_vs_threshold.jpg' % (config.TestOutputDir,
level)
with utils.timer('Save visualization for threshold'):
plot_utils.threshold_vs_toxic(labels, predictions, level,
output_image_file)
def main(argv):
parser = argparse.ArgumentParser()
args = parser.parse_args()
# args.input_file = "yyy-1.jpg"
# args.input_file = "no-sexy.jpg"
# args.input_file = "zzpic19597.jpg"
args.input_file = "sexy.jpg" # 输入图像
print('[Info] 测试图像: {}'.format(args.input_file))
args.image_loader = IMAGE_LOADER_YAHOO
args.input_type = InputType.TENSOR.name.lower()
args.model_weights = "data/open_nsfw-weights.npy"
model = OpenNsfwModel()
fn_load_image = None
input_type = InputType[args.input_type.upper()]
if input_type == InputType.TENSOR:
if args.image_loader == IMAGE_LOADER_TENSORFLOW:
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:
fn_load_image = lambda filename: np.array([base64.urlsafe_b64encode(open(filename, "rb").read())])
with tf.Session() as sess:
model.build(weights_path=args.model_weights, input_type=input_type)
sess.run(tf.global_variables_initializer())
print('\n[Info] 原始版本')
image = fn_load_image(args.input_file) # 源图像处理格式
model_predict(sess, model, image, args.input_file) # 第2个版本
print('\n[Info] 重写OpenCV版本')
img_np = cv2.imread(args.input_file)
img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
image_v2 = process_img_opencv(img_np)
model_predict(sess, model, image_v2, args.input_file) # 第2个版本
# 存储模型的逻辑
# print('\n[Info] 存储模型')
# save_model(sess, model)
print('\n[Info] base64模型版本')
img_np = cv2.imread(args.input_file)
img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
img_np = process_img_opencv(img_np)
print('[Info] Img: {}'.format(img_np.shape))
img_b64 = base64.urlsafe_b64encode(img_np) # 转换base64
img_tf = image_decode(img_b64)
print('[Info] tf shape: {}'.format(img_tf.shape))
img_np = tf.Session().run(img_tf)
print('[Info] tf->np shape: {}'.format(img_np.shape))
export_path = "data/model-tf" # 模型文件
with tf.Session(graph=tf.Graph()) as sess:
tf.saved_model.loader.load(sess, ["serve"], export_path)
graph = tf.get_default_graph()
print(graph.get_operations())
res = sess.run('predictions:0',
feed_dict={'input:0': img_np})
print('[Info] 最终结果: {}'.format(res))
print('[Info] 性感值: {}'.format(res[0][1] * 100.0))
help=
"Input type. Warning: base64_jpeg does not work with the standard TFLite runtime since a lot of operations are not supported",
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")
args = parser.parse_args()
model = OpenNsfwModel()
export_path = args.target
input_type = InputType[args.input_type.upper()]
with tf.Session() as sess:
model.build(weights_path=args.model_weights, input_type=input_type)
sess.run(tf.global_variables_initializer())
converter = tf.contrib.lite.TFLiteConverter.from_session(
sess, [model.input], [model.predictions])
tflite_model = converter.convert()
with open(export_path, "wb") as f:
f.write(tflite_model)
parser.add_argument("-m",
"--model_weights",
required=True,
help="Path to trained model weights file")
args = parser.parse_args()
model = OpenNsfwModel()
export_base_path = args.target
export_version = args.export_version
export_path = os.path.join(export_base_path, export_version)
with tf.Session() as sess:
model.build(weights_path=args.model_weights,
input_type=InputType.BASE64_JPEG)
sess.run(tf.global_variables_initializer())
builder = saved_model_builder.SavedModelBuilder(export_path)
builder.add_meta_graph_and_variables(
sess, [SERVING],
signature_def_map={
DEFAULT_SERVING_SIGNATURE_DEF_KEY:
predict_signature_def(
inputs={PREDICT_INPUTS: model.input},
outputs={PREDICT_OUTPUTS: model.predictions})
})
builder.save()
def main(argv):
parser = argparse.ArgumentParser()
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("-t",
"--input_type",
default=InputType.TENSOR.name.lower(),
help="input type",
choices=[
InputType.TENSOR.name.lower(),
InputType.BASE64_JPEG.name.lower()
])
args = parser.parse_args()
model = OpenNsfwModel()
with tf.Session() as sess:
input_type = InputType[args.input_type.upper()]
model.build(weights_path=args.model_weights, input_type=input_type)
fn_load_image = None
if input_type == InputType.TENSOR:
if args.image_loader == IMAGE_LOADER_TENSORFLOW:
fn_load_image = create_tensorflow_image_loader(sess)
else:
fn_load_image = create_yahoo_image_loader()
elif input_type == InputType.BASE64_JPEG:
import base64
fn_load_image = lambda filename: np.array(
[base64.urlsafe_b64encode(open(filename, "rb").read())])
sess.run(tf.global_variables_initializer())
images = []
images_names = []
for i in os.listdir(args.input_file):
images_names.append(i)
image_path = os.path.join(args.input_file, i)
image = fn_load_image(image_path)
if images == []:
images = image
print(image_path)
else:
images = np.concatenate((images, image), axis=0)
image = images
predictions = \
sess.run(model.predictions,
feed_dict={model.input: image})
classify_to_folder(args, images_names, predictions)
def classify_nsfw_lambda(imgs):
model = OpenNsfwModel()
with tf.Session() as sess:
def classity_nsfw(url):
if '/a/' in url:
return -1
if not (url.endswith(".gif") or url.endswith(".jpg") or url.endswith.("jpeg") or url.endswith(".png")):
return -1
print("Downloading from '{}'".format(url))
local_filename = url.split('/')[-1]
try:
r = requests.get(url, stream=True)
with open(local_filename, 'wb') as f:
for chunk in r.iter_content(chunk_size=1024):
if chunk: # filter out keep-alive new chunks
f.write(chunk)
if '.gif' in local_filename:
print("gif found")
local_filename_gif = local_filename
local_filename = local_filename.replace('.gif', '.jpg')
print('calling convert {gif}[0] {jpg}'.format(gif=local_filename_gif, jpg=local_filename))
ret = subprocess.call('convert {gif}[0] {jpg}'.format(gif=local_filename_gif, jpg=local_filename), shell=True)
print('output: {}'.format(ret))
elif '.png' in local_filename:
print("png found")
local_filename_png = local_filename
local_filename = local_filename.replace('.png', '.jpg')
print('convert {png} {jpg}'.format(png=local_filename_png, jpg=local_filename))
ret = subprocess.call('convert {png} {jpg}'.format(png=local_filename_png, jpg=local_filename), shell=True)
print('output: {}'.format(ret))
image = fn_load_image(local_filename)
predictions = \
sess.run(model.predictions,
feed_dict={model.input: image})
print("Results for '{}'".format(local_filename))
print("\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0]))
return predictions[0][1]
except requests.exceptions.ConnectionError:
print("Connection Err")
return -1
except tf.errors.InvalidArgumentError:
print("Argument Err")
return -1
except (OSError, IOError) as e:
print("FIle Not Found Err")
return -1
finally:
dir_name = os.getcwd()
test = os.listdir(dir_name)
for item in test:
if item.endswith(".gif") or item.endswith(".jpg") or item.endswith(".png") or item.endswith("jpeg"):
os.remove(os.path.join(dir_name, item))
model.build(weights_path='data/open_nsfw-weights.npy', input_type=InputType[InputType.TENSOR.name.upper()])
fn_load_image = create_tensorflow_image_loader(sess)
sess.run(tf.global_variables_initializer())
outputObj = {'values': []}
for img in imgs:
p = classity_nsfw(img)
outputObj['values'].append(p)
return str(outputObj)