HiveCoreRD/hivecore_robot_ctrl_gui
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gen images and skill info

Browse Source
This commit is contained in:
NuoDaJia02
2026-01-27 19:49:27 +08:00
parent 1cd009fd6b
commit ec279ec154
1 changed files with 137 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

153
ctrlgui/ctrlgui_node.py
View File

@@ -87,20 +87,32 @@ class CtrlGuiNode(Node):
# Top Camera subscription
self.bridge = CvBridge()
self.last_top_image = None
self.last_top_depth = None
self.last_left_image = None
self.last_left_depth = None
self.last_right_image = None
self.last_right_depth = None
self._sub_top_image = self.create_subscription(
Image, '/camera/color/image_raw', self._on_top_image, 10
)
self._sub_top_depth = self.create_subscription(
Image, '/camera/depth/image_raw', self._on_top_depth, 10
)
# Left Camera subscription
self._sub_left_image = self.create_subscription(
Image, '/camera1/camera1/color/image_raw', self._on_left_image, 10
)
self._sub_left_depth = self.create_subscription(
Image, '/camera1/camera1/depth/image_rect_raw', self._on_left_depth, 10
)
# Right Camera subscription
self._sub_right_image = self.create_subscription(
Image, '/camera2/camera2/color/image_raw', self._on_right_image, 10
)
self._sub_right_depth = self.create_subscription(
Image, '/camera2/camera2/depth/image_rect_raw', self._on_right_depth, 10
)
# Data collection state
self.is_collecting = False
@@ -358,6 +370,13 @@ class CtrlGuiNode(Node):
except Exception as e:
self.file_logger.error(f"Error converting top image: {str(e)}")
def _on_top_depth(self, msg: Image) -> None:
"""Callback for /camera/depth/image_raw subscription: store latest top depth image."""
try:
self.last_top_depth = self.bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
except Exception as e:
self.file_logger.error(f"Error converting top depth image: {str(e)}")
def _on_left_image(self, msg: Image) -> None:
"""Callback for /camera1/camera1/color/image_raw subscription: store latest left image."""
try:
@@ -365,6 +384,13 @@ class CtrlGuiNode(Node):
except Exception as e:
self.file_logger.error(f"Error converting left image: {str(e)}")
def _on_left_depth(self, msg: Image) -> None:
"""Callback for /camera1/camera1/depth/image_rect_raw subscription: store latest left depth image."""
try:
self.last_left_depth = self.bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
except Exception as e:
self.file_logger.error(f"Error converting left depth image: {str(e)}")
def _on_right_image(self, msg: Image) -> None:
"""Callback for /camera2/camera1/color/image_raw subscription: store latest right image."""
try:
@@ -372,6 +398,13 @@ class CtrlGuiNode(Node):
except Exception as e:
self.file_logger.error(f"Error converting right image: {str(e)}")
def _on_right_depth(self, msg: Image) -> None:
"""Callback for /camera2/camera2/depth/image_rect_raw subscription: store latest right depth image."""
try:
self.last_right_depth = self.bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
except Exception as e:
self.file_logger.error(f"Error converting right depth image: {str(e)}")
def publish_info(self, text: str) -> None:
msg = String()
msg.data = text
@@ -471,6 +504,7 @@ def build_ui(node: CtrlGuiNode) -> None:
node.data_collection_dir = os.path.join(data_base_dir, str(timestamp_s))
try:
os.makedirs(node.data_collection_dir, exist_ok=True)
os.makedirs(os.path.join(node.data_collection_dir, 'images'), exist_ok=True)
skill_info_path = os.path.join(node.data_collection_dir, 'skill_info.txt')
# Create empty file
with open(skill_info_path, 'w') as f:
@@ -931,41 +965,69 @@ def build_ui(node: CtrlGuiNode) -> None:
instance_params_data = str(msg.instance_params)
# Save images if available
timestamp = time.time()
# timestamp = msg.stamp.sec + msg.stamp.nanosec * 1e-9
timestamp = f"{msg.stamp.sec}.{msg.stamp.nanosec:09d}"
image_filename_top = ""
image_filename_top_depth = ""
image_filename_left = ""
image_filename_left_depth = ""
image_filename_right = ""
image_filename_right_depth = ""
if node.last_top_image is not None:
try:
image_filename_top = f"top_cam_{timestamp:.6f}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, image_filename_top), node.last_top_image)
image_filename_top = f"top_color_{timestamp}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, 'images', image_filename_top), node.last_top_image)
except Exception as e:
node.file_logger.error(f"Failed to save top image: {str(e)}")
if node.last_top_depth is not None:
try:
image_filename_top_depth = f"top_depth_{timestamp}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, 'images', image_filename_top_depth), node.last_top_depth)
except Exception as e:
node.file_logger.error(f"Failed to save top depth image: {str(e)}")
if node.last_left_image is not None:
try:
image_filename_left = f"left_cam_{timestamp:.6f}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, image_filename_left), node.last_left_image)
image_filename_left = f"left_color_{timestamp}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, 'images', image_filename_left), node.last_left_image)
except Exception as e:
node.file_logger.error(f"Failed to save left image: {str(e)}")
if node.last_left_depth is not None:
try:
image_filename_left_depth = f"left_depth_{timestamp}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, 'images', image_filename_left_depth), node.last_left_depth)
except Exception as e:
node.file_logger.error(f"Failed to save left depth image: {str(e)}")
if node.last_right_image is not None:
try:
image_filename_right = f"right_cam_{timestamp:.6f}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, image_filename_right), node.last_right_image)
image_filename_right = f"right_color_{timestamp}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, 'images', image_filename_right), node.last_right_image)
except Exception as e:
node.file_logger.error(f"Failed to save right image: {str(e)}")
if node.last_right_depth is not None:
try:
image_filename_right_depth = f"right_depth_{timestamp}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, 'images', image_filename_right_depth), node.last_right_depth)
except Exception as e:
node.file_logger.error(f"Failed to save right depth image: {str(e)}")
data = {
"action_name": current_action_name,
"timestamp": timestamp, #f"{msg.stamp.sec}.{msg.stamp.nanosec:09d}",
"skill": str(msg.skill),
"stamp": f"{msg.stamp.sec}.{msg.stamp.nanosec:09d}",
"instance_params": instance_params_data,
"bt_node_status": current_bt_status,
"image_path_top": image_filename_top,
"image_path_left": image_filename_left,
"image_path_right": image_filename_right
"action_name": current_action_name,
"action_status": current_bt_status,
"parameters": instance_params_data,
"color_path_top": f"images/{image_filename_top}" if image_filename_top else "",
"depth_path_top": f"images/{image_filename_top_depth}" if image_filename_top_depth else "",
"color_path_left": f"images/{image_filename_left}" if image_filename_left else "",
"depth_path_left": f"images/{image_filename_left_depth}" if image_filename_left_depth else "",
"color_path_right": f"images/{image_filename_right}" if image_filename_right else "",
"depth_path_right": f"images/{image_filename_right_depth}" if image_filename_right_depth else ""
}
with open(skill_info_path, 'a') as f:
f.write(json.dumps(data, ensure_ascii=False) + '\n')
@@ -986,7 +1048,7 @@ def build_ui(node: CtrlGuiNode) -> None:
last_entry = json.loads(lines[-1].strip())
# 只有当 action_name 匹配时才更新状态
if last_entry.get("action_name") == current_action_name:
last_entry["bt_node_status"] = current_bt_status
last_entry["action_status"] = current_bt_status
lines[-1] = json.dumps(last_entry, ensure_ascii=False) + '\n'
with open(skill_info_path, 'w') as f:
@@ -995,9 +1057,68 @@ def build_ui(node: CtrlGuiNode) -> None:
# 如果不匹配（理论上不应发生），可以视为新动作或记录异常
node.file_logger.warning(f'Action name mismatch during status update: {last_entry.get("action_name")} vs {current_action_name}')
# 当检测到 action_status 不等于 RUNNING 时，额外触发一次图片存储并记录
if current_bt_status != "RUNNING" and current_bt_status.strip():
# Save images
timestamp = time.time()
img_top = ""
dep_top = ""
img_left = ""
dep_left = ""
img_right = ""
dep_right = ""
if node.last_top_image is not None:
try:
img_top = f"top_color_{timestamp:.9f}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, 'images', img_top), node.last_top_image)
except: pass
if node.last_top_depth is not None:
try:
dep_top = f"top_depth_{timestamp:.9f}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, 'images', dep_top), node.last_top_depth)
except: pass
if node.last_left_image is not None:
try:
img_left = f"left_color_{timestamp:.9f}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, 'images', img_left), node.last_left_image)
except: pass
if node.last_left_depth is not None:
try:
dep_left = f"left_depth_{timestamp:.9f}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, 'images', dep_left), node.last_left_depth)
except: pass
if node.last_right_image is not None:
try:
img_right = f"right_color_{timestamp:.9f}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, 'images', img_right), node.last_right_image)
except: pass
if node.last_right_depth is not None:
try:
dep_right = f"right_depth_{timestamp:.9f}.png"
cv2.imwrite(os.path.join(node.data_collection_dir, 'images', dep_right), node.last_right_depth)
except: pass
tasks = list(msg.task)
last_task_val = str(tasks[0]) if tasks else "N/A"
result_data = {
"timestamp": f"{timestamp:.9f}",
"schedule": last_task_val,
"action_status": current_bt_status,
"color_path_top": f"images/{img_top}" if img_top else "",
"depth_path_top": f"images/{dep_top}" if dep_top else "",
"color_path_left": f"images/{img_left}" if img_left else "",
"depth_path_left": f"images/{dep_left}" if dep_left else "",
"color_path_right": f"images/{img_right}" if img_right else "",
"depth_path_right": f"images/{dep_right}" if dep_right else ""
}
with open(skill_info_path, 'a') as f:
f.write(json.dumps(result_data, ensure_ascii=False) + '\n')
node.last_bt_status = current_bt_status
except Exception as e:
node.file_logger.error(f'Failed to update bt_node_status: {str(e)}')
node.file_logger.error(f'Failed to update action_status: {str(e)}')
# Basic Information
msg_id_label.set_text(f'Message ID: {str(msg.msg_id)}')