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

痛点:TensorRT 10.13 的 Python API 没有元数据读写,engine 一转就把类名丢了;
推理侧只能退化成 str(cid)("0"、"1"),界面上看不出是什么类,属于静默失效。
同时旧的 export_yolo_trt.sh 只能手工搬文件、手工改 site.yaml、标签无处安放。

标签子系统(model_labels.py)
- 旁挂标签文件 <模型名>.labels.yaml,放在权重旁边,文本、可手改、进 git
  (二进制权重不进 git,标签是配置所以进)
- 读取优先级:旁挂文件 > ONNX 元数据 > 模型元数据 > 占位符 class_N
  占位符是故意的:缺标签时要一眼看出来,而不是悄悄退回 cid
- 解析器容忍多种写法:dict / list / 逗号串 / 逐行文本 / Python repr 串

布局与转换(weights_layout.py + weights.py + weights.sh)
- models/<分组>/<模型名>.{pt,onnx,engine} + 同名 .labels.yaml
- convert:PT→ONNX→TensorRT,支持自定义名字、自动归位、固化标签
- staging/:非标准格式先进这里,转换后源文件归档到 staging/done/
- list / labels / doctor / show 四个子命令
- doctor --check-engines 能查出 engine 与 TensorRT 版本不匹配(反序列化返回空)

推理侧接入
- yolo_detector / onnx_workpiece / hole_detect_hough 统一走 read_labels()
- 画框标签不再硬编码 "gongjian",改用实际类名

顺手修掉的真实问题
- setup.py 撞上断掉的符号链接会让整包构建失败,报错还指向 build/ 里的幽灵文件
  (源文件被挪走后 build/ 里留下的断链被 os.walk 列出来)。加 _existing() 过滤。
- P3 遗漏:site.yaml / export_site_env.py / run_handeye.sh / handeye_teach.py /
  handeye_calibrate.py 里的 handeye 工作目录仍是 /tmp/handeye_calib*,
  与已迁到 handeye/ 的其余路径不一致。统一到 handeye/tmp/<臂号>。
- site.yaml 的 hole_weights / hole_detect.weights 原指向不存在的 best.pt,
  改为实际存在的 detect_best_1.onnx;active.pt 符号链接已废弃并清理。
- weights.sh labels 原先不接受 "分组/模型名" 写法,而 list 输出正是这个形式,
  照着敲会报"找不到模型"。补上,并让报错直接给出可复制的全名。

测试:225 项通过(新增 test_model_labels.py / test_weights_layout.py)。
真实 ONNX→engine 端到端转换已验过。

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
arm_dev
2026-09-23 12:39:18 +08:00
parent adc72955c7
commit 2b77ec127d
33 changed files with 3797 additions and 152 deletions
+5 -1
View File
@@ -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
View File
@@ -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
+7 -7
View File
@@ -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"),
+1 -1
View File
@@ -75,7 +75,7 @@ WP="${SITE_HANDEYE_WAYPOINTS:-${HE_DIR}/poses/89_eye_in_hand.yaml}"
CAL="${SITE_HANDEYE_CALIBRATE:-${HE_DIR}/calibrate.yaml}"
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}"
+8 -5
View File
@@ -13,8 +13,11 @@ parse_arm_args "$@"
python3 "${ROOT}/scripts/lib/export_site_env.py" --print ${ROBOT_SN:+--sn "$ROBOT_SN"}
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。)"
+842
View File
@@ -0,0 +1,842 @@
#!/usr/bin/env python3
"""权重模块的统一命令行入口:转换 / 标签 / 列表 / 自检。
通常不直接调用本文件,而是走薄封装(它会先把环境准备好):
bash ~/ros2/scripts/weights.sh list
bash ~/ros2/scripts/weights.sh labels detect_best_0
bash ~/ros2/scripts/weights.sh labels detect_best_0 --set "工件,通止规"
bash ~/ros2/scripts/weights.sh convert staging/best_v3.pt --name workpiece --group yolo
把"纯逻辑"放在 ``arm_application.weights_layout`` / ``arm_application.model_labels``
里,本文件只负责:解析参数、打印人话、真正调用 PyTorch/Ultralytics/TensorRT。
为什么单独做这个模块:原先只有 ``export_yolo_trt.sh``,只能"就地"把
``<同名>.onnx`` 变成 ``<同名>.engine``,既不能改名、不能归位、也不管标签。
实际用起来必须手工搬文件 + 手改 site.yaml,而**标签在 engine 那一环会静默丢失**
(原因见 model_labels 模块头)。这里把三件事一次做完:
转换 → 规范改名 → 放到权重目录 → 顺带把标签落成可编辑的文件。
"""
from __future__ import annotations
import argparse
import json
import os
import shutil
import sys
from pathlib import Path
from typing import Dict, List, Optional, Sequence
# ---------------------------------------------------------------------------
# 让脚本能在没 colcon build / 没 source install 的情况下直接 import
# ---------------------------------------------------------------------------
def _bootstrap_import_path() -> Optional[Path]:
"""把 ``src/arm_application`` 插进 sys.path,返回仓库根(找不到返回 None)。
``arm_application/__init__.py`` 里只有文档字符串、没有 ROS 依赖,
所以这样 import 不会把 ROS 拖进来。
"""
here = Path(__file__).resolve()
for cand in [here, *here.parents]:
pkg = cand / "src" / "arm_application"
if (pkg / "arm_application" / "weights_layout.py").is_file():
if str(pkg) not in sys.path:
sys.path.insert(0, str(pkg))
return cand
return None
_REPO_ROOT = _bootstrap_import_path()
from arm_application.model_labels import ( # noqa: E402
LABELS_SUFFIX,
SOURCE_LABELS_ZH,
SOURCE_PLACEHOLDER,
SOURCE_SIDECAR,
describe,
format_labels_yaml,
labels_from_model,
labels_from_onnx,
labels_path_for,
parse_label_list,
read_label_text_file,
read_labels,
write_labels,
)
from arm_application.weights_layout import ( # noqa: E402
CONVERTIBLE_INPUT_EXTS,
STAGING_DIRNAME,
WEIGHT_EXTS,
ModelEntry,
canonical_path,
conversion_chain,
default_weights_root,
discover,
find_model,
groups,
model_labels,
plan_conversion,
staging_dir,
staging_files,
validate_name,
)
# 默认转换输入尺寸。site.yaml 的 yolo.imgsz 优先(由 weights.sh 通过环境变量传入)。
DEFAULT_IMGSZ = 640
# ---------------------------------------------------------------------------
# 小工具
# ---------------------------------------------------------------------------
def _root(args) -> Path:
return Path(args.root).expanduser() if getattr(args, "root", None) else default_weights_root()
def _imgsz(args) -> int:
"""输入尺寸:命令行 > site.yaml(环境变量) > 默认 640。"""
if getattr(args, "imgsz", None):
return int(args.imgsz)
env = os.environ.get("SITE_YOLO_IMGSZ")
if env:
try:
return int(env)
except ValueError:
pass
return DEFAULT_IMGSZ
def _err(msg: str) -> None:
print(f"错误: {msg}", file=sys.stderr)
def _ok(msg: str) -> None:
print(f"✓ {msg}")
def _warn(msg: str) -> None:
print(f"! {msg}")
def _hr(char: str = "─", n: int = 74) -> None:
print(char * n)
def _print_table(headers: Sequence[str], rows: Sequence[Sequence[str]]) -> None:
"""左对齐打印一张表。宽度按显示宽度(东亚字符按 2 算)估算。"""
def w(s: str) -> int:
return sum(2 if ord(c) > 0x2E80 else 1 for c in s)
widths = [w(h) for h in headers]
for r in rows:
for i, c in enumerate(r):
widths[i] = max(widths[i], w(c))
def pad(s: str, width: int) -> str:
return s + " " * max(0, width - w(s))
print(" ".join(pad(h, widths[i]) for i, h in enumerate(headers)))
print(" ".join("-" * widths[i] for i in range(len(headers))))
for r in rows:
print(" ".join(pad(c, widths[i]) for i, c in enumerate(r)))
def _mode_of(p: Path) -> str:
"""文件可读性/大小的人话描述,用于列表显示。"""
try:
return f"{p.stat().st_size / 1e6:.1f}MB"
except OSError:
return "?"
# ---------------------------------------------------------------------------
# list
# ---------------------------------------------------------------------------
def cmd_list(args) -> int:
root = _root(args)
if not root.is_dir():
_err(f"权重目录不存在: {root}")
return 2
entries = discover(root)
if not entries:
_warn(f"{root} 下没找到任何权重文件(后缀: {', '.join(WEIGHT_EXTS)})")
return 0
payload = []
rows = []
for e in entries:
ls = model_labels(e, deep=args.deep)
if ls.source == SOURCE_PLACEHOLDER:
label_cell = "— 缺标签"
else:
label_cell = f"{', '.join(ls.sorted_names())}"
rows.append(
[
e.group,
e.name,
" ".join(e.exts) or "—",
label_cell,
SOURCE_LABELS_ZH.get(ls.source, ls.source),
]
)
payload.append(
{
"group": e.group,
"name": e.name,
"formats": e.exts,
"files": {k: str(v) for k, v in e.files.items()},
"labels": ls.names,
"labels_source": ls.source,
"labels_path": str(ls.path) if ls.path else None,
}
)
if args.json:
print(json.dumps(payload, ensure_ascii=False, indent=2))
return 0
print(f"权重目录: {root}")
_hr()
_print_table(["分组", "模型名", "已有格式", "标签", "标签来源"], rows)
_hr()
missing = [e for e in entries if model_labels(e).source == SOURCE_PLACEHOLDER]
print(f"共 {len(entries)} 个模型,其中 {len(missing)} 个缺标签。")
if missing and not args.deep:
print("提示: 加 --deep 可再试着从 .pt/.engine 自身读标签(会加载模型,慢)。")
if missing:
print("补标签: bash scripts/weights.sh labels <模型名> --from-model")
return 0
# ---------------------------------------------------------------------------
# labels
# ---------------------------------------------------------------------------
def _resolve_entry(args, root: Path) -> Optional[ModelEntry]:
entry = find_model(root, args.model, getattr(args, "group", None))
if entry is None:
# 提示写成 ``分组/模型名``:这正是 ``list`` 里显示的样子,可以直接复制粘贴,
# 比只列裸名字更不容易再踩一次坑。
avail = ", ".join(f"{e.group}/{e.name}" for e in discover(root)) or "(无)"
_err(
f"找不到模型 '{args.model}'。现有模型(可整串复制): {avail}\n"
f" 也可以只写裸名字,或加 --group 指定分组。全部列出: "
f"bash ~/ros2/scripts/weights.sh list"
)
return entry
def cmd_labels(args) -> int:
root = _root(args)
entry = _resolve_entry(args, root)
if entry is None:
return 2
# 修改类操作 ------------------------------------------------------------
if args.clear:
p = entry.labels_path
if p.is_file():
p.unlink()
_ok(f"已删除标签文件 {p}")
else:
_warn(f"本来就没有标签文件 {p}")
return 0
if args.set is not None and args.set_file:
_err("--set 与 --set-file 不能同时用")
return 2
new_names: Optional[Dict[int, str]] = None
if args.set is not None:
new_names = parse_label_list(args.set)
if not new_names:
_err("--set 解析不出任何类别名")
return 2
elif args.set_file:
src = Path(args.set_file).expanduser()
if not src.is_file():
_err(f"标签文件不存在: {src}")
return 2
new_names = read_label_text_file(src)
if not new_names:
_err(f"{src} 里没读到任何类别名(一行一个)")
return 2
if new_names is None and args.from_model:
# 从模型内嵌 metadata 抄一份出来,目的是"变成可编辑的旁挂文件"。
for p in entry.sorted_files():
got = None
if p.suffix.lower() == ".onnx":
from arm_application.model_labels import labels_from_onnx
got = labels_from_onnx(p)
if got is None:
got = labels_from_model(p)
if got:
new_names = got[0]
print(f"从 {p.name} 读到 {len(new_names)} 个类别名")
break
if new_names is None:
_err(f"{entry.name} 的模型文件里读不到标签,只能用 --set 手工指定")
return 2
if new_names is not None:
dst = write_labels(entry.path(entry.exts[0]) or entry.labels_path, new_names)
_ok(f"已写入 {dst}")
print(f" 类别({len(new_names)}): {', '.join(new_names[k] for k in sorted(new_names))}")
return 0
# 读取(默认动作)------------------------------------------------------
ls = read_labels(entry.path(entry.exts[0]), allow_model_metadata=args.deep) if entry.exts else None
if ls is None or ls.is_placeholder:
ls = model_labels(entry, deep=args.deep)
if args.json:
print(
json.dumps(
{
"model": entry.name,
"group": entry.group,
"labels": ls.names,
"source": ls.source,
"source_zh": SOURCE_LABELS_ZH.get(ls.source, ls.source),
"source_file": str(ls.path) if ls.path else None,
"expected_labels_path": str(entry.labels_path),
},
ensure_ascii=False,
indent=2,
)
)
return 0
print(f"模型: {entry.group}/{entry.name}")
print(f"格式: {' '.join(entry.exts) or '—'}")
_hr()
if ls.is_placeholder:
_warn(f"没有标签。{describe(ls)}")
print(f" 期望的标签文件位置: {entry.labels_path}")
print(" 修复方式(三选一):")
print(f" 1) 从模型内嵌 metadata 抄一份出来再改: weights.sh labels {entry.name} --from-model")
print(f' 2) 直接指定类别名: weights.sh labels {entry.name} --set "工件,通止规"')
print(f" 3) 手工创建 {entry.labels_path.name}")
return 1
print(f"标签来源: {SOURCE_LABELS_ZH.get(ls.source, ls.source)}")
if ls.path:
print(f"来源文件: {ls.path}")
_hr()
for k in sorted(ls.names):
print(f" {k}: {ls.names[k]}")
return 0
# ---------------------------------------------------------------------------
# convert
# ---------------------------------------------------------------------------
def _export_onnx(src: Path, imgsz: int, dst: Path, *, simplify: bool = True) -> Path:
"""YOLO .pt/.pth → ONNX,直接产出到 ``dst``。
**关键:先把 .pt 复制到临时目录再导出。**
Ultralytics 会把 onnx 写在**源文件旁边**(同名同目录),我们对源 .pt 直接导出时,
会把它旁边已有的 ``<同名>.onnx`` 静默覆盖掉 —— 那是用户自己的产物,
可能还是不同 imgsz 导出的。更糟的是随后的搬运会把它彻底删掉。
实测就发生过:转换 detect_hc90.pt 时,models/yolo/detect_hc90.onnx 被覆盖后又被搬走。
在临时目录里导出,源目录一个字节都不会动。
"""
import tempfile
os.environ.setdefault("YOLO_AUTOINSTALL", "false")
from ultralytics import YOLO
dst.parent.mkdir(parents=True, exist_ok=True)
with tempfile.TemporaryDirectory(prefix="weights_export_") as td:
work = Path(td) / src.name
shutil.copy2(src, work)
model = YOLO(str(work))
out = Path(
str(
model.export(
format="onnx",
imgsz=int(imgsz),
device=0,
simplify=simplify,
dynamic=False,
)
)
).expanduser()
if not out.is_file():
raise RuntimeError(f"Ultralytics 说导出了 {out},但文件不存在")
# 临时目录马上会被删,所以这里必须是 move。
if out.resolve() != dst.resolve():
shutil.move(str(out), str(dst))
return dst
def _build_engine(onnx_src: Path, dst: Path, *, fp16: bool = True, workspace_gb: int = 4) -> Path:
"""ONNX → TensorRT engine,直接写到 dst。
这里刻意绕开 Ultralytics 的 engine 导出:它对 .onnx 会按发行包名检查
onnxruntime,环境装的是 onnxruntime-gpu 时会触发 pip AutoUpdate(可能卡死)。
直接用 TensorRT 的 Python API,行为可控、错误信息也更清楚。
"""
import tensorrt as trt
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
# TRT 10 里 EXPLICIT_BATCH 已是默认行为,但这个 flag 仍然接受;
# 保持和 export_yolo_trt.sh 一致,避免不同版本行为差异。
flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
network = builder.create_network(flag)
parser = trt.OnnxParser(network, logger)
if not parser.parse(onnx_src.read_bytes()):
msgs = [str(parser.get_error(i)) for i in range(parser.num_errors)]
raise RuntimeError("ONNX 解析失败:\n " + "\n ".join(msgs))
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, int(workspace_gb) << 30)
if fp16 and builder.platform_has_fast_fp16:
config.set_flag(trt.BuilderFlag.FP16)
serialized = builder.build_serialized_network(network, config)
if serialized is None:
raise RuntimeError("TensorRT 构建失败(返回空)。可试着 --no-fp16 或换 imgsz。")
dst.parent.mkdir(parents=True, exist_ok=True)
dst.write_bytes(bytes(serialized))
return dst
def _materialize_labels(src: Path, dest_weights: Path, explicit: Optional[str]) -> Optional[Path]:
"""转换后把标签落到目标权重旁边。
这是防止"engine 丢标签"再次发生的**关键一步**:转换时就把标签固化成
可编辑的旁挂文件,而不是等到运行时才发现读不到名字。
优先级:命令行 --labels > 源文件/中间产物的内嵌 metadata。
已经有旁挂文件时不覆盖(避免把用户手改的内容冲掉)。
"""
dst_labels = labels_path_for(dest_weights)
if dst_labels.is_file():
_warn(f"标签文件已存在,保持不变: {dst_labels.name}")
return dst_labels
if explicit:
names = parse_label_list(explicit)
if names:
return write_labels(dest_weights, names)
for cand in (src, src.with_suffix(".onnx"), dest_weights):
if not cand.is_file():
continue
if cand.suffix.lower() == ".onnx":
from arm_application.model_labels import labels_from_onnx
got = labels_from_onnx(cand)
else:
got = labels_from_model(cand)
if got:
_ok(f"标签来自 {cand.name}: {', '.join(got[0][k] for k in sorted(got[0]))}")
return write_labels(dest_weights, got[0])
return None
def cmd_convert(args) -> int:
root = _root(args)
src = Path(args.input).expanduser()
if not src.is_file():
# 相对路径优先按**权重目录**解析,而不是当前工作目录。
# 否则 `weights.sh convert staging/x.pt` 会去 <cwd>/staging 找,
# 而用户心里想的是 models/staging/x.pt —— 报"文件不存在"很费解。
cands = []
if not src.is_absolute():
cands.append(root / args.input) # models/staging/x.pt
cands.append(staging_dir(root) / args.input) # 只写文件名: models/staging/x.pt
cands.append(Path(args.input)) # 退一步:相对 cwd
found = next((c for c in cands if c.is_file()), None)
if found is None:
_err(
f"输入文件不存在: {args.input}\n"
f" 找过: " + ", ".join(str(c) for c in cands) +
f"\n 把待转换的文件放进 {staging_dir(root)} 再试。"
)
return 2
src = found
name = (args.name or src.stem).strip()
bad = validate_name(name)
if bad:
_err(bad)
return 2
group = args.group
chain, _outs = plan_conversion(root, src, name=name, group=group, target=args.to)
print(f"输入: {src}")
print(f"输出: {root / group}/{name}.<ext>")
print(f"计划: {chain.describe()}")
print(f"尺寸: imgsz={_imgsz(args)}")
if args.to != "all":
print(f"目标: 只产出 {args.to}")
in_staging = src.parent.resolve() == staging_dir(root).resolve()
# 归位路径会把源文件本身搬到目标去;那种情况下它已经不在 staging 了,
# 下面不能再"归档"一次(否则会 FileNotFoundError)。
src_consumed = False
if chain.is_noop:
ext = src.suffix.lower().lstrip(".")
if ext in WEIGHT_EXTS:
dst = canonical_path(root, group, name, ext)
if dst.resolve() == src.resolve():
_warn(f"{src.name} 已经是最终产物,且就在目标位置,无需处理")
return 0
print(f"{ext} 无需转换,直接归位 → {dst}")
if args.dry_run:
return 0
dst.parent.mkdir(parents=True, exist_ok=True)
# 归位本身就是"移动源文件"。若用户要求保留 staging 原件,
# 就改成复制,否则源文件会连同 staging 一起消失。
if in_staging and args.keep_staging:
shutil.copy2(str(src), str(dst))
else:
shutil.move(str(src), str(dst))
src_consumed = True
_ok(f"已归位 {dst}")
_materialize_labels(dst, dst, args.labels)
if in_staging and not src_consumed:
_ok(f"staging 原件保留在 {src}(--keep-staging)")
elif in_staging:
_ok(f"staging 已清空该文件(它已变成 {dst.name})")
return 0
_err(
f"不认识的输入格式 '{src.suffix}'。"
f"可转换的输入: {', '.join('.' + e for e in CONVERTIBLE_INPUT_EXTS)};"
f"已是最终产物的: .engine/.trt/.plan"
)
return 2
if args.dry_run:
print("(--dry-run,未实际执行)")
return 0
# 真正执行 --------------------------------------------------------------
produced: Dict[str, Path] = {}
try:
for ext in chain.produce:
dst = canonical_path(root, group, name, ext)
if dst.exists() and not args.force:
raise FileExistsError(f"{dst} 已存在(加 --force 覆盖)")
if ext == "onnx":
if src.suffix.lower() == ".onnx":
# 源本身就是 onnx:复制而不是移动 —— 用户给的文件不该消失。
dst.parent.mkdir(parents=True, exist_ok=True)
shutil.copy2(str(src), str(dst))
else:
_export_onnx(src, _imgsz(args), dst)
produced["onnx"] = dst
_ok(f"ONNX → {dst} ({_mode_of(dst)})")
elif ext == "engine":
onnx_in = produced.get("onnx") or src
_build_engine(onnx_in, dst, fp16=not args.no_fp16)
produced["engine"] = dst
_ok(f"engine→ {dst} ({_mode_of(dst)})")
except Exception as exc: # noqa: BLE001
_err(f"转换失败: {type(exc).__name__}: {exc}")
return 1
# 标签 ------------------------------------------------------------------
anchor = produced.get("onnx") or produced.get("engine") or src
lp = _materialize_labels(src, anchor, args.labels)
if lp:
_ok(f"标签 → {lp}")
else:
_warn("没有可用的标签(源模型里也没有)。请手动补:")
print(f" bash scripts/weights.sh labels {name} --set \"类别1,类别2\"")
# staging 收尾 ----------------------------------------------------------
if in_staging and not args.keep_staging and not src_consumed:
done = staging_dir(root) / "done"
done.mkdir(parents=True, exist_ok=True)
target = done / src.name
# 同名已存在时加个序号,避免静默覆盖上一次的源文件。
n = 1
while target.exists():
target = done / f"{src.stem}.{n}{src.suffix}"
n += 1
shutil.move(str(src), str(target))
_ok(f"已把 staging 里的源文件挪到 {target}(加 --keep-staging 可保留)")
print()
_ok(f"完成。共产出 {len(produced)} 个文件。")
print(f"下一步: 把 config/site.yaml 的 paths.yolo_weights 指向 {anchor}")
print(" 或直接跑 bash scripts/weights.sh list 确认标签读得出来")
return 0
# ---------------------------------------------------------------------------
# doctor
# ---------------------------------------------------------------------------
def _read_site_yaml() -> Dict:
if _REPO_ROOT is None:
return {}
p = _REPO_ROOT / "config" / "site.yaml"
if not p.is_file():
return {}
try:
import yaml
return yaml.safe_load(p.read_text(encoding="utf-8")) or {}
except Exception: # noqa: BLE001
return {}
def _engine_usable(p: Path) -> Optional[str]:
"""试着反序列化 engine。返回 None 表示可用,否则返回错误摘要。
引擎和 TensorRT 版本是绑死的:换过 TensorRT 之后旧引擎会直接加载失败。
这个检查就是为了把这种"静默失效"提前暴露出来。
"""
try:
import tensorrt as trt
except Exception as exc: # noqa: BLE001
return f"无法 import tensorrt: {exc}"
try:
rt = trt.Runtime(trt.Logger(trt.Logger.ERROR))
eng = rt.deserialize_cuda_engine(p.read_bytes())
if eng is None:
return "反序列化返回空(多半是 TensorRT 版本不匹配)"
return None
except Exception as exc: # noqa: BLE001
msg = str(exc)
if "incompatible serialization version" in msg or "magicTag" in msg:
return "TensorRT 版本不匹配(需重新导出:weights.sh convert <onnx> --name ...)"
return f"{type(exc).__name__}: {msg[:160]}"
def cmd_doctor(args) -> int:
root = _root(args)
problems = 0
print(f"权重目录: {root}")
_hr("═")
# 1) 逐模型体检 --------------------------------------------------------
entries = discover(root)
print(f"[1/4] 模型({len(entries)} 个)")
for e in entries:
ls = model_labels(e, deep=args.deep)
tag = "OK " if ls.source != SOURCE_PLACEHOLDER else "缺标签"
if ls.source == SOURCE_PLACEHOLDER:
problems += 1
print(f" [{tag}] {e.group}/{e.name} 格式: {' '.join(e.exts)}")
if ls.source == SOURCE_PLACEHOLDER:
print(f" 期望: {e.labels_path.name}")
# 只有 .pt 没有 onnx/engine 时,two_stage/onnx 后端会用不了
if "onnx" not in e.files and any(x in e.files for x in ("engine", "trt")):
print(" 注意: 只有 engine 没有 onnx,CPU/ORT 后端无处可用")
# 2) engine 版本 -------------------------------------------------------
print()
print("[2/4] engine 可用性(TensorRT 版本绑定检查)")
engines = [
p for e in entries for ext in ("engine", "trt") if (p := e.path(ext)) is not None
]
if not engines:
print(" (没有 engine)")
for p in engines:
if args.check_engines:
err = _engine_usable(p)
if err:
problems += 1
print(f" [失效] {p.name}: {err}")
else:
print(f" [可用] {p.name}")
else:
print(f" [跳过] {p.name} (加 --check-engines 实测)")
# 3) 配置引用 ----------------------------------------------------------
print()
print("[3/4] config/site.yaml 的引用")
site = _read_site_yaml()
paths = site.get("paths") or {}
refs: List[tuple] = []
if paths.get("yolo_weights"):
refs.append(("paths.yolo_weights", paths["yolo_weights"]))
if paths.get("hole_weights"):
refs.append(("paths.hole_weights", paths["hole_weights"]))
if paths.get("hole_workpiece_weights"):
refs.append(("paths.hole_workpiece_weights", paths["hole_workpiece_weights"]))
hd = site.get("hole_detect") or {}
if hd.get("weights"):
refs.append(("hole_detect.weights", hd["weights"]))
if hd.get("workpiece_weights"):
refs.append(("hole_detect.workpiece_weights", hd["workpiece_weights"]))
if not refs:
print(" (没读到引用)")
for key, val in refs:
p = Path(str(val)).expanduser()
if not p.is_file():
if p.suffix == ".pt" and p.with_suffix(".onnx").is_file() is False:
problems += 1
print(f" [悬空] {key} → {p} (文件不存在)")
else:
problems += 1
print(f" [悬空] {key} → {p} (文件不存在)")
continue
ls = read_labels(p, allow_model_metadata=False)
origin = SOURCE_LABELS_ZH.get(ls.source, ls.source)
print(f" [OK ] {key} → {p.name} 标签: {origin}")
if ls.is_placeholder:
problems += 1
print(f" → 该模型读不到标签,运行时类名会变成 class_0")
backend = (site.get("yolo") or {}).get("backend")
if backend == "tensorrt":
yw = Path(str(paths.get("yolo_weights") or "")).expanduser()
if yw.name and not yw.with_suffix(".engine").is_file():
problems += 1
print(
f" [风险] yolo.backend=tensorrt 但 {yw.with_suffix('.engine').name} 不存在,"
"启动会直接失败"
)
# 4) staging -----------------------------------------------------------
print()
print("[4/4] staging(待转换)")
st = staging_files(root)
if not st:
print(" (空)")
for p in st:
ext = p.suffix.lower().lstrip(".")
if ext in CONVERTIBLE_INPUT_EXTS:
print(f" [待转换] {p.name} → weights.sh convert {STAGING_DIRNAME}/{p.name} --name <模型名>")
else:
print(f" [需人工] {p.name} 格式 .{ext} 不在支持列表({', '.join(CONVERTIBLE_INPUT_EXTS)})")
_hr("═")
if problems:
print(f"发现 {problems} 处问题。")
return 1
print("没有发现问题。")
return 0
# ---------------------------------------------------------------------------
# show:把模块约定直接打出来,省得翻文档
# ---------------------------------------------------------------------------
def cmd_show(args) -> int:
root = _root(args)
print("权重模块约定")
_hr()
print(f"权重根目录 : {root}")
print(f"仓库根 : {_REPO_ROOT or '(未找到)'}")
print(f"待转换目录 : {staging_dir(root)}")
print(f"标签文件后缀 : {LABELS_SUFFIX} (<模型名>{LABELS_SUFFIX},放权重旁边,进 git)")
print(f"权重后缀 : {', '.join('.' + e for e in WEIGHT_EXTS)} (二进制,不进 git)")
print(f"已存在的分组 : {', '.join(groups(root)) or '(无)'}")
print()
print("命令")
_hr()
print(" weights.sh list [--deep] [--json]")
print(" weights.sh labels <模型名> [--group G] [--set a,b,c] [--set-file f.txt]")
print(" [--from-model] [--clear] [--deep] [--json]")
print(" weights.sh convert <文件> --name <模型名> [--group yolo] [--to all|onnx|engine]")
print(" [--imgsz 640] [--labels a,b] [--force] [--dry-run]")
print(" [--no-fp16] [--keep-staging]")
print(" weights.sh doctor [--deep] [--check-engines]")
print()
print("标签优先级: 旁挂文件 > ONNX 内嵌 metadata > 模型内嵌 metadata > 占位符(class_0)")
return 0
# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------
def build_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(
prog="weights.sh",
description="模型权重模块:转换(PT→ONNX→TensorRT)、标签读写、列表、自检",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=(
"示例:\n"
' weights.sh list\n'
' weights.sh labels detect_best_0\n'
' weights.sh labels detect_best_0 --set "工件,通止规"\n'
' weights.sh convert staging/best_v3.pt --name workpiece\n'
' weights.sh convert models/yolo/detect_best_0.onnx --name detect_best_0 --to engine\n'
),
)
p.add_argument("--root", help="权重根目录(默认自动探测 models/)")
sub = p.add_subparsers(dest="cmd", required=True)
s = sub.add_parser("list", help="列出所有权重及其标签")
s.add_argument("--deep", action="store_true", help="也尝试从 .pt/.engine 自身读标签(慢)")
s.add_argument("--json", action="store_true", help="以 JSON 输出")
s.set_defaults(func=cmd_list)
s = sub.add_parser("labels", help="查看/修改模型标签")
s.add_argument("model", help="模型名(权重目录里的文件名主干)")
s.add_argument("--group", help="限定分组")
s.add_argument("--set", help='直接指定类别名,如 "工件,通止规"(改了立刻生效)')
s.add_argument("--set-file", dest="set_file", help="从纯文本文件导入(一行一个类名)")
s.add_argument("--from-model", dest="from_model", action="store_true",
help="从模型内嵌 metadata 抄一份到旁挂文件(变成可编辑)")
s.add_argument("--clear", action="store_true", help="删除旁挂标签文件")
s.add_argument("--deep", action="store_true", help="读取时也加载 .pt/.engine 自身")
s.add_argument("--json", action="store_true", help="以 JSON 输出")
s.set_defaults(func=cmd_labels)
s = sub.add_parser("convert", help="转换权重并归位")
s.add_argument("input", help="输入文件(可以是 staging/ 里的相对名)")
s.add_argument("--name", help="输出模型名(默认沿用输入文件名主干)")
s.add_argument("--group", default="yolo", help="输出分组子目录(默认 yolo)")
s.add_argument("--to", default="all", choices=["all", "onnx", "engine"], help="产出目标")
s.add_argument("--imgsz", type=int, help="输入尺寸(默认取 site.yaml 的 yolo.imgsz)")
s.add_argument("--labels", help='顺带写标签,如 "工件,通止规"')
s.add_argument("--force", action="store_true", help="覆盖已存在的产物")
s.add_argument("--no-fp16", dest="no_fp16", action="store_true", help="engine 不用 FP16")
s.add_argument("--dry-run", dest="dry_run", action="store_true", help="只打印计划")
s.add_argument("--keep-staging", dest="keep_staging", action="store_true",
help="转换后不把 staging 里的源文件挪走")
s.set_defaults(func=cmd_convert)
s = sub.add_parser("doctor", help="自检:缺标签 / 悬空引用 / engine 失效")
s.add_argument("--deep", action="store_true", help="也加载 .pt/.engine 读标签")
s.add_argument("--check-engines", dest="check_engines", action="store_true",
help="实测反序列化每个 engine(探测 TensorRT 版本不匹配)")
s.set_defaults(func=cmd_doctor)
s = sub.add_parser("show", help="打印模块约定与命令")
s.set_defaults(func=cmd_show)
return p
def main(argv: Optional[Sequence[str]] = None) -> int:
args = build_parser().parse_args(argv)
try:
return int(args.func(args))
except KeyboardInterrupt:
_err("被中断")
return 130
if __name__ == "__main__":
sys.exit(main())
+69
View File
@@ -0,0 +1,69 @@
#!/usr/bin/env bash
# 权重模块的统一入口(薄封装)。
#
# bash ~/ros2/scripts/weights.sh list
# bash ~/ros2/scripts/weights.sh labels detect_best_0
# bash ~/ros2/scripts/weights.sh labels detect_best_0 --set "工件,通止规"
# bash ~/ros2/scripts/weights.sh convert staging/best_v3.pt --name workpiece
# bash ~/ros2/scripts/weights.sh doctor --check-engines
#
# 本脚本只做三件事:准备环境(ROS + TensorRT 动态库)→ 挑一个能用的 python → 转发。
# 真正的逻辑在 scripts/weights.py。
#
# 完全等价的老命令 export_yolo_trt.sh 仍然可用(它是本脚本的转发壳)。
set -eo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
# 用户显式指定解释器时优先用它;否则试 ros_flexiv。
# 本模块需要的重依赖(ultralytics / onnx / tensorrt / onnxruntime)都在那里。
pick_python() {
if [[ -n "${WEIGHTS_PY:-}" ]]; then
echo "${WEIGHTS_PY}"; return 0
fi
local cands=(
"${HOME}/anaconda3/envs/ros_flexiv/bin/python"
"$(command -v python3 || true)"
)
for c in "${cands[@]}"; do
[[ -n "${c}" && -x "${c}" ]] || continue
# list/labels/show 只需要 yaml;convert/doctor 才需要 tensorrt。
# 这里统一要求 yaml,避免半路才发现解释器不对。
if "${c}" -c 'import yaml' >/dev/null 2>&1; then
echo "${c}"; return 0
fi
done
return 1
}
PY="$(pick_python)" || {
echo "错误: 找不到可用的 python(需要能 import yaml)。" >&2
echo " 可用 WEIGHTS_PY=/path/to/python 指定。" >&2
exit 1
}
# TensorRT 的 pip 版把 .so 放在非标准位置,TensorRT 的 Python API 和运行时都要用。
# export_yolo_trt.sh 里也是这么做的;这里保持同样做法。
TRT_LIBS="$("${PY}" - <<'PY' 2>/dev/null || true
import pathlib
try:
import tensorrt
except Exception:
raise SystemExit(0)
p = pathlib.Path(tensorrt.__file__).resolve().parent
cands = [p, p.parent / "tensorrt_libs", p.parent / "nvidia" / "tensorrt" / "lib"]
print(":".join(str(x) for x in cands if x.exists()))
PY
)"
[[ -n "${TRT_LIBS}" ]] && export LD_LIBRARY_PATH="${TRT_LIBS}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
# 不让 Ultralytics 自己去 pip 装东西(会卡住或破坏环境)。
export YOLO_AUTOINSTALL=false
# site.yaml 里的 yolo.imgsz 作为 convert 的默认尺寸。
if [[ -f "${ROOT}/scripts/lib/load_site_config.sh" ]]; then
# shellcheck disable=SC1091
source "${ROOT}/scripts/lib/load_site_config.sh" >/dev/null 2>&1 || true
fi
exec "${PY}" "${ROOT}/scripts/weights.py" "$@"