ai-generated-image-classification

An image classification model to classify ai-generated and human-generated images

This project implements transfer learning from the RegNet800 regnet_y_800mf, and Google vision transformer, vit_b_16 model to detect AI-generated images using PyTorch. The notebook contains code for training and evaluating various models on a dataset of real and AI-generated images.

Overview

The notebook is organized into 4 main sections:

1. Setup

   • Installation of dependencies
   • Data loading and preprocessing
   • Configuration of constants and environment

2. Testing Pre-trained Model

   • Loading saved model weights
   • Evaluating model performance on test data
   • Generating accuracy metrics

3. Training Pipeline

   • Model architecture configuration
   • Training loop implementation
   • Validation and metrics tracking
   • Model checkpointing

4. Parameter Optimization

   • Testing different model architectures
   • Experimenting with trainable layers
   • Optimizing learning rates and scheduler parameters
   • Comparative analysis of model performance

Requirements

• Python 3.x
• PyTorch
• torchvision
• numpy
• matplotlib
• scikit-learn
• PIL
• seaborn

Dataset Structure

The dataset should be organized as follows:

dataset_name/
├── train/
│   ├── 0_real/
│   └── 1_fake/
├── val/
│   ├── 0_real/
│   └── 1_fake/
└── test/
    ├── 0_real/
    └── 1_fake/

Usage

1. Set up the constants in the first section:

FILE_NAME = 'your_dataset_name'
SMALL_FILE_NAME = 'your_small_dataset_name'
MODEL_WEIGHTS_PATH = 'path_to_model_weights'

2. Run the setup section to install dependencies and prepare the environment

3. To test a pre-trained model:

   • Upload model weights to the specified path
   • Run the testing section

4. To train a new model:

   • Configure model parameters
   • Run the training section
   • Monitor training progress and metrics

5. For parameter optimization:

   • Use the smaller dataset (SMALL_FILE_NAME)
   • Run different experiments in the optimization section
   • Compare results using generated plots

Features

• Support for multiple model architectures (RegNet, ResNet, ViT)
• Learning rate scheduling
• Early stopping
• Model checkpointing
• Comprehensive metrics tracking (accuracy, F1 score, precision, recall)
• Visualization tools for training progress
• Parameter optimization experiments

Results

The notebook includes visualization tools to display:

• Training and validation loss curves
• Accuracy metrics
• Confusion matrices
• Comparative analysis of different parameters

Authors

• Gaétan FINE - 40159598
• Syed Ayan Ali - 57278309
• Aday Yarulin - 57278358