README.md

metadata

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

♻️ AI 垃圾分类助手

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

GitHub 仓库

技术栈

模块	技术
深度学习框架	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. 解压项目

unzip Garbage_classification.zip
cd Garbage_classification

2. 安装 Python 环境

方式一：uv（推荐）

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

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

方式二：pip

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

3. 准备数据集

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

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

python split_dataset.py

4. 训练模型

# 使用 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 界面

uv run python main.py webui

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

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

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 命令

# 训练模型
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