PCA visualization collapses to positional encoding instead of semantic features on nuScenes

[niklas030]

I am currently working with the nuScenes dataset and trying to visualize the features using PCA, similar to what is provided in the vanilla demo. I am feeding the raw point cloud into the Utonia model.

However, instead of the PCA clustering around the semantic meaning of the points (as shown in the demos), the output seems to collapse entirely to positional encodings.

Is this expected behavior for raw nuScenes data, or is there a specific preprocessing step / feature requirement I am missing to get semantic-level PCA features?

Expected behavior
The PCA of the extracted features should group points by their semantic meaning (e.g., distinguishing vehicles, roads, buildings) and colorize them accordingly.

Actual behavior
The PCA colors seem strictly tied to the spatial coordinates/distance from the sensor, effectively acting like a positional gradient.

Screenshots

Code to reproduce
Here is the minimal script I am using to load the data, run the model, and extract the PCA colors:

import argparse
import random
from pathlib import Path

import numpy as np
import open3d as o3d
import torch
import utonia
from nuscenes.nuscenes import NuScenes

try:
    import flash_attn  # noqa: F401
except ImportError:
    flash_attn = None

device = "cuda" if torch.cuda.is_available() else "cpu"


def get_pca_color(feat, brightness=1.25, center=True):
    u, s, v = torch.pca_lowrank(feat, center=center, q=12, niter=5)
    projection = feat @ v
    projection = (
        projection[:, :3] * 0.4 + projection[:, 3:6] * 0.2 + projection[:, 9:12] * 0.4
    )
    min_val = projection.min(dim=-2, keepdim=True)[0]
    max_val = projection.max(dim=-2, keepdim=True)[0]
    div = torch.clamp(max_val - min_val, min=1e-6)
    color = (projection - min_val) / div * brightness
    color = color.clamp(0.0, 1.0)
    return color

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dataroot",
        type=str,
        default="/mnt/niklas_storage/datasets/raws",
        help="nuScenes dataset root.",
    )
    parser.add_argument(
        "--version",
        type=str,
        default="v1.0-mini",
        help="nuScenes metadata version.",
    )
    parser.add_argument(
        "--sweep_seed",
        type=int,
        default=135161,
        help="Random seed for selecting one nuScenes sample.",
    )
    parser.add_argument(
        "--out_pc",
        type=str,
        default="pc_nuscenes_simple.ply",
        help="Output path for original point cloud (gray).",
    )
    parser.add_argument(
        "--out_pca",
        type=str,
        default="pca_nuscenes_simple.ply",
        help="Output path for PCA-colored point cloud.",
    )
    return parser.parse_args()


def _resolve_data_path(dataroot: str, relative_path: str) -> Path:
    p = Path(dataroot) / relative_path
    if p.exists():
        return p
    # Some setups store files under dataroot/data.
    p2 = Path(dataroot) / "data" / relative_path
    if p2.exists():
        return p2
    raise FileNotFoundError(
        f"Could not find nuScenes file for relative path '{relative_path}' "
        f"under '{dataroot}' or '{dataroot}/data'."
    )


def load_random_nuscenes_lidar_point_dict(nusc: NuScenes, sweep_seed: int | None):
    if sweep_seed is not None:
        random.seed(sweep_seed)
    sample = random.choice(nusc.sample)
    lidar_sd = nusc.get("sample_data", sample["data"]["LIDAR_TOP"])
    lidar_path = _resolve_data_path(nusc.dataroot, lidar_sd["filename"])
    points = np.fromfile(str(lidar_path), dtype=np.float32, count=-1).reshape([-1, 5])
    coord = points[:, :3].astype(np.float32)
    color = np.zeros_like(coord, dtype=np.float32)
    normal = np.zeros_like(coord, dtype=np.float32)
    point = {"coord": coord, "color": color, "normal": normal}
    meta = (
        f"sample_token={sample['token']} lidar_token={lidar_sd['token']} "
        f"path={lidar_path}"
    )
    return point, meta


if __name__ == "__main__":
    args = parse_args()

    nusc = NuScenes(version=args.version, dataroot=args.dataroot, verbose=False)
    point, meta = load_random_nuscenes_lidar_point_dict(nusc, args.sweep_seed)
    print(f"Loaded {point['coord'].shape[0]} points. {meta}")

    # Keep this fixed seed like the vanilla demo style.
    utonia.utils.set_seed(6985480)

    if flash_attn is not None:
        model = utonia.load("utonia", repo_id="Pointcept/Utonia").to(device)
    else:
        custom_config = dict(enc_patch_size=[1024 for _ in range(5)], enable_flash=False)
        model = utonia.load(
            "utonia", repo_id="Pointcept/Utonia", custom_config=custom_config
        ).to(device)
    model.eval()

    # Same transform function as outdoor_vanilia_demo.py
    transform = utonia.transform.default(0.2, apply_z_positive=False)
    original_coord = point["coord"].copy()
    point = transform(point)

    with torch.inference_mode():
        for key in point.keys():
            if isinstance(point[key], torch.Tensor) and device == "cuda":
                point[key] = point[key].cuda(non_blocking=True)
        point = model(point)
        for _ in range(2):
            assert "pooling_parent" in point.keys()
            assert "pooling_inverse" in point.keys()
            parent = point.pop("pooling_parent")
            inverse = point.pop("pooling_inverse")
            parent.feat = torch.cat([parent.feat, point.feat[inverse]], dim=-1)
            point = parent
        while "pooling_parent" in point.keys():
            assert "pooling_inverse" in point.keys()
            parent = point.pop("pooling_parent")
            inverse = point.pop("pooling_inverse")
            parent.feat = point.feat[inverse]
            point = parent
        pca_color = get_pca_color(point.feat, brightness=1, center=True)

    original_pca_color = pca_color[point.inverse].detach().cpu().numpy()

    # Save original cloud (gray)
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(original_coord)
    pcd.colors = o3d.utility.Vector3dVector(
        np.full_like(original_coord, 0.5, dtype=np.float32)
    )
    o3d.io.write_point_cloud(args.out_pc, pcd)
    print(f"Wrote {args.out_pc}")

    # Save PCA-colored cloud
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(original_coord)
    pcd.colors = o3d.utility.Vector3dVector(original_pca_color)
    o3d.io.write_point_cloud(args.out_pca, pcd)
    print(f"Wrote {args.out_pca}")

Any guidance on whether the PCA script needs adjusting for nuScenes, or if I'm missing a data normalization step, would be greatly appreciated!

[Yujia-Zhang0913]

Hi,
From your last question, I think you're currently finding the best features for distillation.

1. For visualization, you can replace projection[:, :3] * 0.4 + projection[:, 3:6] * 0.2 + projection[:, 9:12] * 0.4 with projection[:, :3] to only see the first three main components, or you can try UMAP for better visualization.
2. For better features, I think you can try using multiple frames as input.

I tried your code on the test split. Below are the visualizations of single-frame and 10-frame point clouds in the first three main components.

single-frame:

10-frame:

And you can also try visualizing different dimensions of PCA results, like projection[:, 3:6]:

[niklas030]

thanks for u efforts i really appreciate!!!!