README.md

---
title: AI 垃圾分类助手
emoji: ♻️
colorFrom: green
colorTo: green
sdk: gradio
sdk_version: 5.0.0
app_file: app.py
pinned: false
---

# ♻️ AI 垃圾分类助手

基于深度学习的垃圾分类识别系统，支持 6 类垃圾（塑料、纸板、纸张、玻璃、金属、其他垃圾）的图像识别，并提供投放指南、环保积分和排行榜功能。

**[GitHub 仓库](https://github.com/hutiger9/garbage-classification)**

## 技术栈

| 模块 | 技术 |
|------|------|
| 深度学习框架 | PyTorch + torchvision |
| 模型架构 | MobileNetV3-Small（迁移学习，ImageNet 预训练） |
| 硬件加速 | Apple Silicon MPS / NVIDIA CUDA / CPU 自动切换 |
| Web 界面 | Gradio |
| REST API | FastAPI + uvicorn |
| 数据库 | SQLite（用户积分、分类记录） |
| 包管理 | uv / pip |
| 数据可视化 | Matplotlib（训练曲线） |
| 推理优化 | Top-K 置信度输出、混淆矩阵分析 |

## 项目结构

```
Garbage_classification/
├── main.py              # CLI 入口（训练/预测/webui/api）
├── train.py             # 模型训练 + 评估报告
├── predict.py           # 推理预测
├── webui.py             # Gradio 网页界面
├── app/
│   └── api.py           # FastAPI REST 接口
├── config.py            # 全局配置
├── database.py          # SQLite 数据库
├── knowledge.py         # 垃圾分类知识库
├── split_dataset.py     # 数据集划分工具
├── download_trashnet.py # 数据集下载
├── models/              # 模型保存目录
├── dataset/             # 数据集目录
├── pyproject.toml       # 项目配置
└── requirements.txt     # 依赖列表
```

## 快速开始

### 1. 解压项目

```bash
unzip Garbage_classification.zip
cd Garbage_classification
```

### 2. 安装 Python 环境

**方式一：uv（推荐）**

```bash
# 安装 uv（如未安装）
curl -LsSf https://astral.sh/uv/install.sh | sh

# 创建虚拟环境并安装依赖
uv venv
uv sync
```

**方式二：pip**

```bash
python -m venv .venv
source .venv/bin/activate  # Windows: .venv\Scripts\activate
pip install -r requirements.txt
```

### 3. 准备数据集

```bash
# 自动下载 TrashNet 数据集并划分训练集/验证集/测试集
python download_trashnet.py
```

或手动下载后将图片放入 `dataset/` 目录，运行：

```bash
python split_dataset.py
```

### 4. 训练模型

```bash
# 使用 uv
uv run python main.py train

# 或使用 pip 虚拟环境
python main.py train
```

训练完成后会在 `models/` 目录生成：
- `garbage_model.pth` — 最佳模型权重
- `training_curves.png` — 训练 Loss/Accuracy 曲线图
- `evaluation_report.txt` — 详细评估报告（各类别准确率 + 混淆矩阵）

### 5. 启动 Web 界面

```bash
uv run python main.py webui
```

打开浏览器访问 `http://localhost:7860`

### 6. 启动 API 服务（可选，供小程序调用）

```bash
uv run python main.py api
```

API 文档访问 `http://localhost:8000/docs`

## 支持的垃圾类别

| 类别 | 英文 | 举例 |
|------|------|------|
| 🥤 塑料 | plastic | 饮料瓶、塑料袋、塑料容器 |
| 📦 纸板 | cardboard | 快递纸箱、包装盒 |
| 📄 纸张 | paper | 办公用纸、报纸、杂志 |
| 🍾 玻璃 | glass | 玻璃瓶、玻璃制品 |
| 🥫 金属 | metal | 易拉罐、金属罐、铁盒 |
| 🍂 其他垃圾 | trash | 杂物、不可回收物 |

## 模型性能

在 TrashNet 测试集上（基于 MobileNetV3-Small）：

| 类别 | 准确率 |
|------|--------|
| 纸板 | 96.77% |
| 纸张 | 96.67% |
| 金属 | 93.65% |
| 塑料 | 91.67% |
| 玻璃 | 90.79% |
| 其他垃圾 | 77.27% |
| **总体** | **92.99%** |

## CLI 命令

```bash
# 训练模型
python main.py train

# 预测单张图片
python main.py predict <图片路径>

# 查询垃圾分类知识
python main.py query <垃圾名称>

# 查看用户环保记录
python main.py record <用户名>

# 查看用户统计
python main.py stats <用户名>

# 查看排行榜
python main.py leaderboard

# 启动 Web 界面
python main.py webui

# 启动 API 服务
python main.py api
```