File size: 2,237 Bytes
4394cfc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
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
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) |