import gradio as gr
from transformers import pipeline

pipe = pipeline("image-classification", "umm-maybe/AI-image-detector")

def image_classifier(image):
    outputs = pipe(image)
    results = {}
    for result in outputs:
        results[result['label']] = result['score']
    return results

title = "BurstPixel AI Image Detector"
description = """
This app is a proof-of-concept demonstration of using a ViT model to predict whether an artistic image was generated using AI.\n
For more information please see the blog post describing it at:
https://medium.com/@matthewmaybe/can-an-ai-learn-to-identify-ai-art-545d9d6af226
"""

demo = gr.Interface(fn=image_classifier, inputs=gr.Image(type="pil"), outputs="label", title=title, description=description)
demo.launch(show_api=False)