-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathbag2image.py
179 lines (153 loc) · 6.72 KB
/
bag2image.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
import argparse
import os
import numpy as np
import cv2
import pyrealsense2 as rs
def streaming_from_file(filepath):
"""Helper function to read rosbag file from realsense.
"""
pipe = rs.pipeline()
cfg = rs.config()
cfg.enable_device_from_file(filepath, repeat_playback=False)
return pipe, cfg
def post_processing(frame,
enable_spatial=True,
enable_temporal=True,
enable_hole=True,
spatial_params=[(rs.option.filter_magnitude, 5),
(rs.option.filter_smooth_alpha, 1),
(rs.option.filter_smooth_delta, 50),
(rs.option.holes_fill, 3)],
temporal_params=[],
hole_params=[]):
"""Filters to cleanup depth maps.
"""
# Filters and settings
depth_to_disparity = rs.disparity_transform(True)
disparity_to_depth = rs.disparity_transform(False)
# Depth to disparity before spatial and temporal filters
frame = depth_to_disparity.process(frame)
# Spatial filter
if enable_spatial:
# Settings
spatial = rs.spatial_filter()
for spatial_param in spatial_params:
spatial.set_option(spatial_param[0], spatial_param[1])
# Apply on frame
frame = spatial.process(frame)
# Temporal filter
if enable_temporal:
temporal = rs.temporal_filter()
for temporal_param in temporal_params:
temporal.set_option(temporal_param[0], temporal_param[1])
frame = temporal.process(frame)
# Back to depth
frame = disparity_to_depth.process(frame)
# Hole filling
if enable_hole:
hole_filling = rs.hole_filling_filter()
for hole_param in hole_params:
hole_filling.set_option(hole_param[0], hole_param[1])
frame = hole_filling.process(frame)
return frame
def extract_frames(pipe,
cfg,
save_path,
post_processing=False,
save_colorize=True,
save_pc=False,
visualize=True):
"""Helper function to align and extract rgb-d image from bagfile.
Check more saving options in arguments.
Args:
pipe: pyrealsense2 pipeline.
cfg: pyrealsense2 pipeline configuration.
save_path: Path to save the extracted frames.
save_colorize: Save colorized depth maps visualization.
save_pc: Save point cloud data.
visualize: Visualize the video frames during saving.
"""
# Configurations
if save_colorize:
colorizer = rs.colorizer()
if save_pc:
pc = rs.pointcloud()
points = rs.points()
# Save path
if not os.path.exists(save_path):
os.makedirs(save_path)
# Start the pipe
i = 0
profile = pipe.start(cfg)
device = profile.get_device()
playback = device.as_playback()
playback.set_real_time(False) # Make sure this is False or frames get dropped
while True:
try:
# Wait for a conherent pairs of frames: (rgb, depth)
pairs = pipe.wait_for_frames()
# Align depth image to rgb image first
align = rs.align(rs.stream.color)
pairs = align.process(pairs)
color_frame = pairs.get_color_frame()
depth_frame = pairs.get_depth_frame()
if not depth_frame or not color_frame:
continue
# Post-processing
if post_processing:
depth_frame = post_processing(depth_frame)
# Get rgb-d images
color_img = np.asanyarray(color_frame.get_data())
color_img = cv2.cvtColor(color_img, cv2.COLOR_RGB2BGR)
depth_img = np.asanyarray(depth_frame.get_data())
print('Frame {}, Depth Image {}, Color Image {}'.format(i+1, depth_img.shape, color_img.shape))
# Save as loseless formats
cv2.imwrite(os.path.join(save_path, '{}_rgb.png'.format(i)), color_img, [cv2.IMWRITE_PNG_COMPRESSION, 0])
np.save(os.path.join(save_path, '{}_depth.npy'.format(i)), depth_img)
if save_colorize:
# Save colorized depth map
depth_img_colorized = np.asanyarray(colorizer.colorize(depth_frame).get_data())
cv2.imwrite(os.path.join(save_path, '{}_depth_colorized.jpg'.format(i)), depth_img_colorized) # No need for lossless here
if save_pc:
# NOTE: Point cloud calculation takes longer time.
pc.map_to(color_frame)
points = pc.calculate(depth_frame)
points.export_to_ply(os.path.join(save_path, '{}_pc.ply'.format(i)), color_frame)
i += 1
if visualize:
# Stack both images horizontally
images = np.vstack((color_img, depth_img_colorized))
# Show images
cv2.namedWindow('RealSense', cv2.WINDOW_NORMAL)
cv2.imshow('RealSense', images)
cv2.waitKey(1)
except Exception as e:
print(e)
break
# Clean pipeline
pipe.stop()
print('{} frames saved in total.'.format(i))
return
if __name__ == '__main__':
# Arguments
parser = argparse.ArgumentParser(description='Convert bagfile to images.')
parser.add_argument('bagfile', help='Path to bagfile.')
parser.add_argument('-post_processing','--post_processing', nargs='?', type=bool, default=False,
help='Enable depth post processing.')
parser.add_argument('-save_colorize','--save_colorize', nargs='?', type=bool, default=True,
help='Save colorized depth map.')
parser.add_argument('-save_pc','--save_pc', nargs='?', type=bool, default=False,
help='Save point cloud data.')
parser.add_argument('-visualize','--visualize', nargs='?', type=bool, default=True,
help='Visualize while saving.')
args = parser.parse_args()
save_path = os.path.join(os.path.curdir, args.bagfile.split('.')[0])
print('Extracting {} and saving to {}...'.format(args.bagfile, save_path))
# Save things here
pipe, cfg = streaming_from_file(args.bagfile)
extract_frames(pipe, cfg,
save_path=save_path,
post_processing=args.post_processing,
save_colorize=args.save_colorize,
save_pc=args.save_pc,
visualize=args.visualize)