Upload 5 files

README.md

@@ -1,3 +1,57 @@
----
-license: apache-2.0
----
+---
+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)
+```

config.json

@@ -0,0 +1,14 @@
+{
+  "base_model": "Falconsai/nsfw_image_detection",
+  "hidden_size": 768,
+  "num_violence_labels": 2,
+  "violence_id2label": {
+    "0": "safe",
+    "1": "violence"
+  },
+  "nsfw_id2label": {
+    "0": "normal",
+    "1": "nsfw"
+  },
+  "model_type": "MultiHeadContentModerator"
+}

model_class.py

@@ -0,0 +1,59 @@
+
+import torch
+import torch.nn as nn
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+
+class MultiHeadContentModerator(nn.Module):
+    """
+    Multi-task model with two classification heads:
+    - Head 1: NSFW detection (frozen, pretrained)
+    - Head 2: Violence detection (trainable)
+    """
+    def __init__(self, base_model_name="Falconsai/nsfw_image_detection", num_violence_labels=2):
+        super().__init__()
+        
+        # Load base model
+        original_model = AutoModelForImageClassification.from_pretrained(base_model_name)
+        hidden_size = original_model.config.hidden_size
+        
+        # ViT backbone (shared)
+        self.vit = original_model.vit
+        
+        # Head 1: Original NSFW classifier
+        self.nsfw_classifier = original_model.classifier
+        
+        # Head 2: Violence classifier
+        self.violence_classifier = nn.Linear(hidden_size, num_violence_labels)
+        
+        # Label mappings - use actual Falconsai config
+        self.nsfw_id2label = original_model.config.id2label  # {0: 'normal', 1: 'nsfw'}
+        self.violence_id2label = {0: 'safe', 1: 'violence'}  # Will be overwritten from checkpoint
+        
+    def forward(self, pixel_values, task='both'):
+        outputs = self.vit(pixel_values=pixel_values)
+        pooled_output = outputs.last_hidden_state[:, 0]
+        
+        if task == 'nsfw':
+            return self.nsfw_classifier(pooled_output)
+        elif task == 'violence':
+            return self.violence_classifier(pooled_output)
+        elif task == 'both':
+            return {
+                'nsfw': self.nsfw_classifier(pooled_output),
+                'violence': self.violence_classifier(pooled_output)
+            }
+        return self.violence_classifier(pooled_output)
+
+def load_multihead_model(checkpoint_path, device='cuda'):
+    """Load trained multi-head model"""
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    
+    model = MultiHeadContentModerator(
+        base_model_name=checkpoint['base_model'],
+        num_violence_labels=checkpoint['num_violence_labels']
+    )
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.violence_id2label = checkpoint['violence_id2label']
+    model.nsfw_id2label = checkpoint['nsfw_id2label']
+    
+    return model.to(device)

multihead_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e3ac359e78361b9a1e71071f61550c3617376d41104a9dba9cebf7fe2ad26bf
+size 343290062

preprocessor_config.json

@@ -0,0 +1,23 @@
+{
+  "do_convert_rgb": null,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "ViTImageProcessor",
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "resample": 2,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 224,
+    "width": 224
+  }
+}