AI Generated Image Identification

With the rise of photo-realistic AI generated images, it is becoming a challenge to differentiate between an AI generated image and an image that was created by a Human Creator. On surface, it is may seem a challenging problem but again Machine Learning can make this problem much simpler. I have a fine tuned an YOLOv8 classifier model foor this purporse. This model takes input an image an outputs the probabilites of two classes: AI Generated and Human Created.

Training Dataset

The dataset that was used to fine tune the model was taken from Roboflow Universe. Please follow this link.

Training procedure

The YOLOv8 nano model for classificatioon ie, yolov8n-cls.pt was chosen as the base model for fine tuning. The model was trained for 50 epochs with a batch size of 9 images on Google Colab Pro using the NVIDIA V100 GPU but you can also fine tune this model on the free tier using the NVIDIA T4 GPU.

Using the Model

Install the following packages to get started:

1. ultralytics
2. huggingface_hub
3. supervision

from ultralytics import YOLO
from supervision import Classifications
from huggingface_hub import hf_hub_download as hub

# download and load the model
model = YOLO(hub(repo_id="arnabdhar/YOLOv8-ai-image-classifier", filename="model.pt"))

# test out inference
image = "https://i.redd.it/0whbc38ljj011.jpg"
result = Classifications.from_ultralytics(model.predict(image)[0])

# get the probabilities
prob_ai, prob_human = result.confidence

print("Probability (AI):", prob_ai)
print("Probability (Human):", prob_human)