Initial commit
This commit is contained in:
commit
a4e596c3aa
234 changed files with 81101 additions and 0 deletions
185
arm_plane_calib/collect_plane_calib.py
Normal file
185
arm_plane_calib/collect_plane_calib.py
Normal file
|
|
@ -0,0 +1,185 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
import yaml
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
import rclpy
|
||||
from rclpy.node import Node
|
||||
from sensor_msgs.msg import Image, CameraInfo
|
||||
from cv_bridge import CvBridge
|
||||
|
||||
|
||||
class PlaneCalibCollector(Node):
|
||||
def __init__(self):
|
||||
super().__init__('plane_calib_collector')
|
||||
|
||||
self.declare_parameter('color_topic', '/camera/color/image_raw')
|
||||
self.declare_parameter('depth_topic', '/camera/depth/image_raw')
|
||||
self.declare_parameter('camera_info_topic', '/camera/depth/camera_info')
|
||||
self.declare_parameter('save_path', 'samples.yaml')
|
||||
self.declare_parameter('depth_scale', 0.001) # 深度单位转米，常见 mm -> m
|
||||
self.declare_parameter('min_depth_m', 0.05)
|
||||
self.declare_parameter('max_depth_m', 1.50)
|
||||
self.declare_parameter('median_kernel', 5)
|
||||
self.declare_parameter('target_count', 9) # 采满多少个点自动退出
|
||||
|
||||
self.color_topic = self.get_parameter('color_topic').value
|
||||
self.depth_topic = self.get_parameter('depth_topic').value
|
||||
self.camera_info_topic = self.get_parameter('camera_info_topic').value
|
||||
self.save_path = self.get_parameter('save_path').value
|
||||
self.depth_scale = float(self.get_parameter('depth_scale').value)
|
||||
self.min_depth_m = float(self.get_parameter('min_depth_m').value)
|
||||
self.max_depth_m = float(self.get_parameter('max_depth_m').value)
|
||||
self.median_kernel = int(self.get_parameter('median_kernel').value)
|
||||
self.target_count = int(self.get_parameter('target_count').value)
|
||||
|
||||
self.bridge = CvBridge()
|
||||
self.color_img = None
|
||||
self.depth_img = None
|
||||
self.K = None
|
||||
self.samples = []
|
||||
self.should_exit = False
|
||||
|
||||
self.create_subscription(Image, self.color_topic, self.on_color, 10)
|
||||
self.create_subscription(Image, self.depth_topic, self.on_depth, 10)
|
||||
self.create_subscription(CameraInfo, self.camera_info_topic, self.on_info, 10)
|
||||
|
||||
self.window_name = 'plane_calib_click'
|
||||
cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
|
||||
cv2.setMouseCallback(self.window_name, self.on_mouse)
|
||||
|
||||
self.get_logger().info('左键点击标定点，终端输入机械臂命令向量，例如: 500,620,410,530')
|
||||
self.get_logger().info(f'每记录 1 个样本自动保存，采满 {self.target_count} 个样本自动退出')
|
||||
self.get_logger().info('按 q 可手动退出')
|
||||
|
||||
def on_color(self, msg: Image):
|
||||
self.color_img = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
|
||||
|
||||
def on_depth(self, msg: Image):
|
||||
if msg.encoding in ('16UC1', 'mono16'):
|
||||
self.depth_img = self.bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
|
||||
else:
|
||||
self.get_logger().warn(f'不支持的深度编码: {msg.encoding}')
|
||||
|
||||
def on_info(self, msg: CameraInfo):
|
||||
self.K = np.array(msg.k, dtype=np.float64).reshape(3, 3)
|
||||
|
||||
def get_xyz_from_pixel(self, u: int, v: int):
|
||||
if self.depth_img is None or self.K is None:
|
||||
return None
|
||||
|
||||
h, w = self.depth_img.shape[:2]
|
||||
if not (0 <= u < w and 0 <= v < h):
|
||||
return None
|
||||
|
||||
k = self.median_kernel
|
||||
r = k // 2
|
||||
u0, u1 = max(0, u - r), min(w, u + r + 1)
|
||||
v0, v1 = max(0, v - r), min(h, v + r + 1)
|
||||
patch = self.depth_img[v0:v1, u0:u1].astype(np.float32)
|
||||
|
||||
valid = patch[patch > 0]
|
||||
if valid.size == 0:
|
||||
return None
|
||||
|
||||
z = np.median(valid) * self.depth_scale
|
||||
if z < self.min_depth_m or z > self.max_depth_m:
|
||||
return None
|
||||
|
||||
fx = self.K[0, 0]
|
||||
fy = self.K[1, 1]
|
||||
cx = self.K[0, 2]
|
||||
cy = self.K[1, 2]
|
||||
|
||||
x = (u - cx) * z / fx
|
||||
y = (v - cy) * z / fy
|
||||
return float(x), float(y), float(z)
|
||||
|
||||
def save_samples(self):
|
||||
data = {
|
||||
'samples': self.samples,
|
||||
'cmd_dim': len(self.samples[0]['arm_cmd']) if self.samples else 0
|
||||
}
|
||||
with open(self.save_path, 'w', encoding='utf-8') as f:
|
||||
yaml.safe_dump(data, f, allow_unicode=True, sort_keys=False)
|
||||
print(f'[SAVE] 已保存到 {self.save_path}')
|
||||
|
||||
def on_mouse(self, event, x, y, flags, param):
|
||||
if event != cv2.EVENT_LBUTTONDOWN:
|
||||
return
|
||||
|
||||
xyz = self.get_xyz_from_pixel(x, y)
|
||||
if xyz is None:
|
||||
print(f'[WARN] 点击点 ({x}, {y}) 无有效深度')
|
||||
return
|
||||
|
||||
X, Y, Z = xyz
|
||||
print(f'\n点击像素: ({x}, {y})')
|
||||
print(f'相机坐标: X={X:.4f} m, Y={Y:.4f} m, Z={Z:.4f} m')
|
||||
print('请手动把机械臂末端移到这个点正上方，然后输入命令向量，例如 500,620,410,530')
|
||||
cmd_str = input('arm_cmd> ').strip()
|
||||
if not cmd_str:
|
||||
print('[SKIP] 未输入，跳过')
|
||||
return
|
||||
|
||||
try:
|
||||
arm_cmd = [float(s) for s in cmd_str.split(',')]
|
||||
except Exception:
|
||||
print('[ERR] 输入格式错误，应类似 500,620,410,530')
|
||||
return
|
||||
|
||||
item = {
|
||||
'u': int(x),
|
||||
'v': int(y),
|
||||
'x_cam': X,
|
||||
'y_cam': Y,
|
||||
'z_cam': Z,
|
||||
'arm_cmd': arm_cmd
|
||||
}
|
||||
self.samples.append(item)
|
||||
print(f'[OK] 已记录第 {len(self.samples)} 个样本')
|
||||
|
||||
# 自动保存
|
||||
self.save_samples()
|
||||
|
||||
# 采满自动退出
|
||||
if len(self.samples) >= self.target_count:
|
||||
print(f'[DONE] 已采满 {self.target_count} 个样本，自动保存并退出')
|
||||
self.should_exit = True
|
||||
|
||||
def spin_loop(self):
|
||||
while rclpy.ok() and not self.should_exit:
|
||||
rclpy.spin_once(self, timeout_sec=0.03)
|
||||
|
||||
if self.color_img is not None:
|
||||
show = self.color_img.copy()
|
||||
for i, s in enumerate(self.samples):
|
||||
cv2.circle(show, (s['u'], s['v']), 4, (0, 255, 0), -1)
|
||||
cv2.putText(
|
||||
show, str(i + 1), (s['u'] + 6, s['v'] - 6),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1
|
||||
)
|
||||
cv2.imshow(self.window_name, show)
|
||||
|
||||
key = cv2.waitKey(1) & 0xFF
|
||||
if key == ord('q'):
|
||||
print('[EXIT] 用户退出')
|
||||
break
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
def main():
|
||||
rclpy.init()
|
||||
node = PlaneCalibCollector()
|
||||
try:
|
||||
node.spin_loop()
|
||||
finally:
|
||||
node.destroy_node()
|
||||
rclpy.shutdown()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue