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 = "Maybe's AI Art 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.

It was created in October 2022, and as such, the training data did not include any samples generated by Midjourney 5, SDXL, or DALLE-3. It still may be able to correctly identify samples from these more recent models due to being trained on outputs of their predecessors.

Furthermore the intended scope of this tool is artistic images; that is to say, it is not a deepfake photo detector, and general computer imagery (webcams, screenshots, etc.) may throw it off.

In general, this tool can only serve as one of many potential indicators that an image was AI-generated. Images scoring as very probably artificial (e.g. 90% or higher) could be referred to a human expert for further investigation, if needed.

For more information please see the blog post describing this project 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)