nsfw / app.py
tangjicheng
new file: app.py
f0e6347
raw
history blame
905 Bytes
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()