arm_dev
6edb8d39ab
在开始'设备契约层 + 目录重构 + 通止规装配'改造之前，先钉住当前可运行状态， 使后续任何删除/重构都可回退。 本基线包含： - 现有 ROS 2 工作区全部源码（arm_application / arm_control / arm_bringup 等） - 官方 flexiv_ros2 栈（含 flexiv_hardware / flexiv_controllers 等符号链接） - 设备契约层 src/arm_interfaces（新增，46 项单测通过） - devices/ 设备选型与品牌清单（新增） - scripts/gen_site_example.py 与由它生成的 config/site.example.yaml 按 .gitignore 排除：build/ / install/ / log/ / config/site.yaml / 模型权重二进制。
91 lines
2.9 KiB
Bash
Executable File
91 lines
2.9 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# Export YOLO .pt/.onnx → TensorRT .engine (FP16).
|
|
#
|
|
# 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.
|
|
|
|
set -eo pipefail
|
|
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
|
# shellcheck disable=SC1091
|
|
source "${ROOT}/scripts/lib/load_site_config.sh"
|
|
|
|
PT="${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
|
|
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
|
|
|
|
echo "Export TensorRT engine from ${PT} imgsz=${IMGSZ}"
|
|
"${PY}" - <<PY
|
|
from pathlib import Path
|
|
import tensorrt as trt
|
|
|
|
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"
|