Semantic Navigation with Language-Embedded Gaussian Splats

git clone https://github.com/ShijieZhou-UCLA/feature-3dgs.git
cd feature-3dgs
pip install -r requirements.txt

# Install the N-dimensional Gaussian rasterizer
pip install submodules/diff-gaussian-rasterization-feature/
pip install submodules/simple-knn/

# Extract SAM features
python extract_sam_features.py \
  --data_dir /data/captures/house_run \
  --output_dir /data/features/house_run/sam

# Extract CLIP-LSeg features
python extract_clip_features.py \
  --data_dir /data/captures/house_run \
  --output_dir /data/features/house_run/clip

python train.py \
  --source_path /data/captures/house_run \
  --feature_dir /data/features/house_run/clip \
  --model_path /data/results/house_run_semantic \
  --iterations 30000

import torch
from open_clip import create_model_and_transforms, get_tokenizer

# Encode the text query into CLIP embedding space
model, _, _ = create_model_and_transforms('ViT-B-32', pretrained='openai')
tokenizer = get_tokenizer('ViT-B-32')
text = tokenizer(["refrigerator"])
text_features = model.encode_text(text)  # shape: (1, D)
text_features = text_features / text_features.norm(dim=-1, keepdim=True)

# Render the feature map from a chosen viewpoint
# rendered_features shape: (H, W, D) where D is the CLIP embedding dimension
rendered_features = render_features(viewpoint_camera)  # from Feature-3DGS
rendered_features = rendered_features / rendered_features.norm(dim=-1, keepdim=True)

# Compute cosine similarity between the text query and each pixel's features
similarity = torch.nn.functional.cosine_similarity(
    rendered_features, 
    text_features.unsqueeze(0).unsqueeze(0),  # broadcast to (1, 1, D)
    dim=-1
)
# similarity is a (H, W) heatmap — high values indicate the queried object

from gsplat import rasterization

# Standard RGB rendering (D=3)
renders, alphas, info = rasterization(
    means=gaussians.means,         # (N, 3) — positions
    quats=gaussians.quats,         # (N, 4) — rotations
    scales=gaussians.scales,       # (N, 3) — scales
    opacities=gaussians.opacities, # (N,)   — opacities
    colors=gaussians.colors,       # (N, 3) — RGB
    viewmats=viewmats,             # (C, 4, 4)
    Ks=Ks,                         # (C, 3, 3)
    width=width,
    height=height,
)
# renders shape: (C, H, W, 3)

# Feature rendering — same API, just change the colors dimension
feature_renders, alphas, info = rasterization(
    means=gaussians.means,         # (N, 3)
    quats=gaussians.quats,         # (N, 4)
    scales=gaussians.scales,       # (N, 3)
    opacities=gaussians.opacities, # (N,)
    colors=gaussians.features,     # (N, D) — D-dimensional features
    viewmats=viewmats,
    Ks=Ks,
    width=width,
    height=height,
)
# feature_renders shape: (C, H, W, D)

import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")

def extract_dense_clip_features(image_path: str) -> torch.Tensor:
    """Extract dense per-patch CLIP features from an image.
    
    CLIP ViT-B/16 divides a 224x224 input into 14x14 = 196 patches.
    Each patch produces a 768-dim feature vector.
    We extract these patch tokens (excluding CLS) and upsample 
    to the original image resolution.
    """
    image = Image.open(image_path).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model.vision_model(**inputs, output_hidden_states=True)
    
    # Extract patch tokens from the last hidden state
    # Shape: (1, 197, 768) — 1 CLS token + 196 patch tokens
    patch_tokens = outputs.last_hidden_state[:, 1:, :]  # (1, 196, 768)
    
    # Reshape to spatial grid: (1, 14, 14, 768) -> (1, 768, 14, 14)
    B, N, D = patch_tokens.shape
    h = w = int(N ** 0.5)
    features = patch_tokens.reshape(B, h, w, D).permute(0, 3, 1, 2)
    
    # Upsample to original image resolution
    orig_h, orig_w = image.size[1], image.size[0]
    features = torch.nn.functional.interpolate(
        features, size=(orig_h, orig_w), mode='bilinear', align_corners=False
    )
    # Output: (1, 768, H, W) -> (H, W, 768)
    return features.squeeze(0).permute(1, 2, 0)

# In the model initialization, alongside existing parameters:
feature_dim = 64  # or 768 if not compressing
self.features = torch.nn.Parameter(
    torch.randn(N, feature_dim) * 0.01  # small random init
)
# Add to optimizer alongside means, quats, scales, etc.

# In the training loop, after RGB rendering:
feature_renders, _, _ = rasterization(
    means=self.means,
    quats=self.quats,
    scales=self.scales,
    opacities=self.opacities,
    colors=self.features,       # (N, feature_dim)
    viewmats=viewmats,
    Ks=Ks,
    width=width,
    height=height,
)
# feature_renders shape: (C, H, W, feature_dim)

# gt_features: ground-truth CLIP features for this view, shape (H, W, feature_dim)
feature_loss = torch.nn.functional.mse_loss(
    feature_renders.squeeze(0),  # (H, W, feature_dim)
    gt_features                   # (H, W, feature_dim)
)

lambda_feat = 0.1  # tune this hyperparameter
total_loss = rgb_loss + lambda_feat * feature_loss
total_loss.backward()

import torch
from transformers import CLIPModel, CLIPProcessor

model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")

def query_scene(text_query: str, viewpoint, trained_splat) -> torch.Tensor:
    """Query the trained splat with a text string.
    
    Returns a (H, W) relevancy heatmap.
    """
    # Encode text
    inputs = processor(text=[text_query], return_tensors="pt", padding=True)
    with torch.no_grad():
        text_features = model.get_text_features(**inputs)  # (1, 512)
    text_features = text_features / text_features.norm(dim=-1, keepdim=True)
    
    # If you used PCA compression, project text_features through the same PCA
    # text_features_compressed = pca.transform(text_features)
    
    # Render feature map from viewpoint
    feature_map = render_features(viewpoint, trained_splat)  # (H, W, D)
    feature_map = feature_map / feature_map.norm(dim=-1, keepdim=True)
    
    # Cosine similarity
    similarity = torch.nn.functional.cosine_similarity(
        feature_map,
        text_features.unsqueeze(0),  # (1, 1, D) broadcast
        dim=-1
    )
    return similarity  # (H, W) heatmap

from gs_semantic_query import SemanticSplatQuery

query_engine = SemanticSplatQuery(
    model_path="/data/results/house_run_semantic"
)

result = query_engine.query("refrigerator")
# result.position        → np.array([x, y, z]) in world frame
# result.relevancy_image → np.ndarray (H, W, 3) heatmap visualization
# result.confidence      → float (0-1) cosine similarity score

User prompt
  "show me all the top trajectories that reach the refrigerator location"
    |
    v
LLM Agent (Claude / GPT-4)
    |
    | 1. Parse intent -> target object = "refrigerator"
    | 2. Query semantic splat -> get 3D location
    | 3. Get current robot pose from /odom
    | 4. Query Nav2 planner for candidate paths
    | 5. Rank trajectories by criteria
    | 6. Render expected views along top trajectories from splat
    v
Response with trajectory visualizations

def query_object_location(object_name: str) -> dict:
    """Query the semantic splat for an object's 3D position.
    Returns: {position: [x, y, z], confidence: float}
    """

def get_robot_pose() -> dict:
    """Get the robot's current pose from /odom.
    Returns: {x: float, y: float, z: float, yaw: float}
    """

def plan_path(start: list, goal: list) -> list:
    """Plan a path from start to goal using Nav2.
    Returns: list of [x, y, z, yaw] waypoints
    """

def render_view_from_splat(pose: list) -> str:
    """Render what the robot would see from a given pose.
    Returns: path to the rendered image
    """

def send_nav2_goal(x: float, y: float, yaw: float) -> str:
    """Send a navigation goal to the robot via Nav2.
    Returns: status string
    """