refactor(umi): 延迟导入外部依赖，增强 xvlib 设备接口
This commit is contained in:
Vinman
parent
57e57358b6
commit
15d9eae9a1
@ -14,3 +14,12 @@ _lerobot_teleoperators.make_teleoperator_from_config = _uf_make_teleoperator_fro
|
||||
_lerobot_cameras_utils.make_cameras_from_configs = _uf_make_cameras_from_configs
|
||||
_lerobot_robot_utils.make_robot_from_config = _uf_make_robot_from_config
|
||||
_lerobot_teleoperators_utils.make_teleoperator_from_config = _uf_make_teleoperator_from_config
|
||||
# register plugin
|
||||
import ufactory_lerobot.cameras.umi_camera
|
||||
import ufactory_lerobot.robots.uf_robot
|
||||
import ufactory_lerobot.robots.uf_mock_robot
|
||||
import ufactory_lerobot.teleoperators.uf_mock_teleop
|
||||
import ufactory_lerobot.teleoperators.gello_teleop
|
||||
import ufactory_lerobot.teleoperators.pika_teleop
|
||||
import ufactory_lerobot.teleoperators.space_mouse
|
||||
import ufactory_lerobot.teleoperators.umi_teleop
|
||||
|
||||
@ -19,9 +19,7 @@ Provides the RealSenseCamera class for capturing frames from Intel RealSense cam
|
||||
import logging
|
||||
import time
|
||||
from typing import Any
|
||||
import cv2 # type: ignore # TODO: add type stubs for OpenCV
|
||||
from lerobot.cameras.camera import Camera
|
||||
from ufactory_lerobot.devices.umi.xvlib import XVLib
|
||||
from .configuration_umi import UmiCameraConfig
|
||||
from lerobot.cameras.configs import ColorMode
|
||||
from lerobot.cameras.utils import get_cv2_rotation
|
||||
@ -53,6 +51,7 @@ class UmiCamera(Camera):
|
||||
self.rotation: int | None = get_cv2_rotation(config.rotation)
|
||||
|
||||
self.last_frame = None
|
||||
from ufactory_lerobot.devices.umi.xvlib import XVLib
|
||||
self.xvlib = XVLib(self.serial_number)
|
||||
self.xvlib.xv_color_camera_init()
|
||||
|
||||
|
||||
@ -1,8 +1,5 @@
|
||||
import time
|
||||
import logging
|
||||
import threading
|
||||
import serial
|
||||
from serial.tools import list_ports
|
||||
|
||||
# 配置日志
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||||
@ -15,6 +12,7 @@ def get_serial_ports(vidpid='1a86:7522'):
|
||||
vidpid: 指定设备的VID:PID字符串, 默认值为'1a86:7522'
|
||||
返回找到的所有符合的串口号列表
|
||||
"""
|
||||
from serial.tools import list_ports
|
||||
ports = list_ports.comports()
|
||||
pika_ports = []
|
||||
for port in ports:
|
||||
@ -34,6 +32,7 @@ def check_pika_device(port):
|
||||
1: Pika Sense设备
|
||||
2: Pika Gripper设备
|
||||
"""
|
||||
import serial
|
||||
try:
|
||||
ser = serial.Serial(
|
||||
port=port,
|
||||
|
||||
@ -79,6 +79,30 @@ class ViveTracker(metaclass=SingletonMeta):
|
||||
if not self._init():
|
||||
raise RuntimeError("Failed to initialize Vive Tracker: pysurvive context creation failed")
|
||||
|
||||
def stop(self):
|
||||
"""停止采集线程 (不释放 context, 可重新 start)"""
|
||||
self.running = False
|
||||
if self._collector_thread is not None:
|
||||
self._collector_thread.join(timeout=2.0)
|
||||
self._collector_thread = None
|
||||
|
||||
def close(self):
|
||||
"""停止采集并释放 pysurvive context。"""
|
||||
self.stop()
|
||||
if self._context is not None:
|
||||
try:
|
||||
self._context.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._context = None
|
||||
|
||||
def __del__(self):
|
||||
# 只做兜底清理，吞掉所有异常避免 __del__ 抛异常导致解释器崩溃
|
||||
try:
|
||||
self.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def __del__(self):
|
||||
logger.info("Stopping Vive Tracker pose tracking...")
|
||||
self.running = False
|
||||
|
||||
Binary file not shown.
@ -3,27 +3,47 @@ import time
|
||||
import threading
|
||||
import ctypes
|
||||
import ctypes.util
|
||||
import cv2
|
||||
import logging
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger('uf.xvlib')
|
||||
|
||||
# cv2 在 .frame() 方法中延迟导入，避免模块顶层强依赖 OpenCV
|
||||
_cv2 = None
|
||||
_rotate_codes = None
|
||||
|
||||
|
||||
def _get_cv2():
|
||||
"""延迟导入 cv2, 避免未安装 opencv 时阻塞模块加载"""
|
||||
global _cv2
|
||||
if _cv2 is None:
|
||||
import cv2 as _cv2_mod
|
||||
_cv2 = _cv2_mod
|
||||
return _cv2
|
||||
|
||||
|
||||
def _get_rotate_codes():
|
||||
"""延迟初始化旋转代码映射表"""
|
||||
global _rotate_codes
|
||||
if _rotate_codes is None:
|
||||
cv2 = _get_cv2()
|
||||
_rotate_codes = {
|
||||
90: cv2.ROTATE_90_CLOCKWISE,
|
||||
-90: cv2.ROTATE_90_COUNTERCLOCKWISE,
|
||||
180: cv2.ROTATE_180,
|
||||
}
|
||||
return _rotate_codes
|
||||
|
||||
|
||||
def _apply_rotate(frame, rotate):
|
||||
"""如果 rotate 不为 None, 对 frame 应用旋转"""
|
||||
if rotate is None or frame is None or (isinstance(rotate, int) and rotate == 0):
|
||||
return frame
|
||||
code = _rotate_codes.get(rotate) if isinstance(rotate, int) else rotate
|
||||
code = _get_rotate_codes().get(rotate) if isinstance(rotate, int) else rotate
|
||||
if code is None:
|
||||
logger.warning(f"Unknown rotate value: {rotate}, skipping")
|
||||
return frame
|
||||
return cv2.rotate(frame, code)
|
||||
return _get_cv2().rotate(frame, code)
|
||||
|
||||
# ============== 数据缓冲区常量 (与 C++ xv_device.h 保持一致) ==============
|
||||
_MAX_COLOR_BUFFER_SIZE = 1280 * 1280 * 3 # MAX_COLOR_BUFFER_SIZE
|
||||
@ -110,6 +130,7 @@ class ColorImageData(ctypes.Structure):
|
||||
("edgeTimestampUs", ctypes.c_longlong)
|
||||
]
|
||||
def frame(self, rgb=False, rotate=None):
|
||||
cv2 = _get_cv2()
|
||||
np_array = np.frombuffer(self.data, dtype=np.uint8, count=self.dataSize)
|
||||
if self.codec == 0: # YUYV 格式, 重塑为 (h, w, 2)，因为每两个字节包含 Y 和 UV 信息
|
||||
yuv_mat = np_array.reshape((self.height, self.width, 2))
|
||||
@ -150,6 +171,7 @@ class DepthImageData(ctypes.Structure):
|
||||
("edgeTimestampUs", ctypes.c_longlong)
|
||||
]
|
||||
def frame(self, rotate=None):
|
||||
cv2 = _get_cv2()
|
||||
np_array = np.frombuffer(self.data, dtype=np.uint8, count=self.dataSize)
|
||||
if self.type == 0: # Depth_16, 数据大小应为 w * h * 2
|
||||
# 1. 转换为 uint16 类型
|
||||
@ -208,6 +230,7 @@ class GrayScaleImage(ctypes.Structure):
|
||||
]
|
||||
def frame(self, rotate=None):
|
||||
"""返回 BGR 三通道图像，兼容 cv2.imshow 显示"""
|
||||
cv2 = _get_cv2()
|
||||
max_size = _MAX_GRAY_BUFFER_SIZE
|
||||
needed = self.width * self.height
|
||||
if needed > max_size:
|
||||
@ -237,6 +260,7 @@ class FisheyeImagesData(ctypes.Structure):
|
||||
if index_valid:
|
||||
return self.images[index].frame(rotate=rotate)
|
||||
else:
|
||||
cv2 = _get_cv2()
|
||||
frame0 = cv2.resize(self.images[0].frame(), (self._DISPLAY_WIDTH, self._DISPLAY_HEIGHT))
|
||||
frame1 = cv2.resize(self.images[1].frame(), (self._DISPLAY_WIDTH, self._DISPLAY_HEIGHT))
|
||||
frame2 = cv2.resize(self.images[2].frame(), (self._DISPLAY_WIDTH, self._DISPLAY_HEIGHT))
|
||||
@ -263,6 +287,7 @@ class EyetrackingImageData(ctypes.Structure):
|
||||
if index_valid:
|
||||
return self.images[index].frame(rotate=rotate)
|
||||
else:
|
||||
cv2 = _get_cv2()
|
||||
frame0 = cv2.resize(self.images[0].frame(), (self._DISPLAY_WIDTH, self._DISPLAY_HEIGHT))
|
||||
frame1 = cv2.resize(self.images[1].frame(), (self._DISPLAY_WIDTH, self._DISPLAY_HEIGHT))
|
||||
frame2 = cv2.resize(self.images[2].frame(), (self._DISPLAY_WIDTH, self._DISPLAY_HEIGHT))
|
||||
@ -288,6 +313,7 @@ class ColorImageDataRef(ctypes.Structure):
|
||||
"""零拷贝帧解码 — 从 SDK buffer 指针直接构造 numpy 数组"""
|
||||
if self.data_size == 0 or not self.data_ptr:
|
||||
return None
|
||||
cv2 = _get_cv2()
|
||||
buf = (ctypes.c_uint8 * self.data_size).from_address(self.data_ptr)
|
||||
np_array = np.frombuffer(buf, dtype=np.uint8, count=self.data_size)
|
||||
if self.codec == 0: # YUYV
|
||||
@ -324,6 +350,7 @@ class DepthImageDataRef(ctypes.Structure):
|
||||
"""零拷贝深度帧解码"""
|
||||
if self.data_size == 0 or not self.data_ptr:
|
||||
return None
|
||||
cv2 = _get_cv2()
|
||||
buf = (ctypes.c_uint8 * self.data_size).from_address(self.data_ptr)
|
||||
np_array = np.frombuffer(buf, dtype=np.uint8, count=self.data_size)
|
||||
if self.type == 0: # Depth_16
|
||||
@ -362,6 +389,7 @@ class GrayScaleImageRef(ctypes.Structure):
|
||||
def frame(self, rotate=None):
|
||||
if not self.data_ptr or self.width <= 0:
|
||||
return None
|
||||
cv2 = _get_cv2()
|
||||
buf = (ctypes.c_uint8 * (self.width * self.height)).from_address(self.data_ptr)
|
||||
gray = np.frombuffer(buf, dtype=np.uint8, count=self.width * self.height).reshape((self.height, self.width))
|
||||
return _apply_rotate(cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR), rotate)
|
||||
@ -384,6 +412,7 @@ class FisheyeImagesDataRef(ctypes.Structure):
|
||||
return self._four_in_one(self.image_count, rotate)
|
||||
|
||||
def _four_in_one(self, count, rotate=None):
|
||||
cv2 = _get_cv2()
|
||||
frames = []
|
||||
for i in range(count):
|
||||
f = self.images[i].frame()
|
||||
@ -447,6 +476,16 @@ class EventData(ctypes.Structure):
|
||||
]
|
||||
|
||||
|
||||
class OrientationData(ctypes.Structure):
|
||||
"""6-DOF 方向数据：四元数 + 3x3 旋转矩阵"""
|
||||
_fields_ = [
|
||||
("quaternion", Vector4D),
|
||||
("rotation", ctypes.c_double * 9),
|
||||
("hostTimestamp", ctypes.c_double),
|
||||
("edgeTimestampUs", ctypes.c_longlong),
|
||||
]
|
||||
|
||||
|
||||
class XVLib:
|
||||
_xvlib = None
|
||||
_load_lock = threading.Lock()
|
||||
@ -463,6 +502,7 @@ class XVLib:
|
||||
self._slam_data = PoseData()
|
||||
self._external_stream_data = PoseData()
|
||||
self._spheretrack_stream_data = PoseData()
|
||||
self._orientation_data = OrientationData()
|
||||
|
||||
# 零拷贝 Ref 数据 holder
|
||||
self._color_image_data_ref = ColorImageDataRef()
|
||||
@ -499,7 +539,14 @@ class XVLib:
|
||||
@classmethod
|
||||
def _check_xvsdk(cls):
|
||||
"""检查系统是否安装了 XVSDK, 没有则抛异常并提示安装命令."""
|
||||
if ctypes.util.find_library("xvsdk") and os.path.exists('/usr/lib/libxvsdk.so'):
|
||||
if ctypes.util.find_library("xvsdk"):
|
||||
candidate_paths = [
|
||||
'/usr/lib/libxvsdk.so',
|
||||
'/usr/lib/x86_64-linux-gnu/libxvsdk.so',
|
||||
'/usr/lib/aarch64-linux-gnu/libxvsdk.so',
|
||||
'/usr/local/lib/libxvsdk.so',
|
||||
]
|
||||
if any(os.path.exists(p) for p in candidate_paths):
|
||||
return
|
||||
raise RuntimeError(
|
||||
"XVSDK not found. Run:\n"
|
||||
@ -664,6 +711,12 @@ class XVLib:
|
||||
'xv_get_fisheye_cameras_data_ref', 'xv_get_eyetracking_camera_data_ref']:
|
||||
lib[name].restype = ctypes.c_int
|
||||
|
||||
# orientation stream data getter
|
||||
lib.xv_get_orientation_stream_data.restype = ctypes.c_int
|
||||
# metadata getters
|
||||
lib.xv_get_color_image_metadata.restype = ctypes.c_int
|
||||
lib.xv_get_fisheye_metadata.restype = ctypes.c_int
|
||||
|
||||
@classmethod
|
||||
def xv_get_devices(cls, timeout=5.0, max_devices=16, double_query=True):
|
||||
"""Scan for connected XV devices.
|
||||
@ -688,15 +741,26 @@ class XVLib:
|
||||
ctypes.c_double(timeout)
|
||||
)
|
||||
if double_query:
|
||||
# 多设备场景下第1次扫描可能遗漏设备，第2次补扫
|
||||
# 多设备场景下第1次扫描可能遗漏设备，第2次补扫（用独立 buffer 避免覆盖）
|
||||
if timeout > 3 and device_count.value != 0:
|
||||
time.sleep(1)
|
||||
second_devices = (DeviceStruct * max_devices)()
|
||||
second_count = ctypes.c_int(0)
|
||||
cls._xvlib.xv_get_devices(
|
||||
ctypes.byref(devices),
|
||||
ctypes.byref(device_count),
|
||||
ctypes.byref(second_devices),
|
||||
ctypes.byref(second_count),
|
||||
ctypes.c_int(max_devices),
|
||||
ctypes.c_double(2.0)
|
||||
)
|
||||
# 合并：将第2次扫描中不在第1次结果里的设备追加进去
|
||||
first_sns = set()
|
||||
for i in range(device_count.value):
|
||||
first_sns.add(devices[i].serial)
|
||||
for i in range(second_count.value):
|
||||
sn = second_devices[i].serial
|
||||
if sn not in first_sns and device_count.value < max_devices:
|
||||
devices[device_count.value] = second_devices[i]
|
||||
device_count.value += 1
|
||||
return device_count.value, list(devices[:device_count.value])
|
||||
|
||||
def xv_init(self, serial_number, init_slam, init_clamp_stream, init_color_camera, init_fisheye_cameras):
|
||||
@ -994,6 +1058,27 @@ class XVLib:
|
||||
ret = self._xvlib.xv_get_spheretrack_stream_data(self.instance_id, ctypes.byref(self._spheretrack_stream_data))
|
||||
return ret, self._spheretrack_stream_data
|
||||
|
||||
def xv_get_orientation_stream_data(self):
|
||||
"""Get the most recently received 6-DOF orientation data (quaternion + 3x3 rotation matrix).
|
||||
Returns (ret_code, OrientationData)."""
|
||||
ret = self._xvlib.xv_get_orientation_stream_data(self.instance_id, ctypes.byref(self._orientation_data))
|
||||
return ret, self._orientation_data
|
||||
|
||||
def xv_get_color_image_metadata(self):
|
||||
"""Get color camera image width & height (updated by callback). Returns (ret_code, width, height)."""
|
||||
width = ctypes.c_int(0)
|
||||
height = ctypes.c_int(0)
|
||||
ret = self._xvlib.xv_get_color_image_metadata(self.instance_id, ctypes.byref(width), ctypes.byref(height))
|
||||
return ret, width.value, height.value
|
||||
|
||||
def xv_get_fisheye_metadata(self):
|
||||
"""Get fisheye cameras image width & height arrays (updated by callback).
|
||||
Returns (ret_code, widths[4], heights[4])."""
|
||||
widths = (ctypes.c_int * 4)(0, 0, 0, 0)
|
||||
heights = (ctypes.c_int * 4)(0, 0, 0, 0)
|
||||
ret = self._xvlib.xv_get_fisheye_metadata(self.instance_id, widths, heights)
|
||||
return ret, list(widths), list(heights)
|
||||
|
||||
# ============== Zero-copy Ref getters (internal, called by xv_get_*_data(use_ref=True)) ==============
|
||||
|
||||
def _xv_get_color_camera_data_ref(self):
|
||||
|
||||
@ -3,8 +3,6 @@ import logging
|
||||
import time
|
||||
import math
|
||||
import numpy as np
|
||||
from gello.dynamixel.driver import DynamixelDriver
|
||||
from gello.agents.gello_agent import GelloAgent, DynamixelRobotConfig
|
||||
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
|
||||
from ..base_teleop import UFBaseTeleop
|
||||
from .gello_teleop_config import GelloTeleopConfig
|
||||
@ -26,6 +24,9 @@ class GelloTeleop(UFBaseTeleop):
|
||||
self._is_connected = False
|
||||
self._is_calibrated = True # CHECK!!
|
||||
|
||||
from gello.dynamixel.driver import DynamixelDriver
|
||||
from gello.agents.gello_agent import DynamixelRobotConfig
|
||||
|
||||
# auto get joint offset from gello
|
||||
joint_ids = []
|
||||
joint_ids.extend(self.config.joint_ids)
|
||||
@ -91,6 +92,7 @@ class GelloTeleop(UFBaseTeleop):
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
if self._is_connected:
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
from gello.agents.gello_agent import GelloAgent
|
||||
|
||||
self.gello_agent = GelloAgent(port=self.config.port, dynamixel_config=self._dynamixel_robo_config)
|
||||
if not self._is_calibrated and calibrate:
|
||||
|
||||
@ -4,7 +4,6 @@ import time
|
||||
import numpy as np
|
||||
from typing import Any
|
||||
from threading import Thread, Event
|
||||
from spnav import spnav_open, spnav_poll_event, spnav_close, SpnavMotionEvent, SpnavButtonEvent
|
||||
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
|
||||
from collections import defaultdict
|
||||
|
||||
@ -35,6 +34,7 @@ class SpaceMouseTeleop(Teleoperator, Thread):
|
||||
self.deadzone = np.array(deadzone, dtype=self.dtype)
|
||||
assert (self.deadzone >= 0).all()
|
||||
|
||||
from spnav import SpnavMotionEvent
|
||||
self.motion_event = SpnavMotionEvent([0,0,0], [0,0,0], 0)
|
||||
self.button_state = defaultdict(lambda: False)
|
||||
self.tx_zup_spnav = np.array([
|
||||
@ -129,6 +129,7 @@ class SpaceMouseTeleop(Teleoperator, Thread):
|
||||
return self.button_state[button_id]
|
||||
|
||||
def run(self):
|
||||
from spnav import spnav_open, spnav_poll_event, spnav_close, SpnavMotionEvent, SpnavButtonEvent
|
||||
spnav_open()
|
||||
self._is_connected = True
|
||||
try:
|
||||
|
||||
@ -5,7 +5,6 @@ from typing import Any
|
||||
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
|
||||
from ufactory_lerobot.devices.umi.vive_tracker.transformations import Transformations
|
||||
from ufactory_lerobot.devices.umi.vive_tracker import ViveTracker
|
||||
from ufactory_lerobot.devices.umi.xvlib import XVLib
|
||||
from ..base_teleop import UFBaseTeleop
|
||||
from .umi_teleop_config import UmiTeleopConfig
|
||||
|
||||
@ -27,6 +26,7 @@ class UmiTeleop(UFBaseTeleop):
|
||||
self.tracker = None
|
||||
self.xvlib = None
|
||||
|
||||
# from ufactory_lerobot.devices.umi.xvlib import XVLib
|
||||
# self.tracker = ViveTracker() if self.config.use_vive_tracker else None
|
||||
# self.xvlib = XVLib(self.config.serial_number, not self.config.use_vive_tracker, self.config.use_gripper)
|
||||
|
||||
@ -84,6 +84,7 @@ class UmiTeleop(UFBaseTeleop):
|
||||
pass
|
||||
|
||||
def connect(self, calibrate: bool = False) -> None:
|
||||
from ufactory_lerobot.devices.umi.xvlib import XVLib
|
||||
self.tracker = ViveTracker() if self.config.use_vive_tracker else None
|
||||
self.xvlib = XVLib(self.config.serial_number, not self.config.use_vive_tracker, self.config.use_gripper)
|
||||
if not self.config.use_vive_tracker:
|
||||
|
||||
Loading…
Reference in New Issue
Block a user