Multi-Head Content Moderator
A multi-task image moderation model with two classification heads:
- NSFW Detection: Detects explicit/adult content (preserved from Falconsai)
- Violence Detection: Detects violent content (newly trained)
Architecture
- Base: ViT (Vision Transformer) from Falconsai/nsfw_image_detection
- Head 1: NSFW classifier (frozen, pretrained)
- Head 2: Violence classifier (trained on violence dataset)
Categories
NSFW Head
- nsfw
- safe
Violence Head
- safe
- violence
Performance (Violence Detection)
- Accuracy: 0.9075
- F1 Score: 0.9076
Usage
import torch
from transformers import AutoImageProcessor
# Load
checkpoint = torch.load('multihead_model.pt')
processor = AutoImageProcessor.from_pretrained('path/to/model')
# Create model class (see notebook for full class definition)
# model = MultiHeadContentModerator(...)
# model.load_state_dict(checkpoint['model_state_dict'])
# Inference
inputs = processor(images=image, return_tensors='pt')
with torch.no_grad():
# Get both predictions
outputs = model(inputs['pixel_values'], task='both')
nsfw_pred = outputs['nsfw'].argmax(-1)
violence_pred = outputs['violence'].argmax(-1)