feat(scripts): 数据转换、ONNX 导出与推理性能分析

- convert_annotations: labelme / VOC -> YOLO txt，顺带生成 data.yaml
- split_dataset: 切分成 images/{train,val} 与 labels/{train,val}，默认建软链不复制数据
- export_onnx: 导出部署用 ONNX，CPU 可跑，缺依赖时提示 `uv sync --group export`
- profile_inference: torch.profiler 测单帧延迟/FPS、CUDA 显存峰值与算子级明细
- check_env: Python / PyTorch / Nvidia GPU 自检

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

scripts/check_env.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python
+"""环境自检：确认依赖与 Nvidia GPU 是否就绪。
+
+python scripts/check_env.py
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from utils.device import NoNvidiaGpuError, cuda_devices, require_cuda  # noqa: E402
+
+
+def main() -> int:
+    print(f"Python: {sys.version.split()[0]}")
+    try:
+        import torch
+
+        print(f"PyTorch: {torch.__version__} (CUDA build: {torch.version.cuda})")
+    except ImportError:
+        print("PyTorch: 未安装 —— 先执行 `uv sync`")
+        return 1
+
+    devices = cuda_devices()
+    print(f"Nvidia GPU: {devices if devices else '无'}")
+
+    try:
+        require_cuda()
+    except NoNvidiaGpuError as exc:
+        print(f"训练不可用: {exc}")
+        print("（评估/导出仍可在 CPU 上进行）")
+        return 1
+    print("训练环境就绪。")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/convert_annotations.py

@@ -0,0 +1,135 @@
+#!/usr/bin/env python
+"""把 labelme / VOC 标注批量转换成 YOLO 风格的 .txt。
+
+    python scripts/convert_annotations.py --format labelme \
+        --input data/film/labelme --images data/film/images \
+        --output data/film/labels --data-yaml data/film/data.yaml
+
+每张图输出一个同名 .txt，每行 `class_id cx cy w h`（相对宽高归一化）；
+类别顺序默认按扫描到的名字排序，也可以用 `--classes scratch bubble ...` 固定。
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import xml.etree.ElementTree as ET
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from datasets.yolo import IMAGE_SUFFIXES, format_label_lines, write_data_yaml  # noqa: E402
+
+
+def from_labelme(path: Path) -> dict:
+    data = json.loads(path.read_text(encoding="utf-8"))
+    boxes, labels = [], []
+    for shape in data.get("shapes", []):
+        points = shape["points"]
+        xs = [p[0] for p in points]
+        ys = [p[1] for p in points]
+        boxes.append([min(xs), min(ys), max(xs), max(ys)])
+        labels.append(shape["label"])
+    return {
+        "image": Path(data.get("imagePath", path.stem + ".png")).name,
+        "width": data.get("imageWidth"),
+        "height": data.get("imageHeight"),
+        "boxes": boxes,
+        "labels": labels,
+    }
+
+
+def from_voc(path: Path) -> dict:
+    root = ET.parse(path).getroot()
+    size = root.find("size")
+    boxes, labels = [], []
+    for obj in root.findall("object"):
+        bnd = obj.find("bndbox")
+        if bnd is None:
+            continue
+        boxes.append([float(bnd.findtext(key, "0")) for key in ("xmin", "ymin", "xmax", "ymax")])
+        labels.append(obj.findtext("name", "unknown"))
+    return {
+        "image": Path(root.findtext("filename") or (path.stem + ".jpg")).name,
+        "width": int(size.findtext("width", "0")) if size is not None else None,
+        "height": int(size.findtext("height", "0")) if size is not None else None,
+        "boxes": boxes,
+        "labels": labels,
+    }
+
+
+PARSERS = {"labelme": (from_labelme, "*.json"), "voc": (from_voc, "*.xml")}
+
+
+def image_size(record: dict, images_dir: Path | None) -> tuple[int, int]:
+    """标注里没写尺寸就去读图片本身——归一化坐标必须要有宽高。"""
+    if record.get("width") and record.get("height"):
+        return int(record["width"]), int(record["height"])
+    if images_dir is None:
+        raise ValueError(f"{record['image']} 的标注里没有图像尺寸，请用 --images 指定图片目录")
+
+    from PIL import Image
+
+    stem = Path(record["image"]).stem
+    for suffix in IMAGE_SUFFIXES:
+        candidate = images_dir / f"{stem}{suffix}"
+        if candidate.exists():
+            with Image.open(candidate) as image:
+                return image.size
+    raise FileNotFoundError(f"在 {images_dir} 下找不到 {record['image']}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument("--format", choices=sorted(PARSERS), required=True)
+    parser.add_argument("--input", required=True, help="标注文件所在目录")
+    parser.add_argument("--output", required=True, help="输出的 labels 目录")
+    parser.add_argument("--images", default=None, help="图片目录；标注里没有尺寸时用来读宽高")
+    parser.add_argument("--classes", nargs="*", default=None, help="固定类别顺序（缺省按名字排序）")
+    parser.add_argument("--data-yaml", default=None, help="顺带写一份 YOLO data.yaml")
+    args = parser.parse_args()
+
+    convert, pattern = PARSERS[args.format]
+    files = sorted(Path(args.input).rglob(pattern))
+    if not files:
+        print(f"在 {args.input} 下没有找到 {pattern}")
+        return 1
+
+    records = [convert(file) for file in files]
+    classes = args.classes or sorted({label for r in records for label in r["labels"]})
+    class_to_idx = {name: i for i, name in enumerate(classes)}
+
+    unknown = {label for r in records for label in r["labels"] if label not in class_to_idx}
+    if unknown:
+        print(f"标注中出现了 --classes 未列出的类别: {sorted(unknown)}")
+        return 1
+
+    images_dir = Path(args.images) if args.images else None
+    out_dir = Path(args.output)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    total_boxes = 0
+    for record in records:
+        width, height = image_size(record, images_dir)
+        lines = format_label_lines(
+            record["boxes"], [class_to_idx[label] for label in record["labels"]], width, height
+        )
+        total_boxes += len(lines)
+        # 无缺陷图也写空文件，方便核对"标了但没有目标" vs "漏标"
+        (out_dir / f"{Path(record['image']).stem}.txt").write_text(
+            "\n".join(lines) + ("\n" if lines else ""), encoding="utf-8"
+        )
+
+    print(f"已转换 {len(records)} 个文件、{total_boxes} 个框 -> {out_dir}")
+    print(f"类别顺序（class_id 即下标）: {classes}")
+    if args.data_yaml:
+        path = write_data_yaml(args.data_yaml, classes, root=out_dir.parent)
+        print(f"data.yaml 已写入 {path}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/export_onnx.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python
+"""导出 ONNX，供实时设备（TensorRT / OpenVINO）部署。
+
+    python scripts/export_onnx.py --config configs/train/baseline.yaml \
+        --checkpoint runs/baseline/best.pth --output runs/baseline/model.onnx
+
+导出在 CPU 上完成，不需要 GPU；需要额外依赖：
+    uv sync --group export
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+import torch  # noqa: E402
+
+from datasets import load_classes  # noqa: E402
+from engine import load_checkpoint  # noqa: E402
+from models import build_model  # noqa: E402
+from utils import load_config  # noqa: E402
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument("--config", required=True)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--opset", type=int, default=17)
+    parser.add_argument("--dynamic-batch", action="store_true")
+    args = parser.parse_args()
+
+    cfg = load_config(args.config)
+    model_cfg = dict(cfg["model"])
+    model_cfg.setdefault("num_classes", len(load_classes(cfg["classes"])))
+    model_cfg.setdefault("label_mode", str(cfg.get("label_mode", "multilabel")))
+    model = build_model(model_cfg)
+    load_checkpoint(args.checkpoint, model, map_location="cpu")
+    model.eval()
+
+    height, width = cfg.get("image_size", [512, 512])
+    channels = 1 if cfg.get("gray") else 3
+    dummy = torch.randn(1, channels, int(height), int(width))
+
+    output_names = ["heatmap", "wh", "offset"] + (["logits"] if model.cls_head is not None else [])
+    dynamic_axes = {"images": {0: "batch"}} if args.dynamic_batch else None
+    if dynamic_axes:
+        dynamic_axes.update({name: {0: "batch"} for name in output_names})
+
+    try:
+        import onnxscript  # noqa: F401
+    except ImportError:
+        print("缺少导出依赖，请先执行: uv sync --group export", file=sys.stderr)
+        return 1
+
+    out_path = Path(args.output)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    torch.onnx.export(
+        model,
+        dummy,
+        str(out_path),
+        input_names=["images"],
+        output_names=output_names,
+        opset_version=args.opset,
+        dynamic_axes=dynamic_axes,
+    )
+    print(f"已导出 {out_path}（输入 {tuple(dummy.shape)}，输出 stride={model.out_stride}）")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/profile_inference.py

@@ -0,0 +1,204 @@
+#!/usr/bin/env python
+"""推理性能与显存分析（torch.profiler）。
+
+部署在实时设备上，显存占用和单帧耗时是硬指标，因此这个脚本在推理路径上跑
+`torch.profiler`，给出算子级的显存与耗时排行，外加 CUDA 显存峰值统计。
+
+    # 用配置里的输入尺寸，batch=1 跑 50 次
+    python scripts/profile_inference.py --config configs/train/baseline.yaml \
+        --checkpoint runs/baseline/best.pth
+
+    # 扫不同 batch，并导出 chrome trace（chrome://tracing 或 https://ui.perfetto.dev 打开）
+    python scripts/profile_inference.py --config configs/train/baseline.yaml \
+        --batch-size 1 4 8 --trace runs/baseline/trace.json
+
+无 GPU 时自动退回 CPU（只有耗时与 CPU 内存，没有显存数据）。
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+import torch  # noqa: E402
+from torch.profiler import ProfilerActivity, profile, record_function  # noqa: E402
+
+from datasets import load_classes  # noqa: E402
+from models import build_model, decode_heatmap  # noqa: E402
+from utils import load_config, setup_logging  # noqa: E402
+
+logger = setup_logging()
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument("--config", required=True)
+    parser.add_argument("--checkpoint", default=None, help="不给则用随机初始化的权重，只看结构开销")
+    parser.add_argument("--batch-size", type=int, nargs="+", default=[1])
+    parser.add_argument("--image-size", type=int, nargs=2, default=None, help="默认取配置的 image_size")
+    parser.add_argument("--iters", type=int, default=50, help="计时轮数")
+    parser.add_argument("--warmup", type=int, default=10, help="预热轮数（不计入统计）")
+    parser.add_argument("--device", default=None, help="默认有 GPU 用 cuda，否则 cpu")
+    parser.add_argument("--amp", action="store_true", help="用 float16 autocast 推理")
+    parser.add_argument("--postprocess", action="store_true", help="把 heatmap 解码也计入")
+    parser.add_argument("--rows", type=int, default=15, help="算子排行显示条数")
+    parser.add_argument("--trace", default=None, help="导出 chrome trace 的路径")
+    parser.add_argument("--out", default=None, help="把汇总结果写入 json")
+    return parser.parse_args()
+
+
+def build(args: argparse.Namespace, device: torch.device) -> torch.nn.Module:
+    cfg = load_config(args.config)
+    model_cfg = dict(cfg["model"])
+    model_cfg.setdefault("num_classes", len(load_classes(cfg["classes"])))
+    model_cfg.setdefault("label_mode", str(cfg.get("label_mode", "multilabel")))
+    model = build_model(model_cfg)
+
+    if args.checkpoint:
+        from engine import load_checkpoint
+
+        load_checkpoint(args.checkpoint, model, map_location="cpu")
+    model.to(device).eval()
+    model.cfg = cfg
+    return model
+
+
+@torch.no_grad()
+def run_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> None:
+    device_type = images.device.type
+    with torch.autocast(device_type, dtype=torch.float16, enabled=args.amp and device_type == "cuda"):
+        with record_function("forward"):
+            outputs = model(images)
+    if args.postprocess:
+        with record_function("postprocess"):
+            decode_heatmap(
+                outputs["heatmap"].float(),
+                outputs["wh"].float(),
+                outputs["offset"].float(),
+                stride=model.out_stride,
+            )
+
+
+def measure(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> dict:
+    """预热后计时，并统计 CUDA 显存峰值。"""
+    device = images.device
+    is_cuda = device.type == "cuda"
+
+    for _ in range(args.warmup):
+        run_once(model, images, args)
+    if is_cuda:
+        torch.cuda.synchronize(device)
+        torch.cuda.reset_peak_memory_stats(device)
+
+    start = time.perf_counter()
+    for _ in range(args.iters):
+        run_once(model, images, args)
+    if is_cuda:
+        torch.cuda.synchronize(device)
+    elapsed = time.perf_counter() - start
+
+    batch = images.shape[0]
+    result = {
+        "batch_size": batch,
+        "latency_ms": elapsed / args.iters * 1000,
+        "throughput_fps": args.iters * batch / elapsed,
+    }
+    if is_cuda:
+        result.update(
+            {
+                "peak_allocated_mb": torch.cuda.max_memory_allocated(device) / 1024**2,
+                "peak_reserved_mb": torch.cuda.max_memory_reserved(device) / 1024**2,
+                "weights_mb": sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**2,
+            }
+        )
+    return result
+
+
+def profile_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> None:
+    """算子级 profiling：按显存占用与耗时各排一次序。"""
+    is_cuda = images.device.type == "cuda"
+    activities = [ProfilerActivity.CPU] + ([ProfilerActivity.CUDA] if is_cuda else [])
+
+    with profile(activities=activities, profile_memory=True, record_shapes=True) as prof:
+        for _ in range(max(args.iters // 5, 1)):
+            run_once(model, images, args)
+    if is_cuda:
+        torch.cuda.synchronize(images.device)
+
+    memory_key = "self_cuda_memory_usage" if is_cuda else "self_cpu_memory_usage"
+    time_key = "self_cuda_time_total" if is_cuda else "self_cpu_time_total"
+    print(f"\n===== 显存占用 Top {args.rows}（按 {memory_key}）=====")
+    print(prof.key_averages().table(sort_by=memory_key, row_limit=args.rows))
+    print(f"\n===== 耗时 Top {args.rows}（按 {time_key}）=====")
+    print(prof.key_averages().table(sort_by=time_key, row_limit=args.rows))
+
+    if args.trace:
+        trace_path = Path(args.trace)
+        trace_path.parent.mkdir(parents=True, exist_ok=True)
+        prof.export_chrome_trace(str(trace_path))
+        print(f"\nchrome trace 已导出: {trace_path}（用 chrome://tracing 或 ui.perfetto.dev 打开）")
+
+
+def main() -> int:
+    args = parse_args()
+    device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    if device.type != "cuda":
+        logger.warning("当前在 %s 上运行：没有显存数据，只能看耗时与 CPU 内存。", device)
+
+    model = build(args, device)
+    cfg = model.cfg
+    height, width = args.image_size or cfg.get("image_size", [512, 512])
+    channels = 1 if cfg.get("gray") else 3
+
+    logger.info(
+        "模型 %s | 参数量 %.2fM | 输出 stride %d | 输入 %dx%dx%d | device %s | amp=%s",
+        cfg["model"]["name"],
+        sum(p.numel() for p in model.parameters()) / 1e6,
+        model.out_stride,
+        channels,
+        height,
+        width,
+        device,
+        args.amp,
+    )
+
+    results = []
+    for batch_size in args.batch_size:
+        images = torch.randn(batch_size, channels, int(height), int(width), device=device)
+        result = measure(model, images, args)
+        results.append(result)
+        memory = (
+            f" | 显存峰值 allocated {result['peak_allocated_mb']:.1f} MB"
+            f" / reserved {result['peak_reserved_mb']:.1f} MB"
+            if "peak_allocated_mb" in result
+            else ""
+        )
+        logger.info(
+            "batch=%-3d 单次 %.2f ms | %.1f FPS%s",
+            result["batch_size"],
+            result["latency_ms"],
+            result["throughput_fps"],
+            memory,
+        )
+
+    # 算子级明细只对最大的 batch 跑一次，避免输出过长。
+    images = torch.randn(max(args.batch_size), channels, int(height), int(width), device=device)
+    profile_once(model, images, args)
+
+    if args.out:
+        out_path = Path(args.out)
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        out_path.write_text(json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8")
+        logger.info("汇总已写入 %s", out_path)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/split_dataset.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python
+"""按比例切分成 YOLO 的 train / val / test 目录结构。
+
+    python scripts/split_dataset.py --images data/film/images --labels data/film/labels \
+        --val-ratio 0.2 --data-yaml data/film/data.yaml
+
+默认建软链（不复制数据），产出：
+
+```
+images/train/  images/val/
+labels/train/  labels/val/
+```
+
+原始文件保留在 `images/` `labels/` 根目录下不动；重复执行会先清掉旧的划分目录。
+"""
+
+from __future__ import annotations
+
+import argparse
+import random
+import shutil
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from datasets.yolo import IMAGE_SUFFIXES, load_classes, write_data_yaml  # noqa: E402
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument("--images", required=True, help="图片目录（划分前的平铺目录）")
+    parser.add_argument("--labels", default=None, help="标注目录，缺省为 images 同级的 labels")
+    parser.add_argument("--val-ratio", type=float, default=0.2)
+    parser.add_argument("--test-ratio", type=float, default=0.0)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--copy", action="store_true", help="复制文件而不是建软链")
+    parser.add_argument("--data-yaml", default=None, help="顺带更新/生成 data.yaml")
+    parser.add_argument("--classes", nargs="*", default=None, help="写 data.yaml 时的类别顺序")
+    args = parser.parse_args()
+
+    images_dir = Path(args.images)
+    labels_dir = Path(args.labels) if args.labels else images_dir.parent / "labels"
+    splits = ("train", "val", "test")
+
+    files = sorted(
+        p
+        for p in images_dir.iterdir()
+        if p.is_file() and p.suffix.lower() in IMAGE_SUFFIXES and p.parent.name not in splits
+    )
+    if not files:
+        print(f"{images_dir} 下没有待划分的图片（子目录里的已划分文件会被跳过）")
+        return 1
+
+    random.Random(args.seed).shuffle(files)
+    n_val = int(len(files) * args.val_ratio)
+    n_test = int(len(files) * args.test_ratio)
+    partition = {
+        "val": files[:n_val],
+        "test": files[n_val : n_val + n_test],
+        "train": files[n_val + n_test :],
+    }
+
+    for name, items in partition.items():
+        img_out, label_out = images_dir / name, labels_dir / name
+        for directory in (img_out, label_out):
+            if directory.exists():
+                shutil.rmtree(directory)
+        if not items:
+            continue
+        img_out.mkdir(parents=True, exist_ok=True)
+        label_out.mkdir(parents=True, exist_ok=True)
+
+        empty = 0
+        for image in items:
+            label = labels_dir / f"{image.stem}.txt"
+            place(image, img_out / image.name, args.copy)
+            if label.exists():
+                place(label, label_out / label.name, args.copy)
+                empty += int(not label.read_text(encoding="utf-8").strip())
+            else:
+                empty += 1  # 缺 .txt 等价于无缺陷
+        print(f"{name}: {len(items)} 张（其中无缺陷 {empty} 张） -> {img_out}")
+
+    if args.data_yaml:
+        classes = load_classes(args.classes or args.data_yaml)
+        path = write_data_yaml(args.data_yaml, classes, root=images_dir.parent)
+        print(f"data.yaml 已更新 {path}")
+    return 0
+
+
+def place(src: Path, dst: Path, copy: bool) -> None:
+    if dst.exists() or dst.is_symlink():
+        dst.unlink()
+    if copy:
+        shutil.copy2(src, dst)
+    else:
+        dst.symlink_to(src.resolve())
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())