import gradio as gr
import numpy as np
import urllib
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import load_model

# Load the pre-trained model
model = load_model('my_model.h5')
def classify_image(img):
    # Preprocess the input image
    img = image.img_to_array(img)
    img = np.expand_dims(img, axis=0)
    img /= 255.0
    
    # Use the model to make a prediction
    prediction = model.predict(img)[0]
    #print(prediction)
    # Map the predicted class to a label
    dic = {'SFW': np.round(prediction[1],2), 'NSFW': np.round(prediction[0],2)}
    return dic#{'SFW': prediction[0][1], 'NSFW': prediction[0][0]}

def classify_url(url):
    # Load the image from the URL
    response = urllib.request.urlopen(url)
    img = image.load_img(response, target_size=(224, 224))
    
    return classify_image(img)

# Define the GRADIO input interface
#inputs = gr.inputs.Image(shape=(224, 224, 3))

# Define the GRADIO output interface


# Define the GRADIO app
app = gr.Interface(classify_image, gr.Image(shape=(224, 224)), outputs="label", allow_flagging="never", title="NSFW/SFW Classifier")

# Start the GRADIO app
app.launch()