XRAI AI Art Detection

Mon, 24 Jun 2024 00:00:00 +0000

Welcome to the StyleAIzed Detection project! This notebook documents the journey to create a model that can differentiate between human-made and AI-generated artwork. Here’s a glimpse of the accomplishments and abilities of this notebook:

Trained on a dataset consisting of over 180,000+ works of art using both traditional and diffusion methods
Implemented model interpretability using XRAI attributions for better understanding of predictions
Detailed explanations from data preprocessing to model training and evaluation

Note that this is not the final version and I am still revising the resolution of the XRAI attributions at the end as it is blurry. And this project is a work in progress as I continue to collect more AI works I encounter on the web to enhance accuracy!

Install Dependencies

First we’ll install all the necessary libraries and dependencies required for this project. Here we’ll be using both Torch and TensorFlow for its Keras Tuner as well as the Saliency package to make the attributions.

!pip install torch torchvision keras_tuner tf tf-nightly saliency keras-tuner

Download and Extract AI-ArtBench

We then download the AI-ArtBench dataset from Kaggle and extract its contents. This dataset will be crucial for training and evaluating our models.

import os
from urllib.request import urlopen
from zipfile import ZipFile
from io import BytesIO

def download_and_unzip(url, extract_to='.'):
 # Define the expected directory path based on the extraction target
 directory_path = os.path.join(extract_to, '/Real_AI_SD_LD_Dataset')

 # Check if the directory already exists
 if not os.path.exists(directory_path):
 # Directory does not exist, download and extract the file
 http_response = urlopen(url)
 zipfile = ZipFile(BytesIO(http_response.read()))
 zipfile.extractall(path=extract_to)
 print('Done')
 else:
 # Directory exists, no need to download and extract
 print('Directory already exists:', directory_path)

# Example usage
download_and_unzip('https://www.kaggle.com/api/v1/datasets/download/ravidussilva/real-ai-art', '.')

Done

Define Test and Training Data

We split up the datasets into training and testing images for both human-made and AI images. Each image has a streamlined particular suffix and prefix for easier parsing and sorting into the folders.

import os

top_dir = './Real_AI_SD_LD_Dataset'
train_dir = os.path.join(top_dir, 'train')
test_dir = os.path.join(top_dir, 'test')
all_train_dir = os.listdir(train_dir)
all_test_dir = os.listdir(test_dir)

train_human = []
train_ai = []
test_human = []
test_ai = []

for dir in all_train_dir:
 if not dir.startswith('AI_'):
 train_human.append(os.path.join(train_dir, dir))
 else:
 train_ai.append(os.path.join(train_dir, dir))

print("Training directories containing human-drawn art:")
for i, dir in enumerate(train_human):
 print(f"{i}: {dir}")

print("\nTraining directories containing AI-drawn art:")
for i, dir in enumerate(train_ai):
 print(f"{i}: {dir}")

for dir in all_test_dir:
 if not dir.startswith('AI_'):
 test_human.append(os.path.join(test_dir, dir))
 else:
 test_ai.append(os.path.join(test_dir, dir))

print("\nTesting directories containing human-drawn art:")
for i, dir in enumerate(test_human):
 print(f"{i}: {dir}")

print("\nTesting directories containing AI-drawn art:")
for i, dir in enumerate(test_ai):
 print(f"{i}: {dir}")

Plot Images in Dataset

Let’s take a visual tour of our dataset by randomly selecting and displaying 9 sample images from our dataset, both from the human-made and AI sections.

import random
import cv2
from matplotlib import pyplot as plt
import numpy as np

def plot_images(directory, k = 9):
 files = os.listdir(directory)
 im = random.choices(files, k=k)
 fig = plt.figure(figsize=(10, 10))

 for i in range(k):
 im_i_path = os.path.join(directory, im[i])
 im_i = cv2.imread(im_i_path)
 axis = fig.add_subplot(int(np.sqrt(k)), int(np.sqrt(k)), i + 1)
 axis.imshow(im_i)
 axis.axis('on')
 axis.set_title(im[i], fontsize=8, pad=2)

 plt.tight_layout()
 plt.show()

# Sample 9 random human-made images
print("Showing 9 human-made images in training...")
r = random.randint(0, len(train_human) - 1)
real_images = plot_images(directory=train_human[r], k = 9)
plt.show()

# Sample 9 AI-generated images
print("Showing 9 AI-generated images in training...")
r = random.randint(0, len(train_ai) - 1)
real_images = plot_images(directory=train_ai[r], k = 9)
plt.show()

Showing 9 human-made images in training...

Showing 9 AI-generated images in training...

Labelling Dataset

This section involves labeling the dataset images as either for training or testing.

Training

import pandas as pd

file_paths, labels = [], []
train_dataframe = pd.DataFrame(columns=['file_path', 'label'])

for dir in train_human:
 for file in os.listdir(dir):
 filepath = os.path.join(dir, file)
 file_paths.append(filepath)
 labels.append('human')

for dir in train_ai:
 for file in os.listdir(dir):
 filepath = os.path.join(dir, file)
 file_paths.append(filepath)
 labels.append('AI')

data = pd.DataFrame({'file_path': file_paths, 'label': labels})
train_dataframe = pd.concat([train_dataframe, data], ignore_index=True)
file_cts = train_dataframe['label'].value_counts()
print(f"\nNumber of files per label:\n\n{file_cts}")

Number of files per label:
label
AI 105015
human 50000
Name: count, dtype: int64

Reduce and Balance Training Data

To avoid any biases, we will reduce and balance the training data. This step ensures that our model gets a fair representation of all classes during training.

random_seed = 129
np.random.seed(random_seed)
amt_drop = 55015
ai_indices = train_dataframe[train_dataframe['label'] == 'AI'].index
indices_to_drop = np.random.choice(ai_indices, amt_drop, replace=False)
train_dataframe = train_dataframe.drop(indices_to_drop)
train_dataframe.reset_index(drop=True, inplace=True)
file_cts = train_dataframe['label'].value_counts()
print(f"\nNumber of files per training label:\n\n{file_cts}")

Number of files per training label:
label
human 50000
AI 50000
Name: count, dtype: int64

Testing

file_paths, labels = [], []
test_dataframe = pd.DataFrame(columns=['file_path', 'label'])

for dir in test_human:
 for file in os.listdir(dir):
 filepath = os.path.join(dir, file)
 file_paths.append(filepath)
 labels.append('human')

for dir in test_ai:
 for file in os.listdir(dir):
 filepath = os.path.join(dir, file)
 file_paths.append(filepath)
 labels.append('AI')

data = pd.DataFrame({'file_path': file_paths, 'label': labels})
test_dataframe = pd.concat([test_dataframe, data], ignore_index=True)
file_cts = test_dataframe['label'].value_counts()
print(f"\nNumber of files per testing label:\n\n{file_cts}")

Number of files per testing label:
label
AI 20000
human 10000
Name: count, dtype: int64

Putting Together Training Data Set

from keras.src.legacy.preprocessing.image import ImageDataGenerator
from keras.utils import image_dataset_from_directory
from keras.models import Sequential
from keras.layers import Rescaling, Conv2D, MaxPooling2D, Flatten, Dense, BatchNormalization, Activation, Dropout, GlobalAveragePooling2D
from keras.metrics import Precision, Recall
import keras_tuner as kt
from keras.optimizers import Adam
from keras.callbacks import EarlyStopping

train_gen = ImageDataGenerator(rescale=1./255,
 rotation_range=7,
 horizontal_flip=True,
 zoom_range=0.2
 )
train_dataset = train_gen.flow_from_dataframe(dataframe=train_dataframe, x_col='file_path', y_col='label', target_size=(32, 32), class_mode='categorical', batch_size=64, shuffle=True)
train_dataset.class_indices

Found 100000 validated image filenames belonging to 2 classes.
{'AI': 0, 'human': 1}

Putting Together Test Data Set

test_gen = ImageDataGenerator(rescale=1./255)
test_dataset = test_gen.flow_from_dataframe(dataframe=test_dataframe, x_col='file_path', y_col='label', target_size=(32, 32), class_mode='categorical', batch_size=1, shuffle=False)
test_dataset.class_indices

Found 30000 validated image filenames belonging to 2 classes.
{'AI': 0, 'human': 1}

Train CNN Model

In this section, we train a Convolutional Neural Network (CNN) model using our prepared dataset. First the mages are resized, normalized, and augmented (with techniques like rotation and flipping) to create a diverse training set. Then we assemble the layers of the CNN which consists of:

Convolutional layers to detect features such as edges and textures
ReLu layers to introduce non-linearity
Pooling layers to reduce spatial dimensions and computational load
Fully connected layers to combine features for classification
Dropout layers to prevent overfitting by randomly setting input units to zero

The model is then compiled with an Adam optimizer, a categorical cross-entropy loss function, forward and backward propagation, periodical validation, and performance metrics to measure accuracy amongst other traits.

Build Model

import tensorflow as tf
from keras import layers, models
from keras.optimizers import Adam
from keras.optimizers.schedules import ExponentialDecay
from keras.callbacks import EarlyStopping

network = Sequential()
network.add(Conv2D(filters = 64, kernel_size = 3, input_shape = (32,32,3), activation='relu'))
network.add(MaxPooling2D())
network.add(Conv2D(filters = 64, kernel_size = 3, activation='relu'))
network.add(MaxPooling2D())
network.add(Flatten())
network.add(Dense(units = 64, activation='relu'))
network.add(Dropout(0.2))
network.add(Dense(units = 2, activation='softmax'))

# optimizers
e_schedule = EarlyStopping(monitor='val_accuracy', patience=10, restore_best_weights=True, mode='max', verbose=1, start_from_epoch=1)
lr_schedule = ExponentialDecay(initial_learning_rate=0.001, decay_steps=14067, decay_rate=0.99)
custom_adam_op = Adam(learning_rate=lr_schedule, clipnorm=1.0)

/usr/local/lib/python3.10/dist-packages/keras/src/layers/convolutional/base_conv.py:99: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead.
super().__init__(

Save and Train the Model

from keras.callbacks import ModelCheckpoint, EarlyStopping

eschedule = EarlyStopping(monitor='val_accuracy', patience=12, restore_best_weights=True, mode='max', verbose=1, start_from_epoch=1)
model_checkpoint = ModelCheckpoint('best_model.keras', monitor='val_accuracy', save_best_only=True, mode='max', verbose=1)
network.compile(optimizer=custom_adam_op,
 loss='categorical_crossentropy',
 metrics=['accuracy', 'Precision', 'Recall'])
network.summary()
history = network.fit(train_dataset, epochs=30, validation_data=test_dataset, callbacks=[model_checkpoint, e_schedule])

Evaluate Validation Accuracy

# Plot validation accuracy
plt.plot(history.history['val_accuracy'], label='Validation Accuracy')

# Plot validation loss
plt.plot(history.history['val_loss'], label='Validation Loss')

# Adding labels and title
plt.xlabel('Epoch')
plt.ylabel('Metric')
plt.title('Validation Accuracy and Loss over Epochs')

# Adding legend
plt.legend()

# Show plot
plt.show()

Predict Test Data Set

Make Predictions

predictions = network.predict(test_dataset)
predictions = np.argmax(predictions, axis=1)

[1m30000/30000[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m86s[0m 3ms/step

Observe Precision and Accuracy

from sklearn.metrics import classification_report
print(classification_report(test_dataset.classes, predictions))

 precision recall f1-score support
0 0.94 0.95 0.95 20000
1 0.90 0.88 0.89 10000
accuracy 0.93 30000
macro avg 0.92 0.92 0.92 30000
weighted avg 0.93 0.93 0.93 30000

Setup Confusion Matrix

We’ll set up and analyze this matrix to gain deeper insights into the model’s predictions. For this run it misclassified a little over a thousand images from human and AI, which means there’s much work in progress to improve its accuracy classifying!

from sklearn.metrics import confusion_matrix
import seaborn as sns

# Get the confusion matrix
cm = confusion_matrix(test_dataset.classes, predictions)

labels = ['AI', 'Human']
cm_df = pd.DataFrame(cm, index=labels, columns=labels)
plt.figure(figsize=(8, 6))
sns.heatmap(cm_df, annot=True, fmt='d', cmap='Blues')
plt.xlabel('Predicted')
plt.ylabel('Actual')
plt.title('Confusion Matrix')
plt.show()

Predict on Unseen Samples

Testing our model’s generalization capability, we will make predictions on new, unseen samples. This step is vital to see how well the model can adapt to novel data. In this case you can replace the input_path with the folder to your test images in Google Drive for your perusal.

from google.colab import drive
drive.mount('/content/drive')

model = keras.models.load_model('best_model.keras')
input_path = '/content/drive/MyDrive/AA-Test-Inputs/'
files = os.listdir(input_path)
image_files = sorted([file for file in files if file.endswith(('.jpg', '.png'))])

def predict_and_visualize(img_path, ax, target_size=(32, 32), display_size=(256, 256)):
 if ax is None:
 fig, ax = plt.subplots()
 # Load the original image for display
 display_img = Image.open(img_path)
 display_img = display_img.resize(display_size)

 # Load and preprocess the image for prediction
 prediction_img = Image.open(img_path)
 prediction_img = prediction_img.resize(target_size, Image.LANCZOS)
 prediction_img = np.asarray(prediction_img)
 prediction_img = prediction_img / 255.0 # Normalize the image
 prediction_img = np.expand_dims(prediction_img, axis=0) # Add the batch dimension

 # Make the prediction
 prediction = model.predict(prediction_img)
 predicted_class = np.argmax(prediction, axis=1) # Get the class with the highest probability
 class_labels = {0: "AI-generated", 1: "Human-made"}
 pred_label = class_labels[predicted_class[0]]

 # Display the larger image
 ax.imshow(np.array(display_img)) # Display with bicubic interpolation
 ax.set_title(pred_label)
 ax.axis('off')

# Set up the plot grid
ROWS = 2
COLS = 5
fig, axes = plt.subplots(ROWS, COLS, figsize=(COLS * 3, ROWS * 3)) # Increase figsize to improve display size
axes_flat = axes.flatten()

# Loop over the sorted image files
for i, image_file in enumerate(image_files):
 img_path = os.path.join(input_path, image_file) # Get the full image path
 predict_and_visualize(img_path, axes_flat[i], target_size=(32, 32), display_size=(256, 256)) # Call function

plt.tight_layout()
plt.show()

# https://stable-diffusion-art.com/sdxl-styles/

Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount("/content/drive", force_remount=True).
[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 90ms/step
[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 32ms/step
[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 31ms/step
[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 29ms/step
[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 29ms/step
[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 29ms/step
[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 28ms/step
[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 31ms/step
[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 30ms/step
[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 29ms/step

Explainable Attributions using XRAI

Moving towards interpretability, this section explores explainable AI techniques using XRAI (eXtended Region-based Attributions Integrated). XRAI is a method for generating saliency maps that highlight the most important regions of an image that influence the model’s predictions.

Unlike other saliency methods that may focus on pixel-level attributions, XRAI integrates region-based attributions to provide more coherent and interpretable visual explanations. It works by initially generating a fine-grained saliency map and then grouping adjacent pixels into superpixels. These superpixels are iteratively merged based on their importance until the most relevant regions are identified, whcih then allows for a more human-interpretable understanding of which parts of an image are most significant for the model’s decision.

For more detailed information, you can read the XRAI paper here.

Import Necessary Dependencies

import tensorflow as tf
import PIL.Image
from matplotlib import pylab as P
from tensorflow import keras
import saliency.tf1 as saliency
import warnings
from PIL import Image
import json
import tensorflow.python.keras.backend as K
%matplotlib inline
warnings.filterwarnings("ignore")

Utility Functions

We’ll define some utility functions that assist in generating and visualizing XRAI attributions. These functions simplify repetitive tasks and enhance our workflow, such as working with our images and labelling them.

def ShowImage(img, title='', ax=None):
 if ax is None:
 fig, ax = plt.subplots()
 P.axis('off')
 # Remove the batch dimension by selecting the first image in the batch
 img = img[0]
 img = (img * 255).astype(np.uint8)
 P.imshow(img, interpolation='lanczos')
 P.title(title)
 P.show()

def LoadImg(img_path, target_size = (32, 32)):
 img = PIL.Image.open(img_path)
 img = img.resize(target_size)
 img = np.asarray(img, dtype=np.float32)
 img = np.expand_dims(img, axis=0) # Add the batch dimension
 return img/255.0

# For processing file uploads by user
def LoadImgStream(file_stream, target_size=(32, 32)):
 img = Image.open(file_stream)
 img = img.resize(target_size)
 img = np.asarray(img, dtype=np.float32) / 255.0 # Normalize
 img = np.expand_dims(img, axis=0) # Add the batch dimension
 return img/255.0

def LoadAndUpsampleImg(img_path, target_size=(256, 256)):
 img = Image.open(img_path)
 img = img.resize(target_size, Image.LANCZOS) # Upsample for display
 img = np.asarray(img, dtype=np.float32) / 255.0 # Normalize
 return img

def GetLabel(model, img_path):
 img = LoadImg(img_path)
 prediction = model.predict(img)
 predicted_class = np.argmax(prediction, axis=1) # Get the class with the highest probability
 class_labels = {0: "AI-generated", 1: "Human-made"}
 pred_label = class_labels[predicted_class[0]]
 return pred_label

Define Call Model Function for Inference

This function computes the gradients for saliency maps using TensorFlow’s GradientTape. It handles two scenarios: computing input-output gradients, reflecting how changes in the input affect model predictions, and obtaining convolution-output gradients for interpreting intermediary layers. While wrapping gradients as NumPy arrays is non-standard, this is done for compatibility with external saliency libraries. This function is slightly tailored to work with the custom Sequential CNN model, providing both predictions and intermediate layer activations as opposed to only the final output.

import saliency.core as saliency

model = tf.keras.models.load_model('best_model.keras')

def call_model_function(images, call_model_args=None, expected_keys=None):
 target_class_idx = call_model_args[0]
 images = tf.convert_to_tensor(images)
 with tf.GradientTape() as tape:
 if expected_keys==[saliency.base.INPUT_OUTPUT_GRADIENTS]:
 tape.watch(images)
 output_layer = model(images)
 output_layer = output_layer[:, target_class_idx]
 gradients = np.array(tape.gradient(output_layer, images))
 return {saliency.base.INPUT_OUTPUT_GRADIENTS: gradients}
 else:
 conv_layer, output_layer = model(images)
 gradients = np.array(tape.gradient(output_layer, conv_layer))
 return {saliency.base.CONVOLUTION_LAYER_VALUES: conv_layer,
 saliency.base.INPUT_OUTPUT_GRADIENTS: gradients}

Define Fast and Full XRAI Attributions

from scipy.ndimage import zoom

def load_xrai(model, img_path):
 # Load image and infer
 test_img = LoadImg(img_path)
 xrai_object = saliency.XRAI()
 test_img = test_img.squeeze()
 call_model_args = {0: 0, 1: 1}
 xrai_attributions = xrai_object.GetMask(test_img, call_model_function, call_model_args)

 # Load the original high-resolution image for display
 original_high_res_img = LoadAndUpsampleImg(img_path)
 upscale_factor = 256 / 32
 xrai_attributions_upsampled = zoom(xrai_attributions, (upscale_factor, upscale_factor), order=1)

 # Create a high-resolution mask for the top 25%
 mask_upsampled = xrai_attributions_upsampled >= np.percentile(xrai_attributions_upsampled, 75)
 im_mask_high_res = original_high_res_img.copy()
 im_mask_high_res[~mask_upsampled] = 0 # Apply mask to high-resolution image

 # Plotting
 fig, axes = plt.subplots(1, 3, figsize=(15, 5))

 # Original Image
 axes[0].imshow(original_high_res_img)
 axes[0].set_title(f"{GetLabel(model, img_path)}")
 axes[0].axis('off') # Hide axes

 # XRAI Heatmap upsampled
 axes[1].imshow(xrai_attributions_upsampled, cmap='magma')
 axes[1].set_title('XRAI Heatmap')
 axes[1].axis('off')

 # Top 30% mask applied to high-resolution image
 axes[2].imshow(im_mask_high_res)
 axes[2].set_title('Top 25%')
 axes[2].axis('off')

 plt.tight_layout()
 plt.show()

Apply XRAI Attributions to Unseen Images

Finally, we demonstrate the application of XRAI attributions to new, unseen images. This step allows us to interpret the model’s behavior on novel data and ensure its decisions are not only accurate but justifiable.

input_path = '/content/drive/MyDrive/AA-Test-Inputs/'
files = os.listdir(input_path)
image_files = sorted([file for file in files if file.endswith(('.jpg', '.png'))])

# Loop over the sorted image files
for i, image_file in enumerate(image_files):
 img_path = os.path.join(input_path, image_file) # Get the full image path
 load_xrai(model, img_path)

[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 28ms/step

[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 26ms/step

[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 26ms/step

[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 31ms/step

[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 28ms/step

[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 28ms/step

[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 29ms/step

[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 29ms/step

[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 28ms/step

[1m1/1[0m [32m━━━━━━━━━━━━━━━━━━━━[0m[37m[0m [1m0s[0m 28ms/step

Article |