Skip to main content

AI or Not — detect AI-generated photos using AI

Try the classifier at ivan-tolkunov–fake-or-not-run.modal.run (warning: the app could take up to 30s to boot up).

The new AI models for image generations, like Stable Diffusion XL or Midjourney v6, are getting really good at making photorealistic images. Sometimes it’s hard to tell if the images are real even when looking at them for a while.

I wanted to see if it’s possible to train a classifier that’s able to detect AI-generated photos.

To make iteration faster, I managed to configure PyTorch to use my GPU on my M2 MacBook Air. It made things a bit faster, but I wish I had a better GPU, since so many things in AI development are optimized for CUDA.

import torch
from fastai.vision.all import *

print(torch.backends.mps.is_available())
default_device(torch.device("mps"))

I used a simple resnet-based classifier model and the FastAI framework for training on a dataset of images. After playing with parameters a bit, I found that resnet18 is enough to get 99%+ accurancy.

dls = DataBlock(
    blocks=(ImageBlock, CategoryBlock),
    get_items=get_image_files,
    splitter=RandomSplitter(valid_pct=0.2, seed=42),
    get_y=parent_label
).dataloaders(self.path, bs=32)

learn = vision_learner(dls, resnet18, metrics=error_rate)
learn.fine_tune(5)

learn.export('./models/realOrFake_resnet18_224px.pkl')

To allow users to test their own images, I made a simple web app. This time I used Gradio, a Python web app framework that makes it super easy to build simple interfaces that actually look nice.

import gradio as gr
from gradio.routes import mount_gradio_app

mount_gradio_app(
    app=web_app,
    blocks=gr.Interface(fn=predict,
            inputs=gr.Image(type="pil"),
            outputs=gr.Label(num_top_classes=2),
            examples=[
                # list of example URLs
            ]),
    path="/",
)

When the user submits an image, our predict function gets executed on the backend. It returns a list of labels and probabilities. Good news is that FastAI classifiers already return data in the right format, I only needed to convert tensors to floats:

def predict(fake_or_real):
    model = load_learner('/models/realOrFake_resnet18_224px.pkl')
    model.eval()

    result ,_ , probs = model.predict(fake_or_real)
    print(result, probs)
    return {'ai': float(probs[0]), 'real': float(probs[1])}

Full source code is available at https://github.com/ivan-tolkunov/fake_or_not

Hi there! Thanks for reading!

My name is Ivan Tolkunov, I'm a software engineer passionate about programming and AI.

In 2024, I'm looking for a fulltime or an internship position as a software engineer, preferably in AI space based in Seattle, WA (remote is OK). Check out my resume and get in touch via email or LinkedIn.