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

更改深度raw的保存方式为.zip

Browse Source
This commit is contained in:
liangyuxuan
2025-12-27 16:45:55 +08:00
parent 87a5697e75
commit 12fed65d8a
4 changed files with 233 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
vision_tools/vision_tools/collect_data_client.py
View File

@@ -34,6 +34,7 @@ def main(args=None):
client = node.create_client(SetBool, "/collect_data")
while not client.wait_for_service(timeout_sec=1.0):
node.get_logger().info('Service not available, waiting again...')
node.get_logger().info("Service available")
node.declare_parameter("collect_sign", False)
node.declare_parameter("end_sign", False)
node.create_timer(0.5, timer_callback)

253
vision_tools/vision_tools/collect_data_node.py
View File

@@ -1,7 +1,10 @@
import os
import json
import zipfile
import threading
import cv2
import numpy as np
import rclpy
from rclpy.node import Node
@@ -12,9 +15,13 @@ from message_filters import Subscriber, ApproximateTimeSynchronizer
from sensor_msgs.msg import Image, CameraInfo
from std_srvs.srv import SetBool
from vision_tools.utils import zip_tools
share_dir = get_package_share_directory('vision_tools')
class CollectDataNode(Node):
def __init__(self):
super().__init__("collect_data_node")
@@ -79,7 +86,8 @@ class CollectDataNode(Node):
)
self.left_sync_subscriber.registerCallback(self.left_callback)
# Camera info
self.left_info_name = configs["left"]["Info"]
self.left_info_name = configs["left"]["info"]
self.get_logger().info("left camera online")
# right
if self.topic_exists(configs["right"]["color"]) and self.topic_exists(
@@ -96,7 +104,8 @@ class CollectDataNode(Node):
)
self.right_sync_subscriber.registerCallback(self.right_callback)
# Camera info
self.right_info_name = configs["right"]["Info"]
self.right_info_name = configs["right"]["info"]
self.get_logger().info("right camera online")
# head
if self.topic_exists(configs["head"]["color"]) and self.topic_exists(
@@ -113,8 +122,8 @@ class CollectDataNode(Node):
)
self.head_sync_subscriber.registerCallback(self.head_callback)
# Camera info
self.head_info_name = configs["head"]["Info"]
self.head_info_name = configs["head"]["info"]
self.get_logger().info("head camera online")
self.service = self.create_service(SetBool, "/collect_data", self.service_callback)
self.get_logger().info("Success start")
@@ -123,47 +132,61 @@ class CollectDataNode(Node):
# Start collect data
if request.data:
self.left_meta = self.right_meta = self.head_meta = False
if self.save_path:
os.makedirs(self.save_path, exist_ok=True)
while (
os.path.exists(os.path.join(
self.save_path, f"left_color_video_{self.index:04d}.mp4"))
self.save_path, f"{self.index:06d}_left_color_video.mp4"))
or os.path.exists(os.path.join(
self.save_path, f"right_color_video_{self.index:04d}.mp4"))
self.save_path, f"{self.index:06d}_right_color_video.mp4"))
or os.path.exists(os.path.join(
self.save_path, f"head_color_video_{self.index:04d}.mp4"))
self.save_path, f"{self.index:06d}_head_color_video.mp4"))
or os.path.exists(os.path.join(
self.save_path, f"left_depth_video_{self.index:04d}.mp4"))
self.save_path, f"{self.index:06d}_left_depth_video.mp4"))
or os.path.exists(os.path.join(
self.save_path, f"right_depth_video_{self.index:04d}.mp4"))
self.save_path, f"{self.index:06d}_right_depth_video.mp4"))
or os.path.exists(os.path.join(
self.save_path, f"head_depth_video_{self.index:04d}.mp4"))
self.save_path, f"{self.index:06d}_head_depth_video.mp4"))
or os.path.exists(os.path.join(
self.save_path, f"left_depth_raw_{self.index:04d}.raw"))
self.save_path, f"{self.index:06d}_left_depth_raw.raw"))
or os.path.exists(os.path.join(
self.save_path, f"right_depth_raw_{self.index:04d}.raw"))
self.save_path, f"{self.index:06d}_right_depth_raw.raw"))
or os.path.exists(os.path.join(
self.save_path, f"head_depth_raw_{self.index:04d}.raw"))
self.save_path, f"{self.index:06d}_head_depth_raw.raw"))
or os.path.exists(os.path.join(
self.save_path, f"left_timestamp_{self.index:04d}.txt"))
self.save_path, f"{self.index:06d}_left_depth_raw.zip"))
or os.path.exists(os.path.join(
self.save_path, f"right_timestamp_{self.index:04d}.txt"))
self.save_path, f"{self.index:06d}_right_depth_raw.zip"))
or os.path.exists(os.path.join(
self.save_path, f"head_timestamp_{self.index:04d}.txt"))
self.save_path, f"{self.index:06d}_head_depth_raw.zip"))
or os.path.exists(os.path.join(
self.save_path, f"left_camera_info_{self.index:04d}.json"))
self.save_path, f"{self.index:06d}_left_timestamp.txt"))
or os.path.exists(os.path.join(
self.save_path, f"right_camera_info_{self.index:04d}.json"))
self.save_path, f"{self.index:06d}_right_timestamp.txt"))
or os.path.exists(os.path.join(
self.save_path, f"head_camera_info_{self.index:04d}.json"))
self.save_path, f"{self.index:06d}_head_timestamp.txt"))
or os.path.exists(os.path.join(
self.save_path, f"{self.index:06d}_left_depth_meta.txt"))
or os.path.exists(os.path.join(
self.save_path, f"{self.index:06d}_right_depth_meta.txt"))
or os.path.exists(os.path.join(
self.save_path, f"{self.index:06d}_head_depth_meta.txt"))
or os.path.exists(os.path.join(
self.save_path, f"{self.index:06d}_left_camera_info.json"))
or os.path.exists(os.path.join(
self.save_path, f"{self.index:06d}_right_camera_info.json"))
or os.path.exists(os.path.join(
self.save_path, f"{self.index:06d}_head_camera_info.json"))
):
self.index += 1
if self.left_sign:
self.left_depth_raw_file = open(
os.path.join(self.save_path, f"left_depth_raw_{self.index:04d}.raw"), 'ab')
os.path.join(self.save_path, f"{self.index:06d}_left_depth_raw.raw"), 'ab')
self.left_timestamp_file = open(
os.path.join(self.save_path, f"left_timestamp_{self.index:04d}.txt"), 'a')
os.path.join(self.save_path, f"{self.index:06d}_left_timestamp.txt"), 'a')
self.left_camera_info_file = open(
os.path.join(self.save_path, f"left_camera_info_{self.index:04d}.json"), 'w')
os.path.join(self.save_path, f"{self.index:06d}_left_camera_info.json"), 'w')
self.left_info_sub = self.create_subscription(
CameraInfo,
self.left_info_name,
@@ -172,11 +195,11 @@ class CollectDataNode(Node):
)
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')
os.path.join(self.save_path, f"{self.index:06d}_right_depth_raw.raw"), 'ab')
self.right_timestamp_file = open(
os.path.join(self.save_path, f"right_timestamp_{self.index:04d}.txt"), 'a')
os.path.join(self.save_path, f"{self.index:06d}_right_timestamp.txt"), 'a')
self.right_camera_info_file = open(
os.path.join(self.save_path, f"right_camera_info_{self.index:04d}.json"), 'w')
os.path.join(self.save_path, f"{self.index:06d}_right_camera_info.json"), 'w')
self.right_info_sub = self.create_subscription(
CameraInfo,
self.right_info_name,
@@ -185,11 +208,11 @@ class CollectDataNode(Node):
)
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')
os.path.join(self.save_path, f"{self.index:06d}_head_depth_raw.raw"), 'ab')
self.head_timestamp_file = open(
os.path.join(self.save_path, f"head_timestamp_{self.index:04d}.txt"), 'a')
os.path.join(self.save_path, f"{self.index:06d}_head_timestamp.txt"), 'a')
self.head_camera_info_file = open(
os.path.join(self.save_path, f"head_camera_info_{self.index:04d}.json"), 'w')
os.path.join(self.save_path, f"{self.index:06d}_head_camera_info.json"), 'w')
self.head_info_sub = self.create_subscription(
CameraInfo,
self.head_info_name,
@@ -211,65 +234,174 @@ class CollectDataNode(Node):
if self.left_depth_raw_file is not None:
self.left_depth_raw_file.close()
self.left_depth_raw_file = None
threading.Thread(
target=zip_tools.zip_worker,
args=(
f"{self.index:06d}_left_depth_raw.zip",
f"{self.index:06d}_left_depth_raw.raw",
self.save_path
),
daemon=True
).start()
# with zipfile.ZipFile(
# os.path.join(self.save_path, f"{self.index:06d}_left_depth_raw.zip"),
# "w",
# compression=zipfile.ZIP_DEFLATED
# ) as zf:
# zf.write(
# os.path.join(self.save_path, f"{self.index:06d}_left_depth_raw.raw"),
# arcname=os.path.basename(f"{self.index:06d}_left_depth_raw.raw")
# )
# if zip_tools.zip_verification(
# f"{self.index:06d}_left_depth_raw.zip",
# f"{self.index:06d}_left_depth_raw.raw",
# self.save_path
# ):
# os.remove(os.path.join(self.save_path, f"{self.index:06d}_left_depth_raw.raw"))
# else:
# self.get_logger().warning("zip file verification file, keep raw file")
# os.remove(os.path.join(self.save_path, f"{self.index:06d}_left_depth_raw.zip"))
self.get_logger().info(
f'left depth raw saved to {self.save_path}, index: {self.index}')
if self.left_timestamp_file is not None:
self.left_timestamp_file.close()
self.left_timestamp_file = None
self.get_logger().info(
f'left timestamp saved to {self.save_path}, index: {self.index}')
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.destroy_subscription(self.left_info_sub)
self.left_info_sub = None
if self.right_sign:
if self.right_depth_raw_file is not None:
self.right_depth_raw_file.close()
self.right_depth_raw_file = None
threading.Thread(
target=zip_tools.zip_worker,
args=(
f"{self.index:06d}_right_depth_raw.zip",
f"{self.index:06d}_right_depth_raw.raw",
self.save_path
),
daemon=True
).start()
# with zipfile.ZipFile(
# os.path.join(self.save_path, f"{self.index:06d}_right_depth_raw.zip"),
# "w",
# compression=zipfile.ZIP_DEFLATED
# ) as zf:
# zf.write(
# os.path.join(self.save_path, f"{self.index:06d}_right_depth_raw.raw"),
# arcname=os.path.basename(f"{self.index:06d}_right_depth_raw.raw")
# )
# if zip_tools.zip_verification(
# f"{self.index:06d}_right_depth_raw.zip",
# f"{self.index:06d}_right_depth_raw.raw",
# self.save_path
# ):
# os.remove(os.path.join(self.save_path, f"{self.index:06d}_right_depth_raw.raw"))
# else:
# self.get_logger().warning("zip file verification file, keep raw file")
# os.remove(os.path.join(self.save_path, f"{self.index:06d}_right_depth_raw.zip"))
self.get_logger().info(
f'right depth raw saved to {self.save_path}, index: {self.index}')
if self.right_timestamp_file is not None:
self.right_timestamp_file.close()
self.right_timestamp_file = None
self.get_logger().info(
f'right timrstamp saved to {self.save_path}, index: {self.index}')
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.destroy_subscription(self.right_info_sub)
self.right_info_sub = None
if self.head_sign:
if self.head_depth_raw_file is not None:
self.head_depth_raw_file.close()
self.head_depth_raw_file = None
threading.Thread(
target=zip_tools.zip_worker,
args=(
f"{self.index:06d}_head_depth_raw.zip",
f"{self.index:06d}_head_depth_raw.raw",
self.save_path
),
daemon=True
).start()
# with zipfile.ZipFile(
# os.path.join(self.save_path, f"{self.index:06d}_head_depth_raw.zip"),
# "w",
# compression=zipfile.ZIP_DEFLATED
# ) as zf:
# zf.write(
# os.path.join(self.save_path, f"{self.index:06d}_head_depth_raw.raw"),
# arcname=os.path.basename(f"{self.index:06d}_head_depth_raw.raw")
# )
# if zip_tools.zip_verification(
# f"{self.index:06d}_head_depth_raw.zip",
# f"{self.index:06d}_head_depth_raw.raw",
# self.save_path
# ):
# os.remove(os.path.join(self.save_path, f"{self.index:06d}_head_depth_raw.raw"))
# else:
# self.get_logger().warning("zip file verification file, keep raw file")
# os.remove(os.path.join(self.save_path, f"{self.index:06d}_head_depth_raw.zip"))
self.get_logger().info(
f'head depth raw saved to {self.save_path}, index: {self.index}')
if self.head_timestamp_file is not None:
self.head_timestamp_file.close()
self.head_timestamp_file = None
self.get_logger().info(
f'head timestamp saved to {self.save_path}, index: {self.index}')
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.destroy_subscription(self.head_info_sub)
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
self.get_logger().info(
f'left color video saved to {self.save_path}, index: {self.index}')
if self.left_depth_writer is not None:
self.left_depth_writer.release()
self.left_depth_writer = None
self.get_logger().info(
f'left depth video saved to {self.save_path}, index: {self.index}')
if self.right_color_writer is not None:
self.right_color_writer.release()
self.right_color_writer = None
self.get_logger().info(
f'right color video saved to {self.save_path}, index: {self.index}')
if self.right_depth_writer is not None:
self.right_depth_writer.release()
self.right_depth_writer = None
self.get_logger().info(
f'right depth video saved to {self.save_path}, index: {self.index}')
if self.head_color_writer is not None:
self.head_color_writer.release()
self.head_color_writer = None
self.get_logger().info(
f'head color video saved to {self.save_path}, index: {self.index}')
if self.head_depth_writer is not None:
self.head_depth_writer.release()
self.head_depth_writer = None
self.get_logger().info(
f'head depth video saved to {self.save_path}, index: {self.index}')
self.collect_sign = False
self.save_sign = True
response.success = False
response.success = True
response.message = "stop collecting data"
self.get_logger().info("stop collecting data")
return response
@@ -278,7 +410,7 @@ class CollectDataNode(Node):
if self.collect_sign:
try:
color_frame = self.cv_bridge.imgmsg_to_cv2(color_msg, desired_encoding='bgr8')
depth_frame = self.cv_bridge.imgmsg_to_cv2(depth_msg, desired_encoding='uint16')
depth_frame = self.cv_bridge.imgmsg_to_cv2(depth_msg, desired_encoding='16UC1')
except Exception as e:
self.get_logger().error(f'cv_bridge error: {e}')
return
@@ -287,17 +419,17 @@ class CollectDataNode(Node):
ch, cw = color_frame.shape[:2]
dh, dw = depth_frame.shape[:2]
color_fourcc = cv2.VideoWriter_fourcc(*'MP4V')
depth_fourcc = cv2.VideoWriter_fourcc(*'MP4V')
color_fourcc = cv2.VideoWriter_fourcc(*'mp4v')
depth_fourcc = cv2.VideoWriter_fourcc(*'mp4v')
self.left_color_writer = cv2.VideoWriter(
os.path.join(self.save_path, f"left_color_video_{self.index:04d}.mp4"),
os.path.join(self.save_path, f"{self.index:06d}_left_color_video.mp4"),
color_fourcc,
self.fps,
(cw, ch)
)
self.left_depth_writer = cv2.VideoWriter(
os.path.join(self.save_path, f"left_depth_video_{self.index:04d}.mp4"),
os.path.join(self.save_path, f"{self.index:06d}_left_depth_video.mp4"),
depth_fourcc,
self.fps,
(dw, dh)
@@ -329,7 +461,7 @@ class CollectDataNode(Node):
"frame_order": "row-major",
"fps": self.fps
}
with open(os.path.join(self.save_path, f"left_depth_meta_{self.index:04d}.json"), "w") as f:
with open(os.path.join(self.save_path, f"{self.index:06d}_left_depth_meta.json"), "w") as f:
json.dump(meta, f, indent=2)
self.left_meta = True
@@ -346,7 +478,7 @@ class CollectDataNode(Node):
if self.collect_sign:
try:
color_frame = self.cv_bridge.imgmsg_to_cv2(color_msg, desired_encoding='bgr8')
depth_frame = self.cv_bridge.imgmsg_to_cv2(depth_msg, desired_encoding='uint16')
depth_frame = self.cv_bridge.imgmsg_to_cv2(depth_msg, desired_encoding='16UC1')
except Exception as e:
self.get_logger().error(f'cv_bridge error: {e}')
return
@@ -355,17 +487,17 @@ class CollectDataNode(Node):
ch, cw = color_frame.shape[:2]
dh, dw = depth_frame.shape[:2]
color_fourcc = cv2.VideoWriter_fourcc(*'MP4V')
depth_fourcc = cv2.VideoWriter_fourcc(*'MP4V')
color_fourcc = cv2.VideoWriter_fourcc(*'mp4v')
depth_fourcc = cv2.VideoWriter_fourcc(*'mp4v')
self.right_color_writer = cv2.VideoWriter(
os.path.join(self.save_path, f"right_color_video_{self.index:04d}.mp4"),
os.path.join(self.save_path, f"{self.index:06d}_right_color_video.mp4"),
color_fourcc,
self.fps,
(cw, ch)
)
self.right_depth_writer = cv2.VideoWriter(
os.path.join(self.save_path, f"right_depth_video_{self.index:04d}.mp4"),
os.path.join(self.save_path, f"{self.index:06d}_right_depth_video.mp4"),
depth_fourcc,
self.fps,
(dw, dh)
@@ -397,7 +529,7 @@ class CollectDataNode(Node):
"frame_order": "row-major",
"fps": self.fps
}
with open(os.path.join(self.save_path, f"right_depth_meta_{self.index:04d}.json"), "w") as f:
with open(os.path.join(self.save_path, f"{self.index:06d}_right_depth_meta.json"), "w") as f:
json.dump(meta, f, indent=2)
self.right_meta = True
@@ -414,7 +546,7 @@ class CollectDataNode(Node):
if self.collect_sign:
try:
color_frame = self.cv_bridge.imgmsg_to_cv2(color_msg, desired_encoding='bgr8')
depth_frame = self.cv_bridge.imgmsg_to_cv2(depth_msg, desired_encoding='uint16')
depth_frame = self.cv_bridge.imgmsg_to_cv2(depth_msg, desired_encoding='16UC1')
except Exception as e:
self.get_logger().error(f'cv_bridge error: {e}')
return
@@ -423,17 +555,17 @@ class CollectDataNode(Node):
ch, cw = color_frame.shape[:2]
dh, dw = depth_frame.shape[:2]
color_fourcc = cv2.VideoWriter_fourcc(*'MP4V')
depth_fourcc = cv2.VideoWriter_fourcc(*'MP4V')
color_fourcc = cv2.VideoWriter_fourcc(*'mp4v')
depth_fourcc = cv2.VideoWriter_fourcc(*'mp4v')
self.head_color_writer = cv2.VideoWriter(
os.path.join(self.save_path, f"head_color_video_{self.index:04d}.mp4"),
os.path.join(self.save_path, f"{self.index:06d}_head_color_video.mp4"),
color_fourcc,
self.fps,
(cw, ch)
)
self.head_depth_writer = cv2.VideoWriter(
os.path.join(self.save_path, f"head_depth_video_{self.index:04d}.mp4"),
os.path.join(self.save_path, f"{self.index:06d}_head_depth_video.mp4"),
depth_fourcc,
self.fps,
(dw, dh)
@@ -465,7 +597,7 @@ class CollectDataNode(Node):
"frame_order": "row-major",
"fps": self.fps
}
with open(os.path.join(self.save_path, f"head_depth_meta_{self.index:04d}.json"), "w") as f:
with open(os.path.join(self.save_path, f"{self.index:06d}_head_depth_meta.json"), "w") as f:
json.dump(meta, f, indent=2)
self.head_meta = True
@@ -480,36 +612,39 @@ class CollectDataNode(Node):
def left_info_callback(self, msgs):
data = {
"K": msgs.k,
"D": msgs.d,
"K": msgs.k.tolist(),
"D": msgs.d.tolist(),
}
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.get_logger().info(f'left camera info saved to {self.save_path}, index: {self.index}')
self.destroy_subscription(self.left_info_sub)
self.left_info_sub = None
def right_info_callback(self, msgs):
data = {
"K": msgs.k,
"D": msgs.d,
"K": msgs.k.tolist(),
"D": msgs.d.tolist(),
}
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.get_logger().info(f'right camera info saved to {self.save_path}, index: {self.index}')
self.destroy_subscription(self.right_info_sub)
self.right_info_sub = None
def head_info_callback(self, msgs):
data = {
"K": msgs.k,
"D": msgs.d,
"K": msgs.k.tolist(),
"D": msgs.d.tolist(),
}
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.get_logger().info(f'head camera info saved to {self.save_path}, index: {self.index}')
self.destroy_subscription(self.head_info_sub)
self.head_info_sub = None
def topic_exists(self, topic_name: str) -> bool:

0
vision_tools/vision_tools/utils/__init__.py Normal file
View File

38
vision_tools/vision_tools/utils/zip_tools.py Normal file
View File

@@ -0,0 +1,38 @@
import os
import zipfile
import logging
__all__ = ["zip_worker", "zip_verification"]
def zip_worker(zip_name:str, raw_name:str, save_path):
with zipfile.ZipFile(
os.path.join(save_path, zip_name),
"w",
compression=zipfile.ZIP_DEFLATED
) as zf:
zf.write(
os.path.join(save_path, raw_name),
arcname=os.path.basename(raw_name)
)
if zip_verification(zip_name, raw_name, save_path):
os.remove(os.path.join(save_path, raw_name))
else:
logging.warning("zip file verification file, keep raw file")
os.remove(os.path.join(save_path, raw_name))
def zip_verification(zip_name:str, raw_name, save_path):
with zipfile.ZipFile(os.path.join(save_path, zip_name), "r") as zf:
zf.extract(raw_name, path="./_tmp")
if os.path.getsize(os.path.join(save_path, raw_name)) == os.path.getsize(f"./_tmp/{raw_name}"):
os.remove(os.path.join(f"./_tmp/{raw_name}"))
# os.rmdir("./_tmp")
return True
os.remove(os.path.join(f"./_tmp/{raw_name}"))
# os.rmdir("./_tmp")
return False