Update README.md

README.md

@@ -11,57 +11,60 @@ tags:
 - exnrt.com
 ---
 
-# Nsfw Image Detection
+# NSFW Image Detection
 
-This model is a fine-tuned for **nsfw image classification**. It has been trained to classify images into three safety-related categories, making it suitable for content moderation, filtering, or safety-aware applications.
+This model is fine-tuned for **NSFW image classification**. It classifies content into three safety-critical categories, making it useful for moderation, safety filtering, and compliant content handling systems.
 
 <p>
   <a href="https://exnrt.com/blog/ai/fine-tuning-siglip2/" target="_blank">
     <img src="https://img.shields.io/badge/View%20Training%20Code-blue?style=for-the-badge&logo=readthedocs"/>
   </a>
+  <a href="https://exnrt.com/blog/ai/fine-tuning-siglip2/" target="_blank">https://exnrt.com/blog/ai/fine-tuning-siglip2/</a>
 </p>
 
+---
+
 ## 🧠 Model Details
 
 * **Base model**: `google/siglip2-base-patch16-224`
-* **Task**: Image Classification (Safety Filtering)
-* **Framework**: PyTorch
-* **Fine-tuned on**: Custom dataset with 3 safety-related categories
-* **Selected checkpoint**: Epoch 3
+* **Task**: Image Classification (NSFW/Safe detection)
+* **Framework**: PyTorch / Hugging Face Transformers
+* **Fine-tuned on**: Custom dataset with 3 content categories
+* **Selected checkpoint**: Epoch 5
 * **Batch size**: 64
-* **Epochs**: 7
+* **Epochs trained**: 5
 
-### 🏷️ Categories
-
-The model classifies images into the following categories:
+---
 
-| ID | Label                 |
-| -- | --------------------- |
-| 0  | `graphically_violent` |
-| 1  | `nudity_pornography`  |
-| 2  | `safe_normal`         |
+### 📌 Epoch Training Results
 
-### 🧾 Label Mapping
+![Epoch Results](https://huggingface.co/Ateeqq/nsfw-image-detection/resolve/main/all-epochs-results.png)
 
-```python
-label2id = {'graphically_violent': 0, 'nudity_pornography': 1, 'safe_normal': 2}
-id2label = {0: 'graphically_violent', 1: 'nudity_pornography', 2: 'safe_normal'}
-```
+### 📌 Confusion Matrix
 
+![Metrics](https://huggingface.co/Ateeqq/nsfw-image-detection/resolve/main/final-epoch-results.png)
 
-## 📈 Visual Results
+![Confusion Matrix](https://huggingface.co/Ateeqq/nsfw-image-detection/resolve/main/confusion_matrix.png)
 
-### 📌 Epoch Training Results
+---
 
-![Epoch Results](https://huggingface.co/Ateeqq/nsfw-image-detection/resolve/main/nsfw-epochs-results.png)
+### 🏷️ Categories
 
-### 📌 Final Metrics & Confusion Matrix
+| ID | Label                    |
+| -- | ------------------------ |
+| 0  | `gore_bloodshed_violent` |
+| 1  | `nudity_pornography`     |
+| 2  | `safe_normal`            |
 
-![Metrics](https://huggingface.co/Ateeqq/nsfw-image-detection/resolve/main/nsfw-training-results.png)
+### 🧾 Label Mapping
 
+```python
+label2id = {'gore_bloodshed_violent': 0, 'nudity_pornography': 1, 'safe_normal': 2}
+id2label = {0: 'gore_bloodshed_violent', 1: 'nudity_pornography', 2: 'safe_normal'}
+```
 ---
 
-## 🚀 Usage
+## 🚀 Usage Example
 
 ```python
 import torch
@@ -73,46 +76,33 @@ model_path = "Ateeqq/nsfw-image-detection"
 processor = AutoImageProcessor.from_pretrained(model_path)
 model = SiglipForImageClassification.from_pretrained(model_path)
 
-image_path = r"your_image_path.jpg"
-image = Image.open(image_path).convert("RGB")
+image = Image.open("your_image_path.jpg").convert("RGB")
 inputs = processor(images=image, return_tensors="pt")
 
 with torch.no_grad():
     logits = model(**inputs).logits
-probabilities = F.softmax(logits, dim=1)
+    probs = F.softmax(logits, dim=1)
 
 predicted_class_id = logits.argmax().item()
-predicted_class_label = model.config.id2label[predicted_class_id]
-confidence_scores = probabilities[0].tolist()
-
-print(f"Predicted class ID: {predicted_class_id}")
-print(f"Predicted class label: {predicted_class_label}\n")
-for i, score in enumerate(confidence_scores):
-    label = model.config.id2label[i]
-    print(f"Confidence for '{label}': {score:.6f}")
-```
-
-### Output
-
-```
-Predicted class ID: 0
-Predicted class label: graphically_violent
+predicted_class = model.config.id2label[predicted_class_id]
 
-Confidence for 'graphically_violent': 0.999988
-Confidence for 'nudity_pornography': 0.000004
-Confidence for 'safe_normal': 0.000008
+print(f"Predicted class: {predicted_class}")
+for i, score in enumerate(probs[0]):
+    print(f"{model.config.id2label[i]}: {score:.4f}")
 ```
 
 ---
 
-## 📊 Training Metrics (Epoch 3 Selected ✅)
+## 📊 Training Metrics (Epoch 5 Selected ✅)
 
 | Epoch | Training Loss | Validation Loss | Accuracy   |
 | ----- | ------------- | --------------- | ---------- |
-| 1     | 0.1086        | 0.0817          | 97.05%     |
-| 2     | 0.0415        | 0.1233          | 95.50%     |
-| 3 ✅   | 0.0302        | 0.0516          | **98.45%** |
-| 4     | 0.0271        | 0.0799          | 97.89%     |
-| 5     | 0.0222        | 0.1015          | 98.03%     |
-| 6     | 0.0026        | 0.0707          | 98.45%     |
-| 7     | 0.0178        | 0.0665          | 98.59%     |
+| 1     | 0.0765        | 0.1166          | 95.70%     |
+| 2     | 0.0719        | 0.0477          | 98.34%     |
+| 3     | 0.0089        | 0.0634          | 98.05%     |
+| 4     | 0.0109        | 0.0437          | 98.61%     |
+| 5 ✅   | 0.0001        | 0.0389          | **99.02%** |
+
+- **Training runtime**: 1h 21m 40s
+- **Final Training Loss**: 0.0727
+- **Steps/sec**: 0.11 | **Samples/sec**: 6.99