diff --git a/scripts/check_env.py b/scripts/check_env.py new file mode 100755 index 0000000..6a2346d --- /dev/null +++ b/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()) diff --git a/scripts/convert_annotations.py b/scripts/convert_annotations.py new file mode 100755 index 0000000..a852b49 --- /dev/null +++ b/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()) diff --git a/scripts/export_onnx.py b/scripts/export_onnx.py new file mode 100755 index 0000000..d2dd042 --- /dev/null +++ b/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()) diff --git a/scripts/profile_inference.py b/scripts/profile_inference.py new file mode 100755 index 0000000..cbf68dc --- /dev/null +++ b/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()) diff --git a/scripts/split_dataset.py b/scripts/split_dataset.py new file mode 100755 index 0000000..3d467a9 --- /dev/null +++ b/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())