Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 29 additions & 14 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ To capture using the right camera and right LiDAR:
```bash
python3 scripts/capture_emulated_rgbd.py --camera right --lidar right
```

*(Note: This must be run on the Stretch 4 robot. It requires `stretch4_body` and the use of the robot's cameras and LiDARs)*

### 2. Preprocessing
Expand Down Expand Up @@ -149,6 +150,11 @@ To capture the live frames and estimate the masks for the right sensor pair:
```bash
python3 scripts/estimate_validity_masks.py --camera right --lidar right
```

To capture and estimate masks for both sides simultaneously:
```bash
python3 scripts/estimate_validity_masks.py --camera left_right --lidar both
```
*(Note: This must be run on the Stretch 4 robot. The script captures exactly 30 synchronized frames, computes the masks at the highest active resolution, and saves them locally to `data/validity_masks/` for the visualizers to use automatically).*

### 8. Visualize Live RGB-D Imagery with the New Calibration (On Robot)
Expand All @@ -164,32 +170,37 @@ For the right camera/lidar:
python3 scripts/visualize_emulated_rgbd.py --camera right --lidar right
```

For both cameras and lidars simultaneously:
```bash
python3 scripts/visualize_emulated_rgbd.py --camera left_right --lidar both
```

### 10. Use PyZMQ to Send RGB-D Imagery

There are two examples of sending RGB-D images via PyZMQ. They can be used to send RGB-D images from the robot to a desktop computer or between processes on the same computer.

The first example sends RGB-D images, receives them, and then visualizes them using OpenCV.

To run the publisher:
To run the publisher (e.g. for both cameras):
```bash
python3 scripts/send_rgbd_images.py --camera left --lidar left
python3 examples/send_rgbd_images.py --camera left_right --lidar both
```

To run the subscriber:
```bash
python3 scripts/recv_rgbd_images.py
python3 examples/recv_rgbd_images.py
```

The second example send RGB-D images with the synchronized joint state of the robot, receives them, and then visualizes the colored 3D point cloud in Rerun along with the joint states.

To run the publisher:
To run the publisher (e.g. for both cameras):
```bash
python3 scripts/send_rgbd_images_and_joint_states.py --camera left --lidar left
python3 examples/send_rgbd_images_and_joint_states.py --camera left_right --lidar both
```

To run the subscriber:
```bash
python3 scripts/recv_rgbd_images_and_joint_states.py
python3 examples/recv_rgbd_images_and_joint_states.py
```

Both examples can send data over the network. To do so, you will need to use the `--remote` flag for both the publisher and the subscriber and provide IP and port information. Sending RGB-D images with joint states uses the following file for IP and port information:
Expand All @@ -203,7 +214,7 @@ The example code for sending RGB-D images alone uses IP and port information pro
For developers writing custom applications, the repository provides a unified API in `stretch4_emulated_rgbd.api` to stream and process synchronized RGB-D frames.

> [!TIP]
> **Quick Start:** For a complete, runnable demonstration of the API capabilities—including lazy properties, calibration extraction, validity masking, dense depth interpolation, and colored 3D point cloud generation—see [`examples/api_example.py`](file:///home/hello-robot/repos/stretch4_emulated_rgbd/examples/api_example.py).
> **Quick Start:** For a complete, runnable demonstration of the API capabilities—including lazy properties, calibration extraction, validity masking, dense depth interpolation, and colored 3D point cloud generation—see [`examples/api_example.py`](file:///home/hello-robot/repos/stretch4_rgbd/examples/api_example.py).

#### Summary of Processing Steps
When the `FastEmulatedRGBDStreamer` captures and aligns an RGB-D frame, it executes the following steps internally *before* yielding it to you:
Expand All @@ -218,19 +229,22 @@ When the `FastEmulatedRGBDStreamer` captures and aligns an RGB-D frame, it execu

#### Code Examples

**Receiving the Stream**
**Receiving the Stream (Simultaneous Left and Right)**
```python
from stretch4_emulated_rgbd.api import get_emulated_rgbd_stream

# Automatically loads the optimized calibration for the current fleet
# Automatically loads the optimized calibrations for the current fleet
streamer, generator = get_emulated_rgbd_stream(
use_left=True,
use_left_right=True,
use_left_lidar=True,
use_right_lidar=True,
emulated_rgbd_fps=10.0
)

# Fetch a single synchronized frame
frame = next(generator)
# Fetch synchronized frames from both cameras
multi_frame = next(generator)
left_frame = multi_frame.left
right_frame = multi_frame.right
```

**Applying Validity Masks**
Expand All @@ -240,13 +254,14 @@ import cv2

mask_manager = ValidityMaskManager()
# Automatically loads the robot's physical vignetting and LiDAR bounds masks
vig_mask, depth_mask = mask_manager.get_masks("left", "left_lidar", frame.image.shape)
vig_mask, depth_mask = mask_manager.get_masks("left", "left_lidar right_lidar", left_frame.image.shape)

# Apply vignette mask to black out the physical camera housing
masked_rgb = frame.image.copy()
masked_rgb = left_frame.image.copy()
masked_rgb[~vig_mask] = 0
```


**Generating a Dense Depth Image**
```python
from stretch4_emulated_rgbd.api import DenseDepthImage
Expand Down
220 changes: 113 additions & 107 deletions examples/api_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,9 +45,11 @@ def main():
camera_fps=args.camera_fps,
resolution_height=args.resolution,
compress=not args.disable_compression,
oak_buffer_size=args.oak_buffer_size
oak_buffer_size=args.oak_buffer_size,
merge_lidars=args.merge_lidars
)


# 2. Initialize the Validity Mask Manager
print("\n[2] Initializing Validity Mask Manager...")
mask_manager = ValidityMaskManager()
Expand All @@ -63,121 +65,125 @@ def main():
rerun_initialized = False

# The generator yields synchronized frames indefinitely
for frame in generator:
if frame is None:
for frame_data in generator:
if frame_data is None:
continue

# 4. Access Lazy Properties
rgb_image = frame.image
depth_image = frame.depth_image

# 5. Access Calibration Data
cam_matrix = frame.camera_matrix
dist_coeffs = frame.distortion_coefficients
T_base_to_cam = frame.T_base_to_cam

# 6. Apply Validity Masks
c_name = frame.camera_type
lidar_str = frame.lidars_used if frame.lidars_used else "no_lidar"
vig_mask, lidar_mask = mask_manager.get_masks(c_name, lidar_str, rgb_image.shape)

if False:
# Combine the masks and ERODE them slightly.
# Eroding shrinks the mask inward by a few pixels. This elegantly drops the extreme
# boundary pixels of the fisheye lens BEFORE unprojection, preventing the tan(theta) explosion.
combined_mask = vig_mask & lidar_mask
erosion_kernel = np.ones((5, 5), np.uint8)
dense_depth_validity_mask = cv2.erode(combined_mask.astype(np.uint8), erosion_kernel, iterations=12).astype(bool)
else:
# Handle both single frames and multi-frame objects
frames = []
if hasattr(frame_data, "left") or hasattr(frame_data, "right") or hasattr(frame_data, "center"):
if getattr(frame_data, "left", None): frames.append(frame_data.left)
if getattr(frame_data, "right", None): frames.append(frame_data.right)
if getattr(frame_data, "center", None): frames.append(frame_data.center)
else:
frames.append(frame_data)

# For visualization stacking
all_rgb_masked = []
all_depth_vis = []

for frame in frames:
# 4. Access Lazy Properties
rgb_image = frame.image
depth_image = frame.depth_image

# 5. Access Calibration Data
cam_matrix = frame.camera_matrix
dist_coeffs = frame.distortion_coefficients
T_base_to_cam = frame.T_base_to_cam

# 6. Apply Validity Masks
c_name = frame.camera_type
lidar_str = frame.lidars_used if frame.lidars_used else "no_lidar"
vig_mask, lidar_mask = mask_manager.get_masks(c_name, lidar_str, rgb_image.shape)

dense_depth_validity_mask = vig_mask & lidar_mask

# Apply the vignetting mask to remove invalid fisheye edges from the RGB image
masked_rgb = rgb_image.copy()
masked_rgb[~vig_mask] = 0

# 7. Generate Dense Depth Map
dense_processor = DenseDepthImage(rgb_image, depth_image, apply_validity_mask=False)
dense_depth = dense_processor.compute_dense_depth()

# Before creating a point cloud, apply the eroded combined mask to drop unstable boundary pixels
dense_depth[~dense_depth_validity_mask] = 0

# 8. Create Colored Point Cloud
pts_cam, colors = create_point_cloud_from_depth(
dense_depth, masked_rgb, cam_matrix, dist_coeffs
)

# Transform points to the robot's base coordinate frame for correct upright 3D viewing
pts_cam_homog = np.hstack((pts_cam, np.ones((pts_cam.shape[0], 1))))
T_cam_to_base = np.linalg.inv(T_base_to_cam)
pts_base = (T_cam_to_base @ pts_cam_homog.T).T[:, :3]

# 9. Visualization
if not rerun_mode:
# 9a. OpenCV Visualization
# Normalize depth for visualization (cap at 5 meters for better contrast)
max_depth = 5.0
depth_vis = np.clip(dense_depth, 0, max_depth) / max_depth
depth_vis = (depth_vis * 255).astype(np.uint8)
depth_colormap = cv2.applyColorMap(depth_vis, cv2.COLORMAP_JET)
# Set invalid depth (0) to black
depth_colormap[dense_depth == 0] = [0, 0, 0]

cv2.imshow("Masked RGB", masked_rgb)
cv2.imshow("Dense Depth", depth_colormap)
# Apply the vignetting mask to remove invalid fisheye edges from the RGB image
masked_rgb = rgb_image.copy()
masked_rgb[~vig_mask] = 0

key = cv2.waitKey(1) & 0xFF
if key == ord('q') or key == 27: # 27 is ESC
rerun_mode = True
cv2.destroyAllWindows()
print("\n[9b] Switching to Rerun Visualization (press Ctrl+C to exit)...")
else:
# 9b. Rerun Visualization
if not rerun_initialized:
try:
import rerun as rr
import rerun.blueprint as rrb
print(" -> Spawning Rerun viewer...")
rr.init("api_example_visualization", spawn=True)

# Explicitly define the layout blueprint to force stacking of the 2D images
# under the 'camera' origin, while keeping the 3D views separate.
blueprint = rrb.Blueprint(
rrb.Horizontal(
rrb.Spatial3DView(name="Sparse Point Cloud", origin="sparse_view"),
rrb.Spatial3DView(name="Dense Point Cloud", origin="dense_view"),
rrb.Spatial2DView(name="Layered RGB-D", origin="camera"),
),
rrb.BlueprintPanel(expanded=False),
rrb.SelectionPanel(expanded=True),
rrb.TimePanel(expanded=False, play_state="following"),
)
rr.send_blueprint(blueprint)
rerun_initialized = True
except ImportError:
print(" -> Rerun is not installed. Exiting.")
break

# Log the images (OpenCV uses BGR, Rerun expects RGB)
rr.log("camera/rgb", rr.Image(masked_rgb[:, :, ::-1]))
rr.log("camera/dense_depth", rr.DepthImage(dense_depth, meter=1.0, depth_range=[0.0, config.RERUN_COLOR_MAX_DEPTH_M]))
# 7. Generate Dense Depth Map
dense_processor = DenseDepthImage(rgb_image, depth_image, apply_validity_mask=False)
dense_depth = dense_processor.compute_dense_depth()

# Add the sparse depth image overlay as requested
rr.log("camera/sparse_depth", rr.DepthImage(depth_image, meter=1.0, depth_range=[0.0, config.RERUN_COLOR_MAX_DEPTH_M]))
# Before creating a point cloud, apply the eroded combined mask to drop unstable boundary pixels
dense_depth[~dense_depth_validity_mask] = 0

# Log the perfectly aligned SPARSE point cloud directly from the API.
# Logged to a separate root path ("sparse_view") to force a side-by-side window in Rerun
rr.log(
"sparse_view/point_cloud",
rr.Points3D(frame.point_cloud_base, colors=frame.point_colors, radii=[0.01])
# 8. Create Colored Point Cloud
pts_cam, colors = create_point_cloud_from_depth(
dense_depth, masked_rgb, cam_matrix, dist_coeffs
)

# Log the perfectly aligned DENSE point cloud generated from the depth map.
# Logged to a separate root path ("dense_view") to force a side-by-side window in Rerun
rr.log(
"dense_view/point_cloud",
rr.Points3D(pts_base, colors=colors[:, ::-1], radii=[0.01])
)
# Transform points to the robot's base coordinate frame for correct upright 3D viewing
pts_cam_homog = np.hstack((pts_cam, np.ones((pts_cam.shape[0], 1))))
T_cam_to_base = np.linalg.inv(T_base_to_cam)
pts_base = (T_cam_to_base @ pts_cam_homog.T).T[:, :3]

# 9. Visualization
if not rerun_mode:
# 9a. OpenCV Visualization
max_depth = 5.0
depth_vis = np.clip(dense_depth, 0, max_depth) / max_depth
depth_vis = (depth_vis * 255).astype(np.uint8)
depth_colormap = cv2.applyColorMap(depth_vis, cv2.COLORMAP_JET)
depth_colormap[dense_depth == 0] = [0, 0, 0]

all_rgb_masked.append(masked_rgb)
all_depth_vis.append(depth_colormap)
else:
# 9b. Rerun Visualization
if not rerun_initialized:
try:
import rerun as rr
import rerun.blueprint as rrb
print(" -> Spawning Rerun viewer...")
rr.init("api_example_visualization", spawn=True)

blueprint = rrb.Blueprint(
rrb.Horizontal(
rrb.Spatial3DView(name="Sparse Point Cloud", origin="sparse_view"),
rrb.Spatial3DView(name="Dense Point Cloud", origin="dense_view"),
rrb.Spatial2DView(name="Layered RGB-D", origin="camera"),
),
rrb.BlueprintPanel(expanded=False),
rrb.SelectionPanel(expanded=True),
rrb.TimePanel(expanded=False, play_state="following"),
)
rr.send_blueprint(blueprint)
rerun_initialized = True
except ImportError:
print(" -> Rerun is not installed. Exiting.")
break

# Prefix paths with camera name for multi-camera support in Rerun
prefix = f"camera/{c_name}/"
rr.log(prefix + "rgb", rr.Image(masked_rgb[:, :, ::-1]))
rr.log(prefix + "dense_depth", rr.DepthImage(dense_depth, meter=1.0, depth_range=[0.0, config.RERUN_COLOR_MAX_DEPTH_M]))
rr.log(prefix + "sparse_depth", rr.DepthImage(depth_image, meter=1.0, depth_range=[0.0, config.RERUN_COLOR_MAX_DEPTH_M]))

rr.log(
f"sparse_view/{c_name}/point_cloud",
rr.Points3D(frame.point_cloud_base, colors=frame.point_colors, radii=[0.01])
)
rr.log(
f"dense_view/{c_name}/point_cloud",
rr.Points3D(pts_base, colors=colors[:, ::-1], radii=[0.01])
)

if not rerun_mode and all_rgb_masked:
# Tile images if there are multiple
stacked_rgb = np.hstack(all_rgb_masked)
stacked_depth = np.hstack(all_depth_vis)
cv2.imshow("Masked RGB", stacked_rgb)
cv2.imshow("Dense Depth", stacked_depth)

key = cv2.waitKey(1) & 0xFF
if key == ord('q') or key == 27:
rerun_mode = True
cv2.destroyAllWindows()
print("\n[9b] Switching to Rerun Visualization (press Ctrl+C to exit)...")


finally:
print("\nStopping streamer...")
Expand Down
Loading