Merge branch 'feature' into feature_cpp_test

vision_detect/configs/launch_configs/medical_sense.json

@@ -25,6 +25,6 @@
     "PCA_configs": {
         "depth_scale": 1000.0,
         "depth_trunc": 3.0,
-        "voxel_size": 0.010
+        "voxel_size": 0.001
     }
 }

vision_detect/setup.py

@@ -15,7 +15,6 @@ setup(
         ('share/' + package_name + '/configs/flexiv_configs', glob('configs/flexiv_configs/*.json')),
         ('share/' + package_name + '/configs/hand_eye_mat', glob('configs/hand_eye_mat/*.json')),
         ('share/' + package_name + '/configs/launch_configs', glob('configs/launch_configs/*.json')),
-        ('share/' + package_name + '/configs/other_configs', glob('configs/other_configs/*.json')),
 
         ('share/' + package_name + '/checkpoints', glob('checkpoints/*.pt')),
         ('share/' + package_name + '/checkpoints', glob('checkpoints/*.onnx')),
@@ -40,12 +39,8 @@ setup(
 
             'sub_pose_node = vision_detect.sub_pose:main',
             'service_client_node = vision_detect.service_client:main',
-            'get_camera_pose_node = vision_detect.get_camera_pose:main',
             
             'detect_node = vision_detect.detect_node:main',
-
-            'collect_data_node = vision_detect.collect_data_node:main',
-            'collect_data_client = vision_detect.collect_data_client:main',
         ],
     },
 )

vision_detect/vision_detect/VisionDetect/node/detect_node.py

@@ -175,6 +175,8 @@ class DetectNode(InitBase):
         sorted_index = np.lexsort((-y_centers, x_centers))
         masks = masks[sorted_index]
         boxes = boxes[sorted_index]
+        class_ids = class_ids[sorted_index]
+        labels = labels[sorted_index]
 
         time3 = time.time()
 

vision_detect/vision_detect/VisionDetect/utils/calculate_tools.py

@@ -1,7 +1,7 @@
 """计算工具"""
 import logging
 
-import cv2
+# import cv2
 import numpy as np
 import open3d as o3d
 import transforms3d as tfs
@@ -53,7 +53,7 @@ def calculate_pose_pca(
         depth_trunc=kwargs.get("depth_trunc", 3.0),
     )
 
-    point_cloud = point_cloud_denoising(point_cloud, kwargs.get("voxel_size", 0.005))
+    point_cloud = point_cloud_denoising(point_cloud, kwargs.get("voxel_size", 0.002))
     if point_cloud is None:
         return None, [0.0, 0.0, 0.0]
 
@@ -158,11 +158,11 @@ def point_cloud_denoising(point_cloud: o3d.geometry.PointCloud, voxel_size: floa
     """点云去噪"""
     point_cloud = point_cloud.remove_non_finite_points()
     while True:
-        down_pcd = point_cloud.voxel_down_sample(voxel_size=voxel_size * 0.5)
+        down_pcd = point_cloud.voxel_down_sample(voxel_size=voxel_size)
         point_num = len(down_pcd.points)
         if point_num <= 1000:
             break
-        voxel_size *= 2
+        voxel_size *= 4
 
     # logging.fatal("point_cloud_denoising: point_num={}".format(len(point_cloud.points)))
     # logging.fatal("point_cloud_denoising: point_num={}".format(point_num))
@@ -170,7 +170,7 @@ def point_cloud_denoising(point_cloud: o3d.geometry.PointCloud, voxel_size: floa
     # 半径滤波
     clean_pcd, _ = down_pcd.remove_radius_outlier(
         nb_points=10,
-        radius=voxel_size * 5
+        radius=voxel_size * 10
         )
 
     # 统计滤波
@@ -189,7 +189,7 @@ def point_cloud_denoising(point_cloud: o3d.geometry.PointCloud, voxel_size: floa
     # if largest_cluster_ratio < 0.5:
     #     return None
 
-    labels = np.array(clean_pcd.cluster_dbscan(eps=voxel_size * 5, min_points=10))
+    labels = np.array(clean_pcd.cluster_dbscan(eps=voxel_size * 10, min_points=10))
     if len(labels[labels >= 0]) == 0:
         return clean_pcd
 

vision_detect/vision_detect/VisionDetect/utils/draw_tools.py

@@ -25,6 +25,12 @@ def draw_box(
     class_ids = segmentation_result.boxes.cls.cpu().numpy()
     labels = segmentation_result.names
 
+    x_centers, y_centers = boxes[:, 0], boxes[:, 1]
+    sorted_index = np.lexsort((-y_centers, x_centers))
+    boxes = boxes[sorted_index]
+    class_ids = class_ids[sorted_index]
+    labels = labels[sorted_index]
+
     for i, box in enumerate(boxes):
         x_center, y_center, width, height = box[:4]
 

vision_detect/vision_detect/collect_data_node.py

@@ -1,408 +0,0 @@
-import os
-import json
-
-import cv2
-
-import rclpy
-from rclpy.node import Node
-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 std_srvs.srv import SetBool
-
-
-share_dir = get_package_share_directory('vision_detect')
-
-class CollectDataNode(Node):
-    def __init__(self):
-        super().__init__("collect_data_node")
-        self.collect_sign = False
-        self.save_sign = False
-        self.cv_bridge = CvBridge()
-
-        self.left_sign = self.right_sign = self.head_sign = False
-        self.left_meta = self.right_meta = self.head_meta = False
-
-        self.left_depth_raw_file = None
-        self.right_depth_raw_file = None
-        self.head_depth_raw_file = None
-
-        self.left_timestamp_file = None
-        self.right_timestamp_file = None
-        self.head_timestamp_file = None
-
-        self.index = 0
-        self.fps = 30.0
-
-        self.left_color_writer = None
-        self.right_color_writer = None
-        self.head_color_writer = None
-
-        self.left_depth_writer = None
-        self.right_depth_writer = None
-        self.head_depth_writer = None
-
-        self.right_writer_initialized = False
-        self.left_writer_initialized = False
-        self.head_writer_initialized = False
-
-        with open(os.path.join(share_dir, 'configs/other_configs/collect_data_node.json'), 'r')as f:
-            configs = json.load(f)
-
-        home_path = os.path.expanduser("~")
-        self.save_path = os.path.join(home_path, configs['save_path'])
-        if self.save_path:
-            os.makedirs(self.save_path, exist_ok=True)
-
-        # left
-        if self.topic_exists(configs["left"]["color"]) and self.topic_exists(
-                configs["left"]["depth"]):
-            self.left_sign = True
-            self.sub_left_color = Subscriber(self, Image, configs["left"]["color"])
-            self.sub_left_depth = Subscriber(self, Image, configs["left"]["depth"])
-
-            self.left_sync_subscriber = ApproximateTimeSynchronizer(
-                [self.sub_left_color, self.sub_left_depth],
-                queue_size=10,
-                slop=0.1
-            )
-            self.left_sync_subscriber.registerCallback(self.left_callback)
-
-        # right
-        if self.topic_exists(configs["right"]["color"]) and self.topic_exists(
-                configs["right"]["depth"]):
-            self.right_sign = True
-            self.sub_right_color = Subscriber(self, Image, configs["right"]["color"])
-            self.sub_right_depth = Subscriber(self, Image, configs["right"]["depth"])
-
-            self.right_sync_subscriber = ApproximateTimeSynchronizer(
-                [self.sub_right_color, self.sub_right_depth],
-                queue_size=10,
-                slop=0.1
-            )
-            self.right_sync_subscriber.registerCallback(self.right_callback)
-
-        # head
-        if self.topic_exists(configs["head"]["color"]) and self.topic_exists(
-                configs["head"]["depth"]):
-            self.head_sign = True
-            self.sub_head_color = Subscriber(self, Image, configs["head"]["color"])
-            self.sub_head_depth = Subscriber(self, Image, configs["head"]["depth"])
-
-            self.head_sync_subscriber = ApproximateTimeSynchronizer(
-                [self.sub_head_color, self.sub_head_depth],
-                queue_size=10,
-                slop=0.1
-            )
-            self.head_sync_subscriber.registerCallback(self.head_callback)
-
-        self.service = self.create_service(SetBool, "/collect_data", self.service_callback)
-
-    def service_callback(self, request, response):
-        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"))):
-                    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')
-                self.left_timestamp_file = open(
-                    os.path.join(self.save_path, f"left_timestamp_{self.index:04d}.txt"), 'a')
-            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')
-            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.left_writer_initialized = False
-            self.right_writer_initialized = False
-            self.head_writer_initialized = False
-
-            self.collect_sign = True
-            response.success = True
-            response.message = "start collecting data"
-            return response
-
-        if self.left_sign:
-            if self.left_depth_raw_file is not None:
-                self.left_depth_raw_file.close()
-                self.left_depth_raw_file = None
-            if self.left_timestamp_file is not None:
-                self.left_timestamp_file.close()
-                self.left_timestamp_file = 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
-            if self.right_timestamp_file is not None:
-                self.right_timestamp_file.close()
-                self.right_timestamp_file = 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
-            if self.head_timestamp_file is not None:
-                self.head_timestamp_file.close()
-                self.head_timestamp_file = None
-
-        if self.left_color_writer is not None:
-            self.left_color_writer.release()
-            self.left_color_writer = None
-        if self.left_depth_writer is not None:
-            self.left_depth_writer.release()
-            self.left_depth_writer = None
-        if self.right_color_writer is not None:
-            self.right_color_writer.release()
-            self.right_color_writer = None
-        if self.right_depth_writer is not None:
-            self.right_depth_writer.release()
-            self.right_depth_writer = None
-        if self.head_color_writer is not None:
-            self.head_color_writer.release()
-            self.head_color_writer = None
-        if self.head_depth_writer is not None:
-            self.head_depth_writer.release()
-            self.head_depth_writer = None
-
-        self.collect_sign = False
-        self.save_sign = True
-        response.success = False
-        response.message = "stop collecting data"
-        return response
-
-    def left_callback(self, color_msg, depth_msg):
-        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')
-            except Exception as e:
-                self.get_logger().error(f'cv_bridge error: {e}')
-                return
-
-            if not self.left_writer_initialized:
-                ch, cw = color_frame.shape[:2]
-                dh, dw = depth_frame.shape[:2]
-
-                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"),
-                    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"),
-                    depth_fourcc,
-                    self.fps,
-                    (dw, dh)
-                )
-
-                if not self.left_color_writer.isOpened() or not self.left_depth_writer.isOpened():
-                    self.get_logger().error('Failed to open VideoWriter')
-                    return
-
-                self.left_writer_initialized = True
-                self.get_logger().info('VideoWriter initialized')
-
-            # RAW
-            if not depth_frame.flags['C_CONTIGUOUS']:
-                depth_frame = depth_frame.copy()
-            if self.left_depth_raw_file is not None:
-                self.left_depth_raw_file.write(depth_frame.tobytes())
-            if self.left_timestamp_file is not None:
-                ts = depth_msg.header.stamp.sec * 1_000_000_000 + depth_msg.header.stamp.nanosec
-                self.left_timestamp_file.write(f"{ts}\n")
-
-            if not self.left_meta:
-                meta = {
-                    "width": int(depth_frame.shape[1]),
-                    "height": int(depth_frame.shape[0]),
-                    "dtype": "uint16",
-                    "endianness": "little",
-                    "unit": "mm",
-                    "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:
-                    json.dump(meta, f, indent=2)
-                self.left_meta = True
-
-            # MP4V
-            self.left_color_writer.write(color_frame)
-            depth_vis = cv2.convertScaleAbs(
-                depth_frame,
-                alpha=255.0 / 10000.0  # 10m，根据相机改
-            )
-            depth_vis = cv2.cvtColor(depth_vis, cv2.COLOR_GRAY2BGR)
-            self.left_depth_writer.write(depth_vis)
-
-    def right_callback(self, color_msg, depth_msg):
-        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')
-            except Exception as e:
-                self.get_logger().error(f'cv_bridge error: {e}')
-                return
-
-            if not self.right_writer_initialized:
-                ch, cw = color_frame.shape[:2]
-                dh, dw = depth_frame.shape[:2]
-
-                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"),
-                    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"),
-                    depth_fourcc,
-                    self.fps,
-                    (dw, dh)
-                )
-
-                if not self.right_color_writer.isOpened() or not self.right_depth_writer.isOpened():
-                    self.get_logger().error('Failed to open VideoWriter')
-                    return
-
-                self.right_writer_initialized = True
-                self.get_logger().info('VideoWriter initialized')
-
-            # RAW
-            if not depth_frame.flags['C_CONTIGUOUS']:
-                depth_frame = depth_frame.copy()
-            if self.right_depth_raw_file is not None:
-                self.right_depth_raw_file.write(depth_frame.tobytes())
-            if self.right_timestamp_file is not None:
-                ts = depth_msg.header.stamp.sec * 1_000_000_000 + depth_msg.header.stamp.nanosec
-                self.right_timestamp_file.write(f"{ts}\n")
-
-            if not self.right_meta:
-                meta = {
-                    "width": int(depth_frame.shape[1]),
-                    "height": int(depth_frame.shape[0]),
-                    "dtype": "uint16",
-                    "endianness": "little",
-                    "unit": "mm",
-                    "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:
-                    json.dump(meta, f, indent=2)
-                self.right_meta = True
-
-            # MP4V
-            self.right_color_writer.write(color_frame)
-            depth_vis = cv2.convertScaleAbs(
-                depth_frame,
-                alpha=255.0 / 10000.0  # 10m，根据相机改
-            )
-            depth_vis = cv2.cvtColor(depth_vis, cv2.COLOR_GRAY2BGR)
-            self.right_depth_writer.write(depth_vis)
-
-    def head_callback(self, color_msg, depth_msg):
-        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')
-            except Exception as e:
-                self.get_logger().error(f'cv_bridge error: {e}')
-                return
-
-            if not self.head_writer_initialized:
-                ch, cw = color_frame.shape[:2]
-                dh, dw = depth_frame.shape[:2]
-
-                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"),
-                    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"),
-                    depth_fourcc,
-                    self.fps,
-                    (dw, dh)
-                )
-
-                if not self.head_color_writer.isOpened() or not self.head_depth_writer.isOpened():
-                    self.get_logger().error('Failed to open VideoWriter')
-                    return
-
-                self.head_writer_initialized = True
-                self.get_logger().info('VideoWriter initialized')
-
-            # RAW
-            if not depth_frame.flags['C_CONTIGUOUS']:
-                depth_frame = depth_frame.copy()
-            if self.head_depth_raw_file is not None:
-                self.head_depth_raw_file.write(depth_frame.tobytes())
-            if self.head_timestamp_file is not None:
-                ts = depth_msg.header.stamp.sec * 1_000_000_000 + depth_msg.header.stamp.nanosec
-                self.head_timestamp_file.write(f"{ts}\n")
-
-            if not self.head_meta:
-                meta = {
-                    "width": int(depth_frame.shape[1]),
-                    "height": int(depth_frame.shape[0]),
-                    "dtype": "uint16",
-                    "endianness": "little",
-                    "unit": "mm",
-                    "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:
-                    json.dump(meta, f, indent=2)
-                self.head_meta = True
-
-            # MP4V
-            self.head_color_writer.write(color_frame)
-            depth_vis = cv2.convertScaleAbs(
-                depth_frame,
-                alpha=255.0 / 10000.0  # 10m，根据相机改
-            )
-            depth_vis = cv2.cvtColor(depth_vis, cv2.COLOR_GRAY2BGR)
-            self.head_depth_writer.write(depth_vis)
-
-    def topic_exists(self, topic_name: str) -> bool:
-        topics = self.get_topic_names_and_types()
-        return any(name == topic_name for name, _ in topics)
-
-    # def destroy_node(self):
-    #     self.get_logger().info('VideoWriter released')
-    #     super().destroy_node()
-
-
-def main(args=None):
-    rclpy.init(args=args)
-    node = CollectDataNode()
-    try:
-        rclpy.spin(node)
-    except KeyboardInterrupt:
-        pass
-    finally:
-        node.destroy_node()
-        rclpy.shutdown()

vision_tools/configs/collect_data_node.json

@@ -1,5 +1,5 @@
 {
-  "save_path": "/collect_data",
+  "save_path": "collect_data",
   "left": {
     "color": "/camera/camera/color/image_raw",
     "depth": "/camera/camera/aligned_depth_to_color/image_raw",

vision_tools/package.xml

@@ -0,0 +1,20 @@
+<?xml version="1.0"?>
+<?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
+<package format="3">
+  <name>vision_tools</name>
+  <version>0.0.0</version>
+  <description>TODO: Package description</description>
+  <maintainer email="16126883+liangyuxuan123@user.noreply.gitee.com">lyx</maintainer>
+  <license>TODO: License declaration</license>
+
+  <depend>rclpy</depend>
+
+  <test_depend>ament_copyright</test_depend>
+  <test_depend>ament_flake8</test_depend>
+  <test_depend>ament_pep257</test_depend>
+  <test_depend>python3-pytest</test_depend>
+
+  <export>
+    <build_type>ament_python</build_type>
+  </export>
+</package>

vision_tools/setup.cfg

@@ -0,0 +1,4 @@
+[develop]
+script_dir=$base/lib/vision_tools
+[install]
+install_scripts=$base/lib/vision_tools

vision_tools/setup.py

@@ -0,0 +1,30 @@
+from glob import glob
+from setuptools import find_packages, setup
+
+package_name = 'vision_tools'
+
+setup(
+    name=package_name,
+    version='0.0.0',
+    packages=find_packages(exclude=['test']),
+    data_files=[
+        ('share/ament_index/resource_index/packages', ['resource/' + package_name]),
+        ('share/' + package_name, ['package.xml']),
+        ('share/' + package_name + '/configs', glob('configs/*.json')),
+    ],
+    install_requires=['setuptools'],
+    zip_safe=True,
+    maintainer='lyx',
+    maintainer_email='lyx@todo.todo',
+    description='TODO: Package description',
+    license='TODO: License declaration',
+    tests_require=['pytest'],
+    entry_points={
+        'console_scripts': [
+            'collect_data_node = vision_tools.collect_data_node:main',
+            'collect_data_client = vision_tools.collect_data_client:main',
+
+            'get_camera_pose_node = vision_tools.get_camera_pose:main',
+        ],
+    },
+)

vision_tools/test/test_copyright.py

@@ -0,0 +1,25 @@
+# Copyright 2015 Open Source Robotics Foundation, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ament_copyright.main import main
+import pytest
+
+
+# Remove the `skip` decorator once the source file(s) have a copyright header
+@pytest.mark.skip(reason='No copyright header has been placed in the generated source file.')
+@pytest.mark.copyright
+@pytest.mark.linter
+def test_copyright():
+    rc = main(argv=['.', 'test'])
+    assert rc == 0, 'Found errors'

vision_tools/test/test_flake8.py

@@ -0,0 +1,25 @@
+# Copyright 2017 Open Source Robotics Foundation, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ament_flake8.main import main_with_errors
+import pytest
+
+
+@pytest.mark.flake8
+@pytest.mark.linter
+def test_flake8():
+    rc, errors = main_with_errors(argv=[])
+    assert rc == 0, \
+        'Found %d code style errors / warnings:\n' % len(errors) + \
+        '\n'.join(errors)

vision_tools/test/test_pep257.py

@@ -0,0 +1,23 @@
+# Copyright 2015 Open Source Robotics Foundation, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ament_pep257.main import main
+import pytest
+
+
+@pytest.mark.linter
+@pytest.mark.pep257
+def test_pep257():
+    rc = main(argv=['.', 'test'])
+    assert rc == 0, 'Found code style errors / warnings'

vision_tools/vision_tools/collect_data_client.py

@@ -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:
@@ -33,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)

vision_tools/vision_tools/collect_data_node.py

@@ -0,0 +1,668 @@
+import os
+import json
+import zipfile
+import threading
+
+import cv2
+import numpy as np
+
+import rclpy
+from rclpy.node import Node
+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, 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")
+        self.collect_sign = False
+        self.save_sign = False
+        self.cv_bridge = CvBridge()
+
+        self.left_sign = self.right_sign = self.head_sign = False
+        self.left_meta = self.right_meta = self.head_meta = False
+
+        self.left_depth_raw_file = None
+        self.right_depth_raw_file = None
+        self.head_depth_raw_file = None
+
+        self.left_timestamp_file = None
+        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
+
+        self.left_color_writer = None
+        self.right_color_writer = None
+        self.head_color_writer = None
+
+        self.left_depth_writer = None
+        self.right_depth_writer = None
+        self.head_depth_writer = None
+
+        self.right_writer_initialized = False
+        self.left_writer_initialized = False
+        self.head_writer_initialized = False
+
+        with open(os.path.join(share_dir, 'configs/collect_data_node.json'), 'r')as f:
+            configs = json.load(f)
+
+        home_path = os.path.expanduser("~")
+        self.save_path = os.path.join(home_path, configs['save_path'])
+        if self.save_path:
+            os.makedirs(self.save_path, exist_ok=True)
+
+        # left
+        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"])
+
+            self.left_sync_subscriber = ApproximateTimeSynchronizer(
+                [self.sub_left_color, self.sub_left_depth],
+                queue_size=10,
+                slop=0.1
+            )
+            self.left_sync_subscriber.registerCallback(self.left_callback)
+            # Camera 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(
+                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"])
+
+            self.right_sync_subscriber = ApproximateTimeSynchronizer(
+                [self.sub_right_color, self.sub_right_depth],
+                queue_size=10,
+                slop=0.1
+            )
+            self.right_sync_subscriber.registerCallback(self.right_callback)
+            # Camera 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(
+                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"])
+
+            self.head_sync_subscriber = ApproximateTimeSynchronizer(
+                [self.sub_head_color, self.sub_head_depth],
+                queue_size=10,
+                slop=0.1
+            )
+            self.head_sync_subscriber.registerCallback(self.head_callback)
+            # Camera 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")
+
+    def service_callback(self, request, response):
+        # 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"{self.index:06d}_left_color_video.mp4"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_right_color_video.mp4"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_head_color_video.mp4"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_left_depth_video.mp4"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_right_depth_video.mp4"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_head_depth_video.mp4"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_left_depth_raw.raw"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_right_depth_raw.raw"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_head_depth_raw.raw"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_left_depth_raw.zip"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_right_depth_raw.zip"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_head_depth_raw.zip"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_left_timestamp.txt"))
+                    or os.path.exists(os.path.join(
+                        self.save_path, f"{self.index:06d}_right_timestamp.txt"))
+                    or os.path.exists(os.path.join(
+                        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"{self.index:06d}_left_depth_raw.raw"), 'ab')
+                self.left_timestamp_file = open(
+                    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"{self.index:06d}_left_camera_info.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"{self.index:06d}_right_depth_raw.raw"), 'ab')
+                self.right_timestamp_file = open(
+                    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"{self.index:06d}_right_camera_info.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"{self.index:06d}_head_depth_raw.raw"), 'ab')
+                self.head_timestamp_file = open(
+                    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"{self.index:06d}_head_camera_info.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
+            self.head_writer_initialized = False
+
+            self.collect_sign = True
+            response.success = True
+            response.message = "start collecting data"
+            self.get_logger().info("start collecting data")
+            return response
+
+        if self.left_sign:
+            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.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.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.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 = True
+        response.message = "stop collecting data"
+        self.get_logger().info("stop collecting data")
+        return response
+
+    def left_callback(self, color_msg, depth_msg):
+        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='16UC1')
+            except Exception as e:
+                self.get_logger().error(f'cv_bridge error: {e}')
+                return
+
+            if not self.left_writer_initialized:
+                ch, cw = color_frame.shape[:2]
+                dh, dw = depth_frame.shape[:2]
+
+                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"{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"{self.index:06d}_left_depth_video.mp4"),
+                    depth_fourcc,
+                    self.fps,
+                    (dw, dh)
+                )
+
+                if not self.left_color_writer.isOpened() or not self.left_depth_writer.isOpened():
+                    self.get_logger().error('Failed to open VideoWriter')
+                    return
+
+                self.left_writer_initialized = True
+                self.get_logger().info('VideoWriter initialized')
+
+            # RAW
+            if not depth_frame.flags['C_CONTIGUOUS']:
+                depth_frame = depth_frame.copy()
+            if self.left_depth_raw_file is not None:
+                self.left_depth_raw_file.write(depth_frame.tobytes())
+            if self.left_timestamp_file is not None:
+                ts = depth_msg.header.stamp.sec * 1_000_000_000 + depth_msg.header.stamp.nanosec
+                self.left_timestamp_file.write(f"{ts}\n")
+
+            if not self.left_meta:
+                meta = {
+                    "width": int(depth_frame.shape[1]),
+                    "height": int(depth_frame.shape[0]),
+                    "dtype": "uint16",
+                    "endianness": "little",
+                    "unit": "mm",
+                    "frame_order": "row-major",
+                    "fps": self.fps
+                }
+                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
+
+            # MP4V
+            self.left_color_writer.write(color_frame)
+            depth_vis = cv2.convertScaleAbs(
+                depth_frame,
+                alpha=255.0 / 10000.0  # 10m，根据相机改
+            )
+            depth_vis = cv2.cvtColor(depth_vis, cv2.COLOR_GRAY2BGR)
+            self.left_depth_writer.write(depth_vis)
+
+    def right_callback(self, color_msg, depth_msg):
+        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='16UC1')
+            except Exception as e:
+                self.get_logger().error(f'cv_bridge error: {e}')
+                return
+
+            if not self.right_writer_initialized:
+                ch, cw = color_frame.shape[:2]
+                dh, dw = depth_frame.shape[:2]
+
+                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"{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"{self.index:06d}_right_depth_video.mp4"),
+                    depth_fourcc,
+                    self.fps,
+                    (dw, dh)
+                )
+
+                if not self.right_color_writer.isOpened() or not self.right_depth_writer.isOpened():
+                    self.get_logger().error('Failed to open VideoWriter')
+                    return
+
+                self.right_writer_initialized = True
+                self.get_logger().info('VideoWriter initialized')
+
+            # RAW
+            if not depth_frame.flags['C_CONTIGUOUS']:
+                depth_frame = depth_frame.copy()
+            if self.right_depth_raw_file is not None:
+                self.right_depth_raw_file.write(depth_frame.tobytes())
+            if self.right_timestamp_file is not None:
+                ts = depth_msg.header.stamp.sec * 1_000_000_000 + depth_msg.header.stamp.nanosec
+                self.right_timestamp_file.write(f"{ts}\n")
+
+            if not self.right_meta:
+                meta = {
+                    "width": int(depth_frame.shape[1]),
+                    "height": int(depth_frame.shape[0]),
+                    "dtype": "uint16",
+                    "endianness": "little",
+                    "unit": "mm",
+                    "frame_order": "row-major",
+                    "fps": self.fps
+                }
+                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
+
+            # MP4V
+            self.right_color_writer.write(color_frame)
+            depth_vis = cv2.convertScaleAbs(
+                depth_frame,
+                alpha=255.0 / 10000.0  # 10m，根据相机改
+            )
+            depth_vis = cv2.cvtColor(depth_vis, cv2.COLOR_GRAY2BGR)
+            self.right_depth_writer.write(depth_vis)
+
+    def head_callback(self, color_msg, depth_msg):
+        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='16UC1')
+            except Exception as e:
+                self.get_logger().error(f'cv_bridge error: {e}')
+                return
+
+            if not self.head_writer_initialized:
+                ch, cw = color_frame.shape[:2]
+                dh, dw = depth_frame.shape[:2]
+
+                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"{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"{self.index:06d}_head_depth_video.mp4"),
+                    depth_fourcc,
+                    self.fps,
+                    (dw, dh)
+                )
+
+                if not self.head_color_writer.isOpened() or not self.head_depth_writer.isOpened():
+                    self.get_logger().error('Failed to open VideoWriter')
+                    return
+
+                self.head_writer_initialized = True
+                self.get_logger().info('VideoWriter initialized')
+
+            # RAW
+            if not depth_frame.flags['C_CONTIGUOUS']:
+                depth_frame = depth_frame.copy()
+            if self.head_depth_raw_file is not None:
+                self.head_depth_raw_file.write(depth_frame.tobytes())
+            if self.head_timestamp_file is not None:
+                ts = depth_msg.header.stamp.sec * 1_000_000_000 + depth_msg.header.stamp.nanosec
+                self.head_timestamp_file.write(f"{ts}\n")
+
+            if not self.head_meta:
+                meta = {
+                    "width": int(depth_frame.shape[1]),
+                    "height": int(depth_frame.shape[0]),
+                    "dtype": "uint16",
+                    "endianness": "little",
+                    "unit": "mm",
+                    "frame_order": "row-major",
+                    "fps": self.fps
+                }
+                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
+
+            # MP4V
+            self.head_color_writer.write(color_frame)
+            depth_vis = cv2.convertScaleAbs(
+                depth_frame,
+                alpha=255.0 / 10000.0  # 10m，根据相机改
+            )
+            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.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.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.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.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.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.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:
+        topics = self.get_topic_names_and_types()
+        return any(name == topic_name for name, _ in topics)
+
+    # def destroy_node(self):
+    #     self.get_logger().info('VideoWriter released')
+    #     super().destroy_node()
+
+
+def main(args=None):
+    rclpy.init(args=args)
+    node = CollectDataNode()
+    try:
+        rclpy.spin(node)
+    except KeyboardInterrupt:
+        pass
+    finally:
+        node.destroy_node()
+        rclpy.shutdown()

vision_tools/vision_tools/utils/zip_tools.py

@@ -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
+
+