This model is a Vision Transformer (ViT) with a base size and a patch size of 16 (ViT-B/16), pre-trained on a large dataset. It has been fine-tuned for the purpose of NSFW (Not Safe For Work) content detection in images, classifying them into specific categories.

Dataset

The model was fine-tuned using the nsfw_dataset_v1 dataset. This dataset comprises five distinct classes:

drawings
hentai
neutral
porn
sexy

The dataset was split into training, validation, and test sets with the following sizes:

Training Set Size: 19600 images
Validation Set Size: 4200 images
Test Set Size: 4200 images
Classes:
classes: {'drawings': 0, 'hentai': 1, 'neutral': 2, 'porn': 3, 'sexy': 4}

Training Details

The model was trained with the following configurations:

Epochs: 10
Batch Size: 64
Optimizer: AdamW, with an initial learning rate of 0.0005.
Loss Function: CrossEntropyLoss, suitable for multi-class classification.
Early Stopping: Set with patience = 3 and delta = 0.001 on Validation Loss

Evaluation

The fine-tuned model's performance was evaluated on the test set, yielding the following metrics:

Loss: 0.3204
Accuracy: 0.9193
F1 (Weighted): 0.9195
AUC (OVR): 0.9908

Classification Report:

        precision    recall  f1-score   support
  
             0       0.85      0.88      0.87       807
             1       0.94      0.88      0.91       839
             2       0.90      0.93      0.92       870
             3       0.98      0.94      0.96       861
             4       0.92      0.97      0.94       823

      accuracy                           0.92      4200
     macro avg       0.92      0.92      0.92      4200
  weighted avg       0.92      0.92      0.92      4200

Confusion Matrix

Usage Instructions

python

from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import requests
import torch

# Load the model and processor
model_name = "acaciabengo/nsfw_image_detection"
processor = ViTImageProcessor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)

# Mapping IDs to human-readable labels
labels = {'drawings': 0, 'hentai': 1, 'neutral': 2, 'porn': 3, 'sexy': 4}
id_to_label = {v: k for k, v in labels.items()}

# Load an image from a URL or local path
url = "https://example.com/sample_image.jpg"
image = Image.open(requests.get(url, stream=True).raw)

# Preprocess image and perform inference
inputs = processor(images=image, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits

# Calculate probabilities and get the top prediction
probabilities = logits.softmax(dim=-1).squeeze().tolist()
predicted_class_idx = logits.argmax(-1).item()

print(f"Predicted label: {id_to_label[predicted_class_idx]}")
print("\nProbabilities:")
for i, prob in enumerate(probabilities):
    print(f"- {id_to_label[i]}: {prob:.2%}")

Commercial Inference API

This model is available via a high-availability, low-latency API for production use cases on RapidApi

Intended Use

This model is intended for the automatic detection and categorization of NSFW content in images. It can be used in applications requiring content moderation, filtering, or analysis to identify images belonging to 'drawings', 'hentai', 'neutral', 'porn', or 'sexy' categories.

Limitations

Bias: Like all AI models, this model may inherit biases present in its training data, potentially leading to unfair or inaccurate classifications for certain demographics or content styles.
Ethical Considerations: The use of NSFW detection models raises ethical concerns regarding privacy, censorship, and the potential for misuse. It is crucial to deploy this model responsibly and with appropriate human oversight.