Implement 3D Gaussian Splatting for Robot Mapping in 45 Minutes

# Check CUDA availability
python -c "import torch; print(torch.cuda.is_available())"

# Create isolated environment
conda create -n robot_gs python=3.10 -y
conda activate robot_gs

# Install PyTorch with CUDA support
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

# Install 3DGS core libraries
pip install plyfile tqdm opencv-python open3d

# Clone optimized 3DGS implementation
git clone https://github.com/graphdeco-inria/gaussian-splatting --recursive
cd gaussian-splatting
pip install submodules/diff-gaussian-rasterization
pip install submodules/simple-knn

# scripts/capture_robot_data.py
import pyrealsense2 as rs
import numpy as np
import cv2
from pathlib import Path
import json

class RobotDataCapture:
    def __init__(self, output_dir="data/robot_scan"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        # Configure RealSense pipeline
        self.pipeline = rs.pipeline()
        config = rs.config()
        config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
        config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
        
        self.pipeline.start(config)
        self.frame_count = 0
        
        # Store camera intrinsics for 3DGS
        profile = self.pipeline.get_active_profile()
        self.intrinsics = profile.get_stream(rs.stream.color).as_video_stream_profile().get_intrinsics()
        
    def capture_frame(self, robot_pose):
        """Capture RGB-D frame with robot pose from odometry"""
        frames = self.pipeline.wait_for_frames()
        color_frame = frames.get_color_frame()
        depth_frame = frames.get_depth_frame()
        
        if not color_frame or not depth_frame:
            return False
            
        # Convert to numpy arrays
        color_image = np.asanyarray(color_frame.get_data())
        depth_image = np.asanyarray(depth_frame.get_data())
        
        # Save images
        frame_id = f"{self.frame_count:05d}"
        cv2.imwrite(str(self.output_dir / f"{frame_id}.png"), color_image)
        np.save(str(self.output_dir / f"{frame_id}_depth.npy"), depth_image)
        
        # Save camera pose (from robot odometry or SLAM)
        pose_data = {
            "frame_id": frame_id,
            "position": robot_pose[:3].tolist(),  # [x, y, z]
            "rotation": robot_pose[3:].tolist(),  # quaternion [qx, qy, qz, qw]
            "timestamp": color_frame.get_timestamp()
        }
        
        with open(self.output_dir / f"{frame_id}_pose.json", 'w') as f:
            json.dump(pose_data, f)
            
        self.frame_count += 1
        return True
    
    def save_camera_params(self):
        """Save camera intrinsics for 3DGS training"""
        params = {
            "fx": self.intrinsics.fx,
            "fy": self.intrinsics.fy,
            "cx": self.intrinsics.ppx,
            "cy": self.intrinsics.ppy,
            "width": self.intrinsics.width,
            "height": self.intrinsics.height
        }
        
        with open(self.output_dir / "camera_params.json", 'w') as f:
            json.dump(params, f)
            
    def stop(self):
        self.pipeline.stop()

# Usage in robot control loop
if __name__ == "__main__":
    capture = RobotDataCapture()
    
    # Simulate robot moving and capturing
    for i in range(100):
        # Get robot pose from your SLAM/odometry system
        robot_pose = get_robot_pose()  # Replace with actual odometry
        
        capture.capture_frame(robot_pose)
        time.sleep(0.1)  # 10 FPS capture
    
    capture.save_camera_params()
    capture.stop()

# scripts/prepare_training_data.py
import numpy as np
import json
from pathlib import Path
from scipy.spatial.transform import Rotation

def convert_to_colmap_format(data_dir):
    """Convert robot poses to COLMAP format for 3DGS"""
    data_dir = Path(data_dir)
    output_dir = data_dir / "colmap_format"
    output_dir.mkdir(exist_ok=True)
    
    # Load camera params
    with open(data_dir / "camera_params.json") as f:
        cam_params = json.load(f)
    
    images_txt = []
    points3d = []
    
    # Process each frame
    pose_files = sorted(data_dir.glob("*_pose.json"))
    for idx, pose_file in enumerate(pose_files):
        with open(pose_file) as f:
            pose_data = json.load(f)
        
        frame_id = pose_data["frame_id"]
        
        # Convert quaternion to rotation matrix
        quat = pose_data["rotation"]  # [qx, qy, qz, qw]
        rot = Rotation.from_quat(quat)
        
        # COLMAP uses camera-to-world transform
        pos = np.array(pose_data["position"])
        
        # Format: IMAGE_ID QW QX QY QZ TX TY TZ CAMERA_ID NAME
        images_txt.append(
            f"{idx+1} {quat[3]} {quat[0]} {quat[1]} {quat[2]} "
            f"{pos[0]} {pos[1]} {pos[2]} 1 {frame_id}.png\n\n"
        )
        
        # Initialize 3D points from depth map
        depth_path = data_dir / f"{frame_id}_depth.npy"
        if depth_path.exists():
            depth = np.load(depth_path)
            points = depth_to_points(depth, cam_params, pos, rot)
            points3d.extend(points)
    
    # Write COLMAP files
    with open(output_dir / "images.txt", 'w') as f:
        f.writelines(images_txt)
    
    # Write camera intrinsics
    with open(output_dir / "cameras.txt", 'w') as f:
        f.write(f"1 PINHOLE {cam_params['width']} {cam_params['height']} "
                f"{cam_params['fx']} {cam_params['fy']} "
                f"{cam_params['cx']} {cam_params['cy']}\n")
    
    # Write initial point cloud
    write_points3d(points3d, output_dir / "points3D.txt")
    
    print(f"Converted {len(pose_files)} frames to COLMAP format")

def depth_to_points(depth, cam_params, position, rotation):
    """Convert depth map to 3D points in world frame"""
    fx, fy = cam_params['fx'], cam_params['fy']
    cx, cy = cam_params['cx'], cam_params['cy']
    
    points = []
    h, w = depth.shape
    
    # Sample every 10 pixels to avoid millions of points
    for v in range(0, h, 10):
        for u in range(0, w, 10):
            z = depth[v, u] * 0.001  # Convert mm to meters
            
            if z < 0.1 or z > 10.0:  # Filter invalid depth
                continue
            
            # Backproject to camera frame
            x = (u - cx) * z / fx
            y = (v - cy) * z / fy
            
            # Transform to world frame
            point_cam = np.array([x, y, z])
            point_world = rotation.apply(point_cam) + position
            
            points.append(point_world)
    
    return points

def write_points3d(points, output_path):
    """Write point cloud in COLMAP points3D.txt format"""
    with open(output_path, 'w') as f:
        for idx, pt in enumerate(points):
            # Format: POINT3D_ID X Y Z R G B ERROR TRACK[]
            f.write(f"{idx+1} {pt[0]} {pt[1]} {pt[2]} 128 128 128 1.0\n")

if __name__ == "__main__":
    convert_to_colmap_format("data/robot_scan")

# Train with robot-optimized parameters
python train.py \
  -s data/robot_scan/colmap_format \
  -m output/robot_model \
  --iterations 7000 \
  --position_lr_init 0.00016 \
  --position_lr_final 0.0000016 \
  --densify_grad_threshold 0.0002 \
  --eval

# In separate Terminal
tensorboard --logdir output/robot_model

# scripts/realtime_renderer.py
import torch
import numpy as np
from gaussian_splatting.scene import Scene
from gaussian_splatting.gaussian_renderer import render
from gaussian_splatting.utils.camera_utils import Camera

class RealtimeGaussianRenderer:
    def __init__(self, model_path):
        """Load trained 3DGS model for real-time rendering"""
        self.device = torch.device("cuda:0")
        
        # Load scene
        self.scene = Scene(model_path, load_iteration=-1)
        self.gaussians = self.scene.gaussians
        
        # Pre-allocate rendering buffers
        self.background = torch.tensor([0, 0, 0], dtype=torch.float32, device=self.device)
        
    def render_from_pose(self, camera_pose, cam_params, width=640, height=480):
        """Render view from robot's current pose in <100ms"""
        
        # Create camera from robot pose
        camera = self._create_camera(camera_pose, cam_params, width, height)
        
        # Render (this is the fast part!)
        with torch.no_grad():
            rendering = render(camera, self.gaussians, self.background)
        
        # Convert to numpy for display/processing
        image = rendering["render"].permute(1, 2, 0).cpu().numpy()
        image = (image * 255).astype(np.uint8)
        
        return image
    
    def _create_camera(self, pose, cam_params, width, height):
        """Convert robot pose to 3DGS camera"""
        # Extract position and rotation from pose
        position = torch.tensor(pose[:3], device=self.device)
        rotation = pose[3:]  # quaternion
        
        # Convert quaternion to rotation matrix
        R = self._quat_to_matrix(rotation)
        
        # Create camera intrinsics matrix
        fx, fy = cam_params['fx'], cam_params['fy']
        cx, cy = cam_params['cx'], cam_params['cy']
        
        FovX = 2 * np.arctan(width / (2 * fx))
        FovY = 2 * np.arctan(height / (2 * fy))
        
        camera = Camera(
            colmap_id=0,
            R=R.cpu().numpy(),
            T=position.cpu().numpy(),
            FoVx=FovX,
            FoVy=FovY,
            image=None,
            gt_alpha_mask=None,
            image_name="realtime",
            uid=0,
            width=width,
            height=height
        )
        
        return camera
    
    def _quat_to_matrix(self, quat):
        """Convert quaternion [qx, qy, qz, qw] to rotation matrix"""
        qx, qy, qz, qw = quat
        
        R = torch.tensor([
            [1 - 2*(qy**2 + qz**2), 2*(qx*qy - qw*qz), 2*(qx*qz + qw*qy)],
            [2*(qx*qy + qw*qz), 1 - 2*(qx**2 + qz**2), 2*(qy*qz - qw*qx)],
            [2*(qx*qz - qw*qy), 2*(qy*qz + qw*qx), 1 - 2*(qx**2 + qy**2)]
        ], device=self.device)
        
        return R

# Integration with robot control
if __name__ == "__main__":
    renderer = RealtimeGaussianRenderer("output/robot_model")
    
    # Simulate robot moving through mapped environment
    while robot_is_active():
        # Get current robot pose from odometry
        current_pose = get_robot_pose()
        
        # Render what robot "sees" from 3DGS map
        rendered_view = renderer.render_from_pose(
            current_pose,
            cam_params={'fx': 615, 'fy': 615, 'cx': 320, 'cy': 240}
        )
        
        # Use for obstacle detection, path planning, etc.
        obstacles = detect_obstacles(rendered_view)
        
        # Display (optional)
        cv2.imshow("3DGS Map View", rendered_view)
        cv2.waitKey(1)

# ros2_ws/src/gaussian_mapping/gaussian_mapping/gs_node.py
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from geometry_msgs.msg import PoseStamped
from cv_bridge import CvBridge
import numpy as np

from realtime_renderer import RealtimeGaussianRenderer

class GaussianSplattingNode(Node):
    def __init__(self):
        super().__init__('gaussian_splatting_node')
        
        # Load trained 3DGS model
        model_path = self.declare_parameter('model_path', 'output/robot_model').value
        self.renderer = RealtimeGaussianRenderer(model_path)
        
        # Subscribe to robot pose
        self.pose_sub = self.create_subscription(
            PoseStamped,
            '/robot/pose',
            self.pose_callback,
            10
        )
        
        # Publish rendered views
        self.image_pub = self.create_publisher(Image, '/gs_map/image', 10)
        
        self.bridge = CvBridge()
        self.latest_pose = None
        
        # Render at 30 Hz
        self.timer = self.create_timer(0.033, self.render_loop)
        
        self.get_logger().info('3D Gaussian Splatting node started')
    
    def pose_callback(self, msg):
        """Store latest robot pose"""
        position = [msg.pose.position.x, msg.pose.position.y, msg.pose.position.z]
        orientation = [
            msg.pose.orientation.x,
            msg.pose.orientation.y,
            msg.pose.orientation.z,
            msg.pose.orientation.w
        ]
        self.latest_pose = position + orientation
    
    def render_loop(self):
        """Render and publish 3DGS view at 30 Hz"""
        if self.latest_pose is None:
            return
        
        # Render from current robot pose
        cam_params = {'fx': 615, 'fy': 615, 'cx': 320, 'cy': 240}
        rendered = self.renderer.render_from_pose(self.latest_pose, cam_params)
        
        # Publish as ROS2 image
        img_msg = self.bridge.cv2_to_imgmsg(rendered, encoding='rgb8')
        img_msg.header.stamp = self.get_clock().now().to_msg()
        img_msg.header.frame_id = 'gs_map'
        
        self.image_pub.publish(img_msg)

def main(args=None):
    rclpy.init(args=args)
    node = GaussianSplattingNode()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()

cd ros2_ws/src
ros2 pkg create --build-type ament_python gaussian_mapping
# Add node to setup.py entry_points

colcon build --packages-select gaussian_mapping
source install/setup.bash

ros2 run gaussian_mapping gs_node --ros-args -p model_path:=output/robot_model

# scripts/benchmark_rendering.py
import time
import numpy as np
from realtime_renderer import RealtimeGaussianRenderer

renderer = RealtimeGaussianRenderer("output/robot_model")

# Generate 100 random poses
poses = [np.random.randn(7) for _ in range(100)]
cam_params = {'fx': 615, 'fy': 615, 'cx': 320, 'cy': 240}

start = time.time()
for pose in poses:
    _ = renderer.render_from_pose(pose, cam_params)
end = time.time()

fps = 100 / (end - start)
print(f"Rendering speed: {fps:.1f} FPS")
print(f"Latency per frame: {(end-start)/100*1000:.1f}ms")

# Compare 3DGS renders to actual robot camera images
python scripts/validate_quality.py \
  --model output/robot_model \
  --test_images data/robot_scan/test_set

# Reduce Gaussian count by 50% with minimal quality loss
python scripts/compress_model.py \
  --input output/robot_model \
  --output output/robot_model_compressed \
  --target_size 0.5

# Incremental training with new frames
python train.py \
  -s data/robot_scan_extended \
  -m output/robot_model \
  --start_checkpoint output/robot_model/chkpnt7000.pth \
  --iterations 10000

from gaussian_splatting.scene import Scene

# Load models from two robots
scene1 = Scene("robot1_model")
scene2 = Scene("robot2_model")

# Merge Gaussian point clouds
merged_gaussians = merge_gaussian_scenes(scene1, scene2)
merged_gaussians.save("merged_model")

# Add to training pipeline
python train.py \
  -s data/robot_scan \
  -m output/robot_model \
  --enable_loop_closure \
  --loop_closure_threshold 0.95