nsfw_image_detection

ViT-B/16 Fine-tuned for NSFW Detection

Model Description

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

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.