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)