nsfw-detector-small

NSFW Detector

A fine-tuned ConvNeXt model for binary NSFW/SFW image classification. This model distinguishes between safe-for-work and not-safe-for-work content in real-world images, including photographs, illustrations, and animated content.

Model Description

This model classifies images as either SFW (0) or NSFW (1). Built on the ConvNeXt architecture family from timm, pre-trained on ImageNet-22k and fine-tuned on a diverse dataset of real images including photographs, animated content, and various image formats.

Available Variants

Intended Use

This model is designed for automated NSFW content detection in various image types including:

• Photographs and real-world images
• Illustrations and artwork
• Animated content
• User-generated content

Primary Use Cases

• Content moderation systems
• Automated NSFW filtering for platforms
• Safe search implementations
• User upload screening

Out of Scope

• Video content (images only, including first frame of animated formats)
• Real-time detection on edge devices (larger variants)
• Legal/medical content classification
• Fine-grained NSFW category detection

Training Data

• Dataset: Pankaj8922/stickers-binary-v2
• Size: 143,902 images
• Classes:
  • 0: SFW (Safe for Work)
  • 1: NSFW (Not Safe for Work)
• Content Types: Real photographs, animated images, illustrations
• Formats: .jpg, .png, .webp, .webm (first frame extraction)
• Additional Features: The dataset includes object detection annotations (bounding boxes, category labels with 80 COCO classes) and image metadata (dimensions, image IDs)

Training Procedure

Preprocessing

Images are resized to 224×224 and normalized using ImageNet statistics. Training augmentations include:

• Random resized crop (scale: 0.8-1.0)
• Random horizontal flip
• Color jitter (brightness, contrast, saturation, hue)

Training Hyperparameters

• Architecture: ConvNeXt Small (ImageNet-22k pretrained)
• Batch size: 32
• Learning rate: 1e-4
• Optimizer: AdamW
• Loss: Cross-Entropy
• Epochs: 2
• Hardware: GPU (CUDA)

Training Results (Small)

Evaluation

Note: This model was trained on the entire dataset without a validation split. The reported metrics are training metrics and may overestimate real-world performance. For production use, consider evaluating on a held-out test set.

Limitations and Biases

Known Limitations

• Trained on a specific image distribution - may not generalize to all image types
• Binary classification only - cannot detect specific categories of NSFW content
• No validation metrics available for unbiased performance estimation
• May reflect biases present in the training data
• Performance may vary across different cultural contexts and content types

Potential Biases

• Training data composition may not represent all cultural contexts and norms
• Class balance in training data may affect prediction confidence
• Content definition of NSFW can be subjective and culturally dependent

How to Use

Installation

pip install timm torch torchvision pillow

Basic Inference

import torch
import timm
from PIL import Image
from torchvision import transforms
from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from huggingface_hub import hf_hub_download

# Download model from Hub
model_path = hf_hub_download(
    repo_id="Pankaj8922/nsfw-detector-small",
    filename="convnext_small_stickers_final.pth"
)

# Load model
model = timm.create_model("convnext_small.fb_in22k", pretrained=False, num_classes=2)
checkpoint = torch.load(model_path)

# Handle both formats: direct state dict or wrapped in checkpoint dict
if 'model_state_dict' in checkpoint:
    model.load_state_dict(checkpoint['model_state_dict'])
else:
    model.load_state_dict(checkpoint)

model.eval()

# Preprocess
transform = transforms.Compose([
    transforms.Resize(224),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD),
])

# Predict
image = Image.open("image.jpg").convert("RGB")
input_tensor = transform(image).unsqueeze(0)

with torch.no_grad():
    output = model(input_tensor)
    prediction = output.argmax().item()
    probability = torch.softmax(output, dim=1)

labels = {0: "SFW", 1: "NSFW"}
print(f"Prediction: {labels[prediction]}")
print(f"Confidence: {probability[0][prediction]:.2%}")

Acknowledgements

• ConvNeXt architecture by Meta Research
• timm library by Hugging Face
• Training dataset by Pankaj8922

License

This model is licensed under Apache 2.0. See the LICENSE file for details.