基于XGBoost的中文垃圾分类系统实战（TF-IDF + XGBoost）

基于 XGBoost 的中文垃圾分类系统实战

垃圾分类已经成为城市管理和环保的重要课题。本项目通过 XGBoost + TF-IDF 实现了中文垃圾名称的自动分类，支持"可回收垃圾、干垃圾、湿垃圾、有害垃圾"四类预测，并附带完整训练与预测代码。

一、项目介绍

本项目主要流程包括：

数据预处理 ：读取垃圾名称数据，并使用 jieba 分词。
特征工程 ：采用字符级 TF-IDF 提取文本特征。
模型训练 ：使用 XGBoost 进行多分类训练。
模型保存与加载 ：训练后的模型与向量化器通过 joblib 保存，方便后续预测。
预测接口：封装成 Python 类，输入垃圾名称即可预测类别。

二、代码解析

1. 数据预处理与分词

python 复制代码

import jieba

def load_and_cut(data):
    data['garbage_name'] = data['garbage_name'].astype(str)
    # 中文分词
    data['text_cut'] = data['garbage_name'].apply(lambda x: " ".join(jieba.lcut(x)))
    return data

这里我们使用 jieba.lcut 对中文垃圾名称进行分词处理，为后续特征工程做准备。

2. 特征工程（TF-IDF）

python 复制代码

from sklearn.feature_extraction.text import TfidfVectorizer
import joblib

def feature_engineering(data):
    vectorizer = TfidfVectorizer(
        analyzer='char', 
        ngram_range=(1, 3), 
        max_features=10000, 
        min_df=2,
        norm='l2',
        sublinear_tf=True
    )
    x = vectorizer.fit_transform(data['garbage_name'])
    y = data["type"] - 1  # 标签从 0 开始
    joblib.dump(vectorizer, '../model/tfidf_vectorizer.pkl')
    return x, y

特点：

使用字符级 n-gram 覆盖更多语义信息
sublinear_tf=True 对高频词进行对数缩放
保存向量化器，便于部署和预测

3. XGBoost 模型训练

python 复制代码

from xgboost import XGBClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, accuracy_score

def model_train(x, y):
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=23, stratify=y)
    estimator = XGBClassifier(
        n_estimators=150, max_depth=6, learning_rate=0.2,
        subsample=1.0, colsample_bytree=0.8
    )
    estimator.fit(x_train, y_train)
    y_pred = estimator.predict(x_test)
    print('模型准确率：', accuracy_score(y_test, y_pred))
    print('分类报告：\n', classification_report(y_test, y_pred))
    joblib.dump(estimator, '../model/gc_model_v1.pkl')

使用 train_test_split 分层采样，保证每类样本比例一致。

XGBoost 参数为之前 RandomizedSearchCV 调优后的最优参数。

4. 封装预测类

python 复制代码

import os
import joblib
import logging

class GarbageClassifier:
    LABEL_MAP = {0: "可回收垃圾", 1: "干垃圾", 2: "湿垃圾", 3: "有害垃圾"}

    def __init__(self, model_dir='../model'):
        self.vectorizer = None
        self.model = None
        self._load_model(model_dir)

    def _load_model(self, model_dir):
        base_dir = os.path.dirname(os.path.abspath(__file__))
        vec_path = os.path.join(base_dir, model_dir, 'tfidf_vectorizer.pkl')
        model_path = os.path.join(base_dir, model_dir, 'gc_model_v1.pkl')
        self.vectorizer = joblib.load(vec_path)
        self.model = joblib.load(model_path)

    def predict(self, text: str) -> dict:
        text_features = self.vectorizer.transform([text])
        label_id = int(self.model.predict(text_features)[0])
        return {"text": text, "label_id": label_id, "label_name": self.LABEL_MAP.get(label_id, "未知分类")}

特点：

模型和向量化器只加载一次
可多次调用 predict，性能极快
输出包含类别 ID 和名称，方便前端显示

5. 测试预测

python 复制代码

if __name__ == "__main__":
    classifier = GarbageClassifier()
    test_items = ["瓶子", "猫砂", "苹果皮", "电池"]
    for item in test_items:
        result = classifier.predict(item)
        print(f"[{result['text']}] 属于: {result['label_name']} (ID: {result['label_id']})")

输出示例：

复制代码

[瓶子] 属于: 可回收垃圾 (ID: 0)
[猫砂] 属于: 干垃圾 (ID: 1)
[苹果皮] 属于: 湿垃圾 (ID: 2)
[电池] 属于: 有害垃圾 (ID: 3)

三、运行说明

克隆项目：

bash 复制代码

git clone https://github.com/MTQ851/garbage-classification-xgboost.git
cd garbage-classification-xgboost

安装依赖：

bash 复制代码

pip install pandas scikit-learn xgboost jieba joblib

数据准备：

将垃圾分类数据放入 ../data/garbage_sorting.csv，并确保包含 garbage_name 与 type 列。
训练模型：

bash 复制代码

python train_model.py

预测测试：

bash 复制代码

python predict.py

四、项目亮点

中文文本分词 + 字符级 TF-IDF 特征工程
XGBoost 多分类训练，准确率高
模型封装成类，部署方便
支持批量预测，响应快速

五、源码地址

项目地址：https://github.com/MTQ851/garbage-classification-xgboost.git

如果你喜欢这个案例，记得 Star ⭐ 并分享给更多小伙伴！