Hard pruning

Dockerfile

@@ -0,0 +1,34 @@
+FROM ghcr.io/nerfstudio-project/nerfstudio
+
+# Install system dependencies
+RUN apt-get update && \
+    apt-get install -y git gcc-10 g++-10 nvidia-cuda-toolkit ninja-build python3-pip wget && \
+    rm -rf /var/lib/apt/lists/*
+
+# Set environment variables to use gcc-10
+ENV CC=gcc-10
+ENV CXX=g++-10
+ENV CUDA_HOME="/usr/local/cuda"
+ENV CMAKE_PREFIX_PATH="$(python -c 'import torch; print(torch.utils.cmake_prefix_path)')"
+ENV TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6+PTX"
+
+# Clone QED-Splatter and install in editable mode
+RUN git clone https://github.com/leggedrobotics/qed-splatter.git /apps/qed-splatter
+RUN pip install /apps/qed-splatter 
+
+# Register QED-Splatter with nerfstudio
+RUN ns-install-cli 
+RUN mkdir -p /usr/local/cuda/bin && ln -s /usr/bin/nvcc /usr/local/cuda/bin/nvcc
+
+# Install additional Python dependencies
+RUN pip install git+https://github.com/rmbrualla/pycolmap@cc7ea4b7301720ac29287dbe450952511b32125e
+RUN pip install git+https://github.com/rahul-goel/fused-ssim@1272e21a282342e89537159e4bad508b19b34157
+RUN pip install nerfview pyntcloud
+
+# Pre-download AlexNet pretrained weights to prevent runtime downloading
+RUN mkdir -p /root/.cache/torch/hub/checkpoints && \
+    wget https://download.pytorch.org/models/alexnet-owt-7be5be79.pth \
+         -O /root/.cache/torch/hub/checkpoints/alexnet-owt-7be5be79.pth
+
+# Set working directory
+WORKDIR /workspace

README.md

@@ -34,3 +34,32 @@ To train the new method, use the following command:
 ```
 ns-train qed-splatter --data [PATH]
 ```
+
+## Pruning Extension
+
+Post-training hard pruners that reduce gaussian count. They follow the same pattern as
+[dn-splatter's mesh export](https://github.com/maturk/dn-splatter/blob/main/dn_splatter/export_mesh.py):
+
+1. ``eval_setup(--load-config)`` loads the training config, dataparser, and checkpoint
+2. Scoring uses ``SplatfactoModel.get_outputs`` over the train cameras
+3. Export uses nerfstudio's ``ExportGaussianSplat.write_ply`` (or a ``.ckpt``)
+
+```
+pip install -e .
+```
+
+### RGB
+```
+qed-rgb-prune --load-config outputs/park/qed-splatter/*/config.yml --output-dir exports/ --pruning-ratio 0.1
+```
+
+### Depth
+Requires a depth-aware run (e.g. qed-splatter).
+```
+qed-depth-prune --load-config outputs/park/qed-splatter/*/config.yml --output-dir exports/ --pruning-ratio 0.1
+```
+
+Useful flags: ``--output-filename``, ``--output-format {ply,ckpt}``, ``--eval-only``, ``--ssim-lambda``.
+
+#### Known Issues
+For the Park scene it tries to generate black gaussians to cover the sky. The entire scene is encased in these gaussians.

pyproject.toml

@@ -5,15 +5,17 @@ version = "0.1.0"
 
 dependencies = ["nerfstudio >= 1.1.0", "open3d", "tyro"]
 
-[project.scripts]
-qed-init-pc = "qed_splatter.create_init_pointcloud:entrypoint"
-
 [tool.setuptools.packages.find]
 include = ["qed_splatter*"]
 
+[project.scripts]
+qed-init-pc = "qed_splatter.create_init_pointcloud:entrypoint"
+qed-rgb-prune = "qed_splatter.rgb_hard_pruner:entrypoint"
+qed-depth-prune = "qed_splatter.depth_hard_pruner:entrypoint"
+
 # register the entry point of your new method here:
 [project.entry-points.'nerfstudio.method_configs']
 method-template = 'qed_splatter.config:qed_splatter_config'
 
 [project.entry-points.'nerfstudio.dataparser_configs']
-qed-dataparser = 'qed_splatter.parser_config:qed_splatter_dataparser'
+qed-dataparser = 'qed_splatter.parser_config:qed_splatter_dataparser'

qed_splatter/depth_hard_pruner.py

@@ -0,0 +1,57 @@
+"""Depth hard pruning via nerfstudio splatfacto ``get_outputs`` gradients."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from qed_splatter.pruning.base import HardPruneConfig, HardPruner, launch_cli
+
+
+@dataclass
+class DepthHardPruneConfig(HardPruneConfig):
+    """Prune gaussians using depth reconstruction loss gradients."""
+
+    def main(self) -> None:
+        DepthHardPruner(self).run()
+
+
+class DepthHardPruner(HardPruner):
+    def default_output_filename(self) -> str:
+        suffix = str(self.cfg.pruning_ratio).replace("0.", "")
+        ext = "ckpt" if self.cfg.output_format == "ckpt" else "ply"
+        return f"{self.cfg.load_config.parent.name}_depth_pruned_{suffix}.{ext}"
+
+    def compute_scores(self) -> torch.Tensor:
+        scores = torch.zeros(len(self.model.means), device=self.device)
+        hits = torch.zeros_like(scores)
+
+        for _, camera, batch in self.iter_train_views():
+            if "depth_image" not in batch:
+                raise KeyError(
+                    "Train batch has no 'depth_image'. Use a depth-aware method "
+                    "(e.g. qed-splatter) when running depth pruning."
+                )
+            outputs = self.render(camera)
+            assert outputs["depth"] is not None, "Model did not return depth"
+            loss = self.depth_loss(outputs["depth"], batch["depth_image"])
+            if loss.numel() == 0 or not torch.isfinite(loss):
+                continue
+            loss.backward()
+
+            means_grad = self.model.gauss_params["means"].grad.abs().mean(dim=-1)
+            with torch.no_grad():
+                active = means_grad > 1e-8
+                scores += means_grad
+                hits += active.float()
+
+        return scores / (hits + 1e-8)
+
+
+def entrypoint() -> None:
+    launch_cli(DepthHardPruneConfig)
+
+
+if __name__ == "__main__":
+    entrypoint()

qed_splatter/pruning/__init__.py

@@ -0,0 +1,9 @@
+"""Hard-pruning utilities for QED-Splatter."""
+
+from qed_splatter.pruning.base import HardPruneConfig, HardPruner, export_gaussian_ply
+
+__all__ = [
+    "HardPruneConfig",
+    "HardPruner",
+    "export_gaussian_ply",
+]

qed_splatter/pruning/base.py

@@ -0,0 +1,233 @@
+"""Shared hard-pruning config and runner (dn-splatter / nerfstudio style).
+
+Loads a past splatfacto / qed-splatter training via ``eval_setup``, scores
+gaussians with ``model.get_outputs``, prunes ``model.gauss_params``, and
+exports with nerfstudio's ``ExportGaussianSplat.write_ply``.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Literal, Optional, Type
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import tqdm
+import tyro
+from nerfstudio.cameras.cameras import Cameras
+from nerfstudio.models.splatfacto import SplatfactoModel
+from nerfstudio.pipelines.base_pipeline import Pipeline
+from nerfstudio.scripts.exporter import ExportGaussianSplat
+from nerfstudio.utils.eval_utils import eval_setup
+from nerfstudio.utils.rich_utils import CONSOLE
+from torchmetrics.image import StructuralSimilarityIndexMeasure
+
+
+@dataclass
+class HardPruneConfig:
+    """Base config matching nerfstudio's ``Exporter`` argument style."""
+
+    load_config: Path
+    """Path to the config YAML file."""
+    output_dir: Path
+    """Path to the output directory."""
+    output_filename: Optional[str] = None
+    """Name of the output file. If unset, derived from the config path and pruning ratio."""
+    pruning_ratio: float = 0.1
+    """Fraction of gaussians to remove (lowest scores)."""
+    output_format: Literal["ply", "ckpt"] = "ply"
+    """Export format: nerfstudio gaussian-splat PLY, or a checkpoint."""
+    ssim_lambda: float = 0.2
+    """Weight of SSIM vs L1 in the scoring loss."""
+    eval_only: bool = False
+    """If True, skip pruning and only report gaussian count."""
+
+    def main(self) -> None:
+        raise NotImplementedError
+
+
+class HardPruner(ABC):
+    """Base runner: eval_setup → score → prune → export."""
+
+    def __init__(self, cfg: HardPruneConfig) -> None:
+        self.cfg = cfg
+        self.cfg.output_dir.mkdir(parents=True, exist_ok=True)
+
+        self.config, self.pipeline, self.checkpoint_path, self.step = eval_setup(
+            cfg.load_config, test_mode="val"
+        )
+        assert isinstance(self.pipeline.model, SplatfactoModel), (
+            f"Expected SplatfactoModel (or subclass), got {type(self.pipeline.model).__name__}"
+        )
+        self.model: SplatfactoModel = self.pipeline.model
+        self.device = self.pipeline.device
+        self.train_dataset = self.pipeline.datamanager.train_dataset
+        assert self.train_dataset is not None
+
+        # Gradients through rasterization; keep eval path (no densify strategy hooks).
+        self.model.eval()
+        for p in self.model.gauss_params.parameters():
+            p.requires_grad_(True)
+
+        self.ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(self.device)
+        CONSOLE.print(
+            f"Loaded {len(self.model.means)} gaussians from {self.checkpoint_path} (step {self.step})"
+        )
+
+    def run(self) -> None:
+        if self.cfg.eval_only or self.cfg.pruning_ratio <= 0:
+            CONSOLE.print(f"Skipping prune (eval_only={self.cfg.eval_only}, ratio={self.cfg.pruning_ratio})")
+            CONSOLE.print(f"Number of GS: {len(self.model.means)}")
+            if not self.cfg.eval_only:
+                self.save()
+            return
+
+        CONSOLE.print("Computing pruning scores...")
+        scores = self.compute_scores()
+        self.prune(scores)
+        CONSOLE.print(f"Number of GS after prune: {len(self.model.means)}")
+        self.save()
+
+    @abstractmethod
+    def compute_scores(self) -> torch.Tensor:
+        """Return a per-gaussian score (higher = keep)."""
+
+    def iter_train_views(self):
+        """Yield ``(image_idx, camera, batch)`` like dn-splatter's dataset loop."""
+        cameras: Cameras = self.train_dataset.cameras  # type: ignore[attr-defined]
+        for image_idx in tqdm.trange(len(self.train_dataset), desc="Views"):
+            batch = self.train_dataset[image_idx]
+            camera = cameras[image_idx : image_idx + 1].to(self.device)
+            yield image_idx, camera, batch
+
+    def render(self, camera: Cameras) -> Dict[str, torch.Tensor]:
+        self.model.zero_grad(set_to_none=True)
+        return self.model.get_outputs(camera)  # type: ignore[return-value]
+
+    def rgb_loss(self, pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
+        gt = self.model.get_gt_img(gt)
+        if gt.shape[-1] == 4:
+            gt = self.model.composite_with_background(gt, self.model._get_background_color())
+        l1 = F.l1_loss(pred, gt)
+        ssim = self.ssim(pred.permute(2, 0, 1)[None], gt.permute(2, 0, 1)[None])
+        return l1 * (1.0 - self.cfg.ssim_lambda) + (1.0 - ssim) * self.cfg.ssim_lambda
+
+    def depth_loss(self, pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
+        if gt.ndim == 2:
+            gt = gt.unsqueeze(-1)
+        gt = gt.to(pred.device, dtype=pred.dtype)
+        # Match shapes if needed
+        if pred.shape[:2] != gt.shape[:2]:
+            gt = F.interpolate(
+                gt.permute(2, 0, 1)[None],
+                size=pred.shape[:2],
+                mode="nearest",
+            )[0].permute(1, 2, 0)
+        valid = torch.isfinite(pred) & torch.isfinite(gt) & (gt > 0)
+        if valid.sum() == 0:
+            return pred.new_zeros(())
+        pred_v, gt_v = pred[valid], gt[valid]
+        l1 = F.l1_loss(pred_v, gt_v)
+        # SSIM on dense maps (fill invalid with 0 for structural term)
+        pred_f = torch.where(valid, pred, torch.zeros_like(pred))
+        gt_f = torch.where(valid, gt, torch.zeros_like(gt))
+        ssim = self.ssim(pred_f.permute(2, 0, 1)[None], gt_f.permute(2, 0, 1)[None])
+        return l1 * (1.0 - self.cfg.ssim_lambda) + (1.0 - ssim) * self.cfg.ssim_lambda
+
+    @torch.no_grad()
+    def prune(self, scores: torch.Tensor) -> None:
+        n = scores.numel() if scores.ndim == 1 else scores.shape[0]
+        num_prune = int(n * self.cfg.pruning_ratio)
+        if num_prune <= 0:
+            return
+        flat = scores.reshape(n, -1).mean(dim=1)
+        _, idx = torch.topk(flat, k=num_prune, largest=False)
+        keep = torch.ones(n, dtype=torch.bool, device=flat.device)
+        keep[idx] = False
+        self._apply_keep_mask(keep)
+
+    def _apply_keep_mask(self, keep: torch.Tensor) -> None:
+        with torch.no_grad():
+            for name, param in list(self.model.gauss_params.items()):
+                self.model.gauss_params[name] = torch.nn.Parameter(param.data[keep].clone())
+                self.model.gauss_params[name].requires_grad_(True)
+
+    def save(self) -> None:
+        path = self.output_path()
+        if self.cfg.output_format == "ply":
+            if path.suffix != ".ply":
+                path = path.with_suffix(".ply")
+            export_gaussian_ply(self.model, path)
+        elif self.cfg.output_format == "ckpt":
+            if path.suffix != ".ckpt":
+                path = path.with_suffix(".ckpt")
+            save_pruned_ckpt(self.pipeline, self.step, path)
+        else:
+            raise ValueError(f"Unknown output_format: {self.cfg.output_format}")
+        CONSOLE.print(f"Wrote {path}")
+
+    def output_path(self) -> Path:
+        if self.cfg.output_filename is not None:
+            return self.cfg.output_dir / self.cfg.output_filename
+        return self.cfg.output_dir / self.default_output_filename()
+
+    @abstractmethod
+    def default_output_filename(self) -> str:
+        """Default filename when ``output_filename`` is not set."""
+        ...
+
+
+def export_gaussian_ply(model: SplatfactoModel, filename: Path) -> None:
+    """Export gaussians using nerfstudio's ``ExportGaussianSplat.write_ply``."""
+    map_to_tensors: OrderedDict[str, np.ndarray] = OrderedDict()
+    with torch.no_grad():
+        positions = model.means.detach().cpu().numpy()
+        count = positions.shape[0]
+        n = count
+        map_to_tensors["x"] = positions[:, 0]
+        map_to_tensors["y"] = positions[:, 1]
+        map_to_tensors["z"] = positions[:, 2]
+        map_to_tensors["nx"] = np.zeros(n, dtype=np.float32)
+        map_to_tensors["ny"] = np.zeros(n, dtype=np.float32)
+        map_to_tensors["nz"] = np.zeros(n, dtype=np.float32)
+
+        shs_0 = model.shs_0.contiguous().cpu().numpy()
+        for i in range(shs_0.shape[1]):
+            map_to_tensors[f"f_dc_{i}"] = shs_0[:, i, None]
+
+        if model.config.sh_degree > 0:
+            shs_rest = model.shs_rest.transpose(1, 2).contiguous().cpu().numpy().reshape((n, -1))
+            for i in range(shs_rest.shape[-1]):
+                map_to_tensors[f"f_rest_{i}"] = shs_rest[:, i, None]
+
+        map_to_tensors["opacity"] = model.opacities.data.cpu().numpy()
+        scales = model.scales.data.cpu().numpy()
+        for i in range(3):
+            map_to_tensors[f"scale_{i}"] = scales[:, i, None]
+        quats = model.quats.data.cpu().numpy()
+        for i in range(4):
+            map_to_tensors[f"rot_{i}"] = quats[:, i, None]
+
+        select = np.ones(n, dtype=bool)
+        for k, t in map_to_tensors.items():
+            select = np.logical_and(select, np.isfinite(t).all(axis=-1))
+        if np.sum(select) < n:
+            for k in map_to_tensors:
+                map_to_tensors[k] = map_to_tensors[k][select]
+            count = int(np.sum(select))
+
+    ExportGaussianSplat.write_ply(str(filename), count, map_to_tensors)
+
+
+def save_pruned_ckpt(pipeline: Pipeline, step: int, path: Path) -> None:
+    """Save a nerfstudio-style checkpoint after in-place pruning."""
+    torch.save({"step": step, "pipeline": pipeline.state_dict()}, path)
+
+
+def launch_cli(config_cls: Type[HardPruneConfig]) -> None:
+    tyro.extras.set_accent_color("bright_yellow")
+    tyro.cli(config_cls).main()