|
|
--- |
|
|
license: apache-2.0 |
|
|
base_model: Falconsai/nsfw_image_detection |
|
|
tags: |
|
|
- image-classification |
|
|
- content-moderation |
|
|
- violence-detection |
|
|
- nsfw-detection |
|
|
- multi-task-learning |
|
|
--- |
|
|
|
|
|
# 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 |
|
|
```python |
|
|
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) |
|
|
``` |
|
|
|
