config.py

"""
AI 垃圾分类助手 - 全局配置
集中管理路径和参数
"""

from pathlib import Path

# 项目根目录
ROOT = Path(__file__).parent

# 原始数据集 (手动下载)
RAW_DATASET_DIR = ROOT / "dataset" / "trashnet"

# 划分后的数据集 (split_dataset.py 生成)
DATASET_DIR = ROOT / "dataset"
TRAIN_DIR = DATASET_DIR / "train"
VAL_DIR = DATASET_DIR / "val"
TEST_DIR = DATASET_DIR / "test"

# 模型
MODEL_DIR = ROOT / "models"
MODEL_PATH = MODEL_DIR / "garbage_model.pth"

# 上传目录 (API)
UPLOAD_DIR = ROOT / "uploads"

# 数据库
DB_PATH = ROOT / "garbage_assistant.db"

# TrashNet 类别 (6类)
CLASS_NAMES = ["cardboard", "glass", "metal", "paper", "plastic", "trash"]
CLASS_NAMES_CN = ["纸板", "玻璃", "金属", "纸张", "塑料", "其他垃圾"]

# 训练参数
TRAIN_PARAMS = {
    "epochs": 30,
    "batch_size": 32,
    "lr": 0.001,
    "input_size": 224,
    "resize_size": 256,
}

# 推理参数
INFER_PARAMS = {
    "top_k": 3,
}