|
import argparse |
|
import json |
|
from os import listdir |
|
from os.path import isfile, join, exists, isdir, abspath |
|
import gradio as gr |
|
|
|
import numpy as np |
|
import tensorflow as tf |
|
from tensorflow import keras |
|
import tensorflow_hub as hub |
|
|
|
|
|
IMAGE_DIM = 299 |
|
|
|
def load_images(image_paths, image_size, verbose=True): |
|
loaded_images = [] |
|
loaded_image_paths = [] |
|
|
|
if isdir(image_paths): |
|
parent = abspath(image_paths) |
|
image_paths = [join(parent, f) for f in listdir(image_paths) if isfile(join(parent, f))] |
|
elif isfile(image_paths): |
|
image_paths = [image_paths] |
|
|
|
for img_path in image_paths: |
|
try: |
|
if verbose: |
|
print(img_path, "size:", image_size) |
|
image = keras.preprocessing.image.load_img(img_path, target_size=image_size) |
|
image = keras.preprocessing.image.img_to_array(image) |
|
image /= 255 |
|
loaded_images.append(image) |
|
loaded_image_paths.append(img_path) |
|
except Exception as ex: |
|
print("Image Load Failure: ", img_path, ex) |
|
|
|
return np.asarray(loaded_images), loaded_image_paths |
|
|
|
def load_model(model_path): |
|
if model_path is None or not exists(model_path): |
|
raise ValueError("saved_model_path must be the valid directory of a saved model to load.") |
|
|
|
model = tf.keras.models.load_model(model_path, custom_objects={'KerasLayer': hub.KerasLayer},compile=False) |
|
return model |
|
|
|
|
|
def classify_nd(model, nd_images, predict_args={}): |
|
model_preds = model.predict(nd_images, **predict_args) |
|
categories = ['drawings', 'hentai', 'neutral', 'porn', 'sexy'] |
|
|
|
probs = [] |
|
for i, single_preds in enumerate(model_preds): |
|
single_probs = {} |
|
for j, pred in enumerate(single_preds): |
|
single_probs[categories[j]] = float(pred) |
|
probs.append(single_probs) |
|
return probs |
|
|
|
|
|
def nsfw(image): |
|
model = load_model("nsfw.299x299.h5") |
|
image_preds = classify_nd(model, image) |
|
return json.dumps(image_preds, indent=2) |
|
|
|
|
|
demo = gr.Interface(fn=nsfw, |
|
inputs= gr.Image(type="pil"), |
|
outputs=["text"], |
|
title="") |
|
demo.launch(share=False) |
