Upload app.py

app.py

@@ -0,0 +1,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)