# NSFW-detector

**Repository Path**: C12bit/nsfw-detector

## Basic Information

- **Project Name**: NSFW-detector
- **Description**: 自用，使用Pytorch+ResNet-50模型将本地保存的pixiv色图按4档色度值分类
- **Primary Language**: Python
- **License**: Apache-2.0
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2025-12-14
- **Last Updated**: 2025-12-14

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 图像分类系统

基于PyTorch实现的图像分类系统，支持四类别图像分类：
- A: 涩情人物图片
- B: 暴露人物图片  
- C: 普通人物图片
- D: 风景图片

## 项目结构

```
NSFW-detector/
├── train.py              # 模型训练脚本
├── inference.py          # 批量推理脚本
├── split_data.py         # 数据集分割脚本
├── requirements.txt      # 依赖包列表
└── README.md            # 说明文档
```

## 环境要求

- Python 3.7+
- PyTorch 1.9.0+
- torchvision 0.10.0+
- Pillow 8.0.0+

## 安装依赖

```bash
pip install -r requirements.txt
```

## 我的模型
本人使用9000张照片，验证集为从训练集随机抽取30%的图片，经20轮训练得到的结果，推荐置信度参数0.7
模型下载链接：
通过网盘分享的文件：model.pth
链接: https://pan.baidu.com/s/1I_bFFrdAIRTK40M_KaFuoA?pwd=t95g 提取码: t95g 
--来自百度网盘超级会员v1的分享

## 数据集分割

### split_data.py 功能介绍

`split_data.py` 脚本用于将原始数据集自动分割为训练集和验证集。

#### 主要功能
- 从训练集中随机抽取指定比例的图片到验证集
- 支持自定义验证集比例（默认30%）
- 自动创建验证集目录结构
- 处理文件名冲突，避免覆盖现有文件
- 提供详细的分割统计信息

#### 使用方法

```bash
# 基本用法（使用默认参数）
python split_data.py

# 自定义数据目录和验证集比例
python split_data.py --data_dir ./my_data --val_ratio 0.2
```

#### 参数说明
- `--data_dir`: 数据目录路径（默认: ./data）
- `--val_ratio`: 验证集比例，0-1之间的小数（默认: 0.3）

#### 输入数据要求
原始数据应按照以下结构组织：
```
data/
└── train/
    ├── A/           # 涩情人物图片
    ├── B/           # 暴露人物图片
    ├── C/           # 普通人物图片
    └── D/           # 风景图片
```

#### 输出结构
分割完成后，数据目录结构变为：
```
data/
├── train/
│   ├── A/           # 训练集 - 涩情人物图片
│   ├── B/           # 训练集 - 暴露人物图片
│   ├── C/           # 训练集 - 普通人物图片
│   └── D/           # 训练集 - 风景图片
└── val/
    ├── A/           # 验证集 - 涩情人物图片
    ├── B/           # 验证集 - 暴露人物图片
    ├── C/           # 验证集 - 普通人物图片
    └── D/           # 验证集 - 风景图片
```

## 数据准备

训练数据需要按照以下目录结构组织：

```
data/
├── train/
│   ├── A/           # 涩情人物图片
│   ├── B/           # 暴露人物图片
│   ├── C/           # 普通人物图片
│   └── D/           # 风景图片
└── val/
    ├── A/
    ├── B/
    ├── C/
    └── D/
```

## 模型训练

### 基本用法

```bash
python train.py --data_dir ./data
```

### 完整参数

```bash
python train.py 
    --data_dir ./data --model_save_path model.pth --batch_size 32 --learning_rate 0.001 --num_epochs 20 --num_classes 4
```

### 参数说明

- `--data_dir`: 训练数据目录路径（必需）
- `--model_save_path`: 模型保存路径（默认: model.pth）
- `--batch_size`: 批次大小（默认: 32）
- `--learning_rate`: 学习率（默认: 0.001）
- `--num_epochs`: 训练轮数（默认: 20）
- `--num_classes`: 类别数量（默认: 4）

## 批量推理

### 基本用法

```bash
python inference.py --model_path model.pth --input_dir ./E
```

### 完整参数

```bash
python inference.py \
    --model_path model.pth \
    --input_dir ./E \
    --output_dir output \
    --confidence_threshold 0.5
```

### 参数说明

- `--model_path`: 训练好的模型路径（必需）
- `--input_dir`: 待分类图像目录路径（必需）
- `--output_dir`: 分类结果输出目录（默认: output）
- `--confidence_threshold`: 置信度阈值（默认: 0.5）

## 输出结构

推理完成后，图像将按照类别分类到以下目录：

```
output/
├── A/           # 色情人物图片
├── B/           # 暴露人物图片
├── C/           # 普通人物图片
└── D/           # 风景图片
```


## 技术特点

### 模型架构
- 使用ResNet-50预训练模型进行迁移学习
- 自定义分类头：全连接层 + Dropout + ReLU激活
- 支持GPU加速训练和推理

### 数据增强
- 随机水平翻转
- 随机旋转
- 颜色抖动
- 标准化处理

### 训练优化
- 交叉熵损失函数
- Adam优化器
- 学习率调度器（StepLR）
- 早停机制（保存最佳模型）

### 推理功能
- 批量图像处理
- 置信度阈值过滤
- 自动创建输出目录
- 错误处理和日志记录

## 注意事项

1. **数据质量**: 确保训练数据的质量和标签准确性
2. **类别平衡**: 尽量保持各类别样本数量均衡
3. **图像格式**: 支持常见图像格式（JPG, PNG, BMP等）
4. **模型性能**: 可根据需要调整模型结构和超参数
5. **硬件要求**: 建议使用GPU进行训练以获得更好的性能

## 故障排除

### 常见问题

1. **内存不足**: 减小batch_size或使用更小的图像尺寸
2. **训练不收敛**: 调整学习率或使用预训练权重
3. **推理错误**: 检查模型路径和图像文件格式

### 错误处理

- 自动检测GPU可用性
- 详细的错误日志输出
- 文件操作异常处理

## 扩展建议

1. **模型改进**: 尝试不同的预训练模型（如EfficientNet, Vision Transformer）
2. **数据增强**: 添加更多数据增强策略
3. **部署优化**: 模型量化、ONNX导出等
4. **监控工具**: 添加TensorBoard可视化训练过程