AI Anime Image Detector ViT

🚀 AI动漫图像检测器ViT

这是一个用于检测动漫风格AI图像的概念验证模型。该模型使用视觉变换器（Vision Transformer），在100万张人工绘制的真实动漫图像和21.7万张AI生成的动漫图像上进行训练。在训练过程中，两种类型的图像数量相等，以避免偏差。该模型在单张RTX 3090 GPU上训练了约40小时，约35个轮次。

训练日志可在我的 wandb 上查看。

✨ 主要特性

• 利用视觉变换器（Vision Transformer）架构，适用于动漫风格AI图像的检测。
• 训练数据平衡，避免因数据类型比例不均导致的偏差。
• 训练日志公开可查，方便了解训练过程。

📦 安装指南

文档未提供具体安装步骤，故跳过此章节。

💻 使用示例

基础用法

from transformers import AutoModelForImageClassification, AutoFeatureExtractor
import torch
from PIL import Image

model = AutoModelForImageClassification.from_pretrained("legekka/AI-Anime-Image-Detector-ViT")
feature_extractor = AutoFeatureExtractor.from_pretrained("legekka/AI-Anime-Image-Detector-ViT")

model.eval()

image = Image.open("example.jpg")
inputs = feature_extractor(images=image, return_tensors="pt")

outputs = model(**inputs)
logits = outputs.logits

label = model.config.id2label[torch.argmax(logits).item()]
confidence = torch.nn.functional.softmax(logits, dim=1)[0][torch.argmax(logits)].item()

print(f"Prediction: {label} ({round(confidence * 100)}%)")

📚 详细文档

评估

每个检查点都在500张真实图像和500张AI生成图像上进行评估。

最终结果：

• 训练损失：0.1009
• 评估损失：0.1386

似乎使用随机裁剪有助于模型更好地泛化，然而，训练数据集仅包含512x512的图像，这意味着每个裁剪后的图像都进行了双线性插值。在1024x1024的图像上训练模型可能会进一步提高其性能。(也许我稍后会这样做)

性能比较

我们使用约5000张图像对当前可用的AI图像检测器进行了小范围的评估测试。请注意，这些模型并非专门针对动漫图像进行训练。

📄 许可证

本项目采用Apache-2.0许可证。