From 2b77ec127d24379c660c6eb2450b9c58b8efd3a3 Mon Sep 17 00:00:00 2001 From: arm_dev Date: Wed, 23 Sep 2026 12:39:18 +0800 Subject: [PATCH] =?UTF-8?q?P4=EF=BC=9A=E6=9D=83=E9=87=8D/=E6=A0=87?= =?UTF-8?q?=E7=AD=BE=E6=A8=A1=E5=9D=97=20=E2=80=94=E2=80=94=20=E6=97=81?= =?UTF-8?q?=E6=8C=82=E6=A0=87=E7=AD=BE=E6=96=87=E4=BB=B6=20+=20PT=E2=86=92?= =?UTF-8?q?ONNX=E2=86=92TRT=20=E8=BD=AC=E6=8D=A2=E5=B7=A5=E5=85=B7=20+=20?= =?UTF-8?q?=E5=BD=92=E6=A1=A3=E5=BD=92=E4=BD=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 痛点：TensorRT 10.13 的 Python API 没有元数据读写，engine 一转就把类名丢了； 推理侧只能退化成 str(cid)（"0"、"1"），界面上看不出是什么类，属于静默失效。 同时旧的 export_yolo_trt.sh 只能手工搬文件、手工改 site.yaml、标签无处安放。 标签子系统（model_labels.py） - 旁挂标签文件 <模型名>.labels.yaml，放在权重旁边，文本、可手改、进 git （二进制权重不进 git，标签是配置所以进） - 读取优先级：旁挂文件 > ONNX 元数据 > 模型元数据 > 占位符 class_N 占位符是故意的：缺标签时要一眼看出来，而不是悄悄退回 cid - 解析器容忍多种写法：dict / list / 逗号串 / 逐行文本 / Python repr 串 布局与转换（weights_layout.py + weights.py + weights.sh） - models/<分组>/<模型名>.{pt,onnx,engine} + 同名 .labels.yaml - convert：PT→ONNX→TensorRT，支持自定义名字、自动归位、固化标签 - staging/：非标准格式先进这里，转换后源文件归档到 staging/done/ - list / labels / doctor / show 四个子命令 - doctor --check-engines 能查出 engine 与 TensorRT 版本不匹配（反序列化返回空） 推理侧接入 - yolo_detector / onnx_workpiece / hole_detect_hough 统一走 read_labels() - 画框标签不再硬编码 "gongjian"，改用实际类名 顺手修掉的真实问题 - setup.py 撞上断掉的符号链接会让整包构建失败，报错还指向 build/ 里的幽灵文件 （源文件被挪走后 build/ 里留下的断链被 os.walk 列出来）。加 _existing() 过滤。 - P3 遗漏：site.yaml / export_site_env.py / run_handeye.sh / handeye_teach.py / handeye_calibrate.py 里的 handeye 工作目录仍是 /tmp/handeye_calib*， 与已迁到 handeye/ 的其余路径不一致。统一到 handeye/tmp/<臂号>。 - site.yaml 的 hole_weights / hole_detect.weights 原指向不存在的 best.pt， 改为实际存在的 detect_best_1.onnx；active.pt 符号链接已废弃并清理。 - weights.sh labels 原先不接受 "分组/模型名" 写法，而 list 输出正是这个形式， 照着敲会报"找不到模型"。补上，并让报错直接给出可复制的全名。 测试：225 项通过（新增 test_model_labels.py / test_weights_layout.py）。 真实 ONNX→engine 端到端转换已验过。 Co-authored-by: Cursor --- .gitignore | 26 +- config/site.example.yaml | 10 +- handeye/waypoints.yaml | 2 +- models/README.md | 273 ++++++ models/hole_insert/README.md | 35 +- models/hole_insert/detect_best_1.labels.yaml | 16 + models/staging/.gitkeep | 0 models/staging/README.md | 26 + models/yolo/README.md | 49 +- models/yolo/detect_best.labels.yaml | 16 + models/yolo/detect_best_0.labels.yaml | 16 + models/yolo/detect_hc90.labels.yaml | 16 + .../yolo/detect_yolo26_gongjian.labels.yaml | 16 + scripts/README.md | 6 +- scripts/export_yolo_trt.sh | 107 +-- scripts/lib/export_site_env.py | 14 +- scripts/run_handeye.sh | 2 +- scripts/show_site_config.sh | 13 +- scripts/weights.py | 842 ++++++++++++++++++ scripts/weights.sh | 69 ++ .../arm_application/handeye_calibrate.py | 2 +- .../arm_application/handeye_teach.py | 2 +- .../arm_application/hole_detect_hough.py | 53 +- .../arm_application/model_labels.py | 555 ++++++++++++ .../arm_application/onnx_workpiece.py | 29 +- .../arm_application/perception_node.py | 6 +- .../arm_application/weights_layout.py | 403 +++++++++ .../arm_application/yolo_detector.py | 70 +- src/arm_application/config/perception.yaml | 16 +- .../launch/hole_detect.launch.py | 2 +- src/arm_application/setup.py | 24 +- src/arm_application/test/test_model_labels.py | 479 ++++++++++ .../test/test_weights_layout.py | 754 ++++++++++++++++ 33 files changed, 3797 insertions(+), 152 deletions(-) create mode 100644 models/README.md create mode 100644 models/hole_insert/detect_best_1.labels.yaml create mode 100644 models/staging/.gitkeep create mode 100644 models/staging/README.md create mode 100644 models/yolo/detect_best.labels.yaml create mode 100644 models/yolo/detect_best_0.labels.yaml create mode 100644 models/yolo/detect_hc90.labels.yaml create mode 100644 models/yolo/detect_yolo26_gongjian.labels.yaml create mode 100755 scripts/weights.py create mode 100755 scripts/weights.sh create mode 100644 src/arm_application/arm_application/model_labels.py create mode 100644 src/arm_application/arm_application/weights_layout.py create mode 100644 src/arm_application/test/test_model_labels.py create mode 100644 src/arm_application/test/test_weights_layout.py diff --git a/.gitignore b/.gitignore index cd471c4..d7f122c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,16 +1,22 @@ # 本机现场配置（含 IP / 序列号；从 site.example.yaml 复制） config/site.yaml -# YOLO 二进制 -models/yolo/*.pt -models/yolo/*.onnx -models/yolo/*.engine -models/yolo/*.trt -models/yolo/active.pt -models/hole_insert/*.pt -models/hole_insert/*.onnx -models/hole_insert/*.engine -models/hole_insert/*.trt +# 模型权重二进制 —— 一律不进 git（体积大、每台机器重新转换即可） +# 用通配模式而不是逐个目录列举，否则新增一个分组就会有人误提交二进制。 +# 注意：*.labels.yaml 是**文本标签配置**，不属于二进制，应该进 git。 +models/**/*.pt +models/**/*.pth +models/**/*.onnx +models/**/*.engine +models/**/*.trt +models/**/*.plan +models/**/*.tflite + +# staging = 待转换的临时区，内容不进 git（含转换后归档的 done/） +# 但 .gitkeep 和说明文档要留着，否则新克隆的仓库里这个目录不存在/看不懂。 +models/staging/** +!models/staging/.gitkeep +!models/staging/README.md # 编译产物 build/ diff --git a/config/site.example.yaml b/config/site.example.yaml index 4ff2094..db47932 100644 --- a/config/site.example.yaml +++ b/config/site.example.yaml @@ -50,7 +50,7 @@ robot: paths: handeye_file: "/home/USER/ros2/handeye/results/89_eye_in_hand.yaml" handeye_waypoints: "/home/USER/ros2/handeye/poses/89_eye_in_hand.yaml" - handeye_workdir: "/tmp/handeye_calib_89" + handeye_workdir: "/home/USER/ros2/handeye/tmp/89" # 左臂跟踪高度（只在 --arm 89 时生效，改 91 的 servo 不影响这里） servo: ee_z_m: 0.35 # 跟踪锁高：法兰相对基座 Z [m] @@ -87,7 +87,7 @@ robot: paths: handeye_file: "/home/USER/ros2/handeye/results/91_eye_in_hand.yaml" handeye_waypoints: "/home/USER/ros2/handeye/poses/91_eye_in_hand.yaml" - handeye_workdir: "/tmp/handeye_calib_91" + handeye_workdir: "/home/USER/ros2/handeye/tmp/91" # >>> 右臂跟踪高度：只改这里，89 仍用上面那段 <<< # 法兰 Z；RGI 比平行爪高，通常要比 89 略抬一点。改完重启： # bash ~/ros2/scripts/run_visual_servo_rt.sh --arm 91 --control @@ -117,12 +117,12 @@ paths: # yolo_weights: "/home/USER/ros2/models/yolo/detect_hc90.pt" # yolo_weights: "/home/USER/ros2/models/yolo/detect_yolo26_gongjian.pt" # 入孔 YOLO（backend=yolo 时把孔当框）。two_stage 用上面的工件权重。 - hole_weights: "/home/USER/ros2/models/hole_insert/best.pt" + hole_weights: "/home/USER/ros2/models/hole_insert/detect_best_1.onnx" # 下面是缺省（未匹配到 robots[] 时）。真机以 robots[].paths 为准：89 / 91 各一份。 handeye_file: "/home/USER/ros2/handeye/results/89_eye_in_hand.yaml" handeye_waypoints: "/home/USER/ros2/handeye/poses/89_eye_in_hand.yaml" handeye_calibrate: "/home/USER/ros2/handeye/calibrate.yaml" - handeye_workdir: "/tmp/handeye_calib_89" + handeye_workdir: "/home/USER/ros2/handeye/tmp/89" # 空=默认 DDS(SHM)；勿强制 UDP-only，否则相机预览收不到 /workpiece/detection fastdds_profile: "" @@ -340,7 +340,7 @@ gripper: # bash ~/ros2/scripts/run_hole_detect.sh --arm 89 --preview hole_detect: backend: "two_stage" # two_stage | yolo | hough - weights: "/home/USER/ros2/models/hole_insert/best.pt" + weights: "/home/USER/ros2/models/hole_insert/detect_best_1.onnx" # 二阶段工件检测权重（YOLO 先框出工件，再 Hough 找 k1/k2 螺纹孔） workpiece_weights: "/home/USER/ros2/models/hole_insert/detect_best_1.onnx" config: "/home/USER/ros2/src/arm_application/config/hole_detect.yaml" diff --git a/handeye/waypoints.yaml b/handeye/waypoints.yaml index d6453d1..d391adb 100644 --- a/handeye/waypoints.yaml +++ b/handeye/waypoints.yaml @@ -20,6 +20,6 @@ handeye_waypoints: camera_namespace: "" image_topic: "" preview_window: true - append_poses_file: "/tmp/handeye_calib/taught_poses.txt" + append_poses_file: "/home/cc/ros2/handeye/tmp/taught_poses.txt" waypoint_times: [4.0, 4.0, 4.0, 4.0, 4.0, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.0, 4.0] diff --git a/models/README.md b/models/README.md new file mode 100644 index 0000000..06ab931 --- /dev/null +++ b/models/README.md @@ -0,0 +1,273 @@ +# 算法权重模块（models/） + +这个目录放**算法模型权重**和它们的**类别标签**。机器人软件里"算法"和"机械/电气" +变化节奏完全不同：机械改一次要几个月，模型换一版可能一周几次。所以这一块单独 +成模块，有自己的一套约定和工具，改模型不需要动任何 C++/节点代码。 + +> **要改这个模块的人，先读完本文。** 下面是规则，不是建议。 + +--- + +## 1. 目录长什么样 + +``` +models/ +├── README.md ← 本文（模块规则） +├── staging/ ← 非标准格式先丢这里，等转换 +│ ├── .gitkeep +│ └── done/ ← 转换完成后源文件自动归档到这里 +├── yolo/ ← 一个"分组"（就是个子目录，名字随意） +│ ├── detect_best_0.onnx +│ ├── detect_best_0.engine +│ └── detect_best_0.labels.yaml +└── hole_insert/ + ├── detect_best_1.onnx + └── detect_best_1.labels.yaml +``` + +**分组**只是一个子目录名，用来分门别类，没有程序语义。默认用 `yolo`。 + +--- + +## 2. 命名规则（重要） + +同一个模型的**所有格式 + 标签**必须用**同一个名字主干**、放在**同一个目录**： + +| 文件 | 作用 | 进 git？ | +|---|---|---| +| `<名字>.pt` | PyTorch 原始权重 | 否 | +| `<名字>.onnx` | 中间格式，CPU 推理后端也用 | 否 | +| `<名字>.engine` | TensorRT 引擎，GPU 推理用 | 否 | +| `<名字>.labels.yaml` | **类别标签** | **是** | + +因为推理代码是按"同目录同主干"去找配套文件的（`resolve_weights` 找 `.engine`， +标签找 `.labels.yaml`）。所以： + +- ✅ `detect_best_0.onnx` + `detect_best_0.engine` + `detect_best_0.labels.yaml` +- ❌ `detect_best_0.onnx` + `detect_best_0_fp16.engine`（名字对不上，找不到） +- ❌ `detect_best_0.onnx` + `labels_detect_best_0.yaml`（标签名对不上） + +**一份标签服务同一个模型的所有格式。** 换 `.pt`→`.engine` 后端时类名不会变。 + +--- + +## 3. 标签（class names） + +### 为什么要有单独的标签文件 + +模型输出的是**类别索引**（0/1/2…），要翻译成人看得懂的名字。这件事原先没有统一 +做法，三种后端各来一套，而且**TensorRT 那一环会把标签弄丢**。 + +丢的原因是硬的：TensorRT 引擎要用 `trt.OnnxParser` 重新构建，而 +**当前安装的 TensorRT（10.13）Python API 根本没有 metadata 读写接口** —— +`INetworkDefinition` / `IBuilderConfig` / `ICudaEngine` 上都没有任何 metadata 方法。 +实测二进制里搜一下就能看出来： + +``` +detect_best_0.engine 搜不到 "workpiece" → 标签已丢 +``` + +丢了的后果很隐蔽：读不到名字时代码会退化成把索引当名字用（类名变成 `"0"`）， +而预览界面又不画类名 —— **没有任何地方看得出标签丢了**。 + +所以标签放在**旁边的文本文件**里。它不是"退而求其次"，而是唯一可行且有好处： + +1. **可修改** —— 改标签不用重新转换模型，存盘即生效； +2. **跨格式统一** —— `.pt`/`.onnx`/`.engine` 共用一份，换后端名字不变。 + +### 写法 + +`<名字>.labels.yaml`： + +```yaml +model: detect_best_0 # 可选，写对了能帮你发现"标签文件拷错了模型" +names: + - workpiece # 第 1 行 = 类别 0 + - tongzhigui # 第 2 行 = 类别 1 +``` + +第 1 行对应类别 0，第 2 行对应类别 1，依次类推。也认这些写法（从别处拷配置常见）： + +```yaml +names: {0: workpiece, 1: tongzhigui} # 字典形式（Ultralytics data.yaml 风格） +``` +```yaml +names: "workpiece, tongzhigui" # 逗号串 +``` +```yaml +workpiece # 纯文本，一行一个（classes.txt 风格，注释行以 # 开头） +``` + +### 读取优先级 + +``` +<名字>.labels.yaml > ONNX 内嵌 metadata > 模型内嵌 metadata > 占位符 class_0 +``` + +旁挂文件**必须**赢，因为它是人能改的那一份。否则"我明明改了标签为什么不生效" +会变成一个很难查的问题。 + +全都读不到时会用占位名 `class_0`、`class_1`…，并且**打一条 warning**。 +刻意不用裸数字 `"0"` —— 那个和"真的有个叫 0 的类"分不出来。 + +### 怎么改 / 怎么看 + +```bash +bash scripts/weights.sh labels detect_best_0 # 看 +bash scripts/weights.sh labels detect_best_0 --set "工件,通止规" # 直接改 +bash scripts/weights.sh labels detect_best_0 --set-file classes.txt # 从文本文件导入 +bash scripts/weights.sh labels detect_best_0 --from-model # 把模型内嵌标签抄出来变成可编辑 +bash scripts/weights.sh labels detect_best_0 --json # 给程序读 +bash scripts/weights.sh labels detect_best_0 --clear # 删掉，回到模型内嵌 +``` + +改完**重跑一下节点**就生效，不用重新转换模型、不用重建 engine。 + +--- + +## 4. 加一个新模型（完整流程） + +### 步骤 1：把原始文件丢进 staging + +不管它原来叫什么、什么格式，先扔进 `models/staging/`： + +```bash +cp ~/下载/best_v3_final\(1\).pt ~/ros2/models/staging/ +``` + +### 步骤 2：一条命令转换 + 改名 + 归位 + +```bash +bash ~/ros2/scripts/weights.sh convert staging/best_v3_final\(1\).pt --name workpiece --group yolo +``` + +这条命令会： + +1. `--name workpiece` → 产出统一叫 `workpiece.onnx` / `workpiece.engine`（**自动改名**） +2. 依次跑 `PT → ONNX → engine`（TensorRT 只能由 ONNX 构建，所以中间产物必出） +3. 把产出放到 `models/yolo/`（**自动归位**） +4. **顺带把标签落成 `workpiece.labels.yaml`**（从源模型内嵌 metadata 抄） +5. 把 staging 里的源文件挪到 `models/staging/done/`（staging 永远是一份"待办清单"） + +常用选项： + +| 选项 | 用途 | +|---|---| +| `--to onnx` | 只要 ONNX，不建 engine（不需要 GPU） | +| `--to engine` | 只要 engine（仍会先产出 onnx 做中间产物） | +| `--imgsz 640` | 输入尺寸（默认取 `config/site.yaml` 的 `yolo.imgsz`） | +| `--labels "工件,通止规"` | 顺手写标签 | +| `--group <分组>` | 输出到别的子目录 | +| `--force` | 覆盖已存在的产出（不加则拒绝覆盖，防误删） | +| `--dry-run` | 只看计划，不动文件 | +| `--no-fp16` | engine 不用 FP16 | +| `--keep-staging` | 保留 staging 里的源文件 | + +### 步骤 3：让配置指向它 + +改 `config/site.yaml`： + +```yaml +paths: + yolo_weights: "/home/USER/ros2/models/yolo/workpiece.onnx" # 写 onnx；engine 会自动找同名的 +yolo: + backend: "tensorrt" # tensorrt / onnx / auto + imgsz: 640 # 必须等于该模型的输入尺寸 +``` + +> 写 `.onnx` 路径即可。`backend: tensorrt` 时程序会自动找同名的 `.engine`； +> 写 `.pt` 也一样能找到 `.engine`。所以换后端不用改配置。 + +### 步骤 4：确认 + +```bash +bash ~/ros2/scripts/weights.sh list # 模型在不在、标签读不读得出来 +bash ~/ros2/scripts/weights.sh doctor # 体检 +``` + +--- + +## 5. 体检与排错 + +```bash +bash ~/ros2/scripts/weights.sh doctor --check-engines +``` + +会检查： + +- 每个模型的标签来源；缺标签的会指出期望的文件名 +- **engine 能不能被当前 TensorRT 加载**（见下） +- `config/site.yaml` 里指向不存在的文件的**悬空引用** +- `backend: tensorrt` 但同名 `.engine` 不存在（启动会直接失败） +- staging 里有没有待转换的文件 + +### 常见问题 + +**「engine 明明在，启动却说加载失败」** + +**引擎和 TensorRT 版本是绑死的。** 换了 TensorRT（重装环境、升级 CUDA） +之后旧引擎会直接反序列化失败，报 `incompatible serialization version` 之类。 +表现是"文件在、大小也对，就是不能用"。用 `doctor --check-engines` 能查出来， +重导一次即可： + +```bash +bash scripts/weights.sh convert models/yolo/<名字>.onnx --name <名字> --to engine --force +``` + +> 实测本仓库里的 `detect_hc90.engine` 就是这样失效的（用旧版 TensorRT 构建， +> 当前 10.13 无法加载）。另外两个 engine 正常。 + +**「类名显示成 `class_0`」** + +标签没配上。`weights.sh labels <名字>` 会告诉你缺什么，然后 +`--from-model` 或 `--set` 补一份。 + +**「同一个模型，换个后端类名就变了」** + +不该发生。如果出现了，说明少了 `.labels.yaml`，各后端的回退来源不一致。 +补上旁挂文件即可。 + +**「`convert` 说目标已存在」** + +默认不覆盖，防止把别人的模型冲掉。确认要覆盖加 `--force`。 + +--- + +## 6. 规则（改动本模块时必须遵守） + +1. **二进制不进 git。** 靠 `.gitignore` 的通配规则挡住。新加分组不用改 `.gitignore`。 +2. **`.labels.yaml` 要进 git。** 它是配置，体积小，是团队间的唯一事实来源。 +3. **不要把类名写死在代码里。** 一律走 `arm_application.model_labels`。 + 已清理过的反例：`onnx_workpiece.py` 里曾硬编码 `{0: "workpiece"}`， + `hole_detect_hough.py` 里曾硬编码框上画的字 `"gongjian"`（而真类名是 `workpiece`）。 +4. **不要在"就地"改权重名字来换模型。** 用 `weights.sh convert --name`， + 让工具保证"名字主干一致 + 同目录"这条不变量。 +5. **不要手工搬 engine 到别的目录或改名。** 改名会切断和 onnx/标签的关联。 +6. **新增推理后端时，标签必须走 `model_labels.read_labels()`**， + 否则又会造出一个"某个后端读不到标签"的坑。 +7. **`models/staging/` 不是永久存储。** 它只放"待转换"的东西。 + +--- + +## 7. 代码在哪 + +| 位置 | 作用 | +|---|---| +| `arm_application/model_labels.py` | 标签读写、来源优先级、占位符（纯逻辑，无 ROS） | +| `arm_application/weights_layout.py` | 目录布局、模型发现、转换链规划（纯逻辑，无 ROS） | +| `scripts/weights.py` | 命令行实现（转换真正干活的地方） | +| `scripts/weights.sh` | 薄封装：准备环境 → 转发给 `weights.py` | +| `scripts/export_yolo_trt.sh` | **旧命令**，现在转发到 `weights.sh`，保持兼容 | +| `arm_application/yolo_detector.py` | 运行时加载模型 + `resolve_class_labels()` | +| `arm_application/onnx_workpiece.py` | ONNX Runtime 后端，读标签同样走 `model_labels` | + +测试：`src/arm_application/test/test_model_labels.py`（标签）、 +`test_weights_layout.py`（布局/发现/CLI）。 + +--- + +## 8. 相关文档 + +- `docs/SITE_CONFIG.md` —— `paths.yolo_weights` / `yolo.backend` 等配置项 +- `docs/WORKSPACE.md` —— 工作区总览 +- `docs/TRACKING.md` —— 换模型要同时改哪几处 diff --git a/models/hole_insert/README.md b/models/hole_insert/README.md index 381227c..65b9158 100644 --- a/models/hole_insert/README.md +++ b/models/hole_insert/README.md @@ -1,23 +1,36 @@ -# 入孔 / 孔定位权重 +# 入孔检测权重 -二进制（`.pt` / `.onnx` / `.engine`）**不进 git**。 +> **改这个目录之前先读 [`../README.md`](../README.md)**（权重模块的规则）。 -`hole_detect.backend: yolo` 时把本目录的检测权重绝对路径写进 `config/site.yaml`： +用于"把孔当框来检测"的模型（`hole_detect.backend: yolo`）。 +当前现场用的是 `two_stage`（两阶段）后端，它用的是**工件**权重 +（`hole_detect.workpiece_weights`，指向 `models/yolo/` 里的模型）， +本目录的模型并不是每次都会用到。 + +## 现状 + +- `detect_best_1.onnx` —— 工件检测模型，被 `hole_detect.workpiece_weights` 使用。 + +## 历史遗留（已知问题） + +`config/site.yaml` 的这两处仍指向一个**不存在**的文件 `models/hole_insert/best.pt`： ```yaml paths: - hole_weights: "/home/USER/ros2/models/hole_insert/best.pt" + hole_weights: ".../models/hole_insert/best.pt" # ← 文件不存在 hole_detect: - backend: "yolo" - weights: "/home/USER/ros2/models/hole_insert/best.pt" + weights: ".../models/hole_insert/best.pt" # ← 文件不存在 ``` -`backend: two_stage` 用的是工件 YOLO（`paths.yolo_weights`），本目录可空。 +`backend: two_stage` 时这两处不会被读，所以现在不影响运行；但一旦切成 +`backend: yolo` 就会启动失败。用 `bash ~/ros2/scripts/weights.sh doctor` +能查出来。要么把真实权重按命名规则放进本目录，要么把这两处配置删掉。 -从旧工程拷贝示例： +## 加新模型 / 看标签 + +和 `models/yolo/` 完全一样，只是分组名换成 `hole_insert`： ```bash -mkdir -p ~/ros2/models/hole_insert -# 若有独立孔模型： -# cp /path/to/best.pt ~/ros2/models/hole_insert/best.pt +bash ~/ros2/scripts/weights.sh convert staging/xxx.onnx --name hole_detect --group hole_insert +bash ~/ros2/scripts/weights.sh labels detect_best_1 --group hole_insert ``` diff --git a/models/hole_insert/detect_best_1.labels.yaml b/models/hole_insert/detect_best_1.labels.yaml new file mode 100644 index 0000000..6f58900 --- /dev/null +++ b/models/hole_insert/detect_best_1.labels.yaml @@ -0,0 +1,16 @@ +# 模型类别标签（class names） +# +# 一个模型对应一份标签，放在权重旁边、同名前缀： +# detect_best_0.engine ← 权重 +# detect_best_0.labels.yaml ← 本文件 +# +# 写法：names 下面一行一个类名，第 1 行对应类别 0，第 2 行对应类别 1，依次类推。 +# 改完保存即生效，不需要重新转换模型、也不用重启环境（重跑节点即可）。 +# +# 命令行查看/修改（推荐，不用手改本文件）： +# bash ~/ros2/scripts/weights.sh labels <模型名> +# bash ~/ros2/scripts/weights.sh labels <模型名> --set "工件,通止规" + +model: detect_best_1 +names: + - workpiece diff --git a/models/staging/.gitkeep b/models/staging/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/models/staging/README.md b/models/staging/README.md new file mode 100644 index 0000000..944b83b --- /dev/null +++ b/models/staging/README.md @@ -0,0 +1,26 @@ +# staging —— 待转换区 + +把**非标准格式**的模型文件先丢进这个目录，然后用一条命令转换、改名、归位： + +```bash +bash ~/ros2/scripts/weights.sh convert staging/<文件名> --name <模型名> --group yolo +``` + +支持的输入：`.pt` / `.pth`（PyTorch）、`.onnx`。 +已经是最终产物的 `.engine` / `.trt` / `.plan` 不需要转换，命令会直接把它归位。 + +其它格式（`.pdparams`、`.caffemodel`、`.tflite`、`.zip`…）工具不认识， +会明确报错并提示需要人工处理 —— 这是有意的，不要让它猜。 + +`convert` 跑完后，源文件会被自动挪到 `done/` 归档，所以本目录 +**始终是一份"待办清单"**：空了就说明都处理完了。 + +```bash +bash ~/ros2/scripts/weights.sh doctor # 看本目录还有哪些待处理 +``` + +## 注意 + +- 本目录内容**不进 git**（`.gitkeep` 除外）。 +- 这不是永久存储。转换完的文件在 `models/<分组>/` 里。 +- `done/` 里的东西可以随时删，只是留个念想。 diff --git a/models/yolo/README.md b/models/yolo/README.md index 4234c7c..3aca446 100644 --- a/models/yolo/README.md +++ b/models/yolo/README.md @@ -1,19 +1,54 @@ -# YOLO 权重 +# YOLO 工件检测权重 -二进制（`.pt` / `.onnx` / `.engine`）**不进 git**。放到本目录后改 `config/site.yaml`： +> **改这个目录之前先读 [`../README.md`](../README.md)**（权重模块的规则）。 + +## 放什么 + +同名的 `.pt` / `.onnx` / `.engine` / `.labels.yaml` 放一起，名字主干必须一致： + +``` +detect_best_0.onnx ← 写进 config/site.yaml 的那个 +detect_best_0.engine ← TensorRT 引擎，backend: tensorrt 时自动使用 +detect_best_0.labels.yaml ← 类别标签（进 git，可手改） +``` + +二进制不进 git；`.labels.yaml` 进 git。 + +## 加新模型 + +```bash +cp ~/下载/xxx.pt ~/ros2/models/staging/ +bash ~/ros2/scripts/weights.sh convert staging/xxx.pt --name workpiece --group yolo +``` + +然后改 `config/site.yaml`： ```yaml paths: - yolo_weights: "/home/USER/ros2/models/yolo/YOUR.onnx" + yolo_weights: "/home/USER/ros2/models/yolo/workpiece.onnx" yolo: backend: "tensorrt" - imgsz: 640 # 必须等于该模型输入 + imgsz: 640 # 必须等于该模型输入尺寸 ``` -TensorRT： +## 查看 / 修改标签 ```bash -bash ~/ros2/scripts/export_yolo_trt.sh /home/USER/ros2/models/yolo/YOUR.onnx 640 +bash ~/ros2/scripts/weights.sh labels detect_best_0 +bash ~/ros2/scripts/weights.sh labels detect_best_0 --set "工件,通止规" ``` -会生成同名 `.engine`。换 ONNX 或 imgsz 后重导。然后重启 `run_visual_servo_rt.sh`。 +## 排错 + +```bash +bash ~/ros2/scripts/weights.sh doctor --check-engines +``` + +`engine` 和 TensorRT 版本绑死：换过 TensorRT 后旧 engine 会加载失败，重导即可 +（`convert <同名>.onnx --name <名字> --to engine --force`）。 + +## 历史遗留 + +`active.pt` 是个指向 `detect_hc90.pt` 的符号链接，早期用来"切换当前模型"， +现在没有任何节点会读它 —— 换模型请改 `config/site.yaml` 的 `paths.yolo_weights`。 +留着只是不想动别人的文件，可以安全删除。 diff --git a/models/yolo/detect_best.labels.yaml b/models/yolo/detect_best.labels.yaml new file mode 100644 index 0000000..9e68252 --- /dev/null +++ b/models/yolo/detect_best.labels.yaml @@ -0,0 +1,16 @@ +# 模型类别标签（class names） +# +# 一个模型对应一份标签，放在权重旁边、同名前缀： +# detect_best_0.engine ← 权重 +# detect_best_0.labels.yaml ← 本文件 +# +# 写法：names 下面一行一个类名，第 1 行对应类别 0，第 2 行对应类别 1，依次类推。 +# 改完保存即生效，不需要重新转换模型、也不用重启环境（重跑节点即可）。 +# +# 命令行查看/修改（推荐，不用手改本文件）： +# bash ~/ros2/scripts/weights.sh labels <模型名> +# bash ~/ros2/scripts/weights.sh labels <模型名> --set "工件,通止规" + +model: detect_best +names: + - workpiece diff --git a/models/yolo/detect_best_0.labels.yaml b/models/yolo/detect_best_0.labels.yaml new file mode 100644 index 0000000..a09e5ca --- /dev/null +++ b/models/yolo/detect_best_0.labels.yaml @@ -0,0 +1,16 @@ +# 模型类别标签（class names） +# +# 一个模型对应一份标签，放在权重旁边、同名前缀： +# detect_best_0.engine ← 权重 +# detect_best_0.labels.yaml ← 本文件 +# +# 写法：names 下面一行一个类名，第 1 行对应类别 0，第 2 行对应类别 1，依次类推。 +# 改完保存即生效，不需要重新转换模型、也不用重启环境（重跑节点即可）。 +# +# 命令行查看/修改（推荐，不用手改本文件）： +# bash ~/ros2/scripts/weights.sh labels <模型名> +# bash ~/ros2/scripts/weights.sh labels <模型名> --set "工件,通止规" + +model: detect_best_0 +names: + - workpiece diff --git a/models/yolo/detect_hc90.labels.yaml b/models/yolo/detect_hc90.labels.yaml new file mode 100644 index 0000000..1df098d --- /dev/null +++ b/models/yolo/detect_hc90.labels.yaml @@ -0,0 +1,16 @@ +# 模型类别标签（class names） +# +# 一个模型对应一份标签，放在权重旁边、同名前缀： +# detect_best_0.engine ← 权重 +# detect_best_0.labels.yaml ← 本文件 +# +# 写法：names 下面一行一个类名，第 1 行对应类别 0，第 2 行对应类别 1，依次类推。 +# 改完保存即生效，不需要重新转换模型、也不用重启环境（重跑节点即可）。 +# +# 命令行查看/修改（推荐，不用手改本文件）： +# bash ~/ros2/scripts/weights.sh labels <模型名> +# bash ~/ros2/scripts/weights.sh labels <模型名> --set "工件,通止规" + +model: detect_hc90 +names: + - workpiece diff --git a/models/yolo/detect_yolo26_gongjian.labels.yaml b/models/yolo/detect_yolo26_gongjian.labels.yaml new file mode 100644 index 0000000..051055a --- /dev/null +++ b/models/yolo/detect_yolo26_gongjian.labels.yaml @@ -0,0 +1,16 @@ +# 模型类别标签（class names） +# +# 一个模型对应一份标签，放在权重旁边、同名前缀： +# detect_best_0.engine ← 权重 +# detect_best_0.labels.yaml ← 本文件 +# +# 写法：names 下面一行一个类名，第 1 行对应类别 0，第 2 行对应类别 1，依次类推。 +# 改完保存即生效，不需要重新转换模型、也不用重启环境（重跑节点即可）。 +# +# 命令行查看/修改（推荐，不用手改本文件）： +# bash ~/ros2/scripts/weights.sh labels <模型名> +# bash ~/ros2/scripts/weights.sh labels <模型名> --set "工件,通止规" + +model: detect_yolo26_gongjian +names: + - gongjian diff --git a/scripts/README.md b/scripts/README.md index f8cf63e..81d53b7 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -8,7 +8,11 @@ **夹爪：** `run_gripper.sh --arm 89|91` · `run_dual_gripper.sh` **入孔：** `run_hole_detect.sh --arm … [--preview]` **触觉：** `run_tactile.sh [--cli]` -**权重：** `export_yolo_trt.sh` +**权重：** `weights.sh`（`list` / `labels` / `convert` / `doctor` / `show`） + - 转换：`weights.sh convert <文件> --name <模型名>` —— PT→ONNX→TensorRT，自动改名归位 + 固化标签 + - 标签：`weights.sh labels <模型名>` 看，`--set "工件,通止规"` 改 + - 体检：`weights.sh doctor --check-engines`（能查出 engine 与 TensorRT 版本不匹配这类静默失效） + - 规则写在 [models/README.md](../models/README.md)。`export_yolo_trt.sh` 现在是它的兼容转发壳。 **配置：** `show_site_config.sh` · `lib/export_site_env.py` **环境：** `env.sh` · `build_all.sh` **RDK：** `test_rdk_python.sh` · `rdk_switch_tool.sh` · `rdk_drag_bridge.py` diff --git a/scripts/export_yolo_trt.sh b/scripts/export_yolo_trt.sh index 46e3dee..5135308 100755 --- a/scripts/export_yolo_trt.sh +++ b/scripts/export_yolo_trt.sh @@ -1,90 +1,47 @@ #!/usr/bin/env bash -# Export YOLO .pt/.onnx → TensorRT .engine (FP16). +# 【旧命令，已改为转发壳】YOLO 权重 → TensorRT engine。 # +# 请改用新命令（多了改名、归位、标签固化、staging 等功能）： +# +# bash ~/ros2/scripts/weights.sh convert <文件> --name <模型名> +# bash ~/ros2/scripts/weights.sh list +# bash ~/ros2/scripts/weights.sh labels <模型名> +# +# 本脚本保留只为兼容旧文档/旧习惯，行为等价于： +# weights.sh convert --name <同名主干> --to engine [--imgsz N] +# +# 原用法（仍然可用）： # bash ~/ros2/scripts/export_yolo_trt.sh -# bash ~/ros2/scripts/export_yolo_trt.sh /home/cc/ros2/models/yolo/detect_best.onnx 640 -# -# Needs: tensorrt-cu12 + onnx + onnxslim in ros_flexiv (do NOT pip install -# the meta package `tensorrt`, it pulls CUDA 13 wheels). -# Then set yolo.backend: tensorrt in site.yaml and restart visual servo. +# bash ~/ros2/scripts/export_yolo_trt.sh /path/to/xxx.onnx 640 set -eo pipefail ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" + # shellcheck disable=SC1091 -source "${ROOT}/scripts/lib/load_site_config.sh" +source "${ROOT}/scripts/lib/load_site_config.sh" >/dev/null 2>&1 || true -PT="${1:-${SITE_YOLO_WEIGHTS:-${ROOT}/models/yolo/detect_best.onnx}}" +WEIGHTS="${1:-${SITE_YOLO_WEIGHTS:-${ROOT}/models/yolo/detect_best.onnx}}" IMGSZ="${2:-${SITE_YOLO_IMGSZ:-640}}" -PY="${HOME}/anaconda3/envs/ros_flexiv/bin/python" -if [[ ! -f "${PT}" ]]; then - echo "ERROR: weights not found: ${PT}" >&2 +if [[ ! -f "${WEIGHTS}" ]]; then + echo "ERROR: weights not found: ${WEIGHTS}" >&2 exit 1 fi -# Pip TensorRT shared libs (needed at export + runtime). -TRT_LIBS="$("${PY}" - <<'PY' -import pathlib, tensorrt -p = pathlib.Path(tensorrt.__file__).resolve().parent -cands = [p, p.parent / "tensorrt_libs", p.parent / "nvidia" / "tensorrt" / "lib"] -print(":".join(str(x) for x in cands if x.exists())) -PY -)" -export LD_LIBRARY_PATH="${TRT_LIBS}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" -export YOLO_AUTOINSTALL=false +# 取文件名主干做模型名：xxx.onnx → xxx +BASE="$(basename "${WEIGHTS}")" +NAME="${BASE%.*}" +DIR="$(dirname "${WEIGHTS}")" +# 分组 = 权重所在目录名（相对 models/ 的那一层） +GROUP="$(basename "${DIR}")" -echo "Export TensorRT engine from ${PT} imgsz=${IMGSZ}" -"${PY}" - < dict: ) weights = str(paths.get("yolo_weights") or "") - # Resolve active.pt symlink to real path for display + # 解析符号链接，方便日志里看到真实文件（历史上有人用 active.pt 做切换，已废弃） try: weights_resolved = str(Path(weights).expanduser().resolve()) except Exception: @@ -231,7 +231,7 @@ def flatten(data: dict) -> dict: or str(_root() / "handeye" / "waypoints.yaml") ), "SITE_HANDEYE_WORKDIR": str( - paths.get("handeye_workdir") or "/tmp/handeye_calib" + paths.get("handeye_workdir") or str(_root() / "handeye" / "tmp") ), "SITE_HANDEYE_FILE": str( paths.get("handeye_file") @@ -280,10 +280,10 @@ def flatten(data: dict) -> dict: "SITE_D435_FPS": str(int(((data.get("cameras") or {}).get("d435") or {}).get("fps") or 30)), "SITE_D405_LEFT_SERIAL": str(((data.get("cameras") or {}).get("d405_left") or {}).get("serial") or "352122274408"), "SITE_D405_RIGHT_SERIAL": str(((data.get("cameras") or {}).get("d405_right") or {}).get("serial") or "260522274442"), - "SITE_YOLO_CONF": str(yolo.get("conf") or 0.35), - "SITE_YOLO_IMGSZ": str(int(yolo.get("imgsz") or 416)), + "SITE_YOLO_CONF": str(yolo.get("conf") or 0.10), + "SITE_YOLO_IMGSZ": str(int(yolo.get("imgsz") or 640)), "SITE_YOLO_DEVICE": str(yolo.get("device") or "0"), - "SITE_YOLO_MIN_CONF": str(yolo.get("min_conf") or 0.35), + "SITE_YOLO_MIN_CONF": str(yolo.get("min_conf") or 0.12), "SITE_GO_HOME_POSE": str(home.get("pose") or "home"), "SITE_GO_HOME_SPEED": str(home.get("speed") or 0.25), "SITE_NET_IFACE": str(robot.get("iface") or net.get("prefer_iface") or ""), @@ -312,7 +312,7 @@ def flatten(data: dict) -> dict: "SITE_HOLE_WEIGHTS": str( (data.get("hole_detect") or {}).get("weights") or paths.get("hole_weights") - or str(_root() / "models" / "hole_insert" / "best.pt") + or str(_root() / "models" / "hole_insert" / "detect_best_1.onnx") ), "SITE_HOLE_WORKPIECE_WEIGHTS": str( (data.get("hole_detect") or {}).get("workpiece_weights") @@ -381,7 +381,7 @@ def perception_overrides(data: dict) -> dict: "weights": str(paths.get("yolo_weights") or ""), "conf": float(yolo.get("conf") or 0.10), "iou": float(yolo.get("iou") or 0.45), - "imgsz": int(yolo.get("imgsz") or 416), + "imgsz": int(yolo.get("imgsz") or 640), "device": str(yolo.get("device") or "0"), "half_precision": bool(yolo.get("half_precision", True)), "backend": str(yolo.get("backend") or "auto"), diff --git a/scripts/run_handeye.sh b/scripts/run_handeye.sh index b7ebe7d..56b4ca1 100755 --- a/scripts/run_handeye.sh +++ b/scripts/run_handeye.sh @@ -75,7 +75,7 @@ WP="${SITE_HANDEYE_WAYPOINTS:-${HE_DIR}/poses/89_eye_in_hand.yaml}" CAL="${SITE_HANDEYE_CALIBRATE:-${HE_DIR}/calibrate.yaml}" TOUR="${SITE_HANDEYE_TOUR:-${HE_DIR}/waypoints.yaml}" RES="${SITE_HANDEYE_FILE:-${HE_DIR}/results/89_eye_in_hand.yaml}" -WORKDIR="${SITE_HANDEYE_WORKDIR:-/tmp/handeye_calib}" +WORKDIR="${SITE_HANDEYE_WORKDIR:-${HE_DIR}/tmp}" TAUGHT="${WORKDIR}/taught_poses.txt" PREFIX="${SITE_JOINT_PREFIX:-Rizon4-063989_}" CAM="${SITE_CAM_NAME:-d405}" diff --git a/scripts/show_site_config.sh b/scripts/show_site_config.sh index 2fadd93..d0bd984 100755 --- a/scripts/show_site_config.sh +++ b/scripts/show_site_config.sh @@ -13,8 +13,11 @@ parse_arm_args "$@" python3 "${ROOT}/scripts/lib/export_site_env.py" --print ${ROBOT_SN:+--sn "$ROBOT_SN"} echo echo "Edit: ${SITE_CONFIG:-$ROOT/config/site.yaml}" -echo "Switch weights example:" -echo " paths.yolo_weights: \"/home/cc/ros2/models/yolo/detect_hc90.pt\"" -echo "Or symlink:" -echo " ln -sfn detect_hc90.pt ~/ros2/models/yolo/active.pt" -echo " # then set paths.yolo_weights to .../active.pt" +echo "Switch weights:" +echo " paths.yolo_weights: \"/home/cc/ros2/models/yolo/detect_best_0.onnx\" # 写 onnx 或 pt 均可" +echo " 改完确认标签和 engine 都对：" +echo " bash ~/ros2/scripts/weights.sh list" +echo " bash ~/ros2/scripts/weights.sh doctor --check-engines" +echo +echo "（旧做法: 用 models/yolo/active.pt 符号链接切换模型 —— 已废弃，" +echo " 没有任何节点会读它。请直接改 paths.yolo_weights。）" diff --git a/scripts/weights.py b/scripts/weights.py new file mode 100755 index 0000000..bb615b1 --- /dev/null +++ b/scripts/weights.py @@ -0,0 +1,842 @@ +#!/usr/bin/env python3 +"""权重模块的统一命令行入口：转换 / 标签 / 列表 / 自检。 + +通常不直接调用本文件，而是走薄封装（它会先把环境准备好）： + + bash ~/ros2/scripts/weights.sh list + bash ~/ros2/scripts/weights.sh labels detect_best_0 + bash ~/ros2/scripts/weights.sh labels detect_best_0 --set "工件,通止规" + bash ~/ros2/scripts/weights.sh convert staging/best_v3.pt --name workpiece --group yolo + +把"纯逻辑"放在 ``arm_application.weights_layout`` / ``arm_application.model_labels`` +里，本文件只负责：解析参数、打印人话、真正调用 PyTorch/Ultralytics/TensorRT。 + +为什么单独做这个模块：原先只有 ``export_yolo_trt.sh``，只能"就地"把 +``<同名>.onnx`` 变成 ``<同名>.engine``，既不能改名、不能归位、也不管标签。 +实际用起来必须手工搬文件 + 手改 site.yaml，而**标签在 engine 那一环会静默丢失** +（原因见 model_labels 模块头）。这里把三件事一次做完： +转换 → 规范改名 → 放到权重目录 → 顺带把标签落成可编辑的文件。 +""" +from __future__ import annotations + +import argparse +import json +import os +import shutil +import sys +from pathlib import Path +from typing import Dict, List, Optional, Sequence + +# --------------------------------------------------------------------------- +# 让脚本能在没 colcon build / 没 source install 的情况下直接 import +# --------------------------------------------------------------------------- + + +def _bootstrap_import_path() -> Optional[Path]: + """把 ``src/arm_application`` 插进 sys.path，返回仓库根（找不到返回 None）。 + + ``arm_application/__init__.py`` 里只有文档字符串、没有 ROS 依赖， + 所以这样 import 不会把 ROS 拖进来。 + """ + here = Path(__file__).resolve() + for cand in [here, *here.parents]: + pkg = cand / "src" / "arm_application" + if (pkg / "arm_application" / "weights_layout.py").is_file(): + if str(pkg) not in sys.path: + sys.path.insert(0, str(pkg)) + return cand + return None + + +_REPO_ROOT = _bootstrap_import_path() + +from arm_application.model_labels import ( # noqa: E402 + LABELS_SUFFIX, + SOURCE_LABELS_ZH, + SOURCE_PLACEHOLDER, + SOURCE_SIDECAR, + describe, + format_labels_yaml, + labels_from_model, + labels_from_onnx, + labels_path_for, + parse_label_list, + read_label_text_file, + read_labels, + write_labels, +) +from arm_application.weights_layout import ( # noqa: E402 + CONVERTIBLE_INPUT_EXTS, + STAGING_DIRNAME, + WEIGHT_EXTS, + ModelEntry, + canonical_path, + conversion_chain, + default_weights_root, + discover, + find_model, + groups, + model_labels, + plan_conversion, + staging_dir, + staging_files, + validate_name, +) + +# 默认转换输入尺寸。site.yaml 的 yolo.imgsz 优先（由 weights.sh 通过环境变量传入）。 +DEFAULT_IMGSZ = 640 + + +# --------------------------------------------------------------------------- +# 小工具 +# --------------------------------------------------------------------------- + + +def _root(args) -> Path: + return Path(args.root).expanduser() if getattr(args, "root", None) else default_weights_root() + + +def _imgsz(args) -> int: + """输入尺寸：命令行 > site.yaml（环境变量） > 默认 640。""" + if getattr(args, "imgsz", None): + return int(args.imgsz) + env = os.environ.get("SITE_YOLO_IMGSZ") + if env: + try: + return int(env) + except ValueError: + pass + return DEFAULT_IMGSZ + + +def _err(msg: str) -> None: + print(f"错误: {msg}", file=sys.stderr) + + +def _ok(msg: str) -> None: + print(f"✓ {msg}") + + +def _warn(msg: str) -> None: + print(f"! {msg}") + + +def _hr(char: str = "─", n: int = 74) -> None: + print(char * n) + + +def _print_table(headers: Sequence[str], rows: Sequence[Sequence[str]]) -> None: + """左对齐打印一张表。宽度按显示宽度（东亚字符按 2 算）估算。""" + + def w(s: str) -> int: + return sum(2 if ord(c) > 0x2E80 else 1 for c in s) + + widths = [w(h) for h in headers] + for r in rows: + for i, c in enumerate(r): + widths[i] = max(widths[i], w(c)) + + def pad(s: str, width: int) -> str: + return s + " " * max(0, width - w(s)) + + print(" ".join(pad(h, widths[i]) for i, h in enumerate(headers))) + print(" ".join("-" * widths[i] for i in range(len(headers)))) + for r in rows: + print(" ".join(pad(c, widths[i]) for i, c in enumerate(r))) + + +def _mode_of(p: Path) -> str: + """文件可读性/大小的人话描述，用于列表显示。""" + try: + return f"{p.stat().st_size / 1e6:.1f}MB" + except OSError: + return "?" + + +# --------------------------------------------------------------------------- +# list +# --------------------------------------------------------------------------- + + +def cmd_list(args) -> int: + root = _root(args) + if not root.is_dir(): + _err(f"权重目录不存在: {root}") + return 2 + + entries = discover(root) + if not entries: + _warn(f"{root} 下没找到任何权重文件（后缀: {', '.join(WEIGHT_EXTS)}）") + return 0 + + payload = [] + rows = [] + for e in entries: + ls = model_labels(e, deep=args.deep) + if ls.source == SOURCE_PLACEHOLDER: + label_cell = "— 缺标签" + else: + label_cell = f"{', '.join(ls.sorted_names())}" + rows.append( + [ + e.group, + e.name, + " ".join(e.exts) or "—", + label_cell, + SOURCE_LABELS_ZH.get(ls.source, ls.source), + ] + ) + payload.append( + { + "group": e.group, + "name": e.name, + "formats": e.exts, + "files": {k: str(v) for k, v in e.files.items()}, + "labels": ls.names, + "labels_source": ls.source, + "labels_path": str(ls.path) if ls.path else None, + } + ) + + if args.json: + print(json.dumps(payload, ensure_ascii=False, indent=2)) + return 0 + + print(f"权重目录: {root}") + _hr() + _print_table(["分组", "模型名", "已有格式", "标签", "标签来源"], rows) + _hr() + missing = [e for e in entries if model_labels(e).source == SOURCE_PLACEHOLDER] + print(f"共 {len(entries)} 个模型，其中 {len(missing)} 个缺标签。") + if missing and not args.deep: + print("提示: 加 --deep 可再试着从 .pt/.engine 自身读标签（会加载模型，慢）。") + if missing: + print("补标签: bash scripts/weights.sh labels <模型名> --from-model") + return 0 + + +# --------------------------------------------------------------------------- +# labels +# --------------------------------------------------------------------------- + + +def _resolve_entry(args, root: Path) -> Optional[ModelEntry]: + entry = find_model(root, args.model, getattr(args, "group", None)) + if entry is None: + # 提示写成 ``分组/模型名``：这正是 ``list`` 里显示的样子，可以直接复制粘贴， + # 比只列裸名字更不容易再踩一次坑。 + avail = ", ".join(f"{e.group}/{e.name}" for e in discover(root)) or "(无)" + _err( + f"找不到模型 '{args.model}'。现有模型（可整串复制）: {avail}\n" + f" 也可以只写裸名字，或加 --group 指定分组。全部列出: " + f"bash ~/ros2/scripts/weights.sh list" + ) + return entry + + +def cmd_labels(args) -> int: + root = _root(args) + entry = _resolve_entry(args, root) + if entry is None: + return 2 + + # 修改类操作 ------------------------------------------------------------ + if args.clear: + p = entry.labels_path + if p.is_file(): + p.unlink() + _ok(f"已删除标签文件 {p}") + else: + _warn(f"本来就没有标签文件 {p}") + return 0 + + if args.set is not None and args.set_file: + _err("--set 与 --set-file 不能同时用") + return 2 + + new_names: Optional[Dict[int, str]] = None + if args.set is not None: + new_names = parse_label_list(args.set) + if not new_names: + _err("--set 解析不出任何类别名") + return 2 + elif args.set_file: + src = Path(args.set_file).expanduser() + if not src.is_file(): + _err(f"标签文件不存在: {src}") + return 2 + new_names = read_label_text_file(src) + if not new_names: + _err(f"{src} 里没读到任何类别名（一行一个）") + return 2 + + if new_names is None and args.from_model: + # 从模型内嵌 metadata 抄一份出来，目的是"变成可编辑的旁挂文件"。 + for p in entry.sorted_files(): + got = None + if p.suffix.lower() == ".onnx": + from arm_application.model_labels import labels_from_onnx + + got = labels_from_onnx(p) + if got is None: + got = labels_from_model(p) + if got: + new_names = got[0] + print(f"从 {p.name} 读到 {len(new_names)} 个类别名") + break + if new_names is None: + _err(f"{entry.name} 的模型文件里读不到标签，只能用 --set 手工指定") + return 2 + + if new_names is not None: + dst = write_labels(entry.path(entry.exts[0]) or entry.labels_path, new_names) + _ok(f"已写入 {dst}") + print(f" 类别({len(new_names)}): {', '.join(new_names[k] for k in sorted(new_names))}") + return 0 + + # 读取（默认动作）------------------------------------------------------ + ls = read_labels(entry.path(entry.exts[0]), allow_model_metadata=args.deep) if entry.exts else None + if ls is None or ls.is_placeholder: + ls = model_labels(entry, deep=args.deep) + + if args.json: + print( + json.dumps( + { + "model": entry.name, + "group": entry.group, + "labels": ls.names, + "source": ls.source, + "source_zh": SOURCE_LABELS_ZH.get(ls.source, ls.source), + "source_file": str(ls.path) if ls.path else None, + "expected_labels_path": str(entry.labels_path), + }, + ensure_ascii=False, + indent=2, + ) + ) + return 0 + + print(f"模型: {entry.group}/{entry.name}") + print(f"格式: {' '.join(entry.exts) or '—'}") + _hr() + if ls.is_placeholder: + _warn(f"没有标签。{describe(ls)}") + print(f" 期望的标签文件位置: {entry.labels_path}") + print(" 修复方式（三选一）:") + print(f" 1) 从模型内嵌 metadata 抄一份出来再改: weights.sh labels {entry.name} --from-model") + print(f' 2) 直接指定类别名: weights.sh labels {entry.name} --set "工件,通止规"') + print(f" 3) 手工创建 {entry.labels_path.name}") + return 1 + + print(f"标签来源: {SOURCE_LABELS_ZH.get(ls.source, ls.source)}") + if ls.path: + print(f"来源文件: {ls.path}") + _hr() + for k in sorted(ls.names): + print(f" {k}: {ls.names[k]}") + return 0 + + +# --------------------------------------------------------------------------- +# convert +# --------------------------------------------------------------------------- + + +def _export_onnx(src: Path, imgsz: int, dst: Path, *, simplify: bool = True) -> Path: + """YOLO .pt/.pth → ONNX，直接产出到 ``dst``。 + + **关键：先把 .pt 复制到临时目录再导出。** + Ultralytics 会把 onnx 写在**源文件旁边**（同名同目录），我们对源 .pt 直接导出时， + 会把它旁边已有的 ``<同名>.onnx`` 静默覆盖掉 —— 那是用户自己的产物， + 可能还是不同 imgsz 导出的。更糟的是随后的搬运会把它彻底删掉。 + 实测就发生过：转换 detect_hc90.pt 时，models/yolo/detect_hc90.onnx 被覆盖后又被搬走。 + 在临时目录里导出，源目录一个字节都不会动。 + """ + import tempfile + + os.environ.setdefault("YOLO_AUTOINSTALL", "false") + from ultralytics import YOLO + + dst.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix="weights_export_") as td: + work = Path(td) / src.name + shutil.copy2(src, work) + model = YOLO(str(work)) + out = Path( + str( + model.export( + format="onnx", + imgsz=int(imgsz), + device=0, + simplify=simplify, + dynamic=False, + ) + ) + ).expanduser() + if not out.is_file(): + raise RuntimeError(f"Ultralytics 说导出了 {out}，但文件不存在") + # 临时目录马上会被删，所以这里必须是 move。 + if out.resolve() != dst.resolve(): + shutil.move(str(out), str(dst)) + return dst + + +def _build_engine(onnx_src: Path, dst: Path, *, fp16: bool = True, workspace_gb: int = 4) -> Path: + """ONNX → TensorRT engine，直接写到 dst。 + + 这里刻意绕开 Ultralytics 的 engine 导出：它对 .onnx 会按发行包名检查 + onnxruntime，环境装的是 onnxruntime-gpu 时会触发 pip AutoUpdate（可能卡死）。 + 直接用 TensorRT 的 Python API，行为可控、错误信息也更清楚。 + """ + import tensorrt as trt + + logger = trt.Logger(trt.Logger.WARNING) + builder = trt.Builder(logger) + # TRT 10 里 EXPLICIT_BATCH 已是默认行为，但这个 flag 仍然接受； + # 保持和 export_yolo_trt.sh 一致，避免不同版本行为差异。 + flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) + network = builder.create_network(flag) + parser = trt.OnnxParser(network, logger) + + if not parser.parse(onnx_src.read_bytes()): + msgs = [str(parser.get_error(i)) for i in range(parser.num_errors)] + raise RuntimeError("ONNX 解析失败:\n " + "\n ".join(msgs)) + + config = builder.create_builder_config() + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, int(workspace_gb) << 30) + if fp16 and builder.platform_has_fast_fp16: + config.set_flag(trt.BuilderFlag.FP16) + + serialized = builder.build_serialized_network(network, config) + if serialized is None: + raise RuntimeError("TensorRT 构建失败（返回空）。可试着 --no-fp16 或换 imgsz。") + + dst.parent.mkdir(parents=True, exist_ok=True) + dst.write_bytes(bytes(serialized)) + return dst + + +def _materialize_labels(src: Path, dest_weights: Path, explicit: Optional[str]) -> Optional[Path]: + """转换后把标签落到目标权重旁边。 + + 这是防止"engine 丢标签"再次发生的**关键一步**：转换时就把标签固化成 + 可编辑的旁挂文件，而不是等到运行时才发现读不到名字。 + + 优先级：命令行 --labels > 源文件/中间产物的内嵌 metadata。 + 已经有旁挂文件时不覆盖（避免把用户手改的内容冲掉）。 + """ + dst_labels = labels_path_for(dest_weights) + if dst_labels.is_file(): + _warn(f"标签文件已存在，保持不变: {dst_labels.name}") + return dst_labels + + if explicit: + names = parse_label_list(explicit) + if names: + return write_labels(dest_weights, names) + + for cand in (src, src.with_suffix(".onnx"), dest_weights): + if not cand.is_file(): + continue + if cand.suffix.lower() == ".onnx": + from arm_application.model_labels import labels_from_onnx + + got = labels_from_onnx(cand) + else: + got = labels_from_model(cand) + if got: + _ok(f"标签来自 {cand.name}: {', '.join(got[0][k] for k in sorted(got[0]))}") + return write_labels(dest_weights, got[0]) + return None + + +def cmd_convert(args) -> int: + root = _root(args) + src = Path(args.input).expanduser() + if not src.is_file(): + # 相对路径优先按**权重目录**解析，而不是当前工作目录。 + # 否则 `weights.sh convert staging/x.pt` 会去 /staging 找， + # 而用户心里想的是 models/staging/x.pt —— 报"文件不存在"很费解。 + cands = [] + if not src.is_absolute(): + cands.append(root / args.input) # models/staging/x.pt + cands.append(staging_dir(root) / args.input) # 只写文件名: models/staging/x.pt + cands.append(Path(args.input)) # 退一步：相对 cwd + found = next((c for c in cands if c.is_file()), None) + if found is None: + _err( + f"输入文件不存在: {args.input}\n" + f" 找过: " + ", ".join(str(c) for c in cands) + + f"\n 把待转换的文件放进 {staging_dir(root)} 再试。" + ) + return 2 + src = found + + name = (args.name or src.stem).strip() + bad = validate_name(name) + if bad: + _err(bad) + return 2 + + group = args.group + chain, _outs = plan_conversion(root, src, name=name, group=group, target=args.to) + + print(f"输入: {src}") + print(f"输出: {root / group}/{name}.") + print(f"计划: {chain.describe()}") + print(f"尺寸: imgsz={_imgsz(args)}") + if args.to != "all": + print(f"目标: 只产出 {args.to}") + + in_staging = src.parent.resolve() == staging_dir(root).resolve() + # 归位路径会把源文件本身搬到目标去；那种情况下它已经不在 staging 了， + # 下面不能再"归档"一次（否则会 FileNotFoundError）。 + src_consumed = False + + if chain.is_noop: + ext = src.suffix.lower().lstrip(".") + if ext in WEIGHT_EXTS: + dst = canonical_path(root, group, name, ext) + if dst.resolve() == src.resolve(): + _warn(f"{src.name} 已经是最终产物，且就在目标位置，无需处理") + return 0 + print(f"{ext} 无需转换，直接归位 → {dst}") + if args.dry_run: + return 0 + dst.parent.mkdir(parents=True, exist_ok=True) + # 归位本身就是"移动源文件"。若用户要求保留 staging 原件， + # 就改成复制，否则源文件会连同 staging 一起消失。 + if in_staging and args.keep_staging: + shutil.copy2(str(src), str(dst)) + else: + shutil.move(str(src), str(dst)) + src_consumed = True + _ok(f"已归位 {dst}") + _materialize_labels(dst, dst, args.labels) + if in_staging and not src_consumed: + _ok(f"staging 原件保留在 {src}（--keep-staging）") + elif in_staging: + _ok(f"staging 已清空该文件（它已变成 {dst.name}）") + return 0 + _err( + f"不认识的输入格式 '{src.suffix}'。" + f"可转换的输入: {', '.join('.' + e for e in CONVERTIBLE_INPUT_EXTS)}；" + f"已是最终产物的: .engine/.trt/.plan" + ) + return 2 + + if args.dry_run: + print("(--dry-run，未实际执行)") + return 0 + + # 真正执行 -------------------------------------------------------------- + produced: Dict[str, Path] = {} + try: + for ext in chain.produce: + dst = canonical_path(root, group, name, ext) + if dst.exists() and not args.force: + raise FileExistsError(f"{dst} 已存在（加 --force 覆盖）") + if ext == "onnx": + if src.suffix.lower() == ".onnx": + # 源本身就是 onnx：复制而不是移动 —— 用户给的文件不该消失。 + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(str(src), str(dst)) + else: + _export_onnx(src, _imgsz(args), dst) + produced["onnx"] = dst + _ok(f"ONNX → {dst} ({_mode_of(dst)})") + elif ext == "engine": + onnx_in = produced.get("onnx") or src + _build_engine(onnx_in, dst, fp16=not args.no_fp16) + produced["engine"] = dst + _ok(f"engine→ {dst} ({_mode_of(dst)})") + except Exception as exc: # noqa: BLE001 + _err(f"转换失败: {type(exc).__name__}: {exc}") + return 1 + + # 标签 ------------------------------------------------------------------ + anchor = produced.get("onnx") or produced.get("engine") or src + lp = _materialize_labels(src, anchor, args.labels) + if lp: + _ok(f"标签 → {lp}") + else: + _warn("没有可用的标签（源模型里也没有）。请手动补：") + print(f" bash scripts/weights.sh labels {name} --set \"类别1,类别2\"") + + # staging 收尾 ---------------------------------------------------------- + if in_staging and not args.keep_staging and not src_consumed: + done = staging_dir(root) / "done" + done.mkdir(parents=True, exist_ok=True) + target = done / src.name + # 同名已存在时加个序号，避免静默覆盖上一次的源文件。 + n = 1 + while target.exists(): + target = done / f"{src.stem}.{n}{src.suffix}" + n += 1 + shutil.move(str(src), str(target)) + _ok(f"已把 staging 里的源文件挪到 {target}（加 --keep-staging 可保留）") + + print() + _ok(f"完成。共产出 {len(produced)} 个文件。") + print(f"下一步: 把 config/site.yaml 的 paths.yolo_weights 指向 {anchor}") + print(" 或直接跑 bash scripts/weights.sh list 确认标签读得出来") + return 0 + + +# --------------------------------------------------------------------------- +# doctor +# --------------------------------------------------------------------------- + + +def _read_site_yaml() -> Dict: + if _REPO_ROOT is None: + return {} + p = _REPO_ROOT / "config" / "site.yaml" + if not p.is_file(): + return {} + try: + import yaml + + return yaml.safe_load(p.read_text(encoding="utf-8")) or {} + except Exception: # noqa: BLE001 + return {} + + +def _engine_usable(p: Path) -> Optional[str]: + """试着反序列化 engine。返回 None 表示可用，否则返回错误摘要。 + + 引擎和 TensorRT 版本是绑死的：换过 TensorRT 之后旧引擎会直接加载失败。 + 这个检查就是为了把这种"静默失效"提前暴露出来。 + """ + try: + import tensorrt as trt + except Exception as exc: # noqa: BLE001 + return f"无法 import tensorrt: {exc}" + try: + rt = trt.Runtime(trt.Logger(trt.Logger.ERROR)) + eng = rt.deserialize_cuda_engine(p.read_bytes()) + if eng is None: + return "反序列化返回空（多半是 TensorRT 版本不匹配）" + return None + except Exception as exc: # noqa: BLE001 + msg = str(exc) + if "incompatible serialization version" in msg or "magicTag" in msg: + return "TensorRT 版本不匹配（需重新导出：weights.sh convert --name ...）" + return f"{type(exc).__name__}: {msg[:160]}" + + +def cmd_doctor(args) -> int: + root = _root(args) + problems = 0 + print(f"权重目录: {root}") + _hr("═") + + # 1) 逐模型体检 -------------------------------------------------------- + entries = discover(root) + print(f"[1/4] 模型（{len(entries)} 个）") + for e in entries: + ls = model_labels(e, deep=args.deep) + tag = "OK " if ls.source != SOURCE_PLACEHOLDER else "缺标签" + if ls.source == SOURCE_PLACEHOLDER: + problems += 1 + print(f" [{tag}] {e.group}/{e.name} 格式: {' '.join(e.exts)}") + if ls.source == SOURCE_PLACEHOLDER: + print(f" 期望: {e.labels_path.name}") + # 只有 .pt 没有 onnx/engine 时，two_stage/onnx 后端会用不了 + if "onnx" not in e.files and any(x in e.files for x in ("engine", "trt")): + print(" 注意: 只有 engine 没有 onnx，CPU/ORT 后端无处可用") + + # 2) engine 版本 ------------------------------------------------------- + print() + print("[2/4] engine 可用性（TensorRT 版本绑定检查）") + engines = [ + p for e in entries for ext in ("engine", "trt") if (p := e.path(ext)) is not None + ] + if not engines: + print(" (没有 engine)") + for p in engines: + if args.check_engines: + err = _engine_usable(p) + if err: + problems += 1 + print(f" [失效] {p.name}: {err}") + else: + print(f" [可用] {p.name}") + else: + print(f" [跳过] {p.name} (加 --check-engines 实测)") + + # 3) 配置引用 ---------------------------------------------------------- + print() + print("[3/4] config/site.yaml 的引用") + site = _read_site_yaml() + paths = site.get("paths") or {} + refs: List[tuple] = [] + if paths.get("yolo_weights"): + refs.append(("paths.yolo_weights", paths["yolo_weights"])) + if paths.get("hole_weights"): + refs.append(("paths.hole_weights", paths["hole_weights"])) + if paths.get("hole_workpiece_weights"): + refs.append(("paths.hole_workpiece_weights", paths["hole_workpiece_weights"])) + hd = site.get("hole_detect") or {} + if hd.get("weights"): + refs.append(("hole_detect.weights", hd["weights"])) + if hd.get("workpiece_weights"): + refs.append(("hole_detect.workpiece_weights", hd["workpiece_weights"])) + + if not refs: + print(" (没读到引用)") + for key, val in refs: + p = Path(str(val)).expanduser() + if not p.is_file(): + if p.suffix == ".pt" and p.with_suffix(".onnx").is_file() is False: + problems += 1 + print(f" [悬空] {key} → {p} (文件不存在)") + else: + problems += 1 + print(f" [悬空] {key} → {p} (文件不存在)") + continue + ls = read_labels(p, allow_model_metadata=False) + origin = SOURCE_LABELS_ZH.get(ls.source, ls.source) + print(f" [OK ] {key} → {p.name} 标签: {origin}") + if ls.is_placeholder: + problems += 1 + print(f" → 该模型读不到标签，运行时类名会变成 class_0") + + backend = (site.get("yolo") or {}).get("backend") + if backend == "tensorrt": + yw = Path(str(paths.get("yolo_weights") or "")).expanduser() + if yw.name and not yw.with_suffix(".engine").is_file(): + problems += 1 + print( + f" [风险] yolo.backend=tensorrt 但 {yw.with_suffix('.engine').name} 不存在，" + "启动会直接失败" + ) + + # 4) staging ----------------------------------------------------------- + print() + print("[4/4] staging（待转换）") + st = staging_files(root) + if not st: + print(" (空)") + for p in st: + ext = p.suffix.lower().lstrip(".") + if ext in CONVERTIBLE_INPUT_EXTS: + print(f" [待转换] {p.name} → weights.sh convert {STAGING_DIRNAME}/{p.name} --name <模型名>") + else: + print(f" [需人工] {p.name} 格式 .{ext} 不在支持列表（{', '.join(CONVERTIBLE_INPUT_EXTS)}）") + + _hr("═") + if problems: + print(f"发现 {problems} 处问题。") + return 1 + print("没有发现问题。") + return 0 + + +# --------------------------------------------------------------------------- +# show：把模块约定直接打出来，省得翻文档 +# --------------------------------------------------------------------------- + + +def cmd_show(args) -> int: + root = _root(args) + print("权重模块约定") + _hr() + print(f"权重根目录 : {root}") + print(f"仓库根 : {_REPO_ROOT or '(未找到)'}") + print(f"待转换目录 : {staging_dir(root)}") + print(f"标签文件后缀 : {LABELS_SUFFIX} （<模型名>{LABELS_SUFFIX}，放权重旁边，进 git）") + print(f"权重后缀 : {', '.join('.' + e for e in WEIGHT_EXTS)} （二进制，不进 git）") + print(f"已存在的分组 : {', '.join(groups(root)) or '(无)'}") + print() + print("命令") + _hr() + print(" weights.sh list [--deep] [--json]") + print(" weights.sh labels <模型名> [--group G] [--set a,b,c] [--set-file f.txt]") + print(" [--from-model] [--clear] [--deep] [--json]") + print(" weights.sh convert <文件> --name <模型名> [--group yolo] [--to all|onnx|engine]") + print(" [--imgsz 640] [--labels a,b] [--force] [--dry-run]") + print(" [--no-fp16] [--keep-staging]") + print(" weights.sh doctor [--deep] [--check-engines]") + print() + print("标签优先级: 旁挂文件 > ONNX 内嵌 metadata > 模型内嵌 metadata > 占位符(class_0)") + return 0 + + +# --------------------------------------------------------------------------- +# main +# --------------------------------------------------------------------------- + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + prog="weights.sh", + description="模型权重模块：转换（PT→ONNX→TensorRT）、标签读写、列表、自检", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=( + "示例:\n" + ' weights.sh list\n' + ' weights.sh labels detect_best_0\n' + ' weights.sh labels detect_best_0 --set "工件,通止规"\n' + ' weights.sh convert staging/best_v3.pt --name workpiece\n' + ' weights.sh convert models/yolo/detect_best_0.onnx --name detect_best_0 --to engine\n' + ), + ) + p.add_argument("--root", help="权重根目录（默认自动探测 models/）") + sub = p.add_subparsers(dest="cmd", required=True) + + s = sub.add_parser("list", help="列出所有权重及其标签") + s.add_argument("--deep", action="store_true", help="也尝试从 .pt/.engine 自身读标签（慢）") + s.add_argument("--json", action="store_true", help="以 JSON 输出") + s.set_defaults(func=cmd_list) + + s = sub.add_parser("labels", help="查看/修改模型标签") + s.add_argument("model", help="模型名（权重目录里的文件名主干）") + s.add_argument("--group", help="限定分组") + s.add_argument("--set", help='直接指定类别名，如 "工件,通止规"（改了立刻生效）') + s.add_argument("--set-file", dest="set_file", help="从纯文本文件导入（一行一个类名）") + s.add_argument("--from-model", dest="from_model", action="store_true", + help="从模型内嵌 metadata 抄一份到旁挂文件（变成可编辑）") + s.add_argument("--clear", action="store_true", help="删除旁挂标签文件") + s.add_argument("--deep", action="store_true", help="读取时也加载 .pt/.engine 自身") + s.add_argument("--json", action="store_true", help="以 JSON 输出") + s.set_defaults(func=cmd_labels) + + s = sub.add_parser("convert", help="转换权重并归位") + s.add_argument("input", help="输入文件（可以是 staging/ 里的相对名）") + s.add_argument("--name", help="输出模型名（默认沿用输入文件名主干）") + s.add_argument("--group", default="yolo", help="输出分组子目录（默认 yolo）") + s.add_argument("--to", default="all", choices=["all", "onnx", "engine"], help="产出目标") + s.add_argument("--imgsz", type=int, help="输入尺寸（默认取 site.yaml 的 yolo.imgsz）") + s.add_argument("--labels", help='顺带写标签，如 "工件,通止规"') + s.add_argument("--force", action="store_true", help="覆盖已存在的产物") + s.add_argument("--no-fp16", dest="no_fp16", action="store_true", help="engine 不用 FP16") + s.add_argument("--dry-run", dest="dry_run", action="store_true", help="只打印计划") + s.add_argument("--keep-staging", dest="keep_staging", action="store_true", + help="转换后不把 staging 里的源文件挪走") + s.set_defaults(func=cmd_convert) + + s = sub.add_parser("doctor", help="自检：缺标签 / 悬空引用 / engine 失效") + s.add_argument("--deep", action="store_true", help="也加载 .pt/.engine 读标签") + s.add_argument("--check-engines", dest="check_engines", action="store_true", + help="实测反序列化每个 engine（探测 TensorRT 版本不匹配）") + s.set_defaults(func=cmd_doctor) + + s = sub.add_parser("show", help="打印模块约定与命令") + s.set_defaults(func=cmd_show) + + return p + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = build_parser().parse_args(argv) + try: + return int(args.func(args)) + except KeyboardInterrupt: + _err("被中断") + return 130 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/weights.sh b/scripts/weights.sh new file mode 100755 index 0000000..a779a7f --- /dev/null +++ b/scripts/weights.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +# 权重模块的统一入口（薄封装）。 +# +# bash ~/ros2/scripts/weights.sh list +# bash ~/ros2/scripts/weights.sh labels detect_best_0 +# bash ~/ros2/scripts/weights.sh labels detect_best_0 --set "工件,通止规" +# bash ~/ros2/scripts/weights.sh convert staging/best_v3.pt --name workpiece +# bash ~/ros2/scripts/weights.sh doctor --check-engines +# +# 本脚本只做三件事：准备环境（ROS + TensorRT 动态库）→ 挑一个能用的 python → 转发。 +# 真正的逻辑在 scripts/weights.py。 +# +# 完全等价的老命令 export_yolo_trt.sh 仍然可用（它是本脚本的转发壳）。 + +set -eo pipefail +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" + +# 用户显式指定解释器时优先用它；否则试 ros_flexiv。 +# 本模块需要的重依赖（ultralytics / onnx / tensorrt / onnxruntime）都在那里。 +pick_python() { + if [[ -n "${WEIGHTS_PY:-}" ]]; then + echo "${WEIGHTS_PY}"; return 0 + fi + local cands=( + "${HOME}/anaconda3/envs/ros_flexiv/bin/python" + "$(command -v python3 || true)" + ) + for c in "${cands[@]}"; do + [[ -n "${c}" && -x "${c}" ]] || continue + # list/labels/show 只需要 yaml；convert/doctor 才需要 tensorrt。 + # 这里统一要求 yaml，避免半路才发现解释器不对。 + if "${c}" -c 'import yaml' >/dev/null 2>&1; then + echo "${c}"; return 0 + fi + done + return 1 +} + +PY="$(pick_python)" || { + echo "错误: 找不到可用的 python（需要能 import yaml）。" >&2 + echo " 可用 WEIGHTS_PY=/path/to/python 指定。" >&2 + exit 1 +} + +# TensorRT 的 pip 版把 .so 放在非标准位置，TensorRT 的 Python API 和运行时都要用。 +# export_yolo_trt.sh 里也是这么做的；这里保持同样做法。 +TRT_LIBS="$("${PY}" - <<'PY' 2>/dev/null || true +import pathlib +try: + import tensorrt +except Exception: + raise SystemExit(0) +p = pathlib.Path(tensorrt.__file__).resolve().parent +cands = [p, p.parent / "tensorrt_libs", p.parent / "nvidia" / "tensorrt" / "lib"] +print(":".join(str(x) for x in cands if x.exists())) +PY +)" +[[ -n "${TRT_LIBS}" ]] && export LD_LIBRARY_PATH="${TRT_LIBS}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + +# 不让 Ultralytics 自己去 pip 装东西（会卡住或破坏环境）。 +export YOLO_AUTOINSTALL=false + +# site.yaml 里的 yolo.imgsz 作为 convert 的默认尺寸。 +if [[ -f "${ROOT}/scripts/lib/load_site_config.sh" ]]; then + # shellcheck disable=SC1091 + source "${ROOT}/scripts/lib/load_site_config.sh" >/dev/null 2>&1 || true +fi + +exec "${PY}" "${ROOT}/scripts/weights.py" "$@" diff --git a/src/arm_application/arm_application/handeye_calibrate.py b/src/arm_application/arm_application/handeye_calibrate.py index 081965c..2e985a4 100644 --- a/src/arm_application/arm_application/handeye_calibrate.py +++ b/src/arm_application/arm_application/handeye_calibrate.py @@ -131,7 +131,7 @@ class HandeyeCalibrate(Node): self.declare_parameter('pin_hold_sec', 0.45) # 每拍多长 self.declare_parameter('skip_unstable_poses', True) - self.declare_parameter('output_dir', '/tmp/handeye_calib') + self.declare_parameter('output_dir', '/home/cc/ros2/handeye/tmp') self.declare_parameter('result_file', '') self.declare_parameter('cleanup_images', True) diff --git a/src/arm_application/arm_application/handeye_teach.py b/src/arm_application/arm_application/handeye_teach.py index f75c1e4..fb323cf 100644 --- a/src/arm_application/arm_application/handeye_teach.py +++ b/src/arm_application/arm_application/handeye_teach.py @@ -244,7 +244,7 @@ class RdkBridge: class HandeyeTeach(Node): def __init__(self) -> None: super().__init__('handeye_teach') - self.declare_parameter('output_dir', '/tmp/handeye_calib') + self.declare_parameter('output_dir', '/home/cc/ros2/handeye/tmp') self.declare_parameter('width', 1280) self.declare_parameter('height', 720) self.declare_parameter('fps', 15) diff --git a/src/arm_application/arm_application/hole_detect_hough.py b/src/arm_application/arm_application/hole_detect_hough.py index 5be8e9d..b3581b6 100644 --- a/src/arm_application/arm_application/hole_detect_hough.py +++ b/src/arm_application/arm_application/hole_detect_hough.py @@ -40,10 +40,34 @@ def _fail( } -def _draw_bbox(img: np.ndarray, bbox: tuple[int, int, int, int], color=(0, 255, 0)) -> None: +def _draw_bbox( + img: np.ndarray, + bbox: tuple[int, int, int, int], + color=(0, 255, 0), + label: str = "", +) -> None: + """画工件框。``label`` 是模型给出的类别名（由标签文件决定）。 + + 这里刻意**不写死类名**：以前写死了 "gongjian"，但实际模型的类别是 + "workpiece"，图上的字和模型的类名对不上，排查时会被误导。 + 现在用 ``Detection.cls_name``，也就是 ``weights.sh labels`` 里能读到的那份。 + """ x1, y1, x2, y2 = bbox cv2.rectangle(img, (x1, y1), (x2, y2), color, 2) - cv2.putText(img, "gongjian", (x1, max(20, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2) + text = (label or "").strip() + if text: + cv2.putText(img, text, (x1, max(20, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2) + + +def _best_label(dets) -> str: + """取置信度最高那个检测的类名（由标签文件决定）。没有则返回空串。""" + if not dets: + return "" + try: + best = max(dets, key=lambda d: float(d.conf)) + except Exception: # noqa: BLE001 + return "" + return str(getattr(best, "cls_name", "") or "").strip() def _draw_center_gate(img: np.ndarray, max_dx: float, max_dy: float, ok: bool) -> None: @@ -332,11 +356,12 @@ def _render_success( hole_radii: list[float], dist: float, gate: tuple[float, float] | None = None, + label: str = "", ) -> np.ndarray: draw_img = image.copy() if gate is not None: _draw_center_gate(draw_img, gate[0], gate[1], True) - _draw_bbox(draw_img, bbox) + _draw_bbox(draw_img, bbox, label=label) _draw_holes(draw_img, hole_points, hole_radii) cv2.putText( draw_img, @@ -391,7 +416,9 @@ def detect_holes( if bbox is None: return _fail("未检测到工件", draw_img) - _draw_bbox(draw_img, bbox) + # 类别名来自标签文件（见 model_labels / weights.sh labels），不再写死。 + label = _best_label(dets) + _draw_bbox(draw_img, bbox, label=label) # 1.5) 居中门控：bbox 中心接近画面中心才进入二阶检测。 # 精对准 k1/k2 时工件会离开画面中心，调用方应 skip_center_gate=True。 @@ -428,13 +455,16 @@ def detect_holes( hole_radii = [float(r1), float(r2)] dist = float(np.hypot(k1[0] - k2[0], k1[1] - k2[1])) gate = (max_dx, max_dy) if bool(cfg.get("center_gate_draw", True)) else None - draw_img = _render_success(image, bbox, hole_points, hole_radii, dist, gate=gate) + draw_img = _render_success( + image, bbox, hole_points, hole_radii, dist, gate=gate, label=label + ) return { "status": True, "msg": f"螺纹嵌件对，两孔距离 {dist:.1f}px", "hole_points": hole_points, "hole_radii": hole_radii, "bbox": list(bbox), + "label": label, "draw_img": draw_img, } # ---- 退路：原 Hough 老路径 ---- @@ -561,13 +591,14 @@ def detect_holes( hole_points = [[x1, y1], [x2, y2]] hole_radii = [r1, r2] gate = (max_dx, max_dy) if bool(cfg.get("center_gate_draw", True)) else None - draw_img = _render_success(image, bbox, hole_points, hole_radii, dist, gate=gate) + draw_img = _render_success(image, bbox, hole_points, hole_radii, dist, gate=gate, label=label) return { "status": True, "msg": f"检测成功，两孔距离 {dist:.1f}px", "hole_points": hole_points, "hole_radii": hole_radii, "bbox": list(bbox), + "label": label, "draw_img": draw_img, } except Exception as exc: @@ -675,7 +706,15 @@ class HoleStabilizer: max_dx = float(self.cfg.get("center_max_offset_ratio_x", 0.12)) * iw max_dy = float(self.cfg.get("center_max_offset_ratio_y", 0.12)) * ih gate = (max_dx, max_dy) - draw_img = _render_success(image, bbox_draw, self.points, self.radii, dist, gate=gate) + draw_img = _render_success( + image, + bbox_draw, + self.points, + self.radii, + dist, + gate=gate, + label=str(result.get("label") or ""), + ) held = self.miss > 0 return { diff --git a/src/arm_application/arm_application/model_labels.py b/src/arm_application/arm_application/model_labels.py new file mode 100644 index 0000000..9bd374a --- /dev/null +++ b/src/arm_application/arm_application/model_labels.py @@ -0,0 +1,555 @@ +"""模型类别标签（class names）的统一读写。 + +## 为什么需要这个模块 + +推理代码需要把模型输出的**类别索引**（整数 0/1/2…）翻译成**人能看懂的名字**。 +这件事原先没有统一做法，导致同一个模型在不同后端下名字都不一样： + +- `.pt`：从 Ultralytics 的 `model.names` 读（来源是 checkpoint 内嵌 metadata） +- `.onnx`：ONNX 文件里有 `names` metadata，但代码没用，而是**硬编码** `{0: "workpiece"}` +- `.engine`：**标签直接丢了** + +最后一条是本模块存在的主要原因。TensorRT engine 要用 `trt.OnnxParser` 重新构建， +而**当前安装的 TensorRT（10.13）Python API 根本没有 metadata 读写接口** —— +`INetworkDefinition` / `IBuilderConfig` / `ICudaEngine` 上都没有任何 metadata 方法。 +所以 engine 构建完就只剩张量，标签无从附着。实测： + + detect_best_0.engine 二进制里搜不到 "workpiece" → 标签已丢 + detect_hc90.engine 能搜到 "workpiece" → 旧版 TRT 建的，且已无法加载 + +后果很隐蔽：`yolo_detector` 读不到名字时会退化成 `str(cid)`，也就是类名变成字符串 +`"0"`；而预览界面又不绘制类名，所以**没有任何地方看得出标签丢了**。 + +## 解决办法：旁挂标签文件 + +在每个权重旁边放一个同名文本文件 `<模型名>.labels.yaml`： + + models/yolo/detect_best_0.onnx + models/yolo/detect_best_0.engine + models/yolo/detect_best_0.labels.yaml ← 就是这个 + +选文本文件而不是想办法塞进 engine，理由有二： + +1. **可修改**。标签是配置，不是模型的一部分。塞进 engine 就只能重建 engine 才能改。 +2. **跨格式统一**。engine / onnx / pt 三种格式共用同一份标签，不会出现 + "换后端名字就变了"。engine 丢标签的问题也一并绕过了。 + +标签文件是纯文本，进 git（体积小、属于配置）；权重二进制不进 git。 + +## 读取优先级 + + sidecar（旁挂文件） > onnx（ONNX metadata） > model（模型自身 metadata） > placeholder + +`read_labels()` 会连**来源**一起返回。这点很重要：调用方可以据此提示 +"标签来自占位符，请补一个 .labels.yaml"，而不是悄悄给出一个错的类名。 +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Tuple, Union + +# 标签文件的扩展名。刻意用双后缀：既能被 *.yaml 工具识别， +# 又能和模型自身的配置（如果有）区分开。 +LABELS_SUFFIX = ".labels.yaml" + +# read_labels() 的来源标记。顺序即优先级（高 → 低）。 +SOURCE_SIDECAR = "sidecar" # 旁挂 <模型名>.labels.yaml（权威，用户可直接改） +SOURCE_ONNX = "onnx" # ONNX 文件内嵌 metadata +SOURCE_MODEL = "model" # 模型自身 metadata（Ultralytics 的 .names） +SOURCE_PLACEHOLDER = "placeholder" # 都没有 → 占位符 class_0/class_1… + +SOURCE_LABELS_ZH = { + SOURCE_SIDECAR: "旁挂标签文件", + SOURCE_ONNX: "ONNX 内嵌 metadata", + SOURCE_MODEL: "模型内嵌 metadata", + SOURCE_PLACEHOLDER: "占位符（缺标签）", +} + + +@dataclass(frozen=True) +class LabelSet: + """一次标签读取的结果。 + + ``names`` 是索引 → 名字；``source`` 说明它从哪来（见 SOURCE_* 常量）。 + + 调用方**应该**检查 ``source``：如果是 ``placeholder``，说明这个模型缺标签， + 此时给出的是 ``class_0`` 这样的占位名，只适合占位显示，不该当成真名使用。 + """ + + names: Dict[int, str] + source: str + path: Optional[Path] = None # 实际读到的文件（placeholder 时为 None） + + @property + def is_placeholder(self) -> bool: + return self.source == SOURCE_PLACEHOLDER + + def get(self, cid: int, default: Optional[str] = None) -> str: + if cid in self.names: + return self.names[cid] + return default if default is not None else f"class_{cid}" + + def sorted_names(self) -> List[str]: + return [self.names[k] for k in sorted(self.names)] + + +# --------------------------------------------------------------------------- +# 路径 +# --------------------------------------------------------------------------- + + +def labels_path_for(weights: Union[str, Path]) -> Path: + """给出权重对应的旁挂标签文件路径。 + + 只换后缀，目录不变：``models/yolo/detect_best_0.engine`` → + ``models/yolo/detect_best_0.labels.yaml``。 + + 这样无论 ``resolve_weights()`` 最后选中了 .pt/.onnx/.engine 中的哪一个， + 算出来的标签路径都相同 —— 一份标签服务同一个模型的所有格式。 + """ + p = Path(weights).expanduser() + # 用 stem 手工拼接，而不是 with_suffix('.labels.yaml')： + # 后者对 `a.b.onnx` 这类含点文件名容易踩坑，手工拼没有歧义。 + return p.parent / (p.stem + LABELS_SUFFIX) + + +# --------------------------------------------------------------------------- +# 解析：把各种形状的标签定义统一成 {索引: 名字} +# --------------------------------------------------------------------------- + + +def _clean(text: object) -> str: + # None 必须先拦掉：str(None) == "None" 是 truthy 的， + # 直接 str() 会让空项变成一个字面量为 "None" 的类名。 + if text is None: + return "" + return str(text).strip() + + +def _try_parse_structured(text: str) -> Optional[Dict[int, str]]: + """尝试把字符串当成 dict/list 的字面量解析。 + + 为什么需要：ONNX metadata 里的 ``names`` 存的**不是** JSON，也不是逗号列表， + 而是 Python 字典的 repr（单引号）： + + "{0: 'workpiece'}" + + Ultralytics 是直接 ``str(names)`` 写进去的，所以 ``json.loads`` 会失败。 + 两种都试一遍（JSON 标准、literal_eval 兜底），都不行才交给逗号解析。 + ``ast.literal_eval`` 只解析字面量、不执行代码，是安全的。 + """ + s = text.strip() + if not s: + return None + import ast + import json + + for loader in (json.loads, ast.literal_eval): + try: + obj = loader(s) + except Exception: # noqa: BLE001 + continue + if isinstance(obj, (dict, list, tuple)): + return parse_names_field(obj) + return None + + +def _parse_freeform_text(text: str) -> Dict[int, str]: + """解析"自由文本"形式的标签内容。 + + 两种情况要分开，否则会互相吃掉： + + - 一行一个名字（有人把 ``classes.txt`` 直接改名成 ``.labels.yaml``）： + ``"a\\nb"`` → ``{0:'a', 1:'b'}`` + - 单行逗号串：``"a, b"`` → ``{0:'a', 1:'b'}`` + + 判据是"有没有多行"。只有一行时才按逗号切， + 否则 ``"a, b"`` 这种文件名里带逗号的写法会被切错。 + """ + items = [ + ln.strip() + for ln in text.splitlines() + if ln.strip() and not ln.strip().startswith("#") + ] + if len(items) > 1: + return {i: n for i, n in enumerate(items)} + return parse_label_list(text) + + +def parse_names_field(value: object) -> Dict[int, str]: + """把 YAML/metadata 里的 ``names`` 字段规范成 ``{int: str}``。 + + 容忍三种写法，因为它们各自都很常见： + + - 字典（Ultralytics data.yaml 的写法）：``{0: workpiece, 1: hole}`` + - 列表（本模块写出的规范形式，索引即顺序）：``[workpiece, hole]`` + - 逗号分隔的字符串：``"workpiece, hole"`` + + 跳过空项但**保留其索引位置**，这样 ``[a, "", c]`` 不会让 c 变成 1 号。 + 无法解析成整数键的字典项会被忽略（不抛异常 —— 标签读不出来不该让程序崩）。 + """ + if value is None: + return {} + if isinstance(value, dict): + out: Dict[int, str] = {} + for k, v in value.items(): + try: + idx = int(k) + except (TypeError, ValueError): + continue + name = _clean(v) + if name: + out[idx] = name + return out + if isinstance(value, (list, tuple)): + out = {} + for i, v in enumerate(value): + name = _clean(v) + if name: + out[i] = name + return out + if isinstance(value, str): + # 字符串可能是三种东西：dict/list 的字面量（ONNX metadata）、 + # 逗号列表、或者多行文本。先试结构化的，再退到逗号/多行解析。 + structured = _try_parse_structured(value) + if structured is not None: + return structured + return _parse_freeform_text(value) + return {} + + +def parse_label_list(text: str) -> Dict[int, str]: + """解析逗号（或中文逗号、分号）分隔的类别名。 + + ``"workpiece, hole"`` → ``{0: 'workpiece', 1: 'hole'}`` + + 同时兼容 ``"0:workpiece, 1:hole"`` 这种带显式索引的写法， + 因为用户从别处拷配置过来时经常带着索引。空项被跳过。 + """ + if not text or not text.strip(): + return {} + parts = [p for p in re.split(r"[,，;；\n]", text)] + out: Dict[int, str] = {} + next_idx = 0 + for raw in parts: + item = raw.strip() + if not item: + continue + m = re.match(r"^\s*(\d+)\s*[:：]\s*(.+)$", item) + if m: + idx = int(m.group(1)) + name = m.group(2).strip() + if name: + out[idx] = name + next_idx = max(next_idx, idx + 1) + continue + out[next_idx] = item + next_idx += 1 + return out + + +def read_label_text_file(path: Union[str, Path]) -> Dict[int, str]: + """读纯文本标签文件：一行一个类名（YOLO 的 classes.txt 就是这个格式）。 + + 空行和 ``#`` 开头的注释行被忽略，且**不占索引**。 + 这样用户可以在文件里随意加注释而不破坏索引对应关系。 + """ + p = Path(path).expanduser() + out: Dict[int, str] = {} + idx = 0 + with p.open("r", encoding="utf-8") as f: + for line in f: + s = line.strip() + if not s or s.startswith("#"): + continue + out[idx] = s + idx += 1 + return out + + +def placeholder_labels(nc: int) -> Dict[int, str]: + """没有标签可用时的占位名。 + + 刻意用 ``class_0`` 而不是直接回退成字符串形式的数字 ``"0"``： + 后者和"真的有一个叫 0 的类"无法区分，正是原先难以发现问题所在。 + ``class_`` 前缀一眼就能看出是没配标签。 + """ + n = max(0, int(nc)) + return {i: f"class_{i}" for i in range(n)} + + +# --------------------------------------------------------------------------- +# 旁挂标签文件读写 +# --------------------------------------------------------------------------- + + +def format_labels_yaml(names: Dict[int, str], model: Optional[str] = None) -> str: + """生成旁挂标签文件的内容（带使用说明的注释头）。 + + 注释头是刻意保留的：这个文件会被不熟悉的人直接打开编辑， + 把"第几行 = 第几类"写在文件里比写在文档里有效得多。 + """ + lines: List[str] = [] + lines.append("# 模型类别标签（class names）") + lines.append("#") + lines.append("# 一个模型对应一份标签，放在权重旁边、同名前缀：") + lines.append("# detect_best_0.engine ← 权重") + lines.append("# detect_best_0.labels.yaml ← 本文件") + lines.append("#") + lines.append("# 写法：names 下面一行一个类名，第 1 行对应类别 0，第 2 行对应类别 1，依次类推。") + lines.append("# 改完保存即生效，不需要重新转换模型、也不用重启环境（重跑节点即可）。") + lines.append("#") + lines.append("# 命令行查看/修改（推荐，不用手改本文件）：") + lines.append("# bash ~/ros2/scripts/weights.sh labels <模型名>") + lines.append("# bash ~/ros2/scripts/weights.sh labels <模型名> --set \"工件,通止规\"") + lines.append("") + if model: + # 记下所属模型，便于发现"标签文件被拷给了另一个模型"这种错配。 + lines.append(f"model: {model}") + lines.append("names:") + if names: + for k in sorted(names): + lines.append(f" - {names[k]}") + else: + lines.append(" []") + return "\n".join(lines) + "\n" + + +def write_labels( + weights: Union[str, Path], + names: Dict[int, str], + *, + path: Optional[Union[str, Path]] = None, +) -> Path: + """把标签写到权重旁边的旁挂文件，返回写入路径。""" + dst = Path(path).expanduser() if path else labels_path_for(weights) + dst.parent.mkdir(parents=True, exist_ok=True) + model = Path(weights).stem + dst.write_text(format_labels_yaml(names, model=model), encoding="utf-8") + return dst + + +def read_label_file(path: Union[str, Path]) -> Optional[Dict[int, str]]: + """读旁挂标签文件。文件不存在或内容为空返回 None（不抛异常）。""" + p = Path(path).expanduser() + if not p.is_file(): + return None + try: + raw = p.read_text(encoding="utf-8") + except Exception: # noqa: BLE001 + return None + try: + import yaml + + data = yaml.safe_load(raw) + except Exception: # noqa: BLE001 + # 文件存在但读不动（语法错、编码错、权限）。返回 None 让调用方按 + # "没有标签"处理，而不是让整个推理进程崩掉 —— 标签缺失不该是致命的。 + return None + + if isinstance(data, dict): + # 规范形式：names: [...] + names = parse_names_field(data.get("names")) + if names: + return names + # 容忍 {0: 'a', 1: 'b'} 这种"顶层直接就是索引字典"的写法（Ultralytics data.yaml） + names = { + int(k): str(v).strip() + for k, v in data.items() + if str(k).lstrip("-").isdigit() and str(v).strip() + } + # 是合法 YAML 字典却读不出类别名 → 就是没有标签。 + # 这里**不能**再回落到"按行读文件"，否则会把 `names: []` 里的 + # 'names:' 这几个字当成类名。 + return names or None + + if isinstance(data, (list, tuple)): + return parse_names_field(data) or None + + if isinstance(data, str): + # 整个文件就是一个标量 —— 可能是"改名过来的 classes.txt"。 + # 注意必须用**原始文本**：YAML 会把多行纯标量折成一行 + # （"a\nb\n" → "a b"），用 data 的话换行已经没了， + # "一行一个类名"会被当成一个带空格的长名字。 + return _parse_freeform_text(raw) or None + + # 空文件 / 只有注释 / 其它标量类型：按行读（能正确处理注释行）。 + try: + return read_label_text_file(p) or None + except Exception: # noqa: BLE001 + return None + + +def label_file_model(path: Union[str, Path]) -> Optional[str]: + """读旁挂标签文件里记录的 ``model:`` 字段（用于发现标签错配）。""" + p = Path(path).expanduser() + if not p.is_file(): + return None + try: + import yaml + + data = yaml.safe_load(p.read_text(encoding="utf-8")) + except Exception: # noqa: BLE001 + return None + if isinstance(data, dict): + v = data.get("model") + return _clean(v) if v else None + return None + + +# --------------------------------------------------------------------------- +# 从模型自身读标签 +# --------------------------------------------------------------------------- + + +def onnx_metadata(path: Union[str, Path]) -> Dict[str, str]: + """读 ONNX 的 ``metadata_props``（全是字符串），读不到返回空字典。 + + 用 ``load_external_data=False``：只读图结构不读权重，对 38MB 的模型也是毫秒级。 + Ultralytics 导出时会写入 ``names`` / ``imgsz`` / ``nc`` / ``task`` 等。 + """ + p = Path(path).expanduser() + if not p.is_file(): + return {} + try: + import onnx + + m = onnx.load(str(p), load_external_data=False) + except Exception: # noqa: BLE001 + return {} + return {x.key: x.value for x in m.metadata_props} + + +def labels_from_onnx(path: Union[str, Path]) -> Optional[Tuple[Dict[int, str], int]]: + """从 ONNX 的 ``metadata_props`` 读 ``names`` 与类别数 ``nc``。 + + 返回 ``(names, nc)``；读不到返回 None。 + + 注意 ``names`` 在 ONNX 里存的是 **Python 字典的 repr 字符串**（单引号）， + 形如 ``"{0: 'workpiece'}"`` —— 不是 JSON。解析细节见 ``parse_names_field``。 + """ + meta = onnx_metadata(path) + if not meta: + return None + + names = parse_names_field(meta.get("names")) + nc = 0 + raw_nc = meta.get("nc") + if raw_nc is not None: + try: + nc = int(str(raw_nc).strip()) + except (TypeError, ValueError): + nc = 0 + if not nc and names: + nc = max(names) + 1 + return (names, nc) if names else None + + +def labels_from_model(path: Union[str, Path]) -> Optional[Tuple[Dict[int, str], int]]: + """从模型自身 metadata 读标签（走 Ultralytics，主要给 .pt 用）。 + + 对 .engine / .onnx 也能用（Ultralytics 会自己反序列化），但代价高： + 要加载整个模型。所以 ``read_labels()`` 只在必要时才调它。 + + Ultralytics 的 model 对象除了 ``names`` 还有 ``nc``；拿不到 nc 时用 names 推断。 + """ + p = Path(path).expanduser() + if not p.is_file(): + return None + try: + from ultralytics import YOLO + + model = YOLO(str(p), task="detect") + raw = getattr(model, "names", None) or {} + names = parse_names_field(raw) + nc = 0 + for attr in ("nc",): + try: + nc = int(getattr(model, attr) or 0) + except Exception: # noqa: BLE001 + nc = 0 + if not nc and names: + nc = max(names) + 1 + return (names, nc) if names else None + except Exception: # noqa: BLE001 + return None + + +# --------------------------------------------------------------------------- +# 统一入口 +# --------------------------------------------------------------------------- + + +def read_labels( + weights: Union[str, Path], + *, + allow_model_metadata: bool = True, + nc_hint: int = 0, +) -> LabelSet: + """按优先级读出一个模型该用的标签。 + + 优先级：旁挂文件 > ONNX metadata > 模型自身 metadata > 占位符。 + + 旁挂优先是刻意的：它是用户能直接编辑的那一份，必须覆盖模型内嵌的旧值， + 否则"我明明改了标签为什么不生效"就会变成一个很难查的问题。 + + 绝不抛异常。读不到就返回占位符并把原因记在 ``source`` 里， + 由调用方决定是记一条警告还是直接显示。 + """ + p = Path(weights).expanduser() + + side = labels_path_for(p) + names = read_label_file(side) + if names: + return LabelSet(names=names, source=SOURCE_SIDECAR, path=side) + + suffix = p.suffix.lower() + # .pt 的标签只可能来自模型自身；.onnx 先看内嵌 metadata（快）； + # .engine / .trt 没有可用 metadata（见模块头说明），直接落到占位符或模型加载。 + if suffix == ".onnx": + got = labels_from_onnx(p) + if got: + return LabelSet(names=got[0], source=SOURCE_ONNX, path=p) + elif suffix in (".pt", ".pth"): + got = labels_from_model(p) + if got: + return LabelSet(names=got[0], source=SOURCE_MODEL, path=p) + + if allow_model_metadata and suffix in (".engine", ".trt"): + # engine 一般读不到（当前 TRT 没 metadata API），但也别直接放弃： + # 若用户是用 Ultralytics 导出的 engine，metadata 可能在。 + got = labels_from_model(p) + if got: + return LabelSet(names=got[0], source=SOURCE_MODEL, path=p) + + return LabelSet(names=placeholder_labels(nc_hint), source=SOURCE_PLACEHOLDER, path=None) + + +def describe(labels: LabelSet) -> str: + """一句人能看懂的说明，用于日志：``workpiece (来源: 旁挂标签文件)``。""" + names = ", ".join(labels.sorted_names()) or "(空)" + src = SOURCE_LABELS_ZH.get(labels.source, labels.source) + return f"{names} (来源: {src})" + + +def label_file_status( + weights: Union[str, Path], +) -> Tuple[bool, Optional[Path], Optional[str]]: + """检查旁挂标签文件状态，供 doctor 用。 + + 返回 ``(存在, 路径, 提示)``；``提示`` 为 None 表示没有问题。 + """ + p = Path(weights).expanduser() + side = labels_path_for(p) + if not side.is_file(): + return False, side, f"缺标签文件：{side.name}" + recorded = label_file_model(side) + if recorded and Path(weights).stem != recorded: + return True, side, ( + f"标签文件记录的 model 是 '{recorded}'，但当前权重是 " + f"'{Path(weights).stem}' —— 可能是拷错了标签文件" + ) + return True, side, None diff --git a/src/arm_application/arm_application/onnx_workpiece.py b/src/arm_application/arm_application/onnx_workpiece.py index bf97b23..0404d89 100644 --- a/src/arm_application/arm_application/onnx_workpiece.py +++ b/src/arm_application/arm_application/onnx_workpiece.py @@ -12,13 +12,22 @@ detect_best_1.onnx 满足（metadata: end2end=True, nms=False）。 from __future__ import annotations +import logging from pathlib import Path from typing import List, Optional, Union import numpy as np +from arm_application.model_labels import ( + SOURCE_MODEL, + LabelSet, + parse_names_field, + read_labels, +) from arm_application.yolo_detector import Detection +logger = logging.getLogger(__name__) + class OnnxWorkpieceDetector: """给 detect_holes() 用的最小接口：detect(image, conf, iou, imgsz)。""" @@ -47,7 +56,23 @@ class OnnxWorkpieceDetector: self._in_h, self._in_w = int(shape[-2]), int(shape[-1]) self.conf = float(conf) self.imgsz = int(imgsz) - self._names = class_names or {0: "workpiece"} + # 类别名：显式传入 > 旁挂标签文件 > ONNX 内嵌 metadata > 占位符。 + # 原先这里是硬编码 {0: "workpiece"}，等于把类名写死在代码里 —— + # 换一个多类模型就会展示错误的名字，而且没人改得动。 + if class_names: + names = parse_names_field(class_names) + self._labels = LabelSet(names=names, source=SOURCE_MODEL, path=path) + else: + self._labels = read_labels(path, allow_model_metadata=False) + if self._labels.is_placeholder: + logger.warning( + "ONNX %s 没有类别标签，将用占位名显示。" + "补法: bash ~/ros2/scripts/weights.sh labels %s --from-model", + path.name, + path.stem, + ) + # 兼容旧属性名，避免外部代码引用时报 AttributeError。 + self._names = dict(self._labels.names) def _letterbox(self, img: np.ndarray): h, w = img.shape[:2] @@ -105,7 +130,7 @@ class OnnxWorkpieceDetector: xyxy=xyxy, conf=float(score), cls_id=cid, - cls_name=self._names.get(cid, str(cid)), + cls_name=self._labels.get(cid), ) ) dets.sort(key=lambda d: d.conf, reverse=True) diff --git a/src/arm_application/arm_application/perception_node.py b/src/arm_application/arm_application/perception_node.py index 321f17f..469f603 100644 --- a/src/arm_application/arm_application/perception_node.py +++ b/src/arm_application/arm_application/perception_node.py @@ -74,10 +74,10 @@ class PerceptionNode(Node): self.declare_parameter("image_topic", "") self.declare_parameter("camera_info_topic", "") self.declare_parameter("depth_topic", "/d405/aligned_depth_to_color/image_raw") - self.declare_parameter("weights", "/home/cc/ros2/models/yolo/detect_hc90.pt") - self.declare_parameter("conf", 0.12) + self.declare_parameter("weights", "/home/cc/ros2/models/yolo/detect_best_0.onnx") + self.declare_parameter("conf", 0.10) self.declare_parameter("iou", 0.45) - self.declare_parameter("imgsz", 416) + self.declare_parameter("imgsz", 640) self.declare_parameter("device", "0") self.declare_parameter("min_conf", 0.12) self.declare_parameter("vision_hz", 30.0) diff --git a/src/arm_application/arm_application/weights_layout.py b/src/arm_application/arm_application/weights_layout.py new file mode 100644 index 0000000..ceb99d6 --- /dev/null +++ b/src/arm_application/arm_application/weights_layout.py @@ -0,0 +1,403 @@ +"""权重模块的目录布局、模型发现与标签归属（纯逻辑）。 + +这个模块刻意**只依赖标准库 + model_labels**，不 import ROS、不 import +ultralytics/tensorrt。原因： + +1. 可以在没有 GPU、没有 ROS 的环境里跑单元测试； +2. ``scripts/weights.sh`` 这个命令行工具能直接用上，不必先 ``source`` 一整套环境。 + +真正干重活的转换（PyTorch → ONNX → TensorRT）在 ``scripts/weights.py`` 里， +那部分才需要 ``ros_flexiv`` 环境和 GPU。 + +## 目录约定 + + models/ + README.md ← 模块规则（要改这个模块先读它） + staging/ ← 非标准格式先丢这里，转换工具有序地消费它 + <分组>/ ← 例如 yolo/、hole_insert/ + <模型名>.onnx + <模型名>.engine + <模型名>.labels.yaml ← 标签（旁挂，可手改，进 git） + +"分组"只是一个子目录名，用来分门别类，没有语义约束。默认 ``yolo``。 +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Tuple, Union + +from .model_labels import ( + SOURCE_MODEL, + SOURCE_ONNX, + SOURCE_PLACEHOLDER, + SOURCE_SIDECAR, + LabelSet, + label_file_status, + labels_path_for, + labels_from_model, + labels_from_onnx, + placeholder_labels, + read_label_file, +) + +# 会被当成"权重"的文件后缀（不带点的小写形式）。 +WEIGHT_EXTS: Tuple[str, ...] = ( + "pt", + "pth", + "onnx", + "engine", + "trt", + "plan", + "tflite", +) + +# 转换工具认识的输入格式。其它格式一律归到 staging 等人工处理。 +# 注意 pth 也收：虽然 Ultralytics 主要用 .pt，但用户手里常是 .pth。 +CONVERTIBLE_INPUT_EXTS: Tuple[str, ...] = ("pt", "pth", "onnx") + +# 已经被转换过的产物，重转没有意义。 +FINAL_EXTS: Tuple[str, ...] = ("engine", "trt", "plan") + +# 非标准格式/待转换文件放这里。 +STAGING_DIRNAME = "staging" + +# 读标签时挑格式的顺序：onnx 有内嵌 metadata 且读取代价极低（毫秒级）， +# 所以排在前面；engine 排最后是因为它很可能什么都没有。 +LABEL_READ_ORDER: Tuple[str, ...] = ("onnx", "pt", "pth", "engine", "trt") + + +# --------------------------------------------------------------------------- +# 仓库根 / 权重根 +# --------------------------------------------------------------------------- + + +def find_repo_root(start: Optional[Union[str, Path]] = None) -> Optional[Path]: + """从 ``start`` 往上找仓库根（以存在 ``config/site.yaml`` 为标志）。 + + 这样无论代码是从 ``src/`` 直接跑、还是从 ``install/`` 跑， + 都能定位到同一份 ``models/``，不会出现"两份权重"。 + """ + here = Path(start).resolve() if start else Path(__file__).resolve() + for cand in [here, *here.parents]: + if (cand / "config" / "site.yaml").is_file(): + return cand + return None + + +def default_weights_root() -> Path: + """``models/`` 的默认位置。找不到仓库根时退回当前工作目录下的 models/。""" + import os + + env = os.environ.get("SITE_WEIGHTS_ROOT") + if env: + return Path(env).expanduser() + root = find_repo_root() + if root is not None: + return root / "models" + return Path.cwd() / "models" + + +# --------------------------------------------------------------------------- +# 模型条目 +# --------------------------------------------------------------------------- + + +@dataclass +class ModelEntry: + """一个"模型" = 权重目录里同名前缀的一组文件。""" + + group: str + name: str + files: Dict[str, Path] = field(default_factory=dict) + + @property + def dir(self) -> Path: + if not self.files: + return Path(".") + return next(iter(self.files.values())).parent + + @property + def labels_path(self) -> Path: + return self.dir / (self.name + ".labels.yaml") + + @property + def exts(self) -> List[str]: + """已存在的格式，按固定顺序返回，保证输出稳定。 + + 顺序就是转换流水线的顺序：``pt → onnx → engine``（源头 → 中间 → 最终）。 + 这样列表里一眼能看出这个模型走到哪一步了。 + """ + order = list(WEIGHT_EXTS) + return sorted(self.files, key=lambda e: (order.index(e) if e in order else 99, e)) + + def path(self, ext: str) -> Optional[Path]: + return self.files.get(ext.lower().lstrip(".")) + + def sorted_files(self) -> List[Path]: + """按读标签的推荐顺序给出文件列表。""" + out: List[Path] = [] + for ext in LABEL_READ_ORDER: + p = self.files.get(ext) + if p: + out.append(p) + for ext in self.exts: + p = self.files.get(ext) + if p and p not in out: + out.append(p) + return out + + def missing_of(self, wanted: Sequence[str]) -> List[str]: + return [e for e in wanted if e not in self.files] + + +# --------------------------------------------------------------------------- +# 发现 +# --------------------------------------------------------------------------- + + +def is_staging(path: Path) -> bool: + return path.name == STAGING_DIRNAME + + +def groups(root: Union[str, Path]) -> List[str]: + """列出所有权重分组目录（跳过 staging 与隐藏目录）。""" + r = Path(root) + if not r.is_dir(): + return [] + return sorted( + p.name + for p in r.iterdir() + if p.is_dir() and not p.name.startswith(".") and not is_staging(p) + ) + + +def staging_dir(root: Union[str, Path]) -> Path: + return Path(root) / STAGING_DIRNAME + + +def staging_files(root: Union[str, Path]) -> List[Path]: + """staging 里的待处理文件。 + + 跳过隐藏文件（``.gitkeep`` 等）和 ``.md`` 说明文档 —— 那些是给人和 git 看的， + 不是待转换的模型。否则 ``doctor`` 会把 README 报成"格式不支持"。 + """ + d = staging_dir(root) + if not d.is_dir(): + return [] + return sorted( + p + for p in d.iterdir() + if p.is_file() and not p.name.startswith(".") and p.suffix.lower() != ".md" + ) + + +def discover(root: Optional[Union[str, Path]] = None) -> List[ModelEntry]: + """扫描权重根，按"分组 + 模型名"归并出所有模型。 + + 同名不同格式（m.onnx / m.engine / m.pt）会被合成一个条目 —— 因为它们 + 共用一个模型名、一份标签。``.labels.yaml`` 本身不算权重。 + """ + r = Path(root) if root is not None else default_weights_root() + out: List[ModelEntry] = [] + if not r.is_dir(): + return out + for g in groups(r): + gdir = r / g + by_name: Dict[str, ModelEntry] = {} + for p in sorted(gdir.iterdir()): + if not p.is_file() or p.name.startswith("."): + continue + ext = p.suffix.lower().lstrip(".") + if ext not in WEIGHT_EXTS: + continue + # m.labels.yaml 的后缀是 .yaml，不会落进 WEIGHT_EXTS，天然被排除。 + stem = p.name[: -(len(ext) + 1)] + entry = by_name.get(stem) + if entry is None: + entry = ModelEntry(group=g, name=stem, files={}) + by_name[stem] = entry + entry.files[ext] = p + out.extend(by_name.values()) + return out + + +def find_model( + root: Optional[Union[str, Path]] = None, + name: str = "", + group: Optional[str] = None, +) -> Optional[ModelEntry]: + """按模型名找模型。 + + - ``name`` 是完整路径 → 直接按路径构造条目（不要求它在权重目录里） + - ``name`` 是 ``分组/模型名`` → 按分组+名字精确查找 + - ``name`` 是裸名字 → 在所有分组里找同名模型；``group`` 可缩小范围 + + 裸名字跨分组同名时会优先返回 ``group`` 指定的那个；都没指定就返回找到的第一个， + 并由调用方提示存在歧义。 + """ + p = Path(name).expanduser() + if p.suffix and (p.is_file() or p.parent.is_dir() and str(name).count("/")): + ext = p.suffix.lower().lstrip(".") + if ext in WEIGHT_EXTS: + entry = ModelEntry(group=group or p.parent.name, name=p.stem, files={}) + entry.files[ext] = p + return entry + + for e in discover(root): + if e.name != name: + continue + if group and e.group != group: + continue + return e + + # 裸名字没匹配上，再试 ``分组/模型名``。 + # ``list`` 的输出里模型是带分组一起显示的，用户照着敲 ``yolo/detect_best`` + # 非常自然；以前这里会直接报"找不到模型"，而报错又只列出裸名字，很别扭。 + # 放在裸名字查找之后，是为了不把真正的相对路径（``a/b.onnx``）误判成分组名。 + if group is None and "/" in name: + head, _, tail = name.rpartition("/") + if head and tail: + return find_model(root, tail, head) + return None + + +# --------------------------------------------------------------------------- +# 标签归属 +# --------------------------------------------------------------------------- + + +def model_labels(entry: ModelEntry, *, deep: bool = False) -> LabelSet: + """给出某个模型当前生效的标签，附带来源。 + + 优先级：旁挂文件 > ONNX 内嵌 > 模型内嵌 > 占位符。 + + ``deep=False``（默认）只读旁挂文件和 ONNX metadata —— 都快。 + ``deep=True`` 才会去加载模型本身（``.pt`` / ``.engine``）， + 代价是要起 Ultralytics、加载权重，秒级，所以默认不开。 + 这样 ``weights.sh list`` 能秒出结果，不因为几个 .pt 卡十几秒。 + """ + files = entry.sorted_files() + if not files: + return LabelSet(names={}, source=SOURCE_PLACEHOLDER, path=None) + + side = label_file_status(files[0])[1] + names = read_label_file(side) if side else None + if names: + return LabelSet(names=names, source=SOURCE_SIDECAR, path=side) + + for p in files: + ext = p.suffix.lower().lstrip(".") + if ext == "onnx": + got = labels_from_onnx(p) + if got: + return LabelSet(names=got[0], source=SOURCE_ONNX, path=p) + elif deep and ext in ("pt", "pth", "engine", "trt", "plan"): + # engine 在当前的 TensorRT 下读不到（没有 metadata API）， + # 但万一用户是拿 Ultralytics 导出的，还是值得试一次。 + got = labels_from_model(p) + if got: + return LabelSet(names=got[0], source=SOURCE_MODEL, path=p) + + return LabelSet(names=placeholder_labels(0), source=SOURCE_PLACEHOLDER, path=None) + + +# --------------------------------------------------------------------------- +# 转换规划 +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ConversionChain: + """从某个输入格式到目标格式，需要依次产出哪些文件。""" + + source_ext: str + target: str + produce: Tuple[str, ...] + + @property + def is_noop(self) -> bool: + return not self.produce + + def describe(self) -> str: + if self.is_noop: + return f"{self.source_ext} → 无需转换" + return f"{self.source_ext} → " + " → ".join(self.produce) + + +def conversion_chain(source_ext: str, target: str = "all") -> ConversionChain: + """算出转换链。 + + TensorRT 的 engine 只能由 ONNX 构建，所以任何通往 engine 的路都必须 + 先得到 ONNX。这点在这里集中表达，避免调用方各写一遍。 + + - ``target="all"``：能产出多少产出多少（.pt → ONNX + engine） + - ``target="onnx"``：只要 ONNX（engine 无法反向拆回 ONNX，返回空） + - ``target="engine"``：只要 engine（.pt 会顺带产出 ONNX 作为中间产物） + """ + src = source_ext.lower().lstrip(".") + tgt = (target or "all").lower().strip() + if tgt not in ("all", "onnx", "engine"): + raise ValueError(f"未知转换目标: {target}（可选 all/onnx/engine）") + + if src not in CONVERTIBLE_INPUT_EXTS: + # 已经是最终产物，或干脆不认识 —— 都交给调用方去报错/跳过。 + return ConversionChain(src, tgt, ()) + + if src in ("pt", "pth"): + if tgt == "onnx": + return ConversionChain(src, tgt, ("onnx",)) + # engine 需要 ONNX 做中间产物，两条路都带上它。 + return ConversionChain(src, tgt, ("onnx", "engine")) + + # src == "onnx" + if tgt == "onnx": + return ConversionChain(src, tgt, ()) + return ConversionChain(src, tgt, ("engine",)) + + +def canonical_path(root: Union[str, Path], group: str, name: str, ext: str) -> Path: + """转换产物该落在哪：``models//.``。 + + 这是"自动改名并归位"的落点。无论输入文件叫什么、放在哪 + （哪怕在 /tmp 或 staging 里叫 ``best_v3_final(1).pt``）， + 产出都用调用方给的 ``name`` 和 ``group``。 + """ + return Path(root) / group / f"{name}.{ext.lstrip('.')}" + + +def validate_name(name: str) -> Optional[str]: + """检查模型名是否安全可用作文件名。返回错误信息或 None。""" + n = (name or "").strip() + if not n: + return "模型名不能为空" + if n != name: + return f"模型名首尾不能有空格: {name!r}" + if any(c.isspace() for c in n): + # 名字里有空格会让命令行、site.yaml、ROS 参数到处都要加引号， + # 早晚有人忘。直接禁掉，改用下划线或短横线。 + return "模型名不能含有空格（建议用 _ 或 - 代替）" + bad = set('/\\:*?"<>|') + hit = sorted(bad & set(n)) + if hit: + return f"模型名不能含有这些字符: {''.join(hit)}" + if n.startswith("."): + return "模型名不能以点开头" + if n.endswith(".labels"): + return "不要带 .labels 后缀，标签文件会自动生成" + return None + + +def plan_conversion( + root: Union[str, Path], + source: Union[str, Path], + *, + name: str, + group: str = "yolo", + target: str = "all", +) -> Tuple[ConversionChain, List[Path]]: + """算出一条完整的转换计划：产出链 + 每个产物的最终落点。""" + src = Path(source).expanduser() + chain = conversion_chain(src.suffix, target) + outs = [canonical_path(root, group, name, ext) for ext in chain.produce] + return chain, outs diff --git a/src/arm_application/arm_application/yolo_detector.py b/src/arm_application/arm_application/yolo_detector.py index 7197a2e..40ab1da 100644 --- a/src/arm_application/arm_application/yolo_detector.py +++ b/src/arm_application/arm_application/yolo_detector.py @@ -4,10 +4,16 @@ Public: ``resolve_weights``, ``YoloDetector``, ``Detection``, ``pick_best``, ``MpYoloDetector``. TensorRT loads the sibling ``.engine`` next to the ``.onnx`` / ``.pt`` path when ``backend='tensorrt'``. + +类别名（标签）由 ``model_labels`` 统一提供：优先读权重旁边的 +``<模型名>.labels.yaml``（用户可随时手改），其次读模型内嵌 metadata。 +这样做是因为 TensorRT engine 在当前的 TensorRT Python API 下**无法携带标签**， +只能靠旁挂文件；详见 ``model_labels`` 模块头。 """ from __future__ import annotations +import logging import multiprocessing as mp import time from dataclasses import dataclass @@ -16,6 +22,17 @@ from typing import List, Optional, Sequence, Union import numpy as np +from .model_labels import ( + SOURCE_PLACEHOLDER, + LabelSet, + describe, + parse_names_field, + placeholder_labels, + read_labels, +) + +logger = logging.getLogger(__name__) + # Sentinel: worker still on the previous frame — caller must NOT count a miss. INFER_BUSY = object() @@ -62,7 +79,7 @@ def resolve_weights(weights: Union[str, Path], backend: str = "auto") -> Path: return engine raise FileNotFoundError( f"TensorRT engine missing: {engine}. " - f"Export: bash ~/ros2/scripts/export_yolo_trt.sh" + f"Export: bash ~/ros2/scripts/weights.sh convert {p} --name {p.stem}" ) if b == "onnx": if onnx.is_file(): @@ -78,6 +95,50 @@ def resolve_weights(weights: Union[str, Path], backend: str = "auto") -> Path: return p +def resolve_class_labels(path: Union[str, Path], model=None) -> LabelSet: + """给出一个权重该用的类别标签（旁挂文件优先）。 + + ``model`` 是已经加载好的 Ultralytics 模型对象（可为 None）。顺序： + + 1. ``<模型名>.labels.yaml`` —— 用户可编辑的权威来源 + 2. 传入的 ``model.names`` —— 模型内嵌 metadata + 3. 占位符 ``class_0 / class_1 …`` + + 第 3 步刻意不用 ``str(cid)``。原先读不到名字时会退化成类名 `"0"`， + 和"真的有个叫 0 的类"无法区分，标签丢了都看不出来。 + 占位符 + 一条 warning 让这件事在日志里可见。 + """ + p = Path(path).expanduser() + + # 先只看旁挂文件（快，不加载模型）：它必须能覆盖模型内嵌的旧值， + # 否则"我改了标签为什么不生效"会变成很难查的问题。 + labels = read_labels(p, allow_model_metadata=False) + if not labels.is_placeholder: + return labels + + raw = getattr(model, "names", None) if model is not None else None + names = parse_names_field(raw or {}) + if names: + from .model_labels import SOURCE_MODEL + + return LabelSet(names=names, source=SOURCE_MODEL, path=p) + + nc = 0 + try: + nc = int(getattr(model, "nc", 0) or 0) + except (TypeError, ValueError): + nc = 0 + labels = LabelSet(names=placeholder_labels(nc), source=SOURCE_PLACEHOLDER, path=None) + logger.warning( + "模型 %s 没有类别标签，暂时用占位名显示。补一个 %s 即可（改完重跑节点生效）。" + "命令行: bash ~/ros2/scripts/weights.sh labels %s --from-model", + p.name, + p.stem + ".labels.yaml", + p.stem, + ) + return labels + + class YoloDetector: def __init__( self, @@ -132,8 +193,8 @@ class YoloDetector: if class_names else None ) - names = getattr(self._model, "names", None) or {} - self._id_to_name = {int(k): str(v) for k, v in dict(names).items()} + self._labels = resolve_class_labels(path, self._model) + self._id_to_name = dict(self._labels.names) self._warmed = False self._pending = None @@ -209,7 +270,8 @@ class YoloDetector: clss = boxes.cls.detach().cpu().numpy().astype(int) for i in range(len(xyxy)): cid = int(clss[i]) - name = self._id_to_name.get(cid, str(cid)) + # 用 LabelSet.get：读不到名字时给 class_N 占位，而不是退化成裸数字 "0"。 + name = self._labels.get(cid) if self.class_ids is not None and cid not in self.class_ids: continue if self.class_names is not None and name.lower() not in self.class_names: diff --git a/src/arm_application/config/perception.yaml b/src/arm_application/config/perception.yaml index dcf4b5f..7b9ab6d 100644 --- a/src/arm_application/config/perception.yaml +++ b/src/arm_application/config/perception.yaml @@ -1,5 +1,9 @@ -# perception_node defaults. Launch overlays ~/ros2/config/site.yaml → yolo: / paths: -# Do not edit this file for day-to-day tuning. +# perception_node 的默认值。launch 会用 ~/ros2/config/site.yaml 的 yolo: / paths: 覆盖。 +# 日常调参**不要改这里**，改 site.yaml。 +# +# 注意：下面 weights / conf / iou / imgsz / device / half_precision / min_conf 这几项 +# 一定被 site.yaml 覆盖，这里只是"没有 site.yaml 时的兜底"。所以它们必须和 site.yaml +# 保持一致 —— 曾经这里写死 imgsz=416 而实际模型是 640，一旦覆盖失效就会用错尺寸。 perception_node: ros__parameters: @@ -8,13 +12,13 @@ perception_node: image_topic: "" camera_info_topic: "" depth_topic: "/d405/aligned_depth_to_color/image_raw" - weights: "/home/cc/ros2/models/yolo/detect_yolo26_gongjian.pt" - conf: 0.35 + weights: "/home/cc/ros2/models/yolo/detect_best_0.onnx" + conf: 0.10 iou: 0.45 - imgsz: 416 + imgsz: 640 device: "0" half_precision: true - min_conf: 0.35 + min_conf: 0.12 vision_hz: 20.0 z_ref_m: 0.30 # Large gate + reacquire: fast motion must NOT freeze track diff --git a/src/arm_application/launch/hole_detect.launch.py b/src/arm_application/launch/hole_detect.launch.py index e74f8c7..9922f2a 100644 --- a/src/arm_application/launch/hole_detect.launch.py +++ b/src/arm_application/launch/hole_detect.launch.py @@ -36,7 +36,7 @@ def _setup(context, *args, **kwargs): if backend in ("two_stage", "hough_roi"): weights = os.environ.get("SITE_HOLE_WORKPIECE_WEIGHTS") or os.environ.get("SITE_YOLO_WEIGHTS") or "" else: - weights = os.environ.get("SITE_HOLE_WEIGHTS") or str(root / "models" / "hole_insert" / "best.pt") + weights = os.environ.get("SITE_HOLE_WEIGHTS") or str(root / "models" / "hole_insert" / "detect_best_1.onnx") return [ LogInfo(msg=f"[hole_detect] cam={cam} backend={backend} weights={weights} cfg={cfg}"), Node( diff --git a/src/arm_application/setup.py b/src/arm_application/setup.py index 8359cd4..8b2dfc6 100644 --- a/src/arm_application/setup.py +++ b/src/arm_application/setup.py @@ -5,18 +5,34 @@ from glob import glob package_name = 'arm_application' +def _existing(paths): + """只保留真实存在的常规文件。 + + 必需：colcon 的 ``build//`` 里放着指向源文件的符号链接， + 源文件被挪走或删掉之后这些链接就断了。``glob`` / ``os.walk`` 仍会把它们 + 列出来，随后 copy 会报 "doesn't exist or not a regular file"， + **整个包构建失败**，而报错信息指向 build/ 下的一个幽灵文件，很难查。 + 这里统一过滤掉。 + """ + out = [] + for p in paths: + if os.path.isfile(p): # 会跟随符号链接，断链返回 False + out.append(p) + return out + + def _data_files(): entries = [ ('share/ament_index/resource_index/packages', ['resource/' + package_name]), ('share/' + package_name, ['package.xml']), - (os.path.join('share', package_name, 'launch'), glob('launch/*.py')), - (os.path.join('share', package_name, 'trajectories'), glob('trajectories/*')), + (os.path.join('share', package_name, 'launch'), _existing(glob('launch/*.py'))), + (os.path.join('share', package_name, 'trajectories'), _existing(glob('trajectories/*'))), ] for root, _dirs, files in os.walk('config'): - if not files: + srcs = _existing([os.path.join(root, name) for name in files]) + if not srcs: continue dest = os.path.join('share', package_name, root) - srcs = [os.path.join(root, name) for name in files] entries.append((dest, srcs)) return entries diff --git a/src/arm_application/test/test_model_labels.py b/src/arm_application/test/test_model_labels.py new file mode 100644 index 0000000..73bd8ae --- /dev/null +++ b/src/arm_application/test/test_model_labels.py @@ -0,0 +1,479 @@ +"""model_labels 的测试。 + +刻意**不依赖仓库里那些模型二进制**（`models/**/*.onnx` 都在 .gitignore 里， +新克隆的仓库根本没有），全部用 tmp_path 自建 fixture。ONNX 用 onnx.helper +现场造一个最小的，这样"从 ONNX metadata 读标签"这条路径也能被测到。 +""" + +from __future__ import annotations + +import numpy as np +import pytest +import yaml + +from arm_application.model_labels import ( + SOURCE_MODEL, + SOURCE_ONNX, + SOURCE_PLACEHOLDER, + SOURCE_SIDECAR, + LabelSet, + describe, + format_labels_yaml, + label_file_model, + label_file_status, + labels_path_for, + parse_label_list, + parse_names_field, + placeholder_labels, + read_label_file, + read_label_text_file, + read_labels, + write_labels, +) + + +# -------------------------------------------------------------------------- +# fixture 工具 +# -------------------------------------------------------------------------- + + +def _make_onnx(path, names_json=None, nc=None, metadata=None): + """造一个最小可读的 ONNX，用于测试 metadata 读取。""" + import onnx + from onnx import TensorProto, helper + + inp = helper.make_tensor_value_info("images", TensorProto.FLOAT, [1, 3, 8, 8]) + out = helper.make_tensor_value_info("output0", TensorProto.FLOAT, [1, 3, 8, 8]) + node = helper.make_node("Identity", ["images"], ["output0"]) + graph = helper.make_graph([node], "g", [inp], [out]) + model = helper.make_model(graph) + model.opset_import[0].version = 13 + if names_json is not None: + model.metadata_props.add(key="names", value=names_json) + if nc is not None: + model.metadata_props.add(key="nc", value=str(nc)) + for k, v in (metadata or {}).items(): + model.metadata_props.add(key=k, value=v) + onnx.save(model, str(path)) + return path + + +def _touch(path): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"not a real model") + return path + + +# -------------------------------------------------------------------------- +# 旁挂路径推导 +# -------------------------------------------------------------------------- + + +def test_labels_path_replaces_suffix_only(): + assert labels_path_for("models/yolo/m.engine") == pytest.approx( + labels_path_for("models/yolo/m.engine") + ) + from pathlib import Path + + assert Path(labels_path_for("models/yolo/m.engine")) == Path( + "models/yolo/m.labels.yaml" + ) + assert Path(labels_path_for("models/yolo/m.onnx")) == Path("models/yolo/m.labels.yaml") + assert Path(labels_path_for("a/b/m.pt")) == Path("a/b/m.labels.yaml") + + +def test_labels_path_same_for_all_formats_of_one_model(): + """关键契约：同一模型的三种格式必须算出同一个标签文件。 + + 否则"换后端名字就变了"这个坑会再次出现。 + """ + from pathlib import Path + + a = Path(labels_path_for("d/m.pt")) + b = Path(labels_path_for("d/m.onnx")) + c = Path(labels_path_for("d/m.engine")) + assert a == b == c + + +def test_labels_path_handles_dotted_stem(): + """含点的文件名（a.b.onnx）不能被截成 a.labels.yaml。 + + 这正是 Path.with_suffix('.labels.yaml') 容易踩的坑，所以手工拼 stem。 + """ + from pathlib import Path + + assert Path(labels_path_for("a.b.onnx")) == Path("a.b.labels.yaml") + + +def test_labels_path_without_extension(): + from pathlib import Path + + assert Path(labels_path_for("models/bare")) == Path("models/bare.labels.yaml") + + +def test_labels_path_expands_user(): + assert "~" not in str(labels_path_for("~/ros2/models/x.onnx")) + + +# -------------------------------------------------------------------------- +# names 字段解析：三种写法都要吃 +# -------------------------------------------------------------------------- + + +def test_parse_names_dict_int_keys(): + assert parse_names_field({0: "workpiece", 1: "hole"}) == {0: "workpiece", 1: "hole"} + + +def test_parse_names_dict_string_keys(): + """YAML 里键常常是字符串形式的数字（'0'），必须吃。""" + assert parse_names_field({"0": "workpiece", "1": "hole"}) == { + 0: "workpiece", + 1: "hole", + } + + +def test_parse_names_list_index_is_position(): + assert parse_names_field(["workpiece", "hole"]) == {0: "workpiece", 1: "hole"} + + +def test_parse_names_list_skips_blank_but_keeps_index(): + """空项被跳过，但后面的项索引不能因此前移 —— 否则类别就整体错位了。""" + assert parse_names_field(["a", "", "c"]) == {0: "a", 2: "c"} + assert parse_names_field(["a", None, "c"]) == {0: "a", 2: "c"} + assert parse_names_field(["a", " ", "c"]) == {0: "a", 2: "c"} + + +def test_parse_names_string(): + assert parse_names_field("workpiece, hole") == {0: "workpiece", 1: "hole"} + + +def test_parse_names_none_and_garbage(): + assert parse_names_field(None) == {} + assert parse_names_field(42) == {} + assert parse_names_field({}) == {} + assert parse_names_field([]) == {} + + +def test_parse_names_ignores_unparsable_dict_keys(): + """键不是整数时跳过这一项，而不是抛异常让推理进程崩掉。""" + assert parse_names_field({"x": "bad", 0: "good"}) == {0: "good"} + + +def test_parse_names_strips_whitespace(): + assert parse_names_field([" workpiece "]) == {0: "workpiece"} + + +# -------------------------------------------------------------------------- +# 逗号串解析（CLI --set 用） +# -------------------------------------------------------------------------- + + +def test_parse_label_list_ascii_comma(): + assert parse_label_list("a,b,c") == {0: "a", 1: "b", 2: "c"} + + +def test_parse_label_list_chinese_comma_and_semicolon(): + assert parse_label_list("工件，通止规") == {0: "工件", 1: "通止规"} + assert parse_label_list("a;b") == {0: "a", 1: "b"} + + +def test_parse_label_list_with_explicit_index(): + """从别处拷配置常带索引，要能认。""" + assert parse_label_list("0:workpiece, 5:hole") == {0: "workpiece", 5: "hole"} + assert parse_label_list("0：workpiece") == {0: "workpiece"} # 全角冒号 + + +def test_parse_label_list_tolerates_sloppy_input(): + assert parse_label_list("") == {} + assert parse_label_list(" ") == {} + assert parse_label_list("a,,b") == {0: "a", 1: "b"} + assert parse_label_list(" a , b ") == {0: "a", 1: "b"} + assert parse_label_list("a\nb") == {0: "a", 1: "b"} + + +def test_parse_label_list_mixed_forms_stay_ordered(): + out = parse_label_list("a, 1:b, c") + assert out[0] == "a" + assert out[1] == "b" + assert out[2] == "c" + + +# -------------------------------------------------------------------------- +# 纯文本标签文件（classes.txt 风格） +# -------------------------------------------------------------------------- + + +def test_read_label_text_file_comments_do_not_consume_index(tmp_path): + p = tmp_path / "classes.txt" + p.write_text("# 这是注释\nworkpiece\n\n# 又一条注释\nhole\n", encoding="utf-8") + assert read_label_text_file(p) == {0: "workpiece", 1: "hole"} + + +def test_read_label_text_file_chinese(tmp_path): + p = tmp_path / "classes.txt" + p.write_text("工件\n通止规\n", encoding="utf-8") + assert read_label_text_file(p) == {0: "工件", 1: "通止规"} + + +# -------------------------------------------------------------------------- +# 占位符 +# -------------------------------------------------------------------------- + + +def test_placeholder_is_labelled_not_bare_digit(): + """占位名必须一眼看出是占位，不能退化成裸数字 '0'。 + + 原先的代码是 _id_to_name.get(cid, str(cid))，也就是类名变成 "0"， + 和"真的有个叫 0 的类"无法区分 —— 标签丢了都看不出来。 + """ + assert placeholder_labels(3) == {0: "class_0", 1: "class_1", 2: "class_2"} + assert placeholder_labels(1)[0] == "class_0" + assert placeholder_labels(0) == {} + assert placeholder_labels(-5) == {} + + +# -------------------------------------------------------------------------- +# 旁挂文件：写 → 读 往返 +# -------------------------------------------------------------------------- + + +def test_format_labels_yaml_is_parseable_and_ordered(): + txt = format_labels_yaml({0: "workpiece", 1: "hole"}, model="m") + data = yaml.safe_load(txt) + assert data["names"] == ["workpiece", "hole"] + assert data["model"] == "m" + + +def test_format_labels_yaml_includes_usage_hint(): + """注释头是给直接打开文件的人看的，必须存在。""" + txt = format_labels_yaml({0: "a"}) + assert "第 1 行对应类别 0" in txt + assert "scripts/weights.sh labels" in txt + + +def test_format_labels_yaml_empty(): + data = yaml.safe_load(format_labels_yaml({})) + assert data["names"] == [] + + +def test_write_then_read_round_trip(tmp_path): + w = _touch(tmp_path / "detect_best_0.engine") + written = write_labels(w, {0: "workpiece", 1: "tongzhigui"}) + assert written.name == "detect_best_0.labels.yaml" + assert read_label_file(written) == {0: "workpiece", 1: "tongzhigui"} + + +def test_write_labels_records_model_name(tmp_path): + w = _touch(tmp_path / "abc.onnx") + p = write_labels(w, {0: "x"}) + assert label_file_model(p) == "abc" + + +def test_read_label_file_tolerates_dict_form(tmp_path): + """用户可能按 Ultralytics data.yaml 的字典写法手写。""" + p = tmp_path / "m.labels.yaml" + p.write_text("names:\n 0: a\n 1: b\n", encoding="utf-8") + assert read_label_file(p) == {0: "a", 1: "b"} + + +def test_read_label_file_tolerates_comma_string(tmp_path): + p = tmp_path / "m.labels.yaml" + p.write_text('names: "a, b"\n', encoding="utf-8") + assert read_label_file(p) == {0: "a", 1: "b"} + + +def test_read_label_file_tolerates_bare_text_file(tmp_path): + """有人直接把手写的 classes.txt 改个名字就用了，也该能读。""" + p = tmp_path / "m.labels.yaml" + p.write_text("a\nb\n", encoding="utf-8") + assert read_label_file(p) == {0: "a", 1: "b"} + + +def test_read_label_file_missing_returns_none(tmp_path): + assert read_label_file(tmp_path / "nope.labels.yaml") is None + + +def test_read_label_file_malformed_yaml_returns_none(tmp_path): + """文件坏了不该让推理进程崩 —— 退回"没有标签"即可。""" + p = tmp_path / "m.labels.yaml" + p.write_text("names: [a, b\n bad: : :\n", encoding="utf-8") + assert read_label_file(p) is None + + +def test_read_label_file_empty_names_returns_none(tmp_path): + p = tmp_path / "m.labels.yaml" + p.write_text("names: []\n", encoding="utf-8") + assert read_label_file(p) is None + + +# -------------------------------------------------------------------------- +# 从 ONNX metadata 读 +# -------------------------------------------------------------------------- + + +def test_labels_from_onnx_via_read_labels(tmp_path): + w = _make_onnx(tmp_path / "m.onnx", names_json="{0: 'workpiece'}", nc=1) + ls = read_labels(w) + assert ls.names == {0: "workpiece"} + assert ls.source == SOURCE_ONNX + assert not ls.is_placeholder + + +def test_onnx_multiclass_names(tmp_path): + w = _make_onnx(tmp_path / "m.onnx", names_json="{0: 'a', 1: 'b', 2: 'c'}") + assert read_labels(w).names == {0: "a", 1: "b", 2: "c"} + + +def test_onnx_without_names_gives_placeholder(tmp_path): + w = _make_onnx(tmp_path / "m.onnx") + ls = read_labels(w) + assert ls.is_placeholder + assert ls.path is None + + +# -------------------------------------------------------------------------- +# 优先级：sidecar 必须赢 +# -------------------------------------------------------------------------- + + +def test_sidecar_beats_onnx_metadata(tmp_path): + """这是本模块最要紧的行为。 + + ONNX 里内嵌的是导出时的旧名字。用户改了旁挂文件却看到旧名字， + 会变成极难排查的问题（"我明明改了为什么不生效"）。 + """ + w = _make_onnx(tmp_path / "m.onnx", names_json="{0: 'OLD_NAME'}") + assert read_labels(w).names == {0: "OLD_NAME"} + + write_labels(w, {0: "NEW_NAME"}) + ls = read_labels(w) + assert ls.names == {0: "NEW_NAME"} + + +def test_sidecar_beats_placeholder_for_engine(tmp_path): + """engine 读不到任何 metadata，旁挂文件是它唯一的标签来源。""" + w = _touch(tmp_path / "m.engine") + assert read_labels(w).is_placeholder + + write_labels(w, {0: "workpiece"}) + ls = read_labels(w) + assert ls.names == {0: "workpiece"} + assert ls.source == SOURCE_SIDECAR + + +def test_sidecar_shared_across_formats_in_same_dir(tmp_path): + """.onnx 和 .engine 同名同目录 → 共用一份标签。""" + onnx = _make_onnx(tmp_path / "m.onnx", names_json="{0: 'x'}") + engine = _touch(tmp_path / "m.engine") + write_labels(engine, {0: "shared", 1: "two"}) + + assert read_labels(onnx).names == {0: "shared", 1: "two"} + assert read_labels(engine).names == {0: "shared", 1: "two"} + assert read_labels(onnx).source == SOURCE_SIDECAR + + +def test_engine_without_sidecar_does_not_guess(tmp_path): + """engine 缺标签时必须显式报占位符，不能假装知道。""" + w = _touch(tmp_path / "m.engine") + ls = read_labels(w) + assert ls.is_placeholder + assert ls.source == SOURCE_PLACEHOLDER + + +def test_read_labels_tolerates_missing_weights_file(tmp_path): + """权重不存在也不该崩（由调用方负责检查存在性）。""" + ls = read_labels(tmp_path / "ghost.engine") + assert isinstance(ls, LabelSet) + assert ls.is_placeholder + + +def test_existing_sidecar_wins_even_if_weights_absent(tmp_path): + """标签文件独立于权重存在，有权重没权重都能读出来。""" + write_labels(tmp_path / "m.engine", {0: "only_labels"}) + ls = read_labels(tmp_path / "m.engine") + assert ls.names == {0: "only_labels"} + assert ls.source == SOURCE_SIDECAR + + +# -------------------------------------------------------------------------- +# LabelSet 行为 +# -------------------------------------------------------------------------- + + +def test_labelset_get_falls_back_to_placeholder_name(): + ls = LabelSet(names={0: "a"}, source=SOURCE_SIDECAR) + assert ls.get(0) == "a" + assert ls.get(7) == "class_7" + assert ls.get(7, "unknown") == "unknown" + + +def test_labelset_sorted_names(): + ls = LabelSet(names={2: "c", 0: "a", 1: "b"}, source=SOURCE_SIDECAR) + assert ls.sorted_names() == ["a", "b", "c"] + + +def test_describe_mentions_source_in_chinese(): + """来源必须出现在日志里 —— 这是"标签丢了"能被发现的关键。""" + ls = LabelSet(names={0: "workpiece"}, source=SOURCE_SIDECAR) + assert "workpiece" in describe(ls) + assert "旁挂标签文件" in describe(ls) + + ls = LabelSet(names={0: "class_0"}, source=SOURCE_PLACEHOLDER) + assert "占位符" in describe(ls) + + +def test_describe_handles_empty(): + assert "(空)" in describe(LabelSet(names={}, source=SOURCE_PLACEHOLDER)) + + +# -------------------------------------------------------------------------- +# 错配检测（doctor 用） +# -------------------------------------------------------------------------- + + +def test_label_file_status_reports_missing(tmp_path): + exists, path, hint = label_file_status(tmp_path / "m.engine") + assert exists is False + assert hint is not None and "缺标签文件" in hint + + +def test_label_file_status_ok(tmp_path): + w = _touch(tmp_path / "m.engine") + write_labels(w, {0: "a"}) + exists, _path, hint = label_file_status(w) + assert exists is True + assert hint is None + + +def test_label_file_status_detects_mismatched_copy(tmp_path): + """把 A 的标签文件拷给 B —— 这种错配必须被指出来。""" + write_labels(tmp_path / "A.onnx", {0: "a"}) + (tmp_path / "A.labels.yaml").rename(tmp_path / "B.labels.yaml") + b = _touch(tmp_path / "B.onnx") + + exists, _path, hint = label_file_status(b) + assert exists is True + assert hint is not None + assert "A" in hint and "B" in hint + + +# -------------------------------------------------------------------------- +# 中文标签（实际会用到的场景） +# -------------------------------------------------------------------------- + + +def test_chinese_labels_round_trip(tmp_path): + w = _touch(tmp_path / "m.onnx") + write_labels(w, {0: "工件", 1: "通止规"}) + ls = read_labels(w) + assert ls.names == {0: "工件", 1: "通止规"} + # 写出的文件必须是 UTF-8，否则中文会乱码 + txt = labels_path_for(w).read_text(encoding="utf-8") + assert "通止规" in txt + + +def test_numpy_int_keys_are_accepted(tmp_path): + """调用方从 numpy 结果里拿到的索引可能是 np.int64，不能因此失败。""" + w = _touch(tmp_path / "m.onnx") + write_labels(w, {np.int64(0): "a", np.int64(1): "b"}) + assert read_labels(w).names == {0: "a", 1: "b"} diff --git a/src/arm_application/test/test_weights_layout.py b/src/arm_application/test/test_weights_layout.py new file mode 100644 index 0000000..e575ba2 --- /dev/null +++ b/src/arm_application/test/test_weights_layout.py @@ -0,0 +1,754 @@ +"""weights_layout / weights.sh CLI 的测试。 + +全部用 tmp_path 自建目录，不依赖仓库里真实的模型二进制 +（`models/**/*.onnx` 都在 .gitignore，新克隆的仓库没有）。 + +需要 GPU 的 ONNX→engine 构建**不在这里测** —— 那条链路由一次性的 +端到端验证覆盖（见 models/README.md）。 +""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +from arm_application.model_labels import ( + SOURCE_ONNX, + SOURCE_PLACEHOLDER, + SOURCE_SIDECAR, + write_labels, +) +from arm_application.weights_layout import ( + STAGING_DIRNAME, + ModelEntry, + canonical_path, + conversion_chain, + default_weights_root, + discover, + find_model, + find_repo_root, + groups, + model_labels, + plan_conversion, + staging_files, + validate_name, +) + +REPO = Path(__file__).resolve().parents[3] + + +# -------------------------------------------------------------------------- +# 工具 +# -------------------------------------------------------------------------- + + +def _make_onnx(path: Path, names_json: str | None = None) -> Path: + import onnx + from onnx import TensorProto, helper + + inp = helper.make_tensor_value_info("images", TensorProto.FLOAT, [1, 3, 8, 8]) + out = helper.make_tensor_value_info("output0", TensorProto.FLOAT, [1, 3, 8, 8]) + node = helper.make_node("Identity", ["images"], ["output0"]) + graph = helper.make_graph([node], "g", [inp], [out]) + model = helper.make_model(graph) + model.opset_import[0].version = 13 + if names_json is not None: + model.metadata_props.add(key="names", value=names_json) + path.parent.mkdir(parents=True, exist_ok=True) + onnx.save(model, str(path)) + return path + + +def _touch(path: Path, data: bytes = b"x") -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + return path + + +def _load_cli(): + """把 scripts/weights.py 当普通模块加载（它自己会修好 sys.path）。""" + spec = importlib.util.spec_from_file_location("weights_cli", REPO / "scripts" / "weights.py") + assert spec and spec.loader + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +# -------------------------------------------------------------------------- +# 转换链 +# -------------------------------------------------------------------------- + + +def test_chain_pt_to_all_goes_through_onnx(): + """engine 只能由 ONNX 构建，所以 pt→engine 必须先把 onnx 造出来。""" + c = conversion_chain(".pt", "all") + assert c.produce == ("onnx", "engine") + assert not c.is_noop + + +def test_chain_pt_to_onnx_only(): + assert conversion_chain(".pt", "onnx").produce == ("onnx",) + + +def test_chain_pt_to_engine_still_produces_onnx_intermediate(): + assert conversion_chain(".pt", "engine").produce == ("onnx", "engine") + + +def test_chain_onnx_to_engine(): + assert conversion_chain(".onnx", "all").produce == ("engine",) + assert conversion_chain(".onnx", "engine").produce == ("engine",) + + +def test_chain_onnx_to_onnx_is_noop(): + assert conversion_chain(".onnx", "onnx").is_noop + + +def test_chain_engine_is_noop(): + """engine 无法反向拆回 ONNX，任何目标都是 noop。""" + for tgt in ("all", "onnx", "engine"): + assert conversion_chain(".engine", tgt).is_noop + + +def test_chain_ext_case_and_dot_insensitive(): + assert conversion_chain("PT", "all").produce == ("onnx", "engine") + assert conversion_chain(".PT", "all").produce == ("onnx", "engine") + + +def test_chain_rejects_unknown_target(): + with pytest.raises(ValueError): + conversion_chain(".pt", "tflite") + + +def test_chain_describe_is_human_readable(): + assert "onnx" in conversion_chain(".pt", "all").describe() + assert "无需转换" in conversion_chain(".engine", "all").describe() + + +def test_chain_pth_treated_like_pt(): + assert conversion_chain(".pth", "all").produce == ("onnx", "engine") + + +# -------------------------------------------------------------------------- +# 命名与归位 +# -------------------------------------------------------------------------- + + +def test_canonical_path_shape(): + p = canonical_path("/m", "yolo", "workpiece", ".engine") + assert str(p) == "/m/yolo/workpiece.engine" + # 带不带点都一样 + assert canonical_path("/m", "yolo", "workpiece", "engine") == p + + +def test_canonical_path_is_what_makes_rename_automatic(): + """这是"自动改名并归位"的落点：输入名再乱，产出都用调用方给的 name。""" + p = canonical_path("/m", "yolo", "workpiece", "onnx") + assert p.name == "workpiece.onnx" + + +@pytest.mark.parametrize( + "bad", + ["", " ", "a b", "a/b", "a\\b", "a:b", "a*b", "a?b", "ab", "a|b", ".hidden", "x.labels"], +) +def test_validate_name_rejects_unsafe(bad): + assert validate_name(bad) is not None + + +@pytest.mark.parametrize("good", ["workpiece", "detect_best_0", "工件", "a-b_c.1", "m2"]) +def test_validate_name_accepts_reasonable(good): + assert validate_name(good) is None + + +def test_validate_name_reports_specific_reason(): + assert "不能为空" in validate_name("") + assert "/" in validate_name("a/b") + assert "点开头" in validate_name(".x") + + +# -------------------------------------------------------------------------- +# 发现 / 归组 +# -------------------------------------------------------------------------- + + +def test_discover_merges_formats_of_same_stem(tmp_path): + """同名不同格式必须合成一个模型 —— 它们共用一份标签。""" + _touch(tmp_path / "yolo" / "m.onnx") + _touch(tmp_path / "yolo" / "m.engine") + _touch(tmp_path / "yolo" / "m.pt") + entries = discover(tmp_path) + assert len(entries) == 1 + assert entries[0].name == "m" + assert set(entries[0].exts) == {"onnx", "engine", "pt"} + + +def test_discover_does_not_treat_labels_file_as_weight(tmp_path): + """m.labels.yaml 的后缀是 .yaml，不能变成第四个格式。""" + _touch(tmp_path / "yolo" / "m.onnx") + (tmp_path / "yolo" / "m.labels.yaml").write_text("names:\n - a\n", encoding="utf-8") + e = discover(tmp_path)[0] + assert e.exts == ["onnx"] + + +def test_discover_skips_staging_and_hidden(tmp_path): + _touch(tmp_path / "yolo" / "m.onnx") + _touch(tmp_path / STAGING_DIRNAME / "raw.pt") + _touch(tmp_path / ".hidden" / "x.onnx") + names = {e.name for e in discover(tmp_path)} + assert names == {"m"} + + +def test_discover_skips_gitkeep_and_hidden_files(tmp_path): + _touch(tmp_path / "yolo" / ".gitkeep", b"") + _touch(tmp_path / "yolo" / "m.onnx") + assert {e.name for e in discover(tmp_path)} == {"m"} + + +def test_discover_empty_and_missing_root(tmp_path): + assert discover(tmp_path) == [] + assert discover(tmp_path / "nope") == [] + + +def test_discover_handles_dotted_model_name(tmp_path): + """名字里带点（m.v2.onnx）不能被截断成 m。""" + _touch(tmp_path / "yolo" / "m.v2.onnx") + assert discover(tmp_path)[0].name == "m.v2" + + +def test_discover_ext_order_follows_pipeline(tmp_path): + """显示顺序 = 转换流水线顺序 pt → onnx → engine，一眼看出走到哪步。""" + _touch(tmp_path / "yolo" / "m.pt") + _touch(tmp_path / "yolo" / "m.onnx") + _touch(tmp_path / "yolo" / "m.engine") + assert discover(tmp_path)[0].exts == ["pt", "onnx", "engine"] + + +def test_groups_lists_subdirs_excluding_staging(tmp_path): + _touch(tmp_path / "yolo" / "a.onnx") + _touch(tmp_path / "hole_insert" / "b.onnx") + _touch(tmp_path / STAGING_DIRNAME / "x.pt") + assert groups(tmp_path) == ["hole_insert", "yolo"] + + +def test_missing_of_reports_absent_formats(tmp_path): + _touch(tmp_path / "yolo" / "m.onnx") + e = discover(tmp_path)[0] + assert e.missing_of(["onnx", "engine"]) == ["engine"] + + +# -------------------------------------------------------------------------- +# 查找 +# -------------------------------------------------------------------------- + + +def test_find_model_by_bare_name(tmp_path): + _touch(tmp_path / "yolo" / "m.onnx") + e = find_model(tmp_path, "m") + assert e is not None and e.name == "m" and e.group == "yolo" + + +def test_find_model_group_filter(tmp_path): + _touch(tmp_path / "yolo" / "m.onnx") + _touch(tmp_path / "other" / "m.onnx") + assert find_model(tmp_path, "m", "other").group == "other" + + +def test_find_model_missing_returns_none(tmp_path): + _touch(tmp_path / "yolo" / "m.onnx") + assert find_model(tmp_path, "nope") is None + + +def test_find_model_accepts_group_slash_name(tmp_path): + """``list`` 里模型是带分组显示的，用户会照着敲 "yolo/m"。""" + _touch(tmp_path / "yolo" / "m.onnx") + e = find_model(tmp_path, "yolo/m") + assert e is not None and e.name == "m" and e.group == "yolo" + + +def test_find_model_group_slash_name_picks_right_group(tmp_path): + """分组/名字写法要能区分跨分组同名模型。""" + _touch(tmp_path / "yolo" / "m.onnx") + _touch(tmp_path / "hole_insert" / "m.onnx") + assert find_model(tmp_path, "hole_insert/m").group == "hole_insert" + assert find_model(tmp_path, "yolo/m").group == "yolo" + + +def test_find_model_group_slash_name_wrong_group_returns_none(tmp_path): + _touch(tmp_path / "yolo" / "m.onnx") + assert find_model(tmp_path, "nosuchgroup/m") is None + + +def test_find_model_bare_name_still_wins_over_slash_form(tmp_path): + """裸名字查找优先，别把相对路径误判成分组名。""" + _touch(tmp_path / "yolo" / "m.onnx") + assert find_model(tmp_path, "m").group == "yolo" + + +def test_find_model_accepts_direct_path(tmp_path): + """用户直接给文件路径时也要能用。""" + w = _make_onnx(tmp_path / "elsewhere" / "custom.onnx", "{0: 'z'}") + e = find_model(tmp_path, str(w)) + assert e is not None and e.name == "custom" + + +# -------------------------------------------------------------------------- +# staging +# -------------------------------------------------------------------------- + + +def test_staging_files_skips_gitkeep(tmp_path): + (tmp_path / STAGING_DIRNAME).mkdir(parents=True) + _touch(tmp_path / STAGING_DIRNAME / ".gitkeep", b"") + _touch(tmp_path / STAGING_DIRNAME / "raw.pt") + _touch(tmp_path / STAGING_DIRNAME / "other.onnx") + assert {p.name for p in staging_files(tmp_path)} == {"raw.pt", "other.onnx"} + + +def test_staging_files_skips_readme(tmp_path): + """staging 里的说明文档不是待转换文件，不能报成"格式不支持"。""" + (tmp_path / STAGING_DIRNAME).mkdir(parents=True) + (tmp_path / STAGING_DIRNAME / "README.md").write_text("hi", encoding="utf-8") + _touch(tmp_path / STAGING_DIRNAME / "raw.pt") + assert {p.name for p in staging_files(tmp_path)} == {"raw.pt"} + + +def test_staging_files_missing_dir(tmp_path): + assert staging_files(tmp_path) == [] + + +# -------------------------------------------------------------------------- +# 标签归属（weights_layout.model_labels） +# -------------------------------------------------------------------------- + + +def test_model_labels_sidecar_wins_over_onnx(tmp_path): + _make_onnx(tmp_path / "yolo" / "m.onnx", "{0: 'OLD'}") + e = discover(tmp_path)[0] + assert model_labels(e).names == {0: "OLD"} + + write_labels(e.path("onnx"), {0: "NEW"}) + ls = model_labels(e) + assert ls.names == {0: "NEW"} + assert ls.source == SOURCE_SIDECAR + + +def test_model_labels_reads_onnx_metadata(tmp_path): + _make_onnx(tmp_path / "yolo" / "m.onnx", "{0: 'workpiece', 1: 'hole'}") + ls = model_labels(discover(tmp_path)[0]) + assert ls.names == {0: "workpiece", 1: "hole"} + assert ls.source == SOURCE_ONNX + + +def test_model_labels_engine_without_sidecar_is_placeholder(tmp_path): + """engine 丢了标签时必须显式报占位符，不能假装知道。""" + _touch(tmp_path / "yolo" / "m.engine") + ls = model_labels(discover(tmp_path)[0]) + assert ls.is_placeholder + assert ls.source == SOURCE_PLACEHOLDER + + +def test_model_labels_pt_needs_deep(tmp_path): + """默认不加载 .pt（慢），所以只给占位符；--deep 才会去读。""" + _touch(tmp_path / "yolo" / "m.pt") + e = discover(tmp_path)[0] + assert model_labels(e, deep=False).is_placeholder + + +def test_model_labels_sidecar_shared_across_formats(tmp_path): + _make_onnx(tmp_path / "yolo" / "m.onnx", "{0: 'x'}") + _touch(tmp_path / "yolo" / "m.engine") + e = discover(tmp_path)[0] + write_labels(e.path("engine"), {0: "shared"}) + ls = model_labels(e) + assert ls.names == {0: "shared"} + assert ls.source == SOURCE_SIDECAR + + +def test_model_labels_empty_entry(): + ls = model_labels(ModelEntry(group="g", name="n", files={})) + assert ls.is_placeholder + + +# -------------------------------------------------------------------------- +# 仓库根探测 +# -------------------------------------------------------------------------- + + +def test_find_repo_root_walks_up(tmp_path): + (tmp_path / "config").mkdir() + (tmp_path / "config" / "site.yaml").write_text("a: 1\n", encoding="utf-8") + deep = tmp_path / "a" / "b" / "c" + deep.mkdir(parents=True) + assert find_repo_root(deep) == tmp_path.resolve() + + +def test_find_repo_root_returns_none_when_absent(tmp_path): + deep = tmp_path / "x" / "y" + deep.mkdir(parents=True) + assert find_repo_root(deep) is None + + +def test_default_weights_root_honours_env(tmp_path, monkeypatch): + monkeypatch.setenv("SITE_WEIGHTS_ROOT", str(tmp_path / "w")) + assert default_weights_root() == tmp_path / "w" + + +def test_plan_conversion_returns_chain_and_targets(tmp_path): + src = _touch(tmp_path / "in.pt") + chain, outs = plan_conversion(tmp_path, src, name="m", group="g") + assert chain.produce == ("onnx", "engine") + assert [p.name for p in outs] == ["m.onnx", "m.engine"] + assert all(p.parent == tmp_path / "g" for p in outs) + + +# -------------------------------------------------------------------------- +# CLI：list / labels +# -------------------------------------------------------------------------- + + +def test_cli_list_empty_root(tmp_path, capsys): + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "list"]) == 0 + assert "没找到" in capsys.readouterr().out + + +def test_cli_list_shows_models_and_missing_labels(tmp_path, capsys): + _make_onnx(tmp_path / "yolo" / "good.onnx", "{0: 'workpiece'}") + _touch(tmp_path / "yolo" / "bad.engine") + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "list"]) == 0 + out = capsys.readouterr().out + assert "good" in out and "bad" in out + assert "workpiece" in out + assert "缺标签" in out + assert "2 个模型" in out + + +def test_cli_list_json_is_machine_readable(tmp_path, capsys): + import json + + _make_onnx(tmp_path / "yolo" / "m.onnx", "{0: 'workpiece'}") + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "list", "--json"]) == 0 + data = json.loads(capsys.readouterr().out) + assert data[0]["name"] == "m" + assert data[0]["labels"] == {"0": "workpiece"} + assert data[0]["labels_source"] == SOURCE_ONNX + + +def test_cli_labels_read(tmp_path, capsys): + _make_onnx(tmp_path / "yolo" / "m.onnx", "{0: 'workpiece'}") + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "labels", "m"]) == 0 + out = capsys.readouterr().out + assert "workpiece" in out + assert "ONNX 内嵌 metadata" in out + + +def test_cli_labels_set_writes_and_reads_back(tmp_path, capsys): + """核心可用性：改了标签，下一次读就是新的。""" + _make_onnx(tmp_path / "yolo" / "m.onnx", "{0: 'OLD'}") + cli = _load_cli() + + assert cli.main(["--root", str(tmp_path), "labels", "m", "--set", "工件,通止规"]) == 0 + capsys.readouterr() + + assert cli.main(["--root", str(tmp_path), "labels", "m"]) == 0 + out = capsys.readouterr().out + assert "工件" in out and "通止规" in out + assert "旁挂标签文件" in out + assert "OLD" not in out + + +def test_cli_labels_set_file(tmp_path, capsys): + _touch(tmp_path / "yolo" / "m.engine") + f = tmp_path / "classes.txt" + f.write_text("# 注释\n工件\n通止规\n", encoding="utf-8") + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "labels", "m", "--set-file", str(f)]) == 0 + capsys.readouterr() + assert cli.main(["--root", str(tmp_path), "labels", "m", "--json"]) == 0 + import json + + data = json.loads(capsys.readouterr().out) + assert data["labels"] == {"0": "工件", "1": "通止规"} + + +def test_cli_labels_from_model_materializes_file(tmp_path, capsys): + """--from-model 把内嵌标签抄成可编辑文件，这是 engine 补标签的路子。""" + _make_onnx(tmp_path / "yolo" / "m.onnx", "{0: 'workpiece'}") + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "labels", "m", "--from-model"]) == 0 + assert (tmp_path / "yolo" / "m.labels.yaml").is_file() + assert "已写入" in capsys.readouterr().out + + +def test_cli_labels_clear(tmp_path, capsys): + _touch(tmp_path / "yolo" / "m.engine") + cli = _load_cli() + cli.main(["--root", str(tmp_path), "labels", "m", "--set", "a"]) + capsys.readouterr() + assert cli.main(["--root", str(tmp_path), "labels", "m", "--clear"]) == 0 + assert not (tmp_path / "yolo" / "m.labels.yaml").exists() + + +def test_cli_labels_unknown_model_exit2(tmp_path, capsys): + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "labels", "ghost"]) == 2 + assert "找不到模型" in capsys.readouterr().err + + +def test_cli_labels_set_and_setfile_conflict(tmp_path, capsys): + _touch(tmp_path / "yolo" / "m.engine") + cli = _load_cli() + rc = cli.main(["--root", str(tmp_path), "labels", "m", "--set", "a", "--set-file", "x.txt"]) + assert rc == 2 + assert "不能同时" in capsys.readouterr().err + + +def test_cli_labels_missing_model_reports_placeholder(tmp_path, capsys): + _touch(tmp_path / "yolo" / "m.engine") + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "labels", "m"]) == 1 + out = capsys.readouterr().out + assert "没有标签" in out + assert "--from-model" in out # 必须给出修复办法 + + +def test_cli_labels_json_exposes_source_file(tmp_path, capsys): + import json + + _touch(tmp_path / "yolo" / "m.engine") + cli = _load_cli() + cli.main(["--root", str(tmp_path), "labels", "m", "--set", "a"]) + capsys.readouterr() + cli.main(["--root", str(tmp_path), "labels", "m", "--json"]) + data = json.loads(capsys.readouterr().out) + assert data["source"] == SOURCE_SIDECAR + assert data["source_file"].endswith("m.labels.yaml") + assert data["expected_labels_path"].endswith("m.labels.yaml") + + +# -------------------------------------------------------------------------- +# CLI：convert（不需要 GPU 的路径） +# -------------------------------------------------------------------------- + + +def test_cli_convert_dry_run_does_not_create_files(tmp_path, capsys): + _touch(tmp_path / "in.pt") + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "convert", "in.pt", "--name", "m", "--dry-run"]) == 0 + out = capsys.readouterr().out + assert "onnx" in out and "engine" in out + assert not (tmp_path / "yolo" / "m.onnx").exists() + assert not (tmp_path / "yolo" / "m.engine").exists() + + +def test_cli_convert_engine_input_just_relocates(tmp_path, capsys): + """已是最终产物 → 不改名转换，直接归位（不需要 GPU）。""" + _touch(tmp_path / "staging" / "raw_engine.engine", b"ENGINE") + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "convert", "raw_engine.engine", "--name", "myengine"]) == 0 + dst = tmp_path / "yolo" / "myengine.engine" + assert dst.is_file() + assert dst.read_bytes() == b"ENGINE" + + +def test_cli_convert_engine_from_staging_is_consumed(tmp_path, capsys): + """归位路径：staging 里的 engine 直接变成模型文件。 + + 这种情况没有"额外产物"要归档 —— 源文件本身就是产出， + 所以 staging 变干净，且不该去动 done/（否则会移动两次而报错）。 + """ + _touch(tmp_path / "staging" / "raw.engine", b"E") + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "convert", "staging/raw.engine", "--name", "m"]) == 0 + assert (tmp_path / "yolo" / "m.engine").is_file() + assert staging_files(tmp_path) == [] + assert not (tmp_path / "staging" / "done").exists() + + +def test_cli_convert_keep_staging(tmp_path): + _touch(tmp_path / "staging" / "raw.engine", b"E") + cli = _load_cli() + cli.main(["--root", str(tmp_path), "convert", "staging/raw.engine", "--name", "m", "--keep-staging"]) + assert (tmp_path / "staging" / "raw.engine").is_file() + + +def test_cli_convert_rejects_bad_name(tmp_path, capsys): + _touch(tmp_path / "staging" / "raw.engine", b"E") + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "convert", "staging/raw.engine", "--name", "bad/name"]) == 2 + assert "错误" in capsys.readouterr().err + + +def test_cli_convert_missing_input_exit2(tmp_path, capsys): + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "convert", "nope.pt", "--name", "m"]) == 2 + assert "不存在" in capsys.readouterr().err + + +def test_cli_convert_unknown_format_reports_clearly(tmp_path, capsys): + """非标准格式要给出明确指引，而不是让用户猜。""" + _touch(tmp_path / "staging" / "weird.pdparams", b"x") + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "convert", "staging/weird.pdparams", "--name", "m"]) == 2 + err = capsys.readouterr().err + assert "不认识的输入格式" in err + + +def test_cli_convert_labels_flag_written_for_engine_relocation(tmp_path): + _touch(tmp_path / "staging" / "raw.engine", b"E") + cli = _load_cli() + cli.main( + ["--root", str(tmp_path), "convert", "staging/raw.engine", "--name", "m", "--labels", "工件,通止规"] + ) + assert (tmp_path / "yolo" / "m.labels.yaml").is_file() + ls = model_labels(find_model(tmp_path, "m")) + assert ls.names == {0: "工件", 1: "通止规"} + + +def test_cli_convert_does_not_overwrite_existing_labels(tmp_path): + """已有旁挂标签时，转换不能把用户手改的内容冲掉。""" + _touch(tmp_path / "staging" / "raw.engine", b"E") + write_labels(tmp_path / "yolo" / "m.engine", {0: "手工改过的"}) + cli = _load_cli() + cli.main(["--root", str(tmp_path), "convert", "staging/raw.engine", "--name", "m"]) + ls = model_labels(find_model(tmp_path, "m")) + assert ls.names == {0: "手工改过的"} + + +# -------------------------------------------------------------------------- +# 回归：PT→ONNX 绝不能碰源文件旁边的同名 .onnx +# -------------------------------------------------------------------------- + + +class _FakeYOLO: + """模拟 Ultralytics：把 onnx 写在**传给它的那个路径旁边**。 + + 真实行为就是这样的，也正是踩坑点：对 models/yolo/a.pt 直接导出会覆盖 + models/yolo/a.onnx（用户自己的产物，可能是别的 imgsz 导的）， + 而紧接着的搬运还会把它彻底删掉。 + """ + + def __init__(self, path): + self.path = Path(path) + + def export(self, **_kw): + out = self.path.with_suffix(".onnx") + out.write_bytes(b"FAKE_ONNX_BYTES") + return str(out) + + +def _patch_ultralytics(monkeypatch): + import sys + import types + + monkeypatch.setitem(sys.modules, "ultralytics", types.SimpleNamespace(YOLO=_FakeYOLO)) + + +def test_export_onnx_never_touches_source_directory(tmp_path, monkeypatch): + """这条是数据丢失事故的回归测试。 + + 事故经过：转换 detect_hc90.pt 时，Ultralytics 把 onnx 写到了 + models/yolo/detect_hc90.onnx（覆盖掉原有文件），随后搬运又把它移走， + 结果用户的 detect_hc90.onnx 直接消失。现在导出在临时目录里做。 + """ + _patch_ultralytics(monkeypatch) + cli = _load_cli() + + src = _touch(tmp_path / "src_dir" / "a.pt", b"PT") + before = sorted(p.name for p in (tmp_path / "src_dir").iterdir()) + + dst = tmp_path / "out" / "renamed.onnx" + cli._export_onnx(src, 640, dst) + + # 目标产出正确 + assert dst.read_bytes() == b"FAKE_ONNX_BYTES" + # 源文件还在，且源目录一个字节都没多出来 + assert src.read_bytes() == b"PT" + assert sorted(p.name for p in (tmp_path / "src_dir").iterdir()) == before + + +def test_export_onnx_does_not_clobber_sibling_onnx(tmp_path, monkeypatch): + """源目录里已有的同名 .onnx（别人的产物）必须原封不动。""" + _patch_ultralytics(monkeypatch) + cli = _load_cli() + + src = _touch(tmp_path / "d" / "a.pt", b"PT") + sibling = _touch(tmp_path / "d" / "a.onnx", b"ORIGINAL_USER_ARTIFACT") + + cli._export_onnx(src, 640, tmp_path / "out" / "other.onnx") + + assert sibling.read_bytes() == b"ORIGINAL_USER_ARTIFACT" + + +def test_cli_convert_pt_to_other_name_keeps_sibling_onnx(tmp_path, monkeypatch): + """端到端（用假 ultralytics）：改名为 other 转换 a.pt，a.onnx 必须还在。""" + _patch_ultralytics(monkeypatch) + cli = _load_cli() + + _touch(tmp_path / "yolo" / "a.pt", b"PT") + sibling = _touch(tmp_path / "yolo" / "a.onnx", b"ORIGINAL") + + rc = cli.main(["--root", str(tmp_path), "convert", "yolo/a.pt", "--name", "other", "--to", "onnx"]) + assert rc == 0 + assert (tmp_path / "yolo" / "other.onnx").is_file() + assert sibling.read_bytes() == b"ORIGINAL" + + +def test_cli_convert_refuses_to_overwrite_without_force(tmp_path, monkeypatch, capsys): + """目标已存在时默认拒绝，避免静默覆盖别人的模型。""" + _patch_ultralytics(monkeypatch) + cli = _load_cli() + + _touch(tmp_path / "yolo" / "a.pt", b"PT") + _touch(tmp_path / "yolo" / "other.onnx", b"EXISTING") + + rc = cli.main(["--root", str(tmp_path), "convert", "yolo/a.pt", "--name", "other", "--to", "onnx"]) + assert rc == 1 + assert "已存在" in capsys.readouterr().err + assert (tmp_path / "yolo" / "other.onnx").read_bytes() == b"EXISTING" + + +def test_cli_convert_force_overwrites(tmp_path, monkeypatch): + _patch_ultralytics(monkeypatch) + cli = _load_cli() + + _touch(tmp_path / "yolo" / "a.pt", b"PT") + _touch(tmp_path / "yolo" / "other.onnx", b"OLD") + + assert cli.main( + ["--root", str(tmp_path), "convert", "yolo/a.pt", "--name", "other", "--to", "onnx", "--force"] + ) == 0 + assert (tmp_path / "yolo" / "other.onnx").read_bytes() == b"FAKE_ONNX_BYTES" + + +# -------------------------------------------------------------------------- +# CLI：doctor / show +# -------------------------------------------------------------------------- + + +def test_cli_doctor_runs_and_reports_problems(tmp_path, capsys): + _touch(tmp_path / "yolo" / "m.engine") + cli = _load_cli() + rc = cli.main(["--root", str(tmp_path), "doctor"]) + out = capsys.readouterr().out + assert rc in (0, 1) + assert "缺标签" in out + assert "staging" in out + + +def test_cli_show_lists_conventions(tmp_path, capsys): + cli = _load_cli() + assert cli.main(["--root", str(tmp_path), "show"]) == 0 + out = capsys.readouterr().out + assert ".labels.yaml" in out + assert "旁挂文件 > ONNX" in out + + +def test_cli_requires_subcommand(): + cli = _load_cli() + with pytest.raises(SystemExit): + cli.main([])