update @huiyu

This commit is contained in:
NuoDaJia02
2026-01-05 09:40:30 +08:00
parent 727b191a2a
commit f60a86b0fd
25 changed files with 1166 additions and 52 deletions

1
.gitignore vendored
View File

@@ -3,3 +3,4 @@
**/build/
**/install/
**/log/
imgs*/*

View File

@@ -0,0 +1,28 @@
{
"T": [
[
0.01868075138315295,
0.8881359544145228,
0.45920099738999437,
-0.19284748723617093
],
[
0.002854261818715398,
-0.4593266423892484,
0.8882628489252997,
0.057154511710361816
],
[
0.9998214254141742,
-0.01528273756969839,
-0.011115539354645598,
-0.04187423675125192
],
[
0.0,
0.0,
0.0,
1.0
]
]
}

View File

@@ -1,22 +1,22 @@
{
"T": [
[
1.0,
0.0,
0.0,
0.0
0.013635791720580838,
0.9992765703636767,
-0.03550212819481636,
-0.08645650535173793
],
[
0.0,
1.0,
0.0,
0.0
-0.9998865276218899,
0.013399556063222661,
-0.006883587549256321,
0.036196137264974
],
[
0.0,
0.0,
1.0,
0.0
-0.006402895000908794,
0.03559196285001416,
0.999345893630474,
0.014407883180676354
],
[
0.0,

View File

@@ -1,22 +1,22 @@
{
"T": [
[
1.0,
0.0,
0.0,
0.0
0.01868075138315295,
0.8881359544145228,
0.45920099738999437,
-0.19284748723617093
],
[
0.0,
1.0,
0.0,
0.0
0.002854261818715398,
-0.4593266423892484,
0.8882628489252997,
0.057154511710361816
],
[
0.0,
0.0,
1.0,
0.0
0.9998214254141742,
-0.01528273756969839,
-0.011115539354645598,
-0.04187423675125192
],
[
0.0,
@@ -25,4 +25,4 @@
1.0
]
]
}
}

View File

@@ -17,7 +17,7 @@
"detect_mode": "Detect",
"Detect_configs": {
"checkpoint_path": "checkpoints/medical_sense-seg.pt",
"confidence": 0.50,
"confidence": 0.70,
"classes": []
},
@@ -25,6 +25,6 @@
"PCA_configs": {
"depth_scale": 1000.0,
"depth_trunc": 3.0,
"voxel_size": 0.010
"voxel_size": 0.005
}
}

View File

@@ -45,7 +45,7 @@ setup(
# 'calibration_node = vision_detect.hand_eye_calibration:main',
# 'crossboard_detect_node = vision_detect.crossboard_detect:main',
'service_client_node = vision_detect.service_client:main',
# 'get_camera_pose_node = vision_detect.get_camera_pose:main',
'get_camera_pose_node = vision_detect.get_camera_pose:main',
# 'calculate_node = vision_detect.calculate:main',
'detect_node = vision_detect.detect_node:main',

View File

@@ -160,6 +160,10 @@ class DetectNode(InitBase):
time2 = time.time()
result = results[0]
home = os.path.expanduser("/home/demo/hivecore_robot_os1/hivecore_robot_vision/imgs")
save_path = os.path.join(home, f"origin_image_{int(time1*1000)}.png")
cv2.imwrite(save_path, rgb_img)
# Get masks
if result.masks is None or len(result.masks) == 0:
self.get_logger().info(f"Detect object num: 0")
@@ -170,6 +174,12 @@ class DetectNode(InitBase):
boxes = result.boxes.xywh.cpu().numpy()
class_ids = result.boxes.cls.cpu().numpy()
labels = result.names
x_centers, y_centers = boxes[:, 0], boxes[:, 1]
sorted_index = np.lexsort((-y_centers, x_centers))
masks = masks[sorted_index]
boxes = boxes[sorted_index]
class_ids = class_ids[sorted_index]
time3 = time.time()
@@ -195,13 +205,16 @@ class DetectNode(InitBase):
mask_crop,
depth_crop,
intrinsics,
calculate_grab_width=True,
calculate_grab_width=False,
**self.calculate_configs
)
if rmat is None:
self.get_logger().warning("Object point cloud have too many noise")
continue
if np.allclose(rmat, np.eye(4)):
continue
self.get_logger().info(f"grab_width: {grab_width}")
rmat = self.hand_eye_mat @ rmat
@@ -232,8 +245,8 @@ class DetectNode(InitBase):
# mask_img and box_img is or not output
if self.output_boxes and not self.output_masks:
home = os.path.expanduser("~")
save_path = os.path.join(home, "detect_image.png")
home = os.path.expanduser("/home/demo/hivecore_robot_os1/hivecore_robot_vision/imgs")
save_path = os.path.join(home, f"detect_image_{int(time1*1000)}.png")
draw_box(rgb_img, result, save_path=save_path)
# draw_box(rgb_img, result, save_path=False)
return self.cv_bridge.cv2_to_imgmsg(rgb_img, "bgr8"), pose_list
@@ -303,7 +316,7 @@ class DetectNode(InitBase):
"class_id": int(98),
"class_name": "red",
"pose": pose,
"grab_width": 0.0
"grab_width": [0.0]
}
)
@@ -372,7 +385,7 @@ class DetectNode(InitBase):
"class_id": int(99),
"class_name": 'crossboard',
"pose": pose,
"grab_width": 0.0
"grab_width": [0.0]
}
)

View File

@@ -62,6 +62,9 @@ def calculate_pose_pca(
return np.eye(4), [0.0, 0.0, 0.0]
if calculate_grab_width:
if np.asarray(point_cloud.points).shape[0] < 4:
logging.warning("点数不足,不能算 OBB")
return np.eye(4), [0.0, 0.0, 0.0]
obb = point_cloud.get_oriented_bounding_box()
x, y, z = obb.center
extent = obb.extent
@@ -154,25 +157,27 @@ def calculate_pose_icp(
def point_cloud_denoising(point_cloud: o3d.geometry.PointCloud, voxel_size: float = 0.005):
"""点云去噪"""
point_cloud = point_cloud.remove_non_finite_points()
while True:
down_pcd = point_cloud.voxel_down_sample(voxel_size=voxel_size * 0.5)
point_num = len(down_pcd.points)
if point_num <= 1000:
break
voxel_size *= 2
# while True:
# down_pcd = point_cloud.voxel_down_sample(voxel_size=voxel_size)
# point_num = len(down_pcd.points)
# if point_num <= 1000:
# break
# voxel_size *= 4
down_pcd = point_cloud.voxel_down_sample(voxel_size=voxel_size)
point_num = len(down_pcd.points)
# logging.fatal("point_cloud_denoising: point_num={}".format(len(point_cloud.points)))
# logging.fatal("point_cloud_denoising: point_num={}".format(point_num))
# 半径滤波
clean_pcd, _ = down_pcd.remove_radius_outlier(
nb_points=10,
radius=voxel_size * 5
nb_points=int(round(10 * voxel_size / 0.005)),
radius=voxel_size * 10
)
# 统计滤波
clean_pcd, _ = clean_pcd.remove_statistical_outlier(
nb_neighbors=10,
nb_neighbors=int(round(10 * voxel_size / 0.005)),
std_ratio=2.0
)
@@ -186,7 +191,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=int(round(10 * voxel_size / 0.005))))
if len(labels[labels >= 0]) == 0:
return clean_pcd
@@ -213,7 +218,7 @@ def point_cloud_denoising(point_cloud: o3d.geometry.PointCloud, voxel_size: floa
# 使用最近簇判断噪音强度
largest_cluster_ratio = len(clean_pcd.points) / point_num
if largest_cluster_ratio < 0.20:
if largest_cluster_ratio < 0.08:
return None
return clean_pcd

View File

@@ -25,6 +25,11 @@ def draw_box(
confidences = segmentation_result.boxes.conf.cpu().numpy()
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]
for i, box in enumerate(boxes):
x_center, y_center, width, height = box[:4]

View File

@@ -44,17 +44,17 @@ class GetCameraPose(Node):
pose_dict = {}
for object in camera_pose.objects:
pose_dict[object.class_name] = object.pose_list
pose_dict[object.class_name] = object.pose
if self.class_name in pose_dict:
_camera = [
pose_dict[self.class_name][-1].position.x,
pose_dict[self.class_name][-1].position.y,
pose_dict[self.class_name][-1].position.z,
pose_dict[self.class_name][-1].orientation.w,
pose_dict[self.class_name][-1].orientation.x,
pose_dict[self.class_name][-1].orientation.y,
pose_dict[self.class_name][-1].orientation.z,
pose_dict[self.class_name].position.x,
pose_dict[self.class_name].position.y,
pose_dict[self.class_name].position.z,
pose_dict[self.class_name].orientation.w,
pose_dict[self.class_name].orientation.x,
pose_dict[self.class_name].orientation.y,
pose_dict[self.class_name].orientation.z,
]
self.get_logger().info(f"add camera: {_camera}")

View File

@@ -0,0 +1,18 @@
{
"save_path": "collect_data",
"left": {
"color": "",
"depth": "",
"info": ""
},
"right": {
"color": "/camera/camera/color/image_raw",
"depth": "/camera/camera/aligned_depth_to_color/image_raw",
"info": "/camera/camera/color/camera_info"
},
"head": {
"color": "/camera/color/image_raw",
"depth": "/camera/depth/image_raw",
"info": "/camera/color/camera_info"
}
}

20
vision_tools/package.xml Normal file
View File

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

View File

4
vision_tools/setup.cfg Normal file
View File

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

30
vision_tools/setup.py Normal file
View File

@@ -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',
],
},
)

View File

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

View File

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

View File

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

View File

View File

@@ -0,0 +1,46 @@
import rclpy
from rclpy.node import Node
from rclpy.parameter import Parameter
from std_srvs.srv import SetBool
def main(args=None):
def timer_callback():
start_sign = node.get_parameter("collect_sign").value
end_sign = node.get_parameter("end_sign").value
if start_sign:
request = SetBool.Request()
request.data = True
client.call_async(request).add_done_callback(result_callback)
node.set_parameters([Parameter('collect_sign', Parameter.Type.BOOL, False)])
if end_sign:
request = SetBool.Request()
request.data = False
client.call_async(request).add_done_callback(result_callback)
node.set_parameters([Parameter('end_sign', Parameter.Type.BOOL, False)])
def result_callback(result):
response = result.result()
if response.success:
node.get_logger().info(response.message)
else:
node.get_logger().error(response.message)
rclpy.init(args=args)
node = Node("collect_data_client")
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)
try:
rclpy.spin(node)
except KeyboardInterrupt:
pass
finally:
node.destroy_node()

View File

@@ -0,0 +1,664 @@
import os
import json
import threading
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, CameraInfo
from std_srvs.srv import SetBool
from vision_tools.utils import zip_folder
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_index = self.right_index = self.head_index = 0
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.isdir(os.path.join(
self.save_path, f"{self.index:06d}_left_depth_raw"))
or os.path.isdir(os.path.join(
self.save_path, f"{self.index:06d}_right_depth_raw"))
or os.path.isdir(os.path.join(
self.save_path, f"{self.index:06d}_head_depth_raw"))
# 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_index = 0
os.mkdir(os.path.join(self.save_path, f"{self.index:06d}_left_depth_raw"))
# 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_index = 0
os.mkdir(os.path.join(self.save_path, f"{self.index:06d}_right_depth_raw"))
# 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_index = 0
os.mkdir(os.path.join(self.save_path, f"{self.index:06d}_head_depth_raw"))
# 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_folder,
args=(
f"{self.index:06d}_left_depth_raw.zip",
f"{self.index:06d}_left_depth_raw",
self.save_path
),
daemon=True
).start()
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_folder,
args=(
f"{self.index:06d}_right_depth_raw.zip",
f"{self.index:06d}_right_depth_raw",
self.save_path
),
daemon=True
).start()
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_folder,
args=(
f"{self.index:06d}_head_depth_raw.zip",
f"{self.index:06d}_head_depth_raw",
self.save_path
),
daemon=True
).start()
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())
cv2.imwrite(
os.path.join(
self.save_path,
f"{self.index:06d}_left_depth_raw",
f"{self.left_index:08d}.png"
),
depth_frame
)
self.left_index += 1
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())
cv2.imwrite(
os.path.join(
self.save_path,
f"{self.index:06d}_right_depth_raw",
f"{self.right_index:08d}.png"
),
depth_frame
)
self.right_index += 1
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())
cv2.imwrite(
os.path.join(
self.save_path,
f"{self.index:06d}_head_depth_raw",
f"{self.head_index:08d}.png"
),
depth_frame
)
self.head_index += 1
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()

View File

@@ -0,0 +1,104 @@
#!/usr/bin/env python
# coding: utf-8
import json
# import os
import rclpy
from rclpy.node import Node
from rclpy.parameter import Parameter
from interfaces.msg import PoseArrayClassAndID
class GetCameraPose(Node):
def __init__(self, name):
super(GetCameraPose, self).__init__(name)
self.sync_subscriber = None
self.sub_camera_pose = None
self.sub_hand_pose = None
self.camera = []
self.calibration_matrix = None
self.save = False
self.collect = False
self.done = False
self.done_collecting = False
# self.filename = "camera_pose_data.json"
self.declare_parameter('class_name', 'crossboard')
self.declare_parameter('start_collect_once', False)
self.declare_parameter('save', False)
# if not os.path.exists(self.filename):
# with open(self.filename, "w", encoding="utf-8") as f:
# json.dump([], f, ensure_ascii=False, indent=2)
self.sub_camera_pose = self.create_subscription(PoseArrayClassAndID, '/pose/cv_detect_pose', self.get_pose_callback, 10)
self.timer_ = self.create_timer(1, self.timer_callback)
def get_pose_callback(self, camera_pose):
if self.collect and not self.done_collecting:
self.done_collecting = True
_camera = None
pose_dict = {}
for object in camera_pose.objects:
pose_dict[object.class_name] = object.pose_list
if self.class_name in pose_dict:
_camera = [
pose_dict[self.class_name][-1].position.x,
pose_dict[self.class_name][-1].position.y,
pose_dict[self.class_name][-1].position.z,
pose_dict[self.class_name][-1].orientation.w,
pose_dict[self.class_name][-1].orientation.x,
pose_dict[self.class_name][-1].orientation.y,
pose_dict[self.class_name][-1].orientation.z,
]
self.get_logger().info(f"add camera: {_camera}")
self.done_collecting = False
self.collect = False
self.camera.extend(_camera)
self.set_parameters([Parameter('start_collect_once', Parameter.Type.BOOL, False)])
else:
self.done_collecting = False
self.collect = False
self.set_parameters([Parameter('start_collect_once', Parameter.Type.BOOL, False)])
self.get_logger().info(f"Have not camera data")
return
def timer_callback(self):
self.collect = self.get_parameter('start_collect_once').value
self.save = self.get_parameter('save').value
self.class_name = self.get_parameter('class_name').value
if self.save:
print(self.camera)
self.get_logger().info(f'Save as camera_pose_data.json')
with open(f"camera_pose_data.json", "w") as f:
json.dump(self.camera, f, indent=4)
self.done = True
def main(args=None):
rclpy.init(args=args)
node = GetCameraPose('get_camera_pose')
try:
while rclpy.ok() and not node.done:
rclpy.spin_once(node)
except KeyboardInterrupt:
pass
finally:
node.destroy_node()
rclpy.shutdown()
if __name__ == '__main__':
main()

View File

@@ -0,0 +1 @@
from .zip_tools import *

View File

@@ -0,0 +1,102 @@
import os
import threading
import time
import zipfile
import logging
import shutil
__all__ = ["zip_verification", "zip_folder"]
# 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, zip_name))
def zip_folder(zip_name:str, folder_name:str, save_path:str):
logging.info("zip folder start")
with zipfile.ZipFile(
os.path.join(save_path, zip_name),
'w',
compression=zipfile.ZIP_DEFLATED
) as zf:
for root, dirs, files in os.walk(os.path.join(save_path, folder_name)):
for file in files:
abs_path = os.path.join(root, file)
rel_path = os.path.relpath(abs_path, os.path.join(save_path, folder_name))
zf.write(abs_path, rel_path)
if zip_verification(zip_name, folder_name, save_path):
shutil.rmtree(os.path.join(save_path, folder_name))
else:
logging.warning("zip file verification file, keep folder")
os.remove(os.path.join(save_path, zip_name))
logging.info("zip folder end")
def zip_verification(zip_name:str, folder_name:str, save_path:str):
i = 0
while os.path.exists(os.path.join(save_path, f"_tmp_{i}")):
i += 1
tmp_dir = os.path.join(save_path, f"_tmp_{i}")
# 清理旧的临时目录
if os.path.exists(tmp_dir):
shutil.rmtree(tmp_dir)
os.makedirs(tmp_dir, exist_ok=True)
zip_path = os.path.join(save_path, zip_name)
try:
# 解压所有内容到临时目录
with zipfile.ZipFile(zip_path, "r") as zf:
zf.extractall(tmp_dir)
# 遍历原始文件夹,逐个文件比对大小
for root, dirs, files in os.walk(os.path.join(save_path, folder_name)):
for file in files:
# 原文件完整路径
abs_path = os.path.join(root, file)
# 计算文件相对 folder_name 的路径
rel_path = os.path.relpath(abs_path, os.path.join(save_path, folder_name))
# 解压后的文件路径
extracted_path = os.path.join(tmp_dir, rel_path)
if not os.path.exists(extracted_path):
return False # 文件不存在
if os.path.getsize(abs_path) != os.path.getsize(extracted_path):
return False # 文件大小不一致
return True
finally:
# 清理临时目录
if os.path.exists(tmp_dir):
shutil.rmtree(tmp_dir)
if __name__ == '__main__':
threading.Thread(
target=zip_folder,
args=(
f"000000_right_depth_raw.zip",
f"000000_right_depth_raw",
"/home/lyx/collect_data"
),
daemon=True
).start()
time.sleep(10)