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

修正数据采集节点

Browse Source
This commit is contained in:
liangyuxuan
2025-12-25 15:39:13 +08:00
parent aff813b107
commit 8279687f4c
2 changed files with 131 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
vision_detect/vision_detect/collect_data_client.py
View File

@@ -22,7 +22,8 @@ def main(args=None):
client.call_async(request).add_done_callback(result_callback)
node.set_parameters([Parameter('end_sign', Parameter.Type.BOOL, False)])
def result_callback(response):
def result_callback(result):
response = result.result()
if response.success:
node.get_logger().info(response.message)
else:

136
vision_detect/vision_detect/collect_data_node.py
View File

@@ -9,7 +9,7 @@ from cv_bridge import CvBridge
from ament_index_python import get_package_share_directory
from message_filters import Subscriber, ApproximateTimeSynchronizer
from sensor_msgs.msg import Image
from sensor_msgs.msg import Image, CameraInfo
from std_srvs.srv import SetBool
@@ -33,6 +33,14 @@ class CollectDataNode(Node):
self.right_timestamp_file = None
self.head_timestamp_file = None
self.left_camera_info_file = None
self.right_camera_info_file = None
self.head_camera_info_file = None
self.left_info_sub = None
self.right_info_sub = None
self.head_info_sub = None
self.index = 0
self.fps = 30.0
@@ -60,6 +68,7 @@ class CollectDataNode(Node):
if self.topic_exists(configs["left"]["color"]) and self.topic_exists(
configs["left"]["depth"]):
self.left_sign = True
# Color and Depth
self.sub_left_color = Subscriber(self, Image, configs["left"]["color"])
self.sub_left_depth = Subscriber(self, Image, configs["left"]["depth"])
@@ -69,11 +78,14 @@ class CollectDataNode(Node):
slop=0.1
)
self.left_sync_subscriber.registerCallback(self.left_callback)
# Camera info
self.left_info_name = configs["left"]["Info"]
# right
if self.topic_exists(configs["right"]["color"]) and self.topic_exists(
configs["right"]["depth"]):
self.right_sign = True
# Color and Depth
self.sub_right_color = Subscriber(self, Image, configs["right"]["color"])
self.sub_right_depth = Subscriber(self, Image, configs["right"]["depth"])
@@ -83,11 +95,14 @@ class CollectDataNode(Node):
slop=0.1
)
self.right_sync_subscriber.registerCallback(self.right_callback)
# Camera info
self.right_info_name = configs["right"]["Info"]
# head
if self.topic_exists(configs["head"]["color"]) and self.topic_exists(
configs["head"]["depth"]):
self.head_sign = True
# Color and Depth
self.sub_head_color = Subscriber(self, Image, configs["head"]["color"])
self.sub_head_depth = Subscriber(self, Image, configs["head"]["depth"])
@@ -97,18 +112,48 @@ class CollectDataNode(Node):
slop=0.1
)
self.head_sync_subscriber.registerCallback(self.head_callback)
# Camera info
self.head_info_name = configs["head"]["Info"]
self.service = self.create_service(SetBool, "/collect_data", self.service_callback)
def service_callback(self, request, response):
# Start collect data
if request.data:
self.left_meta = self.right_meta = self.head_meta = False
while (os.path.exists(os.path.join(self.save_path,
f"left_color_video_{self.index:04d}.mp4"))
or os.path.exists(os.path.join(self.save_path,
f"right_color_video_{self.index:04d}.mp4"))
or os.path.exists(os.path.join(self.save_path,
f"head_color_video_{self.index:04d}.mp4"))):
while (
os.path.exists(os.path.join(
self.save_path, f"left_color_video_{self.index:04d}.mp4"))
or os.path.exists(os.path.join(
self.save_path, f"right_color_video_{self.index:04d}.mp4"))
or os.path.exists(os.path.join(
self.save_path, f"head_color_video_{self.index:04d}.mp4"))
or os.path.exists(os.path.join(
self.save_path, f"left_depth_video_{self.index:04d}.mp4"))
or os.path.exists(os.path.join(
self.save_path, f"right_depth_video_{self.index:04d}.mp4"))
or os.path.exists(os.path.join(
self.save_path, f"head_depth_video_{self.index:04d}.mp4"))
or os.path.exists(os.path.join(
self.save_path, f"left_depth_raw_{self.index:04d}.raw"))
or os.path.exists(os.path.join(
self.save_path, f"right_depth_raw_{self.index:04d}.raw"))
or os.path.exists(os.path.join(
self.save_path, f"head_depth_raw_{self.index:04d}.raw"))
or os.path.exists(os.path.join(
self.save_path, f"left_timestamp_{self.index:04d}.txt"))
or os.path.exists(os.path.join(
self.save_path, f"right_timestamp_{self.index:04d}.txt"))
or os.path.exists(os.path.join(
self.save_path, f"head_timestamp_{self.index:04d}.txt"))
or os.path.exists(os.path.join(
self.save_path, f"left_camera_info_{self.index:04d}.json"))
or os.path.exists(os.path.join(
self.save_path, f"right_camera_info_{self.index:04d}.json"))
or os.path.exists(os.path.join(
self.save_path, f"head_camera_info_{self.index:04d}.json"))
):
self.index += 1
if self.left_sign:
@@ -116,16 +161,40 @@ class CollectDataNode(Node):
os.path.join(self.save_path, f"left_depth_raw_{self.index:04d}.raw"), 'ab')
self.left_timestamp_file = open(
os.path.join(self.save_path, f"left_timestamp_{self.index:04d}.txt"), 'a')
self.left_camera_info_file = open(
os.path.join(self.save_path, f"left_camera_info_{self.index:04d}.json"), 'w')
self.left_info_sub = self.create_subscription(
CameraInfo,
self.left_info_name,
self.left_info_callback,
10
)
if self.right_sign:
self.right_depth_raw_file = open(
os.path.join(self.save_path, f"right_depth_raw_{self.index:04d}.raw"), 'ab')
self.right_timestamp_file = open(
os.path.join(self.save_path, f"right_timestamp_{self.index:04d}.txt"), 'a')
self.right_camera_info_file = open(
os.path.join(self.save_path, f"right_camera_info_{self.index:04d}.json"), 'w')
self.right_info_sub = self.create_subscription(
CameraInfo,
self.right_info_name,
self.right_info_callback,
10
)
if self.head_sign:
self.head_depth_raw_file = open(
os.path.join(self.save_path, f"head_depth_raw_{self.index:04d}.raw"), 'ab')
self.head_timestamp_file = open(
os.path.join(self.save_path, f"head_timestamp_{self.index:04d}.txt"), 'a')
self.head_camera_info_file = open(
os.path.join(self.save_path, f"head_camera_info_{self.index:04d}.json"), 'w')
self.head_info_sub = self.create_subscription(
CameraInfo,
self.head_info_name,
self.head_info_callback,
10
)
self.left_writer_initialized = False
self.right_writer_initialized = False
@@ -143,6 +212,12 @@ class CollectDataNode(Node):
if self.left_timestamp_file is not None:
self.left_timestamp_file.close()
self.left_timestamp_file = None
if self.left_camera_info_file is not None:
self.left_camera_info_file.close()
self.left_camera_info_file = None
if self.left_info_sub is not None:
self.left_info_sub.destroy()
self.left_info_sub = None
if self.right_sign:
if self.right_depth_raw_file is not None:
self.right_depth_raw_file.close()
@@ -150,6 +225,12 @@ class CollectDataNode(Node):
if self.right_timestamp_file is not None:
self.right_timestamp_file.close()
self.right_timestamp_file = None
if self.right_camera_info_file is not None:
self.right_camera_info_file.close()
self.right_camera_info_file = None
if self.right_info_sub is not None:
self.right_info_sub.destroy()
self.right_info_sub = None
if self.head_sign:
if self.head_depth_raw_file is not None:
self.head_depth_raw_file.close()
@@ -157,7 +238,14 @@ class CollectDataNode(Node):
if self.head_timestamp_file is not None:
self.head_timestamp_file.close()
self.head_timestamp_file = None
if self.head_camera_info_file is not None:
self.head_camera_info_file.close()
self.head_camera_info_file = None
if self.head_info_sub is not None:
self.head_info_sub.destroy()
self.head_info_sub = None
# End collect data
if self.left_color_writer is not None:
self.left_color_writer.release()
self.left_color_writer = None
@@ -387,6 +475,40 @@ class CollectDataNode(Node):
depth_vis = cv2.cvtColor(depth_vis, cv2.COLOR_GRAY2BGR)
self.head_depth_writer.write(depth_vis)
def left_info_callback(self, msgs):
data = {
"K": msgs.k,
"D": msgs.d,
}
json.dump(data, self.left_camera_info_file, ensure_ascii=False, indent=4)
self.left_camera_info_file.close()
self.left_camera_info_file = None
self.left_info_sub.destroy()
self.left_info_sub = None
def right_info_callback(self, msgs):
data = {
"K": msgs.k,
"D": msgs.d,
}
json.dump(data, self.right_camera_info_file, ensure_ascii=False, indent=4)
self.right_camera_info_file.close()
self.right_camera_info_file = None
self.right_info_sub.destroy()
self.right_info_sub = None
def head_info_callback(self, msgs):
data = {
"K": msgs.k,
"D": msgs.d,
}
json.dump(data, self.head_camera_info_file, ensure_ascii=False, indent=4)
self.head_camera_info_file.close()
self.head_camera_info_file = None
self.head_info_sub.destroy()
self.head_info_sub = None
def topic_exists(self, topic_name: str) -> bool:
topics = self.get_topic_names_and_types()
return any(name == topic_name for name, _ in topics)