Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """ | |
| AI 垃圾分类助手 - 环保小能手 | |
| 主入口程序 (CLI) | |
| """ | |
| import argparse | |
| from pathlib import Path | |
| from config import DATASET_DIR, MODEL_DIR, MODEL_PATH | |
| def cmd_train(args): | |
| from train import train | |
| train(args) | |
| def cmd_predict(args): | |
| from predict import GarbageClassifier | |
| from knowledge import get_class_info | |
| classifier = GarbageClassifier(model_path=args.model_path) | |
| image_path = Path(args.image) | |
| if not image_path.exists(): | |
| print(f"✗ 图片不存在: {image_path}") | |
| return | |
| print(f"正在识别: {image_path.name}") | |
| results = classifier.predict(str(image_path), top_k=args.top_k) | |
| print(f"\n{'='*40}") | |
| print("识别结果:") | |
| print(f"{'='*40}") | |
| for i, r in enumerate(results, 1): | |
| pct = r["confidence"] * 100 | |
| bar = "█" * int(pct // 5) + "░" * (20 - int(pct // 5)) | |
| print(f" {i}. [{bar}] {r['class_name_cn']} ({pct:.1f}%)") | |
| if args.detail: | |
| info = get_class_info(results[0]["class_name"]) | |
| if info: | |
| print(f"\n{'='*40}") | |
| print(f"垃圾分类指南 - {info['name_cn']}") | |
| print(f"分类: {info['category']}") | |
| print(f"{'='*40}") | |
| print(info["disposal"]) | |
| print(f"\n💡 {info['fun_fact']}") | |
| def cmd_query(args): | |
| from knowledge import search_knowledge | |
| results = search_knowledge(args.keyword) | |
| if not results: | |
| print(f"未找到与 '{args.keyword}' 相关的知识") | |
| return | |
| for _, info in results: | |
| print(f"\n{'='*50}") | |
| print(f"{info['name_cn']} | 类别: {info['category']}") | |
| print(f"{'='*50}") | |
| print(info["description"]) | |
| print(f"\n📋 投放方法:") | |
| print(info["disposal"]) | |
| print(f"\n💡 小贴士:") | |
| for tip in info["tips"]: | |
| print(f" • {tip}") | |
| print(f"\n🎯 {info['fun_fact']}") | |
| print(f"⏱ 降解时间: {info['degradation_time']}") | |
| def cmd_record(args): | |
| from database import Database | |
| from predict import GarbageClassifier | |
| from knowledge import get_class_info | |
| db = Database() | |
| user_id = db.register_user(args.username) | |
| classifier = GarbageClassifier(model_path=args.model_path) | |
| image_path = Path(args.image) | |
| if not image_path.exists(): | |
| print(f"✗ 图片不存在: {image_path}") | |
| return | |
| results = classifier.predict(str(image_path)) | |
| best = results[0] | |
| points = db.add_record(user_id, best["class_name"], best["confidence"]) | |
| print(f"\n✓ 已记录! {best['class_name_cn']} (置信度: {best['confidence']*100:.1f}%)") | |
| print(f" +{points} 环保积分!") | |
| info = get_class_info(best["class_name"]) | |
| if info: | |
| print(f"\n📋 投放提示: {info['disposal'].split(chr(10))[0]}") | |
| def cmd_stats(args): | |
| from database import Database | |
| from knowledge import KNOWLEDGE_BASE | |
| db = Database() | |
| user = db.get_user(args.username) | |
| if not user: | |
| print(f"用户 '{args.username}' 不存在，请先使用 record 命令") | |
| return | |
| stats = db.get_user_stats(user["id"]) | |
| print(f"\n{'='*50}") | |
| print(f" 环保统计 - {stats['username']}") | |
| print(f"{'='*50}") | |
| print(f" 总分类次数: {stats['total_classifications']}") | |
| print(f" 总环保积分: {stats['total_points']}") | |
| print(f" 今日分类: {stats['today']['count']} 次") | |
| print(f" 今日积分: {stats['today']['points']}") | |
| if stats["class_distribution"]: | |
| print(f"\n 各类别分类统计:") | |
| for item in stats["class_distribution"]: | |
| cn = KNOWLEDGE_BASE.get(item["predicted_class"], {}).get("name_cn", item["predicted_class"]) | |
| print(f" • {cn}: {item['count']} 次") | |
| if stats["recent_records"]: | |
| print(f"\n 最近记录:") | |
| for r in stats["recent_records"][:5]: | |
| cn = KNOWLEDGE_BASE.get(r["predicted_class"], {}).get("name_cn", r["predicted_class"]) | |
| print(f" • {cn} | 积分: +{r['points']} | {r['created_at']}") | |
| def cmd_leaderboard(args): | |
| from database import Database | |
| db = Database() | |
| leaders = db.get_leaderboard(args.limit) | |
| if not leaders: | |
| print("暂无环保数据，快去分类吧！") | |
| return | |
| print(f"\n{'='*50}") | |
| print(" 🏆 环保积分排行榜") | |
| print(f"{'='*50}") | |
| print(f" {'排名':>4} {'用户名':<15} {'积分':<8} {'分类次数':<8}") | |
| print(f" {'-'*35}") | |
| badges = ["🥇", "🥈", "🥉"] | |
| for i, u in enumerate(leaders, 1): | |
| badge = badges[i - 1] if i <= 3 else " " | |
| print(f" {badge} {i:<2} {u['username']:<15} {u['total_points']:<8} {u['total_classifications']:<8}") | |
| def cmd_web(args): | |
| from app.api import start_api | |
| start_api(host=args.host, port=args.port) | |
| def cmd_webui(args): | |
| from webui import launch_gradio | |
| launch_gradio(server_port=args.port) | |
| def main(): | |
| parser = argparse.ArgumentParser( | |
| description="AI 垃圾分类助手 - 环保小能手", | |
| formatter_class=argparse.RawDescriptionHelpFormatter, | |
| epilog=""" | |
| 使用示例: | |
| %(prog)s train --data-dir dataset/trashnet | |
| %(prog)s predict image.jpg | |
| %(prog)s query 塑料 | |
| %(prog)s record --username 小明 image.jpg | |
| %(prog)s stats --username 小明 | |
| %(prog)s leaderboard | |
| %(prog)s web # 启动 API 服务 (供小程序调用) | |
| %(prog)s webui # 启动 Gradio 网页界面 | |
| """, | |
| ) | |
| sub = parser.add_subparsers(dest="command", help="可用命令") | |
| p = sub.add_parser("train", help="训练分类模型") | |
| p.add_argument("--data-dir", default=str(DATASET_DIR)) | |
| p.add_argument("--model-dir", default=str(MODEL_DIR)) | |
| p.add_argument("--epochs", type=int, default=30) | |
| p.add_argument("--batch-size", type=int, default=32) | |
| p.add_argument("--lr", type=float, default=0.001) | |
| p = sub.add_parser("predict", help="分类垃圾图片") | |
| p.add_argument("image") | |
| p.add_argument("--model-path", default=str(MODEL_PATH)) | |
| p.add_argument("--top-k", type=int, default=3) | |
| p.add_argument("--no-detail", dest="detail", action="store_false", default=True) | |
| p = sub.add_parser("query", help="查询垃圾分类知识") | |
| p.add_argument("keyword") | |
| p = sub.add_parser("record", help="分类并记录积分") | |
| p.add_argument("image") | |
| p.add_argument("--username", default="default") | |
| p.add_argument("--model-path", default=str(MODEL_PATH)) | |
| p = sub.add_parser("stats", help="查看个人统计") | |
| p.add_argument("--username", default="default") | |
| p = sub.add_parser("leaderboard", help="查看排行榜") | |
| p.add_argument("--limit", type=int, default=10) | |
| p = sub.add_parser("web", help="启动 API 服务") | |
| p.add_argument("--host", default="0.0.0.0") | |
| p.add_argument("--port", type=int, default=8000) | |
| p = sub.add_parser("webui", help="启动 Gradio 网页界面") | |
| p.add_argument("--port", type=int, default=7860) | |
| args = parser.parse_args() | |
| cmds = { | |
| "train": cmd_train, | |
| "predict": cmd_predict, | |
| "query": cmd_query, | |
| "record": cmd_record, | |
| "stats": cmd_stats, | |
| "leaderboard": cmd_leaderboard, | |
| "web": cmd_web, | |
| "webui": cmd_webui, | |
| } | |
| fn = cmds.get(args.command) | |
| if fn: | |
| fn(args) | |
| else: | |
| parser.print_help() | |
| if __name__ == "__main__": | |
| main() | |