| 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 # required/default image dimensionality | |
| model = tf.keras.models.load_model("nsfw.299x299.h5", custom_objects={'KerasLayer': hub.KerasLayer},compile=False) | |
| # def load_images(image_paths, image_size, verbose=True): | |
| # ''' | |
| # Function for loading images into numpy arrays for passing to model.predict | |
| # inputs: | |
| # image_paths: list of image paths to load | |
| # image_size: size into which images should be resized | |
| # verbose: show all of the image path and sizes loaded | |
| # outputs: | |
| # loaded_images: loaded images on which keras model can run predictions | |
| # loaded_image_indexes: paths of images which the function is able to process | |
| # ''' | |
| # 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 | |
| # image | |
| # 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) | |
| def load_images(image): | |
| loaded_images = [] | |
| image = keras.preprocessing.image.array_to_img(image) | |
| image = image.resize((299, 299)) | |
| image = keras.preprocessing.image.img_to_array(image) | |
| image /= 255 | |
| loaded_images.append(image) | |
| return np.asarray(loaded_images) | |
| def classify_nd(model, nd_images, predict_args={}): | |
| # file_path = nd_images.filename | |
| img = load_images(nd_images) | |
| model_preds = model.predict(img, **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): | |
| image_preds = classify_nd(model, image) | |
| return image_preds | |
| demo = gr.Interface(fn=nsfw, | |
| inputs= gr.Image(type="pil"), | |
| outputs=["text"], | |
| title="") | |
| demo.launch(share=False) |