Skip to content

Commit b95ebef

Browse files
AadarshHegde123AadarshHegde123
committed
modified rosbag conversion script
1 parent 483d83f commit b95ebef

1 file changed

Lines changed: 208 additions & 118 deletions

File tree

‎GEMstack/scripts/convert_to_rosbag.py‎

Lines changed: 208 additions & 118 deletions
Original file line numberDiff line numberDiff line change
@@ -11,161 +11,251 @@
1111
from datetime import datetime
1212
import time
1313
import rospy
14-
import std_msgs.msg
14+
import std_msgs.msg
1515
import argparse
1616

17-
def get_matching_files(dir):
18-
"""Get files that have corresponding data in all three formats"""
19-
# Get sorted lists of files
20-
png_files = sorted(glob.glob(os.path.join(dir, 'color*.png')))
21-
tif_files = sorted(glob.glob(os.path.join(dir, 'depth*.tif')))
22-
npz_files = sorted(glob.glob(os.path.join(dir, 'lidar*.npz')))
23-
24-
print(f"Found {len(png_files)} PNG files")
25-
print(f"Found {len(tif_files)} TIF files")
26-
print(f"Found {len(npz_files)} NPZ files")
27-
28-
# Extract numbers from filenames
17+
def get_matching_files(data_dir):
18+
"""
19+
Collect and return files that have matching (common) index numbers:
20+
- camera_fl[index].png
21+
- camera_fr[index].png
22+
- camera_rr[index].png
23+
- camera_rl[index].png
24+
- front_cam[index].png
25+
- lidar_front[index].npz
26+
- lidar_top[index].npz
27+
"""
28+
29+
# 1) Gather each pattern of files
30+
camera_fl_files = sorted(glob.glob(os.path.join(data_dir, 'camera_fl*.png')))
31+
camera_fr_files = sorted(glob.glob(os.path.join(data_dir, 'camera_fr*.png')))
32+
camera_rr_files = sorted(glob.glob(os.path.join(data_dir, 'camera_rr*.png')))
33+
camera_rl_files = sorted(glob.glob(os.path.join(data_dir, 'camera_rl*.png')))
34+
front_cam_files = sorted(glob.glob(os.path.join(data_dir, 'front_cam*.png')))
35+
lidar_front_files = sorted(glob.glob(os.path.join(data_dir, 'lidar_front*.npz')))
36+
lidar_top_files = sorted(glob.glob(os.path.join(data_dir, 'lidar_top*.npz')))
37+
38+
print(f"Found {len(camera_fl_files)} FL camera PNG files")
39+
print(f"Found {len(camera_fr_files)} FR camera PNG files")
40+
print(f"Found {len(camera_rr_files)} RR camera PNG files")
41+
print(f"Found {len(camera_rl_files)} RL camera PNG files")
42+
print(f"Found {len(front_cam_files)} front_cam PNG files")
43+
print(f"Found {len(lidar_front_files)} lidar_front NPZ files")
44+
print(f"Found {len(lidar_top_files)} lidar_top NPZ files")
45+
46+
# 2) Helper function to extract index from filename
2947
def get_number(filename):
30-
return int(''.join(filter(str.isdigit, os.path.basename(filename))))
31-
32-
# Create dictionaries with numbers as keys
33-
png_dict = {get_number(f): f for f in png_files}
34-
tif_dict = {get_number(f): f for f in tif_files}
35-
npz_dict = {get_number(f): f for f in npz_files}
36-
37-
# Find common numbers
38-
common_numbers = set(png_dict.keys()) & set(tif_dict.keys()) & set(npz_dict.keys())
39-
print(f"\nFound {len(common_numbers)} matching frame numbers")
40-
41-
# Get matching files in sorted order
42-
common_numbers = sorted(common_numbers)
43-
matching_pngs = [png_dict[n] for n in common_numbers]
44-
matching_tifs = [tif_dict[n] for n in common_numbers]
45-
matching_npzs = [npz_dict[n] for n in common_numbers]
46-
47-
return matching_pngs, matching_tifs, matching_npzs
48-
49-
def create_rosbag_from_data(
50-
dir,
51-
output_bag_path,
52-
frame_interval
53-
):
54-
# Initialize CV bridge
48+
# Example: "camera_fl123.png" => 123
49+
base = os.path.basename(filename)
50+
# Filter out digits in the basename, then convert to int
51+
return int(''.join(filter(str.isdigit, base)))
52+
53+
# 3) Convert to dictionaries keyed by index number
54+
fl_dict = {get_number(f): f for f in camera_fl_files}
55+
fr_dict = {get_number(f): f for f in camera_fr_files}
56+
rr_dict = {get_number(f): f for f in camera_rr_files}
57+
rl_dict = {get_number(f): f for f in camera_rl_files}
58+
front_dict = {get_number(f): f for f in front_cam_files}
59+
60+
lidar_front_dict = {get_number(f): f for f in lidar_front_files}
61+
lidar_top_dict = {get_number(f): f for f in lidar_top_files}
62+
63+
# 4) Find common index numbers across all 7 dictionaries
64+
common_indices = (
65+
set(fl_dict.keys())
66+
& set(fr_dict.keys())
67+
& set(rr_dict.keys())
68+
& set(rl_dict.keys())
69+
& set(front_dict.keys())
70+
& set(lidar_front_dict.keys())
71+
& set(lidar_top_dict.keys())
72+
)
73+
74+
print(f"\nFound {len(common_indices)} matching indices among all file types")
75+
76+
# 5) For convenience, sort them and pull out matching file paths in the same order
77+
common_indices = sorted(common_indices)
78+
matching_fl = [fl_dict[idx] for idx in common_indices]
79+
matching_fr = [fr_dict[idx] for idx in common_indices]
80+
matching_rr = [rr_dict[idx] for idx in common_indices]
81+
matching_rl = [rl_dict[idx] for idx in common_indices]
82+
matching_front = [front_dict[idx] for idx in common_indices]
83+
84+
matching_lidar_front = [lidar_front_dict[idx] for idx in common_indices]
85+
matching_lidar_top = [lidar_top_dict[idx] for idx in common_indices]
86+
87+
return (
88+
matching_fl,
89+
matching_fr,
90+
matching_rr,
91+
matching_rl,
92+
matching_front,
93+
matching_lidar_front,
94+
matching_lidar_top
95+
)
96+
97+
def create_rosbag_from_data(data_dir, output_bag_path, frame_rate):
98+
"""Create a ROS bag with camera (PNG) and LiDAR (NPZ) data at a given frame rate."""
99+
100+
# Initialize CV bridge and ROS node
55101
bridge = CvBridge()
56-
57-
# Initialize ROS node
58102
rospy.init_node('bag_creator', anonymous=True)
59-
60-
# Get matching files
61-
png_files, tif_files, npz_files = get_matching_files(dir)
62-
63-
print(f"Found {len(png_files)} matching files in all three formats")
64-
65-
if len(png_files) == 0:
103+
104+
# Gather matching files
105+
(
106+
fl_files,
107+
fr_files,
108+
rr_files,
109+
rl_files,
110+
front_files,
111+
lidar_front_files,
112+
lidar_top_files
113+
) = get_matching_files(data_dir)
114+
115+
num_frames = len(fl_files)
116+
print(f"Found {num_frames} total matching sets of files")
117+
118+
if num_frames == 0:
66119
raise ValueError("No matching files found!")
67-
68-
# Create a new bag file
120+
121+
# Create a new bag
69122
with rosbag.Bag(output_bag_path, 'w') as bag:
70123
# Start time for the messages
71124
start_time = rospy.Time.from_sec(time.time())
72125
message_count = 0
73-
74-
for idx, (png_file, tif_file, npz_file) in enumerate(zip(png_files, tif_files, npz_files)):
75-
# Calculate timestamp for this frame
76-
timestamp = rospy.Time.from_sec(start_time.to_sec() + idx * (1/frame_interval))
77-
126+
127+
# Publish each frame
128+
for idx in range(num_frames):
129+
# Prepare a timestamp for this frame
130+
timestamp = rospy.Time.from_sec(start_time.to_sec() + idx * (1.0 / frame_rate))
131+
78132
try:
79-
# Process PNG image
80-
png_img = cv2.imread(png_file)
81-
if png_img is not None:
82-
png_msg = bridge.cv2_to_imgmsg(png_img, encoding="bgr8")
83-
png_msg.header.stamp = timestamp
84-
png_msg.header.frame_id = "camera_rgb"
85-
bag.write('/camera/rgb/image_raw', png_msg, timestamp)
133+
# -------------------------------------------------
134+
# 1) CAMERA: FL, FR, RR, RL, FRONT
135+
# -------------------------------------------------
136+
# Read each .png file
137+
fl_img = cv2.imread(fl_files[idx])
138+
fr_img = cv2.imread(fr_files[idx])
139+
rr_img = cv2.imread(rr_files[idx])
140+
rl_img = cv2.imread(rl_files[idx])
141+
front_img = cv2.imread(front_files[idx])
142+
143+
# Convert to ROS Image messages
144+
if fl_img is not None:
145+
fl_msg = bridge.cv2_to_imgmsg(fl_img, encoding="bgr8")
146+
fl_msg.header.stamp = timestamp
147+
fl_msg.header.frame_id = "camera_fl"
148+
bag.write('/camera/fl/image_raw', fl_msg, timestamp)
86149
message_count += 1
87-
else:
88-
print(f"Warning: Could not read PNG file: {png_file}")
89-
90-
# Process TIF image
91-
tif_img = cv2.imread(tif_file, -1) # -1 to preserve original depth
92-
if tif_img is not None:
93-
tif_msg = bridge.cv2_to_imgmsg(tif_img, encoding="passthrough")
94-
tif_msg.header.stamp = timestamp
95-
tif_msg.header.frame_id = "camera_depth"
96-
bag.write('/camera/depth/image_raw', tif_msg, timestamp)
150+
151+
if fr_img is not None:
152+
fr_msg = bridge.cv2_to_imgmsg(fr_img, encoding="bgr8")
153+
fr_msg.header.stamp = timestamp
154+
fr_msg.header.frame_id = "camera_fr"
155+
bag.write('/camera/fr/image_raw', fr_msg, timestamp)
97156
message_count += 1
98-
else:
99-
print(f"Warning: Could not read TIF file: {tif_file}")
100-
101-
# Process pointcloud NPZ
102-
pc_data = np.load(npz_file)
103-
points = pc_data['arr_0'] # Using 'arr_0' based on the provided files'
104-
105-
# Create pointcloud message
106-
header = std_msgs.msg.Header()
107-
header.stamp = timestamp
108-
header.frame_id = "velodyne"
109-
157+
158+
if rr_img is not None:
159+
rr_msg = bridge.cv2_to_imgmsg(rr_img, encoding="bgr8")
160+
rr_msg.header.stamp = timestamp
161+
rr_msg.header.frame_id = "camera_rr"
162+
bag.write('/camera/rr/image_raw', rr_msg, timestamp)
163+
message_count += 1
164+
165+
if rl_img is not None:
166+
rl_msg = bridge.cv2_to_imgmsg(rl_img, encoding="bgr8")
167+
rl_msg.header.stamp = timestamp
168+
rl_msg.header.frame_id = "camera_rl"
169+
bag.write('/camera/rl/image_raw', rl_msg, timestamp)
170+
message_count += 1
171+
172+
if front_img is not None:
173+
front_msg = bridge.cv2_to_imgmsg(front_img, encoding="bgr8")
174+
front_msg.header.stamp = timestamp
175+
front_msg.header.frame_id = "camera_front"
176+
bag.write('/camera/front/image_raw', front_msg, timestamp)
177+
message_count += 1
178+
179+
# -------------------------------------------------
180+
# 2) LIDAR: FRONT, TOP
181+
# -------------------------------------------------
182+
# For .npz, assume data is in 'arr_0' (x, y, z, or however your data is structured)
183+
front_data = np.load(lidar_front_files[idx])
184+
front_points = front_data['arr_0'] # e.g., Nx3 array
185+
186+
top_data = np.load(lidar_top_files[idx])
187+
top_points = top_data['arr_0'] # e.g., Nx3 array
188+
189+
# A) Lidar FRONT
190+
front_header = std_msgs.msg.Header()
191+
front_header.stamp = timestamp
192+
front_header.frame_id = "lidar_front"
193+
194+
# Example: 3 float32 fields for x,y,z
110195
fields = [
111196
pc2.PointField('x', 0, pc2.PointField.FLOAT32, 1),
112197
pc2.PointField('y', 4, pc2.PointField.FLOAT32, 1),
113-
pc2.PointField('z', 8, pc2.PointField.FLOAT32, 1)
198+
pc2.PointField('z', 8, pc2.PointField.FLOAT32, 1),
114199
]
115-
116-
pc_msg = pc2.create_cloud(header, fields, points)
117-
bag.write('/velodyne_points', pc_msg, timestamp)
200+
front_pc_msg = pc2.create_cloud(front_header, fields, front_points)
201+
bag.write('/lidar/front/points', front_pc_msg, timestamp)
118202
message_count += 1
119-
203+
204+
# B) Lidar TOP
205+
top_header = std_msgs.msg.Header()
206+
top_header.stamp = timestamp
207+
top_header.frame_id = "lidar_top"
208+
209+
top_pc_msg = pc2.create_cloud(top_header, fields, top_points)
210+
bag.write('/lidar/top/points', top_pc_msg, timestamp)
211+
message_count += 1
212+
120213
except Exception as e:
121-
print(f"Error processing frame {idx}: {str(e)}")
214+
print(f"Error processing frame {idx}: {e}")
122215
continue
123-
216+
217+
# Log progress every so often
124218
if idx % 10 == 0:
125-
print(f"Processed {idx} frames")
126-
127-
print(f"Total messages written to bag: {message_count}")
219+
print(f"Processed {idx} frames...")
220+
221+
print(f"Finished writing {message_count} total messages to {output_bag_path}.")
128222

129223
if __name__ == "__main__":
130-
#This initializes the parser and sets the parameters that the user will be asked to provide in the terminal
131224
parser = argparse.ArgumentParser(
132-
description='A script to convert data gathered by cameras and lidar sensors to ros .bag messages'
225+
description='Convert multiple camera (PNG) and LiDAR (NPZ) files into a single ROS bag.'
133226
)
134227

135228
parser.add_argument(
136-
'files_directory', type = str,
137-
help = 'The path to the directory with all the rgb images, depth maps, and point clouds. The file formats must be PNG, TIF, and NPZ, respectively'
229+
'files_directory',
230+
type=str,
231+
help='Path to the directory containing the camera and lidar files.'
138232
)
139233

140234
parser.add_argument(
141-
'output_bag', type = str,
142-
help = 'The path to the directory where the bag file will be saved'
235+
'output_bag',
236+
type=str,
237+
help='Full path (including filename) where the output ROS bag will be saved.'
143238
)
144-
239+
145240
parser.add_argument(
146-
'rate', type = int,
147-
help = 'The rate at which the data is collected in Hz'
241+
'rate',
242+
type=int,
243+
help='Playback (or capture) rate in Hz (frames per second).'
148244
)
245+
149246
args = parser.parse_args()
150-
directory = args.files_directory
151-
output_bag = args.output_bag
152-
rate = args.rate
153-
# Example directories below:
154-
#directory = "/home/username/host/CS588/GEMstack/data/data_sample/data/"
155-
#output_bag = "/home/username/host/CS588/GEMstack/data/data_sample/data/output.bag"
156-
247+
157248
try:
158249
create_rosbag_from_data(
159-
directory,
160-
output_bag,
161-
rate
250+
data_dir=args.files_directory,
251+
output_bag_path=args.output_bag,
252+
frame_rate=args.rate
162253
)
163254
print("Successfully created ROS bag file!")
164-
165-
# Verify bag contents
255+
256+
# Show basic info about the new bag
166257
print("\nBag file contents:")
167-
info_cmd = f"rosbag info {output_bag}"
168-
os.system(info_cmd)
169-
258+
os.system(f"rosbag info {args.output_bag}")
259+
170260
except Exception as e:
171-
print(f"Error creating ROS bag: {str(e)}")
261+
print(f"Error creating ROS bag: {e}")

0 commit comments

Comments
 (0)