P4:权重/标签模块 —— 旁挂标签文件 + PT→ONNX→TRT 转换工具 + 归档归位

痛点: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 <cursoragent@cursor.com>
This commit is contained in:
arm_dev
2026-09-23 12:39:18 +08:00
parent adc72955c7
commit 2b77ec127d
33 changed files with 3797 additions and 152 deletions
+16 -10
View File
@@ -1,16 +1,22 @@
# 本机现场配置(含 IP / 序列号;从 site.example.yaml 复制) # 本机现场配置(含 IP / 序列号;从 site.example.yaml 复制)
config/site.yaml config/site.yaml
# YOLO 二进制 # 模型权重二进制 —— 一律不进 git(体积大、每台机器重新转换即可)
models/yolo/*.pt # 用通配模式而不是逐个目录列举,否则新增一个分组就会有人误提交二进制。
models/yolo/*.onnx # 注意:*.labels.yaml 是**文本标签配置**,不属于二进制,应该进 git。
models/yolo/*.engine models/**/*.pt
models/yolo/*.trt models/**/*.pth
models/yolo/active.pt models/**/*.onnx
models/hole_insert/*.pt models/**/*.engine
models/hole_insert/*.onnx models/**/*.trt
models/hole_insert/*.engine models/**/*.plan
models/hole_insert/*.trt models/**/*.tflite
# staging = 待转换的临时区,内容不进 git(含转换后归档的 done/)
# 但 .gitkeep 和说明文档要留着,否则新克隆的仓库里这个目录不存在/看不懂。
models/staging/**
!models/staging/.gitkeep
!models/staging/README.md
# 编译产物 # 编译产物
build/ build/
+5 -5
View File
@@ -50,7 +50,7 @@ robot:
paths: paths:
handeye_file: "/home/USER/ros2/handeye/results/89_eye_in_hand.yaml" 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_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 不影响这里) # 左臂跟踪高度(只在 --arm 89 时生效,改 91 的 servo 不影响这里)
servo: servo:
ee_z_m: 0.35 # 跟踪锁高:法兰相对基座 Z [m] ee_z_m: 0.35 # 跟踪锁高:法兰相对基座 Z [m]
@@ -87,7 +87,7 @@ robot:
paths: paths:
handeye_file: "/home/USER/ros2/handeye/results/91_eye_in_hand.yaml" 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_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 仍用上面那段 <<< # >>> 右臂跟踪高度:只改这里,89 仍用上面那段 <<<
# 法兰 Z;RGI 比平行爪高,通常要比 89 略抬一点。改完重启: # 法兰 Z;RGI 比平行爪高,通常要比 89 略抬一点。改完重启:
# bash ~/ros2/scripts/run_visual_servo_rt.sh --arm 91 --control # 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_hc90.pt"
# yolo_weights: "/home/USER/ros2/models/yolo/detect_yolo26_gongjian.pt" # yolo_weights: "/home/USER/ros2/models/yolo/detect_yolo26_gongjian.pt"
# 入孔 YOLO(backend=yolo 时把孔当框)。two_stage 用上面的工件权重。 # 入孔 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 各一份。 # 下面是缺省(未匹配到 robots[] 时)。真机以 robots[].paths 为准:89 / 91 各一份。
handeye_file: "/home/USER/ros2/handeye/results/89_eye_in_hand.yaml" 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_waypoints: "/home/USER/ros2/handeye/poses/89_eye_in_hand.yaml"
handeye_calibrate: "/home/USER/ros2/handeye/calibrate.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 # 空=默认 DDS(SHM);勿强制 UDP-only,否则相机预览收不到 /workpiece/detection
fastdds_profile: "" fastdds_profile: ""
@@ -340,7 +340,7 @@ gripper:
# bash ~/ros2/scripts/run_hole_detect.sh --arm 89 --preview # bash ~/ros2/scripts/run_hole_detect.sh --arm 89 --preview
hole_detect: hole_detect:
backend: "two_stage" # two_stage | yolo | hough 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 螺纹孔) # 二阶段工件检测权重(YOLO 先框出工件,再 Hough 找 k1/k2 螺纹孔)
workpiece_weights: "/home/USER/ros2/models/hole_insert/detect_best_1.onnx" workpiece_weights: "/home/USER/ros2/models/hole_insert/detect_best_1.onnx"
config: "/home/USER/ros2/src/arm_application/config/hole_detect.yaml" config: "/home/USER/ros2/src/arm_application/config/hole_detect.yaml"
+1 -1
View File
@@ -20,6 +20,6 @@ handeye_waypoints:
camera_namespace: "" camera_namespace: ""
image_topic: "" image_topic: ""
preview_window: true 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] 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]
+273
View File
@@ -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` —— 换模型要同时改哪几处
+24 -11
View File
@@ -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 ```yaml
paths: paths:
hole_weights: "/home/USER/ros2/models/hole_insert/best.pt" hole_weights: ".../models/hole_insert/best.pt" # ← 文件不存在
hole_detect: hole_detect:
backend: "yolo" weights: ".../models/hole_insert/best.pt" # ← 文件不存在
weights: "/home/USER/ros2/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 ```bash
mkdir -p ~/ros2/models/hole_insert 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
# cp /path/to/best.pt ~/ros2/models/hole_insert/best.pt
``` ```
@@ -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
View File
+26
View File
@@ -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/` 里的东西可以随时删,只是留个念想。
+42 -7
View File
@@ -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 ```yaml
paths: paths:
yolo_weights: "/home/USER/ros2/models/yolo/YOUR.onnx" yolo_weights: "/home/USER/ros2/models/yolo/workpiece.onnx"
yolo: yolo:
backend: "tensorrt" backend: "tensorrt"
imgsz: 640 # 必须等于该模型输入 imgsz: 640 # 必须等于该模型输入尺寸
``` ```
TensorRT: ## 查看 / 修改标签
```bash ```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`。
留着只是不想动别人的文件,可以安全删除。
+16
View File
@@ -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
+16
View File
@@ -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
+16
View File
@@ -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
@@ -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
+5 -1
View File
@@ -8,7 +8,11 @@
**夹爪:** `run_gripper.sh --arm 89|91` · `run_dual_gripper.sh` **夹爪:** `run_gripper.sh --arm 89|91` · `run_dual_gripper.sh`
**入孔:** `run_hole_detect.sh --arm … [--preview]` **入孔:** `run_hole_detect.sh --arm … [--preview]`
**触觉:** `run_tactile.sh [--cli]` **触觉:** `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` **配置:** `show_site_config.sh` · `lib/export_site_env.py`
**环境:** `env.sh` · `build_all.sh` **环境:** `env.sh` · `build_all.sh`
**RDK:** `test_rdk_python.sh` · `rdk_switch_tool.sh` · `rdk_drag_bridge.py` **RDK:** `test_rdk_python.sh` · `rdk_switch_tool.sh` · `rdk_drag_bridge.py`
+32 -75
View File
@@ -1,90 +1,47 @@
#!/usr/bin/env bash #!/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 <weights> --name <同名主干> --to engine [--imgsz N]
#
# 原用法(仍然可用):
# bash ~/ros2/scripts/export_yolo_trt.sh # bash ~/ros2/scripts/export_yolo_trt.sh
# bash ~/ros2/scripts/export_yolo_trt.sh /home/cc/ros2/models/yolo/detect_best.onnx 640 # bash ~/ros2/scripts/export_yolo_trt.sh /path/to/xxx.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.
set -eo pipefail set -eo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
# shellcheck disable=SC1091 # 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}}" IMGSZ="${2:-${SITE_YOLO_IMGSZ:-640}}"
PY="${HOME}/anaconda3/envs/ros_flexiv/bin/python"
if [[ ! -f "${PT}" ]]; then if [[ ! -f "${WEIGHTS}" ]]; then
echo "ERROR: weights not found: ${PT}" >&2 echo "ERROR: weights not found: ${WEIGHTS}" >&2
exit 1 exit 1
fi fi
# Pip TensorRT shared libs (needed at export + runtime). # 取文件名主干做模型名:xxx.onnx → xxx
TRT_LIBS="$("${PY}" - <<'PY' BASE="$(basename "${WEIGHTS}")"
import pathlib, tensorrt NAME="${BASE%.*}"
p = pathlib.Path(tensorrt.__file__).resolve().parent DIR="$(dirname "${WEIGHTS}")"
cands = [p, p.parent / "tensorrt_libs", p.parent / "nvidia" / "tensorrt" / "lib"] # 分组 = 权重所在目录名(相对 models/ 的那一层)
print(":".join(str(x) for x in cands if x.exists())) GROUP="$(basename "${DIR}")"
PY
)"
export LD_LIBRARY_PATH="${TRT_LIBS}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
export YOLO_AUTOINSTALL=false
echo "Export TensorRT engine from ${PT} imgsz=${IMGSZ}" echo "提示: export_yolo_trt.sh 已改为转发壳,建议直接用:"
"${PY}" - <<PY echo " bash ~/ros2/scripts/weights.sh convert ${WEIGHTS} --name ${NAME}"
from pathlib import Path echo
import tensorrt as trt
src = Path("${PT}").expanduser() exec bash "${ROOT}/scripts/weights.sh" convert "${WEIGHTS}" \
print("tensorrt", trt.__version__) --name "${NAME}" \
print("source", src) --group "${GROUP}" \
engine_path = src.with_suffix(".engine") --to engine \
--imgsz "${IMGSZ}" \
if src.suffix.lower() == ".onnx": --force
logger = trt.Logger(trt.Logger.INFO)
flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
builder = trt.Builder(logger)
network = builder.create_network(flag)
parser = trt.OnnxParser(network, logger)
raw = src.read_bytes()
if not parser.parse(raw):
for i in range(parser.num_errors):
print("ONNX parse error:", parser.get_error(i))
raise SystemExit("ONNX parse failed")
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30)
if builder.platform_has_fast_fp16:
config.set_flag(trt.BuilderFlag.FP16)
print("FP16 enabled")
serialized = builder.build_serialized_network(network, config)
if serialized is None:
raise SystemExit("TensorRT build failed")
engine_path.write_bytes(bytes(serialized))
print("wrote", engine_path)
else:
from ultralytics import YOLO
m = YOLO(str(src))
try:
out = m.export(
format="engine",
imgsz=int("${IMGSZ}"),
device=0,
quantize=16,
simplify=True,
workspace=4,
)
except TypeError:
out = m.export(
format="engine",
imgsz=int("${IMGSZ}"),
device=0,
half=True,
simplify=True,
workspace=4,
)
print("wrote", out)
PY
echo "Done. Restart: bash ~/ros2/scripts/run_visual_servo_rt.sh --control"
+7 -7
View File
@@ -194,7 +194,7 @@ def flatten(data: dict) -> dict:
) )
weights = str(paths.get("yolo_weights") or "") weights = str(paths.get("yolo_weights") or "")
# Resolve active.pt symlink to real path for display # 解析符号链接,方便日志里看到真实文件(历史上有人用 active.pt 做切换,已废弃)
try: try:
weights_resolved = str(Path(weights).expanduser().resolve()) weights_resolved = str(Path(weights).expanduser().resolve())
except Exception: except Exception:
@@ -231,7 +231,7 @@ def flatten(data: dict) -> dict:
or str(_root() / "handeye" / "waypoints.yaml") or str(_root() / "handeye" / "waypoints.yaml")
), ),
"SITE_HANDEYE_WORKDIR": str( "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( "SITE_HANDEYE_FILE": str(
paths.get("handeye_file") 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_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_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_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_CONF": str(yolo.get("conf") or 0.10),
"SITE_YOLO_IMGSZ": str(int(yolo.get("imgsz") or 416)), "SITE_YOLO_IMGSZ": str(int(yolo.get("imgsz") or 640)),
"SITE_YOLO_DEVICE": str(yolo.get("device") or "0"), "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_POSE": str(home.get("pose") or "home"),
"SITE_GO_HOME_SPEED": str(home.get("speed") or 0.25), "SITE_GO_HOME_SPEED": str(home.get("speed") or 0.25),
"SITE_NET_IFACE": str(robot.get("iface") or net.get("prefer_iface") or ""), "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( "SITE_HOLE_WEIGHTS": str(
(data.get("hole_detect") or {}).get("weights") (data.get("hole_detect") or {}).get("weights")
or paths.get("hole_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( "SITE_HOLE_WORKPIECE_WEIGHTS": str(
(data.get("hole_detect") or {}).get("workpiece_weights") (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 ""), "weights": str(paths.get("yolo_weights") or ""),
"conf": float(yolo.get("conf") or 0.10), "conf": float(yolo.get("conf") or 0.10),
"iou": float(yolo.get("iou") or 0.45), "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"), "device": str(yolo.get("device") or "0"),
"half_precision": bool(yolo.get("half_precision", True)), "half_precision": bool(yolo.get("half_precision", True)),
"backend": str(yolo.get("backend") or "auto"), "backend": str(yolo.get("backend") or "auto"),
+1 -1
View File
@@ -75,7 +75,7 @@ WP="${SITE_HANDEYE_WAYPOINTS:-${HE_DIR}/poses/89_eye_in_hand.yaml}"
CAL="${SITE_HANDEYE_CALIBRATE:-${HE_DIR}/calibrate.yaml}" CAL="${SITE_HANDEYE_CALIBRATE:-${HE_DIR}/calibrate.yaml}"
TOUR="${SITE_HANDEYE_TOUR:-${HE_DIR}/waypoints.yaml}" TOUR="${SITE_HANDEYE_TOUR:-${HE_DIR}/waypoints.yaml}"
RES="${SITE_HANDEYE_FILE:-${HE_DIR}/results/89_eye_in_hand.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" TAUGHT="${WORKDIR}/taught_poses.txt"
PREFIX="${SITE_JOINT_PREFIX:-Rizon4-063989_}" PREFIX="${SITE_JOINT_PREFIX:-Rizon4-063989_}"
CAM="${SITE_CAM_NAME:-d405}" CAM="${SITE_CAM_NAME:-d405}"
+8 -5
View File
@@ -13,8 +13,11 @@ parse_arm_args "$@"
python3 "${ROOT}/scripts/lib/export_site_env.py" --print ${ROBOT_SN:+--sn "$ROBOT_SN"} python3 "${ROOT}/scripts/lib/export_site_env.py" --print ${ROBOT_SN:+--sn "$ROBOT_SN"}
echo echo
echo "Edit: ${SITE_CONFIG:-$ROOT/config/site.yaml}" echo "Edit: ${SITE_CONFIG:-$ROOT/config/site.yaml}"
echo "Switch weights example:" echo "Switch weights:"
echo " paths.yolo_weights: \"/home/cc/ros2/models/yolo/detect_hc90.pt\"" echo " paths.yolo_weights: \"/home/cc/ros2/models/yolo/detect_best_0.onnx\" # 写 onnx 或 pt 均可"
echo "Or symlink:" echo " 改完确认标签和 engine 都对:"
echo " ln -sfn detect_hc90.pt ~/ros2/models/yolo/active.pt" echo " bash ~/ros2/scripts/weights.sh list"
echo " # then set paths.yolo_weights to .../active.pt" echo " bash ~/ros2/scripts/weights.sh doctor --check-engines"
echo
echo "(旧做法: 用 models/yolo/active.pt 符号链接切换模型 —— 已废弃,"
echo " 没有任何节点会读它。请直接改 paths.yolo_weights。)"
+842
View File
@@ -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` 会去 <cwd>/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}.<ext>")
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 <onnx> --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())
+69
View File
@@ -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" "$@"
@@ -131,7 +131,7 @@ class HandeyeCalibrate(Node):
self.declare_parameter('pin_hold_sec', 0.45) # 每拍多长 self.declare_parameter('pin_hold_sec', 0.45) # 每拍多长
self.declare_parameter('skip_unstable_poses', True) 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('result_file', '')
self.declare_parameter('cleanup_images', True) self.declare_parameter('cleanup_images', True)
@@ -244,7 +244,7 @@ class RdkBridge:
class HandeyeTeach(Node): class HandeyeTeach(Node):
def __init__(self) -> None: def __init__(self) -> None:
super().__init__('handeye_teach') 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('width', 1280)
self.declare_parameter('height', 720) self.declare_parameter('height', 720)
self.declare_parameter('fps', 15) self.declare_parameter('fps', 15)
@@ -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 x1, y1, x2, y2 = bbox
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2) 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: 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], hole_radii: list[float],
dist: float, dist: float,
gate: tuple[float, float] | None = None, gate: tuple[float, float] | None = None,
label: str = "",
) -> np.ndarray: ) -> np.ndarray:
draw_img = image.copy() draw_img = image.copy()
if gate is not None: if gate is not None:
_draw_center_gate(draw_img, gate[0], gate[1], True) _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) _draw_holes(draw_img, hole_points, hole_radii)
cv2.putText( cv2.putText(
draw_img, draw_img,
@@ -391,7 +416,9 @@ def detect_holes(
if bbox is None: if bbox is None:
return _fail("未检测到工件", draw_img) 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 中心接近画面中心才进入二阶检测。 # 1.5) 居中门控:bbox 中心接近画面中心才进入二阶检测。
# 精对准 k1/k2 时工件会离开画面中心,调用方应 skip_center_gate=True。 # 精对准 k1/k2 时工件会离开画面中心,调用方应 skip_center_gate=True。
@@ -428,13 +455,16 @@ def detect_holes(
hole_radii = [float(r1), float(r2)] hole_radii = [float(r1), float(r2)]
dist = float(np.hypot(k1[0] - k2[0], k1[1] - k2[1])) 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 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 { return {
"status": True, "status": True,
"msg": f"螺纹嵌件对,两孔距离 {dist:.1f}px", "msg": f"螺纹嵌件对,两孔距离 {dist:.1f}px",
"hole_points": hole_points, "hole_points": hole_points,
"hole_radii": hole_radii, "hole_radii": hole_radii,
"bbox": list(bbox), "bbox": list(bbox),
"label": label,
"draw_img": draw_img, "draw_img": draw_img,
} }
# ---- 退路:原 Hough 老路径 ---- # ---- 退路:原 Hough 老路径 ----
@@ -561,13 +591,14 @@ def detect_holes(
hole_points = [[x1, y1], [x2, y2]] hole_points = [[x1, y1], [x2, y2]]
hole_radii = [r1, r2] hole_radii = [r1, r2]
gate = (max_dx, max_dy) if bool(cfg.get("center_gate_draw", True)) else None 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 { return {
"status": True, "status": True,
"msg": f"检测成功,两孔距离 {dist:.1f}px", "msg": f"检测成功,两孔距离 {dist:.1f}px",
"hole_points": hole_points, "hole_points": hole_points,
"hole_radii": hole_radii, "hole_radii": hole_radii,
"bbox": list(bbox), "bbox": list(bbox),
"label": label,
"draw_img": draw_img, "draw_img": draw_img,
} }
except Exception as exc: 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_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 max_dy = float(self.cfg.get("center_max_offset_ratio_y", 0.12)) * ih
gate = (max_dx, max_dy) 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 held = self.miss > 0
return { return {
@@ -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
@@ -12,13 +12,22 @@ detect_best_1.onnx 满足(metadata: end2end=True, nms=False)。
from __future__ import annotations from __future__ import annotations
import logging
from pathlib import Path from pathlib import Path
from typing import List, Optional, Union from typing import List, Optional, Union
import numpy as np 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 from arm_application.yolo_detector import Detection
logger = logging.getLogger(__name__)
class OnnxWorkpieceDetector: class OnnxWorkpieceDetector:
"""给 detect_holes() 用的最小接口:detect(image, conf, iou, imgsz)。""" """给 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._in_h, self._in_w = int(shape[-2]), int(shape[-1])
self.conf = float(conf) self.conf = float(conf)
self.imgsz = int(imgsz) 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): def _letterbox(self, img: np.ndarray):
h, w = img.shape[:2] h, w = img.shape[:2]
@@ -105,7 +130,7 @@ class OnnxWorkpieceDetector:
xyxy=xyxy, xyxy=xyxy,
conf=float(score), conf=float(score),
cls_id=cid, 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) dets.sort(key=lambda d: d.conf, reverse=True)
@@ -74,10 +74,10 @@ class PerceptionNode(Node):
self.declare_parameter("image_topic", "") self.declare_parameter("image_topic", "")
self.declare_parameter("camera_info_topic", "") self.declare_parameter("camera_info_topic", "")
self.declare_parameter("depth_topic", "/d405/aligned_depth_to_color/image_raw") 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("weights", "/home/cc/ros2/models/yolo/detect_best_0.onnx")
self.declare_parameter("conf", 0.12) self.declare_parameter("conf", 0.10)
self.declare_parameter("iou", 0.45) self.declare_parameter("iou", 0.45)
self.declare_parameter("imgsz", 416) self.declare_parameter("imgsz", 640)
self.declare_parameter("device", "0") self.declare_parameter("device", "0")
self.declare_parameter("min_conf", 0.12) self.declare_parameter("min_conf", 0.12)
self.declare_parameter("vision_hz", 30.0) self.declare_parameter("vision_hz", 30.0)
@@ -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/<group>/<name>.<ext>``。
这是"自动改名并归位"的落点。无论输入文件叫什么、放在哪
(哪怕在 /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
@@ -4,10 +4,16 @@
Public: ``resolve_weights``, ``YoloDetector``, ``Detection``, ``pick_best``, Public: ``resolve_weights``, ``YoloDetector``, ``Detection``, ``pick_best``,
``MpYoloDetector``. TensorRT loads the sibling ``.engine`` next to the ``MpYoloDetector``. TensorRT loads the sibling ``.engine`` next to the
``.onnx`` / ``.pt`` path when ``backend='tensorrt'``. ``.onnx`` / ``.pt`` path when ``backend='tensorrt'``.
类别名(标签)由 ``model_labels`` 统一提供:优先读权重旁边的
``<模型名>.labels.yaml``(用户可随时手改),其次读模型内嵌 metadata。
这样做是因为 TensorRT engine 在当前的 TensorRT Python API 下**无法携带标签**,
只能靠旁挂文件;详见 ``model_labels`` 模块头。
""" """
from __future__ import annotations from __future__ import annotations
import logging
import multiprocessing as mp import multiprocessing as mp
import time import time
from dataclasses import dataclass from dataclasses import dataclass
@@ -16,6 +22,17 @@ from typing import List, Optional, Sequence, Union
import numpy as np 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. # Sentinel: worker still on the previous frame — caller must NOT count a miss.
INFER_BUSY = object() INFER_BUSY = object()
@@ -62,7 +79,7 @@ def resolve_weights(weights: Union[str, Path], backend: str = "auto") -> Path:
return engine return engine
raise FileNotFoundError( raise FileNotFoundError(
f"TensorRT engine missing: {engine}. " 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 b == "onnx":
if onnx.is_file(): if onnx.is_file():
@@ -78,6 +95,50 @@ def resolve_weights(weights: Union[str, Path], backend: str = "auto") -> Path:
return p 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: class YoloDetector:
def __init__( def __init__(
self, self,
@@ -132,8 +193,8 @@ class YoloDetector:
if class_names if class_names
else None else None
) )
names = getattr(self._model, "names", None) or {} self._labels = resolve_class_labels(path, self._model)
self._id_to_name = {int(k): str(v) for k, v in dict(names).items()} self._id_to_name = dict(self._labels.names)
self._warmed = False self._warmed = False
self._pending = None self._pending = None
@@ -209,7 +270,8 @@ class YoloDetector:
clss = boxes.cls.detach().cpu().numpy().astype(int) clss = boxes.cls.detach().cpu().numpy().astype(int)
for i in range(len(xyxy)): for i in range(len(xyxy)):
cid = int(clss[i]) 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: if self.class_ids is not None and cid not in self.class_ids:
continue continue
if self.class_names is not None and name.lower() not in self.class_names: if self.class_names is not None and name.lower() not in self.class_names:
+10 -6
View File
@@ -1,5 +1,9 @@
# perception_node defaults. Launch overlays ~/ros2/config/site.yaml → yolo: / paths: # perception_node 的默认值。launch 会用 ~/ros2/config/site.yaml 的 yolo: / paths: 覆盖。
# Do not edit this file for day-to-day tuning. # 日常调参**不要改这里**,改 site.yaml。
#
# 注意:下面 weights / conf / iou / imgsz / device / half_precision / min_conf 这几项
# 一定被 site.yaml 覆盖,这里只是"没有 site.yaml 时的兜底"。所以它们必须和 site.yaml
# 保持一致 —— 曾经这里写死 imgsz=416 而实际模型是 640,一旦覆盖失效就会用错尺寸。
perception_node: perception_node:
ros__parameters: ros__parameters:
@@ -8,13 +12,13 @@ perception_node:
image_topic: "" image_topic: ""
camera_info_topic: "" camera_info_topic: ""
depth_topic: "/d405/aligned_depth_to_color/image_raw" depth_topic: "/d405/aligned_depth_to_color/image_raw"
weights: "/home/cc/ros2/models/yolo/detect_yolo26_gongjian.pt" weights: "/home/cc/ros2/models/yolo/detect_best_0.onnx"
conf: 0.35 conf: 0.10
iou: 0.45 iou: 0.45
imgsz: 416 imgsz: 640
device: "0" device: "0"
half_precision: true half_precision: true
min_conf: 0.35 min_conf: 0.12
vision_hz: 20.0 vision_hz: 20.0
z_ref_m: 0.30 z_ref_m: 0.30
# Large gate + reacquire: fast motion must NOT freeze track # Large gate + reacquire: fast motion must NOT freeze track
@@ -36,7 +36,7 @@ def _setup(context, *args, **kwargs):
if backend in ("two_stage", "hough_roi"): if backend in ("two_stage", "hough_roi"):
weights = os.environ.get("SITE_HOLE_WORKPIECE_WEIGHTS") or os.environ.get("SITE_YOLO_WEIGHTS") or "" weights = os.environ.get("SITE_HOLE_WORKPIECE_WEIGHTS") or os.environ.get("SITE_YOLO_WEIGHTS") or ""
else: 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 [ return [
LogInfo(msg=f"[hole_detect] cam={cam} backend={backend} weights={weights} cfg={cfg}"), LogInfo(msg=f"[hole_detect] cam={cam} backend={backend} weights={weights} cfg={cfg}"),
Node( Node(
+20 -4
View File
@@ -5,18 +5,34 @@ from glob import glob
package_name = 'arm_application' package_name = 'arm_application'
def _existing(paths):
"""只保留真实存在的常规文件。
必需:colcon 的 ``build/<pkg>/`` 里放着指向源文件的符号链接,
源文件被挪走或删掉之后这些链接就断了。``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(): def _data_files():
entries = [ entries = [
('share/ament_index/resource_index/packages', ['resource/' + package_name]), ('share/ament_index/resource_index/packages', ['resource/' + package_name]),
('share/' + package_name, ['package.xml']), ('share/' + package_name, ['package.xml']),
(os.path.join('share', package_name, 'launch'), glob('launch/*.py')), (os.path.join('share', package_name, 'launch'), _existing(glob('launch/*.py'))),
(os.path.join('share', package_name, 'trajectories'), glob('trajectories/*')), (os.path.join('share', package_name, 'trajectories'), _existing(glob('trajectories/*'))),
] ]
for root, _dirs, files in os.walk('config'): 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 continue
dest = os.path.join('share', package_name, root) dest = os.path.join('share', package_name, root)
srcs = [os.path.join(root, name) for name in files]
entries.append((dest, srcs)) entries.append((dest, srcs))
return entries return entries
@@ -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"}
@@ -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", "a<b", "a>b", "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([])