run.py

"""缺陷合成流水线 CLI 入口。

瓶体褶皱(demo 架构, 配置 demo_pipeline.yaml):
  python run.py demo-gen --products 500多效极润 --per-angle 125 --dry-run  # 离线出计划
  python run.py demo-gen --products 500多效极润 --per-angle 125            # 前台跑完整流程
  python run.py demo-gen --products 500多效极润 --per-angle 125 --detach   # 任务计划后台跑
  python run.py demo-report --run-dir <目录> [--anchors <json>]             # 重算指标、重出打分表
  python run.py demo-job --job <job.json>                                  # 仅供 tools/demo_task.ps1 调用
  python run.py demo-selftest                                              # 离线自检

PCB 缺陷网站(配置 pcb_web.yaml):
  python run.py pcb-web [--host 0.0.0.0] [--port 8000]                     # 启动实验室网站
  python run.py pcb-lib add --image 缺陷.jpg --box x0,y0,x1,y1 --preset bridge  # 参考缺陷入库
  python run.py pcb-lib list                                               # 列出缺陷库

缺陷库:
  python run.py test-api                    # 测试中转站连通性(视觉+编辑模型)
  python run.py inspect                     # 离线体检缺陷库
  python run.py build-catalog [--max-per-class N] [--overwrite] [--classes 类别,类别]
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from src.config import load_config
from src.demo_pipeline import cli as demo_cli


def cmd_test_api(args):
    from src.api_client import RelayClient
    from PIL import Image
    cfg = load_config(args.config)
    relay = RelayClient(cfg)
    print(f"base_url = {relay.base_url}")

    print("\n[1/2] 测试视觉模型 ...")
    img = Image.new("RGB", (64, 64), (120, 120, 120))
    try:
        txt = relay.chat_vision("用一句话描述这张图的主色。", images=[img])
        print("  OK ->", txt[:120])
    except Exception as e:
        print("  失败:", e)

    print("\n[2/2] 测试图像编辑模型 ...")
    try:
        out = relay.edit_image(
            "Add a tiny realistic dark scratch in the center. Keep everything else identical.",
            img)
        print(f"  OK -> 返回图像 size={out.size}")
    except Exception as e:
        print("  失败:", e)


def cmd_inspect(args):
    """离线体检缺陷库: 格式/类型/严重度/字段完整度/裁剪图是否存在。"""
    from src import catalog_report
    cfg = load_config(args.config)
    report = catalog_report.build_report(cfg)
    for line in catalog_report.format_lines(report):
        print(line)


def cmd_build_catalog(args):
    from src.defect_catalog import build_catalog, load_analyzed
    cfg = load_config(args.config)
    only = args.classes.split(",") if args.classes else None

    if args.dry_run:
        import json as _json
        from src.dataset import scan_class_images
        cat = cfg.out_path("catalog") / "catalog.json"
        recs = (_json.loads(cat.read_text(encoding="utf-8"))
                if cat.exists() else [])
        analyzed = load_analyzed(cfg) if recs and not args.overwrite else set()
        imgs = scan_class_images(cfg, cfg.defect_root())
        if only:
            imgs = {k: v for k, v in imgs.items() if k in only}
        todo_total = 0
        print(f"{'类别':8} {'总图数':>7} {'已分析':>7} {'本次待分析':>10}")
        for cls in sorted(imgs):
            paths = imgs[cls]
            done = sum(1 for p in paths if str(p.resolve()) in analyzed)
            todo = len(paths) - done
            if args.max_per_class:
                todo = min(todo, args.max_per_class)
            todo_total += todo
            print(f"{cls:8} {len(paths):>7} {done:>7} {todo:>10}")
        rate = (len(recs) / max(1, len(analyzed))) if analyzed else 1.36
        print(f"\n本次将调用视觉模型 {todo_total} 次(每张图一次)")
        print(f"按当前产出率 {rate:.2f} 条/图估算, 预计新增约 "
              f"{int(todo_total * rate)} 条, 总量约 "
              f"{len(recs) + int(todo_total * rate)} 条")
        workers = max(1, int(cfg.api.get("max_workers", 3)))
        print(f"并发 {workers}, 按单次 15~30 秒估算, 预计耗时 "
              f"{todo_total * 15 / workers / 60:.0f}~"
              f"{todo_total * 30 / workers / 60:.0f} 分钟")
        print("\n(试算, 未调用任何 API。确认后去掉 --dry-run 即开始)")
        return

    build_catalog(cfg, max_per_class=args.max_per_class,
                  overwrite=args.overwrite, only_classes=only)


def cmd_pcb_web(args):
    from src.pcb_web.server import serve
    return serve(args.host, args.port, args.web_config)


def cmd_pcb_lib(args):
    """PCB 缺陷库: 批量入库或列出。和网站共用 pcb_web.yaml 的 data_dir。"""
    from PIL import Image
    from src.pcb_web import generator
    from src.pcb_web.library import Library
    from src.pcb_web.server import load_web_config
    lib = Library(load_web_config(args.web_config).data_path())
    if args.action == "list":
        for e in lib.list():
            print(f"{e['id']}  {e['label']:<8} {e['size'][0]}x{e['size'][1]}  用过 {e['uses']} 张  {e['source']}  {e['note']}")
        print(f"共 {len(lib.list())} 条, 目录 {lib.root}")
        return 0
    if not args.image or args.preset not in generator.PRESETS:
        print(f"add 需要 --image 和 --preset({'/'.join(generator.PRESETS)})")
        return 2
    im = Image.open(args.image).convert("RGB")
    box = None
    if args.box:
        box = tuple(int(v) for v in args.box.split(","))
        if len(box) != 4 or not (0 <= box[0] < box[2] <= im.width and 0 <= box[1] < box[3] <= im.height):
            print(f"--box 需要 x0,y0,x1,y1 且在图内, 图片尺寸 {im.width}x{im.height}")
            return 2
        im = im.crop(box)
    e = lib.add(im, args.preset, source=args.image, box=box, note=args.note)
    print(f"已入库 {e['id']} {e['label']} {e['size'][0]}x{e['size'][1]}")
    return 0


def _configure_stdio() -> None:
    # 控制台默认 GBK 时, 中文和特殊字符不能让日志输出崩掉
    for stream in (sys.stdout, sys.stderr):
        try:
            stream.reconfigure(encoding="utf-8", errors="replace")
        except (AttributeError, ValueError):
            pass


def main():
    ap = argparse.ArgumentParser(description="工业缺陷合成流水线")
    ap.add_argument("--config", default=None, help="config.yaml 路径")
    sub = ap.add_subparsers(dest="cmd", required=True)
    demo_cli.register(sub)

    sub.add_parser("test-api").set_defaults(func=cmd_test_api)
    sub.add_parser("inspect", help="离线体检缺陷库(不调用API)").set_defaults(
        func=cmd_inspect)

    p = sub.add_parser(
        "build-catalog",
        help="编目缺陷库。已有库时默认增量(只分析没分析过的图), 结果追加")
    p.add_argument("--max-per-class", type=int, default=None,
                   help="每类最多分析多少张缺陷图(控制成本)")
    p.add_argument("--dry-run", action="store_true",
                   help="只试算规模/成本/耗时, 不调用 API")
    p.add_argument("--overwrite", action="store_true",
                   help="全量重建(丢弃现有库并重新分析所有图, 慎用)")
    p.add_argument("--classes", type=str, default=None,
                   help="仅(重)建指定类别, 逗号分隔, 如 500水次方_角度1,500水次方_角度2")
    p.set_defaults(func=cmd_build_catalog)

    p = sub.add_parser("pcb-web", help="启动 PCB 缺陷生成网站")
    p.add_argument("--host", default=None, help="监听地址, 默认取 pcb_web.yaml")
    p.add_argument("--port", type=int, default=None, help="监听端口, 默认取 pcb_web.yaml")
    p.add_argument("--web-config", default=None, help="配置文件, 默认 pcb_web.yaml")
    p.set_defaults(func=cmd_pcb_web)

    p = sub.add_parser("pcb-lib", help="PCB 缺陷库: add 入库 / list 列出")
    p.add_argument("action", choices=["add", "list"])
    p.add_argument("--image", default=None, help="参考缺陷图")
    p.add_argument("--box", default=None, help="缺陷框 x0,y0,x1,y1(原图像素); 不给就用整张图")
    p.add_argument("--preset", default=None, help="缺陷类型: debris/flake/bridge/strip/custom")
    p.add_argument("--note", default="", help="备注")
    p.add_argument("--web-config", default=None, help="配置文件, 默认 pcb_web.yaml")
    p.set_defaults(func=cmd_pcb_lib)

    args = ap.parse_args()
    _configure_stdio()
    # 子命令可以返回退出码; 返回 None 视为成功, 保持原有行为
    return args.func(args) or 0


if __name__ == "__main__":
    raise SystemExit(main())