app.py

# prompt: gradio image 分类 not safe for work
from PIL import Image
import gradio as gr
from transformers import pipeline
# Load the image classification pipeline
classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")

# Define a function to classify the image and return the results
def classify_image(img):
    # Convert the Gradio image input to a PIL image
    pil_image = Image.fromarray(img.astype('uint8'), 'RGB')
    # Classify the image using the pipeline
    results = classifier(pil_image)
    # Format the results for display in Gradio
    formatted_results = {result['label']: result['score'] for result in results}
    return formatted_results

# Create the Gradio interface
image_input = gr.inputs.Image(shape=(256, 256))
label_output = gr.outputs.Label(num_top_classes=3)
interface = gr.Interface(fn=classify_image, inputs=image_input, outputs=label_output)

# Launch the interface
interface.launch()