import opennsfw2
import numpy as np
import gradio as gr


def predict_nsfw(pil_image) -> str:
    image = opennsfw2.preprocess_image(
        pil_image, opennsfw2.Preprocessing.YAHOO)

    model = opennsfw2.make_open_nsfw_model()

    # Add batch axis (for single image).
    inputs = np.expand_dims(image, axis=0)

    predictions = model.predict(inputs)

    # The shape of predictions is (num_images, 2).
    sfw_probability, nsfw_probability = predictions[0]
    if nsfw_probability > sfw_probability:
        result_str = f"It is NSFW image.\nnsfw_probability: {nsfw_probability:.5f}"
    else:
        result_str = f"It is NOT NSFW image.\nnsfw_probability: {nsfw_probability:.5f}"
    return result_str


inputs = gr.inputs.Image(label="input_image", type="pil")
outputs = gr.outputs.Textbox(label="output_string")

app = gr.Interface(fn=predict_nsfw, inputs=inputs, outputs=outputs)

app.launch()