145 lines
5.1 KiB
Python
145 lines
5.1 KiB
Python
# coding=utf-8
|
||
"""一键导入模板：文件夹内图片文件名 = 菜名。
|
||
|
||
用法（在项目根目录）::
|
||
|
||
python tools/import_templates_by_name.py "D:\\菜品模板"
|
||
python tools/import_templates_by_name.py "D:\\菜品模板" --dry-run
|
||
python tools/import_templates_by_name.py "D:\\菜品模板" --local-only
|
||
|
||
说明:
|
||
- 旧版逐张 add_template 时，每张都会从 MySQL 全量/增量同步，容易把本地刚写入的冲掉
|
||
- 本脚本改为批量写入，最后只同步一次
|
||
- --local-only：只写本机 SQLite（并入上传队列），不立刻连 MySQL 拉取
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
_ROOT = Path(__file__).resolve().parents[1]
|
||
if str(_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(_ROOT))
|
||
|
||
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
|
||
|
||
|
||
def collect_images(folder: Path, recursive: bool) -> list[Path]:
|
||
files = folder.rglob("*") if recursive else folder.iterdir()
|
||
out = [p for p in files if p.is_file() and p.suffix.lower() in IMAGE_EXTS]
|
||
out.sort(key=lambda x: x.name)
|
||
return out
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="按文件名（菜名）批量导入模板")
|
||
parser.add_argument("folder", help="模板图片文件夹路径")
|
||
parser.add_argument("--recursive", "-r", action="store_true", help="包含子文件夹")
|
||
parser.add_argument("--dry-run", action="store_true", help="只列出，不写入")
|
||
parser.add_argument(
|
||
"--local-only",
|
||
action="store_true",
|
||
help="强制只写本地库（推荐客户机 MySQL 不通时使用）",
|
||
)
|
||
parser.add_argument(
|
||
"--fail-log",
|
||
default="import_fail.log",
|
||
help="失败明细日志路径（相对项目根或绝对路径）",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
folder = Path(args.folder).expanduser().resolve()
|
||
if not folder.is_dir():
|
||
print(f"[ERR] 文件夹不存在: {folder}")
|
||
sys.exit(1)
|
||
|
||
images = collect_images(folder, args.recursive)
|
||
if not images:
|
||
print(f"[ERR] 未找到图片: {folder}")
|
||
sys.exit(1)
|
||
|
||
pairs = []
|
||
for p in images:
|
||
name = p.stem.strip()
|
||
if not name:
|
||
print(f"[WARN] 跳过空文件名: {p}")
|
||
continue
|
||
pairs.append((name, str(p)))
|
||
|
||
print(f"文件夹: {folder}")
|
||
print(f"图片数: {len(images)} | 待导入: {len(pairs)}")
|
||
if args.dry_run:
|
||
from collections import Counter
|
||
c = Counter(n for n, _ in pairs)
|
||
print(f"菜品种类: {len(c)}")
|
||
for name, n in c.most_common(20):
|
||
print(f" {name}: {n} 张")
|
||
if len(c) > 20:
|
||
print(f" ... 其余 {len(c) - 20} 个菜品")
|
||
print("[dry-run] 未写入。")
|
||
return
|
||
|
||
if args.local_only:
|
||
# 让 is_server_available() 在本进程内始终失败 → 走本地插入 + 上传队列
|
||
import dao.database as db
|
||
|
||
db.is_server_available = lambda: False # type: ignore
|
||
print("[INFO] --local-only：只写本机 SQLite，不拉取 MySQL")
|
||
|
||
print("加载特征提取器 / 模板库（首次可能较慢）…")
|
||
from core.common_classifier import PartTemplateManager, TEMPLATE_SAVE_DIR
|
||
from dao.database import local_session_scope
|
||
from models.Global_templates import GlobalTemplate
|
||
from models.Dish_categories import DishCategory
|
||
from sqlalchemy import func
|
||
|
||
manager = PartTemplateManager()
|
||
|
||
def progress(i, total, name, path):
|
||
if i == 1 or i % 20 == 0 or i == total:
|
||
print(f" 提特征 [{i}/{total}] {name}", flush=True)
|
||
|
||
print("开始批量提特征并写入…")
|
||
ok, failures = manager.add_templates_batch(pairs, progress_callback=progress)
|
||
|
||
fail_log = Path(args.fail_log)
|
||
if not fail_log.is_absolute():
|
||
fail_log = _ROOT / fail_log
|
||
if failures:
|
||
with open(fail_log, "w", encoding="utf-8") as f:
|
||
for path, reason in failures:
|
||
f.write(f"{reason}\t{path}\n")
|
||
print(f"[WARN] 失败 {len(failures)} 张，明细: {fail_log}")
|
||
|
||
with local_session_scope() as session:
|
||
n_cat = (
|
||
session.query(func.count(DishCategory.id))
|
||
.filter(DishCategory.is_active == 1)
|
||
.scalar()
|
||
) or 0
|
||
n_tpl = (
|
||
session.query(func.count(GlobalTemplate.id))
|
||
.filter(GlobalTemplate.is_active == 1)
|
||
.scalar()
|
||
) or 0
|
||
|
||
print()
|
||
print("========== 完成 ==========")
|
||
print(f"本次成功写入: {ok}")
|
||
print(f"本次失败: {len(failures)}")
|
||
print(f"本地库当前: {n_cat} 个菜品 / {n_tpl} 张模板")
|
||
print(f"模板图片目录: {TEMPLATE_SAVE_DIR}")
|
||
if n_tpl < len(pairs) * 0.5:
|
||
print()
|
||
print("[提示] 模板数远小于图片数，常见原因：")
|
||
print(" 1) 特征提取失败（看 import_fail.log）")
|
||
print(" 2) 打开识别软件后 MySQL 全量同步把本地冲掉了")
|
||
print(" → 请用本脚本加 --local-only 重新导入，导入后先别连总库同步")
|
||
print("请重启识别程序查看模板库。")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|