# import libraries
from transformers import AutoModel, AutoModelForImageClassification
import torch
from datasets import load_dataset

# load dataset
dataset = load_dataset("competitions/aiornot")

# list of images
images = dataset["test"][10:20]["image"]

# load models
feature_extractor = AutoModel.from_pretrained(
    "RishiDarkDevil/ai-image-det-resnet152", trust_remote_code=True).to('cuda')
classifier = AutoModelForImageClassification.from_pretrained(
    "RishiDarkDevil/ai-image-det-resnet152", trust_remote_code=True).to('cuda')

# extract features from images
inputs = feature_extractor(images)

# classification using extracted features
with torch.no_grad():
    logits = classifier(inputs)['logits']

# model predicts one of the 2 classes
predicted_label = logits.argmax(-1)

# predictions
print(predicted_label) # 0 is Not AI, 1 is AI

Backbone for Feature Extraction: ResNet152

Performance

Trained MLP Fine-tuning layers for 150 epochs.
Accuracy: 0.9250 on validation data (~5% of the training data).