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
+16
-10
@@ -1,16 +1,22 @@
|
||||
# 本机现场配置（含 IP / 序列号；从 site.example.yaml 复制）
|
||||
config/site.yaml
|
||||
|
||||
# YOLO 二进制
|
||||
models/yolo/*.pt
|
||||
models/yolo/*.onnx
|
||||
models/yolo/*.engine
|
||||
models/yolo/*.trt
|
||||
models/yolo/active.pt
|
||||
models/hole_insert/*.pt
|
||||
models/hole_insert/*.onnx
|
||||
models/hole_insert/*.engine
|
||||
models/hole_insert/*.trt
|
||||
# 模型权重二进制 —— 一律不进 git（体积大、每台机器重新转换即可）
|
||||
# 用通配模式而不是逐个目录列举，否则新增一个分组就会有人误提交二进制。
|
||||
# 注意：*.labels.yaml 是**文本标签配置**，不属于二进制，应该进 git。
|
||||
models/**/*.pt
|
||||
models/**/*.pth
|
||||
models/**/*.onnx
|
||||
models/**/*.engine
|
||||
models/**/*.trt
|
||||
models/**/*.plan
|
||||
models/**/*.tflite
|
||||
|
||||
# staging = 待转换的临时区，内容不进 git（含转换后归档的 done/）
|
||||
# 但 .gitkeep 和说明文档要留着，否则新克隆的仓库里这个目录不存在/看不懂。
|
||||
models/staging/**
|
||||
!models/staging/.gitkeep
|
||||
!models/staging/README.md
|
||||
|
||||
# 编译产物
|
||||
build/
|
||||
|
||||
@@ -50,7 +50,7 @@ robot:
|
||||
paths:
|
||||
handeye_file: "/home/USER/ros2/handeye/results/89_eye_in_hand.yaml"
|
||||
handeye_waypoints: "/home/USER/ros2/handeye/poses/89_eye_in_hand.yaml"
|
||||
handeye_workdir: "/tmp/handeye_calib_89"
|
||||
handeye_workdir: "/home/USER/ros2/handeye/tmp/89"
|
||||
# 左臂跟踪高度（只在 --arm 89 时生效，改 91 的 servo 不影响这里）
|
||||
servo:
|
||||
ee_z_m: 0.35 # 跟踪锁高：法兰相对基座 Z [m]
|
||||
@@ -87,7 +87,7 @@ robot:
|
||||
paths:
|
||||
handeye_file: "/home/USER/ros2/handeye/results/91_eye_in_hand.yaml"
|
||||
handeye_waypoints: "/home/USER/ros2/handeye/poses/91_eye_in_hand.yaml"
|
||||
handeye_workdir: "/tmp/handeye_calib_91"
|
||||
handeye_workdir: "/home/USER/ros2/handeye/tmp/91"
|
||||
# >>> 右臂跟踪高度：只改这里，89 仍用上面那段 <<<
|
||||
# 法兰 Z；RGI 比平行爪高，通常要比 89 略抬一点。改完重启：
|
||||
# bash ~/ros2/scripts/run_visual_servo_rt.sh --arm 91 --control
|
||||
@@ -117,12 +117,12 @@ paths:
|
||||
# yolo_weights: "/home/USER/ros2/models/yolo/detect_hc90.pt"
|
||||
# yolo_weights: "/home/USER/ros2/models/yolo/detect_yolo26_gongjian.pt"
|
||||
# 入孔 YOLO（backend=yolo 时把孔当框）。two_stage 用上面的工件权重。
|
||||
hole_weights: "/home/USER/ros2/models/hole_insert/best.pt"
|
||||
hole_weights: "/home/USER/ros2/models/hole_insert/detect_best_1.onnx"
|
||||
# 下面是缺省（未匹配到 robots[] 时）。真机以 robots[].paths 为准：89 / 91 各一份。
|
||||
handeye_file: "/home/USER/ros2/handeye/results/89_eye_in_hand.yaml"
|
||||
handeye_waypoints: "/home/USER/ros2/handeye/poses/89_eye_in_hand.yaml"
|
||||
handeye_calibrate: "/home/USER/ros2/handeye/calibrate.yaml"
|
||||
handeye_workdir: "/tmp/handeye_calib_89"
|
||||
handeye_workdir: "/home/USER/ros2/handeye/tmp/89"
|
||||
# 空=默认 DDS(SHM)；勿强制 UDP-only，否则相机预览收不到 /workpiece/detection
|
||||
fastdds_profile: ""
|
||||
|
||||
@@ -340,7 +340,7 @@ gripper:
|
||||
# bash ~/ros2/scripts/run_hole_detect.sh --arm 89 --preview
|
||||
hole_detect:
|
||||
backend: "two_stage" # two_stage | yolo | hough
|
||||
weights: "/home/USER/ros2/models/hole_insert/best.pt"
|
||||
weights: "/home/USER/ros2/models/hole_insert/detect_best_1.onnx"
|
||||
# 二阶段工件检测权重（YOLO 先框出工件，再 Hough 找 k1/k2 螺纹孔）
|
||||
workpiece_weights: "/home/USER/ros2/models/hole_insert/detect_best_1.onnx"
|
||||
config: "/home/USER/ros2/src/arm_application/config/hole_detect.yaml"
|
||||
|
||||
@@ -20,6 +20,6 @@ handeye_waypoints:
|
||||
camera_namespace: ""
|
||||
image_topic: ""
|
||||
preview_window: true
|
||||
append_poses_file: "/tmp/handeye_calib/taught_poses.txt"
|
||||
append_poses_file: "/home/cc/ros2/handeye/tmp/taught_poses.txt"
|
||||
|
||||
waypoint_times: [4.0, 4.0, 4.0, 4.0, 4.0, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.0, 4.0]
|
||||
|
||||
@@ -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` —— 换模型要同时改哪几处
|
||||
@@ -1,23 +1,36 @@
|
||||
# 入孔 / 孔定位权重
|
||||
# 入孔检测权重
|
||||
|
||||
二进制（`.pt` / `.onnx` / `.engine`）**不进 git**。
|
||||
> **改这个目录之前先读 [`../README.md`](../README.md)**（权重模块的规则）。
|
||||
|
||||
`hole_detect.backend: yolo` 时把本目录的检测权重绝对路径写进 `config/site.yaml`：
|
||||
用于"把孔当框来检测"的模型（`hole_detect.backend: yolo`）。
|
||||
当前现场用的是 `two_stage`（两阶段）后端，它用的是**工件**权重
|
||||
（`hole_detect.workpiece_weights`，指向 `models/yolo/` 里的模型），
|
||||
本目录的模型并不是每次都会用到。
|
||||
|
||||
## 现状
|
||||
|
||||
- `detect_best_1.onnx` —— 工件检测模型，被 `hole_detect.workpiece_weights` 使用。
|
||||
|
||||
## 历史遗留（已知问题）
|
||||
|
||||
`config/site.yaml` 的这两处仍指向一个**不存在**的文件 `models/hole_insert/best.pt`：
|
||||
|
||||
```yaml
|
||||
paths:
|
||||
hole_weights: "/home/USER/ros2/models/hole_insert/best.pt"
|
||||
hole_weights: ".../models/hole_insert/best.pt" # ← 文件不存在
|
||||
hole_detect:
|
||||
backend: "yolo"
|
||||
weights: "/home/USER/ros2/models/hole_insert/best.pt"
|
||||
weights: ".../models/hole_insert/best.pt" # ← 文件不存在
|
||||
```
|
||||
|
||||
`backend: two_stage` 用的是工件 YOLO（`paths.yolo_weights`），本目录可空。
|
||||
`backend: two_stage` 时这两处不会被读，所以现在不影响运行；但一旦切成
|
||||
`backend: yolo` 就会启动失败。用 `bash ~/ros2/scripts/weights.sh doctor`
|
||||
能查出来。要么把真实权重按命名规则放进本目录，要么把这两处配置删掉。
|
||||
|
||||
从旧工程拷贝示例：
|
||||
## 加新模型 / 看标签
|
||||
|
||||
和 `models/yolo/` 完全一样，只是分组名换成 `hole_insert`：
|
||||
|
||||
```bash
|
||||
mkdir -p ~/ros2/models/hole_insert
|
||||
# 若有独立孔模型：
|
||||
# cp /path/to/best.pt ~/ros2/models/hole_insert/best.pt
|
||||
bash ~/ros2/scripts/weights.sh convert staging/xxx.onnx --name hole_detect --group hole_insert
|
||||
bash ~/ros2/scripts/weights.sh labels detect_best_1 --group hole_insert
|
||||
```
|
||||
|
||||
@@ -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
|
||||
@@ -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
@@ -1,19 +1,54 @@
|
||||
# YOLO 权重
|
||||
# YOLO 工件检测权重
|
||||
|
||||
二进制（`.pt` / `.onnx` / `.engine`）**不进 git**。放到本目录后改 `config/site.yaml`：
|
||||
> **改这个目录之前先读 [`../README.md`](../README.md)**（权重模块的规则）。
|
||||
|
||||
## 放什么
|
||||
|
||||
同名的 `.pt` / `.onnx` / `.engine` / `.labels.yaml` 放一起，名字主干必须一致：
|
||||
|
||||
```
|
||||
detect_best_0.onnx ← 写进 config/site.yaml 的那个
|
||||
detect_best_0.engine ← TensorRT 引擎，backend: tensorrt 时自动使用
|
||||
detect_best_0.labels.yaml ← 类别标签（进 git，可手改）
|
||||
```
|
||||
|
||||
二进制不进 git；`.labels.yaml` 进 git。
|
||||
|
||||
## 加新模型
|
||||
|
||||
```bash
|
||||
cp ~/下载/xxx.pt ~/ros2/models/staging/
|
||||
bash ~/ros2/scripts/weights.sh convert staging/xxx.pt --name workpiece --group yolo
|
||||
```
|
||||
|
||||
然后改 `config/site.yaml`：
|
||||
|
||||
```yaml
|
||||
paths:
|
||||
yolo_weights: "/home/USER/ros2/models/yolo/YOUR.onnx"
|
||||
yolo_weights: "/home/USER/ros2/models/yolo/workpiece.onnx"
|
||||
yolo:
|
||||
backend: "tensorrt"
|
||||
imgsz: 640 # 必须等于该模型输入
|
||||
imgsz: 640 # 必须等于该模型输入尺寸
|
||||
```
|
||||
|
||||
TensorRT：
|
||||
## 查看 / 修改标签
|
||||
|
||||
```bash
|
||||
bash ~/ros2/scripts/export_yolo_trt.sh /home/USER/ros2/models/yolo/YOUR.onnx 640
|
||||
bash ~/ros2/scripts/weights.sh labels detect_best_0
|
||||
bash ~/ros2/scripts/weights.sh labels detect_best_0 --set "工件,通止规"
|
||||
```
|
||||
|
||||
会生成同名 `.engine`。换 ONNX 或 imgsz 后重导。然后重启 `run_visual_servo_rt.sh`。
|
||||
## 排错
|
||||
|
||||
```bash
|
||||
bash ~/ros2/scripts/weights.sh doctor --check-engines
|
||||
```
|
||||
|
||||
`engine` 和 TensorRT 版本绑死：换过 TensorRT 后旧 engine 会加载失败，重导即可
|
||||
（`convert <同名>.onnx --name <名字> --to engine --force`）。
|
||||
|
||||
## 历史遗留
|
||||
|
||||
`active.pt` 是个指向 `detect_hc90.pt` 的符号链接，早期用来"切换当前模型"，
|
||||
现在没有任何节点会读它 —— 换模型请改 `config/site.yaml` 的 `paths.yolo_weights`。
|
||||
留着只是不想动别人的文件，可以安全删除。
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
@@ -8,7 +8,11 @@
|
||||
**夹爪：** `run_gripper.sh --arm 89|91` · `run_dual_gripper.sh`
|
||||
**入孔：** `run_hole_detect.sh --arm … [--preview]`
|
||||
**触觉：** `run_tactile.sh [--cli]`
|
||||
**权重：** `export_yolo_trt.sh`
|
||||
**权重：** `weights.sh`（`list` / `labels` / `convert` / `doctor` / `show`）
|
||||
- 转换：`weights.sh convert <文件> --name <模型名>` —— PT→ONNX→TensorRT，自动改名归位 + 固化标签
|
||||
- 标签：`weights.sh labels <模型名>` 看，`--set "工件,通止规"` 改
|
||||
- 体检：`weights.sh doctor --check-engines`（能查出 engine 与 TensorRT 版本不匹配这类静默失效）
|
||||
- 规则写在 [models/README.md](../models/README.md)。`export_yolo_trt.sh` 现在是它的兼容转发壳。
|
||||
**配置：** `show_site_config.sh` · `lib/export_site_env.py`
|
||||
**环境：** `env.sh` · `build_all.sh`
|
||||
**RDK：** `test_rdk_python.sh` · `rdk_switch_tool.sh` · `rdk_drag_bridge.py`
|
||||
|
||||
+32
-75
@@ -1,90 +1,47 @@
|
||||
#!/usr/bin/env bash
|
||||
# Export YOLO .pt/.onnx → TensorRT .engine (FP16).
|
||||
# 【旧命令，已改为转发壳】YOLO 权重 → TensorRT engine。
|
||||
#
|
||||
# 请改用新命令（多了改名、归位、标签固化、staging 等功能）：
|
||||
#
|
||||
# bash ~/ros2/scripts/weights.sh convert <文件> --name <模型名>
|
||||
# bash ~/ros2/scripts/weights.sh list
|
||||
# bash ~/ros2/scripts/weights.sh labels <模型名>
|
||||
#
|
||||
# 本脚本保留只为兼容旧文档/旧习惯，行为等价于：
|
||||
# weights.sh convert <weights> --name <同名主干> --to engine [--imgsz N]
|
||||
#
|
||||
# 原用法（仍然可用）：
|
||||
# bash ~/ros2/scripts/export_yolo_trt.sh
|
||||
# bash ~/ros2/scripts/export_yolo_trt.sh /home/cc/ros2/models/yolo/detect_best.onnx 640
|
||||
#
|
||||
# Needs: tensorrt-cu12 + onnx + onnxslim in ros_flexiv (do NOT pip install
|
||||
# the meta package `tensorrt`, it pulls CUDA 13 wheels).
|
||||
# Then set yolo.backend: tensorrt in site.yaml and restart visual servo.
|
||||
# bash ~/ros2/scripts/export_yolo_trt.sh /path/to/xxx.onnx 640
|
||||
|
||||
set -eo pipefail
|
||||
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
||||
|
||||
# shellcheck disable=SC1091
|
||||
source "${ROOT}/scripts/lib/load_site_config.sh"
|
||||
source "${ROOT}/scripts/lib/load_site_config.sh" >/dev/null 2>&1 || true
|
||||
|
||||
PT="${1:-${SITE_YOLO_WEIGHTS:-${ROOT}/models/yolo/detect_best.onnx}}"
|
||||
WEIGHTS="${1:-${SITE_YOLO_WEIGHTS:-${ROOT}/models/yolo/detect_best.onnx}}"
|
||||
IMGSZ="${2:-${SITE_YOLO_IMGSZ:-640}}"
|
||||
PY="${HOME}/anaconda3/envs/ros_flexiv/bin/python"
|
||||
|
||||
if [[ ! -f "${PT}" ]]; then
|
||||
echo "ERROR: weights not found: ${PT}" >&2
|
||||
if [[ ! -f "${WEIGHTS}" ]]; then
|
||||
echo "ERROR: weights not found: ${WEIGHTS}" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Pip TensorRT shared libs (needed at export + runtime).
|
||||
TRT_LIBS="$("${PY}" - <<'PY'
|
||||
import pathlib, tensorrt
|
||||
p = pathlib.Path(tensorrt.__file__).resolve().parent
|
||||
cands = [p, p.parent / "tensorrt_libs", p.parent / "nvidia" / "tensorrt" / "lib"]
|
||||
print(":".join(str(x) for x in cands if x.exists()))
|
||||
PY
|
||||
)"
|
||||
export LD_LIBRARY_PATH="${TRT_LIBS}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
|
||||
export YOLO_AUTOINSTALL=false
|
||||
# 取文件名主干做模型名：xxx.onnx → xxx
|
||||
BASE="$(basename "${WEIGHTS}")"
|
||||
NAME="${BASE%.*}"
|
||||
DIR="$(dirname "${WEIGHTS}")"
|
||||
# 分组 = 权重所在目录名（相对 models/ 的那一层）
|
||||
GROUP="$(basename "${DIR}")"
|
||||
|
||||
echo "Export TensorRT engine from ${PT} imgsz=${IMGSZ}"
|
||||
"${PY}" - <<PY
|
||||
from pathlib import Path
|
||||
import tensorrt as trt
|
||||
echo "提示: export_yolo_trt.sh 已改为转发壳，建议直接用："
|
||||
echo " bash ~/ros2/scripts/weights.sh convert ${WEIGHTS} --name ${NAME}"
|
||||
echo
|
||||
|
||||
src = Path("${PT}").expanduser()
|
||||
print("tensorrt", trt.__version__)
|
||||
print("source", src)
|
||||
engine_path = src.with_suffix(".engine")
|
||||
|
||||
if src.suffix.lower() == ".onnx":
|
||||
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"
|
||||
exec bash "${ROOT}/scripts/weights.sh" convert "${WEIGHTS}" \
|
||||
--name "${NAME}" \
|
||||
--group "${GROUP}" \
|
||||
--to engine \
|
||||
--imgsz "${IMGSZ}" \
|
||||
--force
|
||||
|
||||
@@ -194,7 +194,7 @@ def flatten(data: dict) -> dict:
|
||||
)
|
||||
|
||||
weights = str(paths.get("yolo_weights") or "")
|
||||
# Resolve active.pt symlink to real path for display
|
||||
# 解析符号链接，方便日志里看到真实文件（历史上有人用 active.pt 做切换，已废弃）
|
||||
try:
|
||||
weights_resolved = str(Path(weights).expanduser().resolve())
|
||||
except Exception:
|
||||
@@ -231,7 +231,7 @@ def flatten(data: dict) -> dict:
|
||||
or str(_root() / "handeye" / "waypoints.yaml")
|
||||
),
|
||||
"SITE_HANDEYE_WORKDIR": str(
|
||||
paths.get("handeye_workdir") or "/tmp/handeye_calib"
|
||||
paths.get("handeye_workdir") or str(_root() / "handeye" / "tmp")
|
||||
),
|
||||
"SITE_HANDEYE_FILE": str(
|
||||
paths.get("handeye_file")
|
||||
@@ -280,10 +280,10 @@ def flatten(data: dict) -> dict:
|
||||
"SITE_D435_FPS": str(int(((data.get("cameras") or {}).get("d435") or {}).get("fps") or 30)),
|
||||
"SITE_D405_LEFT_SERIAL": str(((data.get("cameras") or {}).get("d405_left") or {}).get("serial") or "352122274408"),
|
||||
"SITE_D405_RIGHT_SERIAL": str(((data.get("cameras") or {}).get("d405_right") or {}).get("serial") or "260522274442"),
|
||||
"SITE_YOLO_CONF": str(yolo.get("conf") or 0.35),
|
||||
"SITE_YOLO_IMGSZ": str(int(yolo.get("imgsz") or 416)),
|
||||
"SITE_YOLO_CONF": str(yolo.get("conf") or 0.10),
|
||||
"SITE_YOLO_IMGSZ": str(int(yolo.get("imgsz") or 640)),
|
||||
"SITE_YOLO_DEVICE": str(yolo.get("device") or "0"),
|
||||
"SITE_YOLO_MIN_CONF": str(yolo.get("min_conf") or 0.35),
|
||||
"SITE_YOLO_MIN_CONF": str(yolo.get("min_conf") or 0.12),
|
||||
"SITE_GO_HOME_POSE": str(home.get("pose") or "home"),
|
||||
"SITE_GO_HOME_SPEED": str(home.get("speed") or 0.25),
|
||||
"SITE_NET_IFACE": str(robot.get("iface") or net.get("prefer_iface") or ""),
|
||||
@@ -312,7 +312,7 @@ def flatten(data: dict) -> dict:
|
||||
"SITE_HOLE_WEIGHTS": str(
|
||||
(data.get("hole_detect") or {}).get("weights")
|
||||
or paths.get("hole_weights")
|
||||
or str(_root() / "models" / "hole_insert" / "best.pt")
|
||||
or str(_root() / "models" / "hole_insert" / "detect_best_1.onnx")
|
||||
),
|
||||
"SITE_HOLE_WORKPIECE_WEIGHTS": str(
|
||||
(data.get("hole_detect") or {}).get("workpiece_weights")
|
||||
@@ -381,7 +381,7 @@ def perception_overrides(data: dict) -> dict:
|
||||
"weights": str(paths.get("yolo_weights") or ""),
|
||||
"conf": float(yolo.get("conf") or 0.10),
|
||||
"iou": float(yolo.get("iou") or 0.45),
|
||||
"imgsz": int(yolo.get("imgsz") or 416),
|
||||
"imgsz": int(yolo.get("imgsz") or 640),
|
||||
"device": str(yolo.get("device") or "0"),
|
||||
"half_precision": bool(yolo.get("half_precision", True)),
|
||||
"backend": str(yolo.get("backend") or "auto"),
|
||||
|
||||
@@ -75,7 +75,7 @@ WP="${SITE_HANDEYE_WAYPOINTS:-${HE_DIR}/poses/89_eye_in_hand.yaml}"
|
||||
CAL="${SITE_HANDEYE_CALIBRATE:-${HE_DIR}/calibrate.yaml}"
|
||||
TOUR="${SITE_HANDEYE_TOUR:-${HE_DIR}/waypoints.yaml}"
|
||||
RES="${SITE_HANDEYE_FILE:-${HE_DIR}/results/89_eye_in_hand.yaml}"
|
||||
WORKDIR="${SITE_HANDEYE_WORKDIR:-/tmp/handeye_calib}"
|
||||
WORKDIR="${SITE_HANDEYE_WORKDIR:-${HE_DIR}/tmp}"
|
||||
TAUGHT="${WORKDIR}/taught_poses.txt"
|
||||
PREFIX="${SITE_JOINT_PREFIX:-Rizon4-063989_}"
|
||||
CAM="${SITE_CAM_NAME:-d405}"
|
||||
|
||||
@@ -13,8 +13,11 @@ parse_arm_args "$@"
|
||||
python3 "${ROOT}/scripts/lib/export_site_env.py" --print ${ROBOT_SN:+--sn "$ROBOT_SN"}
|
||||
echo
|
||||
echo "Edit: ${SITE_CONFIG:-$ROOT/config/site.yaml}"
|
||||
echo "Switch weights example:"
|
||||
echo " paths.yolo_weights: \"/home/cc/ros2/models/yolo/detect_hc90.pt\""
|
||||
echo "Or symlink:"
|
||||
echo " ln -sfn detect_hc90.pt ~/ros2/models/yolo/active.pt"
|
||||
echo " # then set paths.yolo_weights to .../active.pt"
|
||||
echo "Switch weights:"
|
||||
echo " paths.yolo_weights: \"/home/cc/ros2/models/yolo/detect_best_0.onnx\" # 写 onnx 或 pt 均可"
|
||||
echo " 改完确认标签和 engine 都对："
|
||||
echo " bash ~/ros2/scripts/weights.sh list"
|
||||
echo " bash ~/ros2/scripts/weights.sh doctor --check-engines"
|
||||
echo
|
||||
echo "（旧做法: 用 models/yolo/active.pt 符号链接切换模型 —— 已废弃，"
|
||||
echo " 没有任何节点会读它。请直接改 paths.yolo_weights。）"
|
||||
|
||||
Executable
+842
@@ -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())
|
||||
Executable
+69
@@ -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('skip_unstable_poses', True)
|
||||
|
||||
self.declare_parameter('output_dir', '/tmp/handeye_calib')
|
||||
self.declare_parameter('output_dir', '/home/cc/ros2/handeye/tmp')
|
||||
self.declare_parameter('result_file', '')
|
||||
self.declare_parameter('cleanup_images', True)
|
||||
|
||||
|
||||
@@ -244,7 +244,7 @@ class RdkBridge:
|
||||
class HandeyeTeach(Node):
|
||||
def __init__(self) -> None:
|
||||
super().__init__('handeye_teach')
|
||||
self.declare_parameter('output_dir', '/tmp/handeye_calib')
|
||||
self.declare_parameter('output_dir', '/home/cc/ros2/handeye/tmp')
|
||||
self.declare_parameter('width', 1280)
|
||||
self.declare_parameter('height', 720)
|
||||
self.declare_parameter('fps', 15)
|
||||
|
||||
@@ -40,10 +40,34 @@ def _fail(
|
||||
}
|
||||
|
||||
|
||||
def _draw_bbox(img: np.ndarray, bbox: tuple[int, int, int, int], color=(0, 255, 0)) -> None:
|
||||
def _draw_bbox(
|
||||
img: np.ndarray,
|
||||
bbox: tuple[int, int, int, int],
|
||||
color=(0, 255, 0),
|
||||
label: str = "",
|
||||
) -> None:
|
||||
"""画工件框。``label`` 是模型给出的类别名（由标签文件决定）。
|
||||
|
||||
这里刻意**不写死类名**：以前写死了 "gongjian"，但实际模型的类别是
|
||||
"workpiece"，图上的字和模型的类名对不上，排查时会被误导。
|
||||
现在用 ``Detection.cls_name``，也就是 ``weights.sh labels`` 里能读到的那份。
|
||||
"""
|
||||
x1, y1, x2, y2 = bbox
|
||||
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
|
||||
cv2.putText(img, "gongjian", (x1, max(20, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
|
||||
text = (label or "").strip()
|
||||
if text:
|
||||
cv2.putText(img, text, (x1, max(20, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
|
||||
|
||||
|
||||
def _best_label(dets) -> str:
|
||||
"""取置信度最高那个检测的类名（由标签文件决定）。没有则返回空串。"""
|
||||
if not dets:
|
||||
return ""
|
||||
try:
|
||||
best = max(dets, key=lambda d: float(d.conf))
|
||||
except Exception: # noqa: BLE001
|
||||
return ""
|
||||
return str(getattr(best, "cls_name", "") or "").strip()
|
||||
|
||||
|
||||
def _draw_center_gate(img: np.ndarray, max_dx: float, max_dy: float, ok: bool) -> None:
|
||||
@@ -332,11 +356,12 @@ def _render_success(
|
||||
hole_radii: list[float],
|
||||
dist: float,
|
||||
gate: tuple[float, float] | None = None,
|
||||
label: str = "",
|
||||
) -> np.ndarray:
|
||||
draw_img = image.copy()
|
||||
if gate is not None:
|
||||
_draw_center_gate(draw_img, gate[0], gate[1], True)
|
||||
_draw_bbox(draw_img, bbox)
|
||||
_draw_bbox(draw_img, bbox, label=label)
|
||||
_draw_holes(draw_img, hole_points, hole_radii)
|
||||
cv2.putText(
|
||||
draw_img,
|
||||
@@ -391,7 +416,9 @@ def detect_holes(
|
||||
if bbox is None:
|
||||
return _fail("未检测到工件", draw_img)
|
||||
|
||||
_draw_bbox(draw_img, bbox)
|
||||
# 类别名来自标签文件（见 model_labels / weights.sh labels），不再写死。
|
||||
label = _best_label(dets)
|
||||
_draw_bbox(draw_img, bbox, label=label)
|
||||
|
||||
# 1.5) 居中门控：bbox 中心接近画面中心才进入二阶检测。
|
||||
# 精对准 k1/k2 时工件会离开画面中心，调用方应 skip_center_gate=True。
|
||||
@@ -428,13 +455,16 @@ def detect_holes(
|
||||
hole_radii = [float(r1), float(r2)]
|
||||
dist = float(np.hypot(k1[0] - k2[0], k1[1] - k2[1]))
|
||||
gate = (max_dx, max_dy) if bool(cfg.get("center_gate_draw", True)) else None
|
||||
draw_img = _render_success(image, bbox, hole_points, hole_radii, dist, gate=gate)
|
||||
draw_img = _render_success(
|
||||
image, bbox, hole_points, hole_radii, dist, gate=gate, label=label
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"msg": f"螺纹嵌件对，两孔距离 {dist:.1f}px",
|
||||
"hole_points": hole_points,
|
||||
"hole_radii": hole_radii,
|
||||
"bbox": list(bbox),
|
||||
"label": label,
|
||||
"draw_img": draw_img,
|
||||
}
|
||||
# ---- 退路：原 Hough 老路径 ----
|
||||
@@ -561,13 +591,14 @@ def detect_holes(
|
||||
hole_points = [[x1, y1], [x2, y2]]
|
||||
hole_radii = [r1, r2]
|
||||
gate = (max_dx, max_dy) if bool(cfg.get("center_gate_draw", True)) else None
|
||||
draw_img = _render_success(image, bbox, hole_points, hole_radii, dist, gate=gate)
|
||||
draw_img = _render_success(image, bbox, hole_points, hole_radii, dist, gate=gate, label=label)
|
||||
return {
|
||||
"status": True,
|
||||
"msg": f"检测成功，两孔距离 {dist:.1f}px",
|
||||
"hole_points": hole_points,
|
||||
"hole_radii": hole_radii,
|
||||
"bbox": list(bbox),
|
||||
"label": label,
|
||||
"draw_img": draw_img,
|
||||
}
|
||||
except Exception as exc:
|
||||
@@ -675,7 +706,15 @@ class HoleStabilizer:
|
||||
max_dx = float(self.cfg.get("center_max_offset_ratio_x", 0.12)) * iw
|
||||
max_dy = float(self.cfg.get("center_max_offset_ratio_y", 0.12)) * ih
|
||||
gate = (max_dx, max_dy)
|
||||
draw_img = _render_success(image, bbox_draw, self.points, self.radii, dist, gate=gate)
|
||||
draw_img = _render_success(
|
||||
image,
|
||||
bbox_draw,
|
||||
self.points,
|
||||
self.radii,
|
||||
dist,
|
||||
gate=gate,
|
||||
label=str(result.get("label") or ""),
|
||||
)
|
||||
|
||||
held = self.miss > 0
|
||||
return {
|
||||
|
||||
@@ -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
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from arm_application.model_labels import (
|
||||
SOURCE_MODEL,
|
||||
LabelSet,
|
||||
parse_names_field,
|
||||
read_labels,
|
||||
)
|
||||
from arm_application.yolo_detector import Detection
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OnnxWorkpieceDetector:
|
||||
"""给 detect_holes() 用的最小接口：detect(image, conf, iou, imgsz)。"""
|
||||
@@ -47,7 +56,23 @@ class OnnxWorkpieceDetector:
|
||||
self._in_h, self._in_w = int(shape[-2]), int(shape[-1])
|
||||
self.conf = float(conf)
|
||||
self.imgsz = int(imgsz)
|
||||
self._names = class_names or {0: "workpiece"}
|
||||
# 类别名：显式传入 > 旁挂标签文件 > ONNX 内嵌 metadata > 占位符。
|
||||
# 原先这里是硬编码 {0: "workpiece"}，等于把类名写死在代码里 ——
|
||||
# 换一个多类模型就会展示错误的名字，而且没人改得动。
|
||||
if class_names:
|
||||
names = parse_names_field(class_names)
|
||||
self._labels = LabelSet(names=names, source=SOURCE_MODEL, path=path)
|
||||
else:
|
||||
self._labels = read_labels(path, allow_model_metadata=False)
|
||||
if self._labels.is_placeholder:
|
||||
logger.warning(
|
||||
"ONNX %s 没有类别标签，将用占位名显示。"
|
||||
"补法: bash ~/ros2/scripts/weights.sh labels %s --from-model",
|
||||
path.name,
|
||||
path.stem,
|
||||
)
|
||||
# 兼容旧属性名，避免外部代码引用时报 AttributeError。
|
||||
self._names = dict(self._labels.names)
|
||||
|
||||
def _letterbox(self, img: np.ndarray):
|
||||
h, w = img.shape[:2]
|
||||
@@ -105,7 +130,7 @@ class OnnxWorkpieceDetector:
|
||||
xyxy=xyxy,
|
||||
conf=float(score),
|
||||
cls_id=cid,
|
||||
cls_name=self._names.get(cid, str(cid)),
|
||||
cls_name=self._labels.get(cid),
|
||||
)
|
||||
)
|
||||
dets.sort(key=lambda d: d.conf, reverse=True)
|
||||
|
||||
@@ -74,10 +74,10 @@ class PerceptionNode(Node):
|
||||
self.declare_parameter("image_topic", "")
|
||||
self.declare_parameter("camera_info_topic", "")
|
||||
self.declare_parameter("depth_topic", "/d405/aligned_depth_to_color/image_raw")
|
||||
self.declare_parameter("weights", "/home/cc/ros2/models/yolo/detect_hc90.pt")
|
||||
self.declare_parameter("conf", 0.12)
|
||||
self.declare_parameter("weights", "/home/cc/ros2/models/yolo/detect_best_0.onnx")
|
||||
self.declare_parameter("conf", 0.10)
|
||||
self.declare_parameter("iou", 0.45)
|
||||
self.declare_parameter("imgsz", 416)
|
||||
self.declare_parameter("imgsz", 640)
|
||||
self.declare_parameter("device", "0")
|
||||
self.declare_parameter("min_conf", 0.12)
|
||||
self.declare_parameter("vision_hz", 30.0)
|
||||
|
||||
@@ -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``,
|
||||
``MpYoloDetector``. TensorRT loads the sibling ``.engine`` next to the
|
||||
``.onnx`` / ``.pt`` path when ``backend='tensorrt'``.
|
||||
|
||||
类别名（标签）由 ``model_labels`` 统一提供：优先读权重旁边的
|
||||
``<模型名>.labels.yaml``（用户可随时手改），其次读模型内嵌 metadata。
|
||||
这样做是因为 TensorRT engine 在当前的 TensorRT Python API 下**无法携带标签**，
|
||||
只能靠旁挂文件；详见 ``model_labels`` 模块头。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import multiprocessing as mp
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
@@ -16,6 +22,17 @@ from typing import List, Optional, Sequence, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .model_labels import (
|
||||
SOURCE_PLACEHOLDER,
|
||||
LabelSet,
|
||||
describe,
|
||||
parse_names_field,
|
||||
placeholder_labels,
|
||||
read_labels,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Sentinel: worker still on the previous frame — caller must NOT count a miss.
|
||||
INFER_BUSY = object()
|
||||
|
||||
@@ -62,7 +79,7 @@ def resolve_weights(weights: Union[str, Path], backend: str = "auto") -> Path:
|
||||
return engine
|
||||
raise FileNotFoundError(
|
||||
f"TensorRT engine missing: {engine}. "
|
||||
f"Export: bash ~/ros2/scripts/export_yolo_trt.sh"
|
||||
f"Export: bash ~/ros2/scripts/weights.sh convert {p} --name {p.stem}"
|
||||
)
|
||||
if b == "onnx":
|
||||
if onnx.is_file():
|
||||
@@ -78,6 +95,50 @@ def resolve_weights(weights: Union[str, Path], backend: str = "auto") -> Path:
|
||||
return p
|
||||
|
||||
|
||||
def resolve_class_labels(path: Union[str, Path], model=None) -> LabelSet:
|
||||
"""给出一个权重该用的类别标签（旁挂文件优先）。
|
||||
|
||||
``model`` 是已经加载好的 Ultralytics 模型对象（可为 None）。顺序：
|
||||
|
||||
1. ``<模型名>.labels.yaml`` —— 用户可编辑的权威来源
|
||||
2. 传入的 ``model.names`` —— 模型内嵌 metadata
|
||||
3. 占位符 ``class_0 / class_1 …``
|
||||
|
||||
第 3 步刻意不用 ``str(cid)``。原先读不到名字时会退化成类名 `"0"`，
|
||||
和"真的有个叫 0 的类"无法区分，标签丢了都看不出来。
|
||||
占位符 + 一条 warning 让这件事在日志里可见。
|
||||
"""
|
||||
p = Path(path).expanduser()
|
||||
|
||||
# 先只看旁挂文件（快，不加载模型）：它必须能覆盖模型内嵌的旧值，
|
||||
# 否则"我改了标签为什么不生效"会变成很难查的问题。
|
||||
labels = read_labels(p, allow_model_metadata=False)
|
||||
if not labels.is_placeholder:
|
||||
return labels
|
||||
|
||||
raw = getattr(model, "names", None) if model is not None else None
|
||||
names = parse_names_field(raw or {})
|
||||
if names:
|
||||
from .model_labels import SOURCE_MODEL
|
||||
|
||||
return LabelSet(names=names, source=SOURCE_MODEL, path=p)
|
||||
|
||||
nc = 0
|
||||
try:
|
||||
nc = int(getattr(model, "nc", 0) or 0)
|
||||
except (TypeError, ValueError):
|
||||
nc = 0
|
||||
labels = LabelSet(names=placeholder_labels(nc), source=SOURCE_PLACEHOLDER, path=None)
|
||||
logger.warning(
|
||||
"模型 %s 没有类别标签，暂时用占位名显示。补一个 %s 即可（改完重跑节点生效）。"
|
||||
"命令行: bash ~/ros2/scripts/weights.sh labels %s --from-model",
|
||||
p.name,
|
||||
p.stem + ".labels.yaml",
|
||||
p.stem,
|
||||
)
|
||||
return labels
|
||||
|
||||
|
||||
class YoloDetector:
|
||||
def __init__(
|
||||
self,
|
||||
@@ -132,8 +193,8 @@ class YoloDetector:
|
||||
if class_names
|
||||
else None
|
||||
)
|
||||
names = getattr(self._model, "names", None) or {}
|
||||
self._id_to_name = {int(k): str(v) for k, v in dict(names).items()}
|
||||
self._labels = resolve_class_labels(path, self._model)
|
||||
self._id_to_name = dict(self._labels.names)
|
||||
self._warmed = False
|
||||
self._pending = None
|
||||
|
||||
@@ -209,7 +270,8 @@ class YoloDetector:
|
||||
clss = boxes.cls.detach().cpu().numpy().astype(int)
|
||||
for i in range(len(xyxy)):
|
||||
cid = int(clss[i])
|
||||
name = self._id_to_name.get(cid, str(cid))
|
||||
# 用 LabelSet.get：读不到名字时给 class_N 占位，而不是退化成裸数字 "0"。
|
||||
name = self._labels.get(cid)
|
||||
if self.class_ids is not None and cid not in self.class_ids:
|
||||
continue
|
||||
if self.class_names is not None and name.lower() not in self.class_names:
|
||||
|
||||
@@ -1,5 +1,9 @@
|
||||
# perception_node defaults. Launch overlays ~/ros2/config/site.yaml → yolo: / paths:
|
||||
# Do not edit this file for day-to-day tuning.
|
||||
# perception_node 的默认值。launch 会用 ~/ros2/config/site.yaml 的 yolo: / paths: 覆盖。
|
||||
# 日常调参**不要改这里**，改 site.yaml。
|
||||
#
|
||||
# 注意：下面 weights / conf / iou / imgsz / device / half_precision / min_conf 这几项
|
||||
# 一定被 site.yaml 覆盖，这里只是"没有 site.yaml 时的兜底"。所以它们必须和 site.yaml
|
||||
# 保持一致 —— 曾经这里写死 imgsz=416 而实际模型是 640，一旦覆盖失效就会用错尺寸。
|
||||
|
||||
perception_node:
|
||||
ros__parameters:
|
||||
@@ -8,13 +12,13 @@ perception_node:
|
||||
image_topic: ""
|
||||
camera_info_topic: ""
|
||||
depth_topic: "/d405/aligned_depth_to_color/image_raw"
|
||||
weights: "/home/cc/ros2/models/yolo/detect_yolo26_gongjian.pt"
|
||||
conf: 0.35
|
||||
weights: "/home/cc/ros2/models/yolo/detect_best_0.onnx"
|
||||
conf: 0.10
|
||||
iou: 0.45
|
||||
imgsz: 416
|
||||
imgsz: 640
|
||||
device: "0"
|
||||
half_precision: true
|
||||
min_conf: 0.35
|
||||
min_conf: 0.12
|
||||
vision_hz: 20.0
|
||||
z_ref_m: 0.30
|
||||
# Large gate + reacquire: fast motion must NOT freeze track
|
||||
|
||||
@@ -36,7 +36,7 @@ def _setup(context, *args, **kwargs):
|
||||
if backend in ("two_stage", "hough_roi"):
|
||||
weights = os.environ.get("SITE_HOLE_WORKPIECE_WEIGHTS") or os.environ.get("SITE_YOLO_WEIGHTS") or ""
|
||||
else:
|
||||
weights = os.environ.get("SITE_HOLE_WEIGHTS") or str(root / "models" / "hole_insert" / "best.pt")
|
||||
weights = os.environ.get("SITE_HOLE_WEIGHTS") or str(root / "models" / "hole_insert" / "detect_best_1.onnx")
|
||||
return [
|
||||
LogInfo(msg=f"[hole_detect] cam={cam} backend={backend} weights={weights} cfg={cfg}"),
|
||||
Node(
|
||||
|
||||
@@ -5,18 +5,34 @@ from glob import glob
|
||||
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():
|
||||
entries = [
|
||||
('share/ament_index/resource_index/packages', ['resource/' + package_name]),
|
||||
('share/' + package_name, ['package.xml']),
|
||||
(os.path.join('share', package_name, 'launch'), glob('launch/*.py')),
|
||||
(os.path.join('share', package_name, 'trajectories'), glob('trajectories/*')),
|
||||
(os.path.join('share', package_name, 'launch'), _existing(glob('launch/*.py'))),
|
||||
(os.path.join('share', package_name, 'trajectories'), _existing(glob('trajectories/*'))),
|
||||
]
|
||||
for root, _dirs, files in os.walk('config'):
|
||||
if not files:
|
||||
srcs = _existing([os.path.join(root, name) for name in files])
|
||||
if not srcs:
|
||||
continue
|
||||
dest = os.path.join('share', package_name, root)
|
||||
srcs = [os.path.join(root, name) for name in files]
|
||||
entries.append((dest, srcs))
|
||||
return entries
|
||||
|
||||
|
||||
@@ -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([])
|
||||
Reference in New Issue
Block a user