# 垃圾分类

**Repository Path**: KillerQueenKING/garbage-classification

## Basic Information

- **Project Name**: 垃圾分类
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 1
- **Forks**: 0
- **Created**: 2026-06-28
- **Last Updated**: 2026-06-28

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 垃圾分类识别系统

基于 ResNet50 迁移学习的图像分类系统，输入一张垃圾图片即可识别其所属四大类别（可回收物 / 其他垃圾 / 厨余垃圾 / 有害垃圾），并通过 FastAPI 提供 Web 推理服务。

## 项目特点

- **模型**：ImageNet 预训练 ResNet50，替换全连接层为 4 类输出，冻结-微调两阶段训练
- **数据**：复用数据集官方 train/validate/test 三段划分，统一分辨率为 256x256
- **训练优化**：判别式学习率、类别加权损失、混合精度（AMP）、早停、AdamW
- **评估**：在独立测试集上输出分类报告与混淆矩阵，测试准确率 **95.57%**
- **服务**：FastAPI 提供 `/api/classify` 图片上传推理接口，返回类别与置信度

## 技术栈

| 组件 | 选型 |
|------|------|
| 预训练模型 | ResNet50（ImageNet 权重） |
| 深度学习框架 | PyTorch 2.6.0 + CUDA 12.4 |
| Web 框架 | FastAPI + Uvicorn |
| 图像处理 | torchvision、Pillow |
| 评估工具 | scikit-learn、matplotlib |

## 项目结构

```
garbage-sorting/
├── main.py                     # FastAPI 应用入口
├── api/
│   ├── __init__.py
│   └── classify.py             # 分类 API 路由
├── module/
│   ├── __init__.py
│   ├── config.py               # 全局配置常量
│   ├── prepare_data.py         # 数据准备：读外部数据集→官方划分→统一分辨率→落盘
│   ├── dataset.py              # 数据集定义与 DataLoader
│   ├── model.py                # ResNet50 模型定义
│   ├── train.py                # 训练脚本（两阶段）
│   ├── evaluate.py             # 评估脚本（测试集报告+混淆矩阵）
│   └── predict.py              # 推理逻辑
├── data/                       # 由 prepare_data.py 生成的预处理数据
│   ├── train/<大类>/<子类>/*.jpg
│   ├── val/<大类>/<子类>/*.jpg
│   └── test/<大类>/<子类>/*.jpg
├── checkpoints/                # 模型检查点与评估产物
│   ├── best_model.pth          # 最佳模型权重（90MB）
│   ├── phase1_model.pth        # Phase 1 冻结阶段权重
│   ├── confusion_matrix.png    # 混淆矩阵图
│   └── train.log / eval.log    # 训练与评估日志
├── docs/
│   ├── 垃圾分类识别项目方案.md  # 项目方案对比文档
│   └── 方案一实现计划.md        # 详细实现计划
├── requirements.txt            # 依赖清单
├── test_main.http              # API 测试请求
└── README.md
```

## 环境要求

- Python 3.12（已验证）
- NVIDIA GPU（推荐，训练约 20 分钟；CPU 也可运行但较慢）
- 外部原始数据集（含 0-39 子类目录、`garbage_classify_rule.json`、`train.txt`/`validate.txt`/`test.txt`）

## 安装

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

若需要 GPU 版 PyTorch，从 PyTorch 官方 CUDA 索引安装：

```bash
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
```

验证 GPU 可用：

```bash
python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))"
```

## 数据准备

原始数据集位于外部目录，结构为 `0-39/<img>.jpg`，搭配 `garbage_classify_rule.json`（0-39 → 大类/子类）与官方划分 txt 文件。

在 `module/config.py` 中设置 `GARBAGE_SOURCE_ROOT` 指向原始数据集根目录，然后运行：

```bash
python -m module.prepare_data
```

脚本会：

1. 读取官方 `train.txt`/`validate.txt`/`test.txt` 划分
2. 按 `garbage_classify_rule.json` 将每张图归入 `data/<split>/<大类>/<子类>/`
3. 将所有图像统一为 256x256 RGB（LANCZOS 重采样，JPEG quality=95）
4. 跳过损坏图片并打印警告（test 与 validate 重复的 1 张自动剔除）
5. 已处理图片跳过，支持断点续跑

生成结果：train 8869 张、val 2978 张、test 2954 张。

## 训练

```bash
python -m module.train
```

两阶段训练策略：

**Phase 1（冻结主干，训练分类头）**
- 冻结 ResNet50 卷积层，仅训练新增的 fc 层
- 2 个 epoch，学习率 1e-3
- 保存权重至 `phase1_model.pth`

**Phase 2（解冻主干，端到端微调）**
- 从 Phase 1 权重加载（复用已训练的分类头，而非随机初始化）
- 解冻全部参数，判别式学习率：fc 组 1e-3、backbone 组 1e-5
- 最多 20 个 epoch，连续 5 轮无提升则早停
- `ReduceLROnPlateau` 动态降学习率（factor=0.5, patience=2）
- 混合精度训练（AMP）加速

关键优化：

- **类别加权损失**：可回收物 5159 张 vs 有害垃圾 689 张（约 7.5:1），按反频率加权 `CrossEntropyLoss`
- **判别式学习率**：分类头用大学习率快速收敛，主干用小学习率保护预训练特征
- **AdamW**：解耦权重衰减，避免 Adam + L2 的正则干扰
- **早停**：防止过拟合，节省训练时间

训练日志写入 `checkpoints/train.log`，最佳模型保存至 `checkpoints/best_model.pth`。

## 评估

```bash
python -m module.evaluate
```

在独立测试集（官方 `test.txt`，未参与训练与模型选择）上评估，输出：

- 分类报告（precision / recall / f1 / support）
- 混淆矩阵（打印到控制台并保存 `checkpoints/confusion_matrix.png`）

### 测试集结果（准确率 95.57%）

| 类别     | Precision | Recall | F1         | Support |
|--------|-----------|--------|------------|---------|
| 可回收物   | 0.9674    | 0.9662 | 0.9668     | 1718    |
| 其他垃圾   | 0.8865    | 0.8758 | 0.8811     | 330     |
| 厨余垃圾   | 0.9693    | 0.9778 | 0.9735     | 677     |
| 有害垃圾   | 0.9258    | 0.9258 | 0.9258     | 229     |
| **总体** |           |        | **0.9557** | 2954    |

macro F1 = 0.9368，weighted F1 = 0.9556。

其他垃圾的 precision/recall 最低，主要误判为可回收物（33 张），符合预期：塑料餐盒（其他垃圾）与塑料瓶（可回收物）外观高度相似。

## API 服务

启动服务：

```bash
uvicorn main:app --host 0.0.0.0 --port 8000
```

访问交互式文档：`http://localhost:8000/docs`

### 推理接口

```
POST /api/classify
Content-Type: multipart/form-data

参数: file (图片文件)
```
响应示例：

```json
{
  "class_id": 0,
  "category": "可回收物",
  "confidence": 0.9995
}
```

测试：

```bash
curl -X POST "http://localhost:8000/api/classify" -F "file=@test_image.jpg"
```

模型首次请求时懒加载，后续推理复用内存中的模型实例。

## 配置说明

核心配置位于 `module/config.py`，主要参数：

| 参数 | 默认值 | 说明 |
|------|--------|------|
| `GARBAGE_SOURCE_ROOT` | 外部数据集路径 | 原始数据集根目录 |
| `UNIFIED_SIZE` | 256 | 数据准备统一分辨率 |
| `IMAGE_SIZE` | 224 | 模型输入尺寸 |
| `BATCH_SIZE` | 32 | 训练批大小 |
| `NUM_EPOCHS_FREEZE` | 2 | Phase 1 轮数 |
| `NUM_EPOCHS_FINETUNE` | 20 | Phase 2 最大轮数 |
| `LEARNING_RATE_FINETUNE_HEAD` | 1e-3 | Phase 2 分类头学习率 |
| `LEARNING_RATE_FINETUNE_BACKBONE` | 1e-5 | Phase 2 主干学习率 |
| `EARLY_STOP_PATIENCE` | 5 | 早停耐心值 |
| `NUM_WORKERS` | 0 | DataLoader 进程数（Windows sandbox 设 0） |
| `USE_AMP` | GPU 时 True | 混合精度开关 |
| `SEED` | 42 | 随机种子 |

## 分类体系

四大类与 40 种子类的映射（见原始数据集 `garbage_classify_rule.json`）：

| 大类 | 子类示例 |
|------|---------|
| 可回收物 | 饮料瓶、易拉罐、纸板箱、玻璃杯、旧衣服、充电宝等（23 种） |
| 其他垃圾 | 一次性快餐盒、污损塑料、烟蒂、牙签、竹筷等（6 种） |
| 厨余垃圾 | 剩饭剩菜、大骨头、水果果皮、茶叶渣、蛋壳、鱼骨等（8 种） |
| 有害垃圾 | 干电池、软膏、过期药物（3 种） |

## 完整运行流程

```bash
# 1. 安装依赖
pip install -r requirements.txt

# 2. 准备数据（设置 config.py 中的 GARBAGE_SOURCE_ROOT 后）
python -m module.prepare_data

# 3. 训练
python -m module.train

# 4. 评估
python -m module.evaluate

# 5. 启动 API 服务
uvicorn main:app --host 0.0.0.0 --port 8000
```