Repository: nerfbaselines/nerfbaselines (language: python)

Issue:
gsplat metrics
Hey @jkulhanek, I am wondering if you are interested in adding the results of gsplat into your baseline comparisons as well? Currently, from my findings, some features in gsplat definitely beat the current metrics reported for gaussian-splatting (Inria based) and rival and possibly even beat the ZiP-NeRF metrics on some scenes.

Let me know... I could help out with this if you would be interested.


Discussion / hints:
Hi @maturk, Thanks! I will integrate it soon. It is very high on my priority list. Sorry, I was on a vacation...

Fix patch:
diff --git a/README.md b/README.md
index 2031b73a..d33a8962 100644
--- a/README.md
+++ b/README.md
@@ -84,6 +84,7 @@ Detailed results are available on the project page: [https://jkulhanek.com/nerfb
 | [Mip-Splatting](https://jkulhanek.com/nerfbaselines/m-mip-splatting)                     |     27.492 |     0.815 |       0.258 |     25m 37s |  11.0 GB |
 | [Gaussian Splatting](https://jkulhanek.com/nerfbaselines/m-gaussian-splatting)           |     27.434 |     0.814 |       0.257 |     23m 25s |  11.1 GB |
 | [Gaussian Opacity Fields](https://jkulhanek.com/nerfbaselines/m-gaussian-opacity-fields) |     27.421 |   *0.826* |     *0.234* |   1h 3m 54s |  28.4 GB |
+| [gsplat](https://jkulhanek.com/nerfbaselines/m-gsplat)                                   |     27.412 |     0.815 |       0.256 |     29m 19s |   8.3 GB |
 | [NerfStudio](https://jkulhanek.com/nerfbaselines/m-nerfacto)                             |     26.388 |     0.731 |       0.343 |   *19m 30s* | *5.9 GB* |
 | [Instant NGP](https://jkulhanek.com/nerfbaselines/m-instant-ngp)                         |     25.507 |     0.684 |       0.398 |  **3m 54s** |   7.8 GB |
 | [COLMAP](https://jkulhanek.com/nerfbaselines/m-colmap)                                   |     16.670 |     0.445 |       0.590 |  2h 52m 55s | **0 MB** |
@@ -103,6 +104,7 @@ Detailed results are available on the project page: [https://jkulhanek.com/nerfb
 | [K-Planes](https://jkulhanek.com/nerfbaselines/m-kplanes)                                |     32.265 |     0.961 |       0.062 |     23m 58s |   4.6 GB |
 | [Instant NGP](https://jkulhanek.com/nerfbaselines/m-instant-ngp)                         |     32.198 |     0.959 |       0.055 |  **2m 23s** | *2.6 GB* |
 | [Tetra-NeRF](https://jkulhanek.com/nerfbaselines/m-tetra-nerf)                           |     31.951 |     0.957 |       0.056 |  6h 53m 20s |  29.6 GB |
+| [gsplat](https://jkulhanek.com/nerfbaselines/m-gsplat)                                   |     31.471 |     0.966 |       0.054 |     14m 45s |   2.8 GB |
 | [Mip-NeRF 360](https://jkulhanek.com/nerfbaselines/m-mipnerf360)                         |     30.345 |     0.951 |       0.060 |  3h 29m 39s | 114.8 GB |
 | [NerfStudio](https://jkulhanek.com/nerfbaselines/m-nerfacto)                             |     29.191 |     0.941 |       0.095 |      9m 38s |   3.6 GB |
 | [NeRF](https://jkulhanek.com/nerfbaselines/m-nerf)                                       |     28.723 |     0.936 |       0.092 | 23h 26m 30s |  10.2 GB |
@@ -145,6 +147,7 @@ Detailed results are available on the project page: [https://jkulhanek.com/nerfb
 | Tetra-NeRF                | 🥈 silver | ❔        | 🥈 silver    | ❔         | ❔            | ❔           | ❔                |
 | WildGaussians             | ❔        | ❔        | ❔           | ❔         | 🥇 gold       | ❔           | ❔                |
 | Zip-NeRF                  | 🥇 gold   | ❌        | 🥇 gold      | 🥇 gold    | ❔            | ❔           | ❔                |
+| gsplat                    | 🥇 gold   | ❔        | 🥇 gold      | ❔         | 🥇 gold       | ❔           | 🥇 gold           |
 
 
 
@@ -177,11 +180,13 @@ For the currently implemented methods, the following licenses apply:
 - Mip-Splatting: [custom, research only](https://raw.githubusercontent.com/autonomousvision/mip-splatting/main/LICENSE.md)
 - NeRF-W (reimplementation): [MIT](https://raw.githubusercontent.com/kwea123/nerf_pl/master/LICENSE)
 - NeRF: [MIT](https://github.com/bmild/nerf/blob/master/LICENSE)
+- NerfStudio: [Apache 2.0](https://raw.githubusercontent.com/nerfstudio-project/nerfstudio/main/LICENSE)
 - SeaThru-NeRF: [Apache 2.0](https://raw.githubusercontent.com/deborahLevy130/seathru_NeRF/master/LICENSE)
 - TensoRF: [MIT](https://github.com/apchenstu/TensoRF/blob/main/LICENSE)
 - Tetra-NeRF: [MIT](https://raw.githubusercontent.com/jkulhanek/tetra-nerf/master/LICENSE)
 - WildGaussians: [MIT](https://raw.githubusercontent.com/jkulhanek/wild-gaussians/main/LICENSE), [custom, research only](https://raw.githubusercontent.com/graphdeco-inria/gaussian-splatting/main/LICENSE.md)
 - Zip-NeRF: [Apache 2.0](https://raw.githubusercontent.com/jonbarron/camp_zipnerf/main/LICENSE)
+- gsplat: [Apache 2.0](https://raw.githubusercontent.com/nerfstudio-project/gsplat/main/LICENSE)
 
 
 ## Acknowledgements
diff --git a/nerfbaselines/__init__.py b/nerfbaselines/__init__.py
index bd345892..c8b792a5 100644
--- a/nerfbaselines/__init__.py
+++ b/nerfbaselines/__init__.py
@@ -5,7 +5,7 @@
     RenderOutput as RenderOutput,
     RenderOutputType as RenderOutputType,
     RenderOptions as RenderOptions,
-    OptimizeEmbeddingsOutput as OptimizeEmbeddingsOutput,
+    OptimizeEmbeddingOutput as OptimizeEmbeddingOutput,
     Cameras as Cameras,
     ColorSpace as ColorSpace,
     CameraModel as CameraModel,
diff --git a/nerfbaselines/_types.py b/nerfbaselines/_types.py
index 202b548f..ff6e240f 100644
--- a/nerfbaselines/_types.py
+++ b/nerfbaselines/_types.py
@@ -372,7 +372,7 @@ def new_dataset(*,
 # accumulation: np.ndarray  # [h w]
 
 
-class OptimizeEmbeddingsOutput(TypedDict):
+class OptimizeEmbeddingOutput(TypedDict):
     embedding: Required[np.ndarray]
     render_output: NotRequired[RenderOutput]
     metrics: NotRequired[Dict[str, Sequence[float]]]
@@ -453,7 +453,7 @@ def get_train_embedding(self, index: int) -> Optional[np.ndarray]:
 
     def optimize_embedding(self, 
                            dataset: Dataset, *,
-                           embedding: Optional[np.ndarray] = None) -> OptimizeEmbeddingsOutput:
+                           embedding: Optional[np.ndarray] = None) -> OptimizeEmbeddingOutput:
         """
         Optimize embedding for a single image (passed as a dataset with a single image).
 
diff --git a/nerfbaselines/methods/gaussian_splatting_wild.py b/nerfbaselines/methods/gaussian_splatting_wild.py
index 2aac62f5..9d854221 100644
--- a/nerfbaselines/methods/gaussian_splatting_wild.py
+++ b/nerfbaselines/methods/gaussian_splatting_wild.py
@@ -24,7 +24,7 @@
 from PIL import Image
 from nerfbaselines import (
     Method, MethodInfo, ModelInfo, 
-    OptimizeEmbeddingsOutput, RenderOutput,
+    OptimizeEmbeddingOutput, RenderOutput,
     Cameras, camera_model_to_int, Dataset,
 )
 from nerfbaselines.utils import convert_image_dtype
@@ -537,7 +537,7 @@ def render(self, camera: Cameras, *, options=None, _gt_image=None, _store_cache=
                 "color": color,
             }
 
-    def optimize_embedding(self, dataset: Dataset, *, embedding: Optional[np.ndarray] = None) -> OptimizeEmbeddingsOutput:
+    def optimize_embedding(self, dataset: Dataset, *, embedding: Optional[np.ndarray] = None) -> OptimizeEmbeddingOutput:
         """
         Optimize embeddings on a single image (passed as a dataset).
 
diff --git a/nerfbaselines/methods/gsplat.py b/nerfbaselines/methods/gsplat.py
new file mode 100644
index 00000000..6b6e556f
--- /dev/null
+++ b/nerfbaselines/methods/gsplat.py
@@ -0,0 +1,678 @@
+import pprint
+import dataclasses
+import json
+import io
+import base64
+import yaml
+from contextlib import contextmanager
+import warnings
+from functools import partial
+import numpy as np
+import argparse
+import os
+import ast
+import importlib.util
+from typing import cast, Optional, List
+from nerfbaselines import Method, MethodInfo, ModelInfo, Dataset, OptimizeEmbeddingOutput
+from nerfbaselines.utils import pad_poses, image_to_srgb, convert_image_dtype
+from typing import Optional, Union
+from typing_extensions import Literal, get_origin, get_args
+
+import torch  # type: ignore
+from torch.nn import functional as F  # type: ignore
+
+
+def numpy_to_base64(array: np.ndarray) -> str:
+    with io.BytesIO() as f:
+        np.save(f, array)
+        return base64.b64encode(f.getvalue()).decode("ascii")
+
+
+def numpy_from_base64(data: str) -> np.ndarray:
+    with io.BytesIO(base64.b64decode(data)) as f:
+        return np.load(f)
+
+
+def cast_value(tp, value):
+    origin = get_origin(tp)
+    if origin is Literal:
+        for val in get_args(tp):
+            try:
+                value_casted = cast_value(type(val), value)
+                if val == value_casted:
+                    return value_casted
+            except ValueError:
+                pass
+            except TypeError:
+                pass
+        raise TypeError(f"Value {value} is not in {get_args(tp)}")
+    if origin is tuple:
+        # NOTE: not supporting ellipsis
+        if Ellipsis in get_args(tp):
+            raise TypeError("Ellipsis not supported")
+        if isinstance(value, str):
+            value = value.split(",")
+        if len(get_args(tp)) != len(value):
+            raise TypeError(f"Length of value {value} is not equal to {tp}")
+        return tuple(cast_value(t, v) for t, v in zip(get_args(tp), value))
+    if origin is Union:
+        for t in get_args(tp):
+            try:
+                return cast_value(t, value)
+            except ValueError:
+                import traceback
+                traceback.print_exc()
+                pass
+            except TypeError:
+                import traceback
+                traceback.print_exc()
+                pass
+        raise TypeError(f"Value {value} is not in {tp}")
+    if tp is type(None):
+        if str(value).lower() == "none":
+            return None
+        else:
+            raise TypeError(f"Value {value} is not None")
+    if tp is bool:
+        if str(value).lower() in {"true", "1", "yes"}:
+            return True
+        elif str(value).lower() in {"false", "0", "no"}:
+            return False
+        else:
+            raise TypeError(f"Value {value} is not a bool")
+    if tp in {int, float, bool, str}:
+        return tp(value)
+    if isinstance(value, tp):
+        return value
+    raise TypeError(f"Cannot cast value {value} to type {tp}")
+
+
+def flatten_hparams(hparams, *, separator: str = "/", _prefix: str = ""):
+    def format_value(v, only_simple_types=True):
+        if isinstance(v, (str, float, int, bool, bytes, type(None))):
+            return v
+        if torch.is_tensor(v) or isinstance(v, np.ndarray):
+            return format_value(v.tolist(), only_simple_types=only_simple_types)
+        if isinstance(v, (list, tuple)):
+            # If list of simple types, convert to string
+            if not only_simple_types:
+                return type(v)([format_value(x, only_simple_types=False) for x in v])
+            formatted = [format_value(x) for x in v]
+            if all(isinstance(x, (str, float, int, bool, bytes, type(None))) for x in formatted):
+                return ",".join(str(x) for x in formatted)
+            return ",".join(pprint.pformat(x) for x in formatted)
+        if isinstance(v, dict):
+            if not only_simple_types:
+                return {k: format_value(v, only_simple_types=False) for k, v in v.items()}
+            return pprint.pformat(format_value(v, only_simple_types=False))
+        if isinstance(v, type):
+            return v.__module__ + ":" + v.__name__
+        if dataclasses.is_dataclass(v):
+            return format_value({
+                f.name: getattr(v, f.name) for f in dataclasses.fields(v)
+            }, only_simple_types=only_simple_types)
+        if callable(v):
+            return v.__module__ + ":" + v.__name__
+        return repr(v)
+
+    flat = {}
+    if dataclasses.is_dataclass(hparams):
+        hparams = {f.name: getattr(hparams, f.name) for f in dataclasses.fields(hparams)}
+    _blacklist = set((
+        "port",
+        "ckpt",
+        "disable_viewer",
+        "render_traj_path",
+        "data_dir",
+        "result_dir",
+        "lpips_net",
+        "tb_every",
+        "tb_save_image",
+    ))
+    for k, v in hparams.items():
+        if _prefix:
+            k = f"{_prefix}{separator}{k}"
+        if k in _blacklist:
+            continue
+        if isinstance(v, dict) or dataclasses.is_dataclass(v):
+            flat.update(flatten_hparams(v, _prefix=k, separator=separator))
+        else:
+            flat[k] = format_value(v)
+    return flat
+
+
+class gs_Parser:
+    def __init__(self, 
+                 data_dir: str,
+                 factor: int = 1,
+                 normalize: bool = False,
+                 test_every: int = 8,
+                 state = None,
+                 dataset: Optional[Dataset] = None):
+        assert factor == 1, "Factor must be 1"
+        del test_every, data_dir
+
+        if state is not None:
+            self.transform = numpy_from_base64(state["transform_base64"])
+            self.scene_scale = state["scene_scale"]
+            self.points = self.points_rgb = None
+            if state["num_points"] is not None:
+                self.points = np.zeros((state["num_points"], 3), dtype=np.float32)
+                self.points_rgb = np.zeros((state["num_points"], 3), dtype=np.uint8)
+            self.num_train_images = state["num_train_images"]
+            self.dataset = None
+            return
+
+        assert dataset is not None, "Dataset must be provided"
+        self.num_train_images = len(dataset.get("images"))
+
+        # Optional normalize
+        from datasets.normalize import (  # type: ignore
+            similarity_from_cameras, transform_cameras, transform_points, align_principle_axes
+        )
+        if normalize:
+            points = dataset.get("points3D_xyz")
+            camtoworlds = pad_poses(dataset.get("cameras").poses)
+            transform = T1 = similarity_from_cameras(camtoworlds)
+            camtoworlds = transform_cameras(T1, camtoworlds)
+            if points is not None:
+                points = transform_points(T1, points)
+
+                T2 = align_principle_axes(points)
+                camtoworlds = transform_cameras(T2, camtoworlds)
+                points = transform_points(T2, points)
+
+                transform = cast(np.ndarray, T2 @ T1)
+
+            # Apply transform to the dataset
+            dataset = dataset.copy()
+            dataset["cameras"] = dataset["cameras"].replace(poses=camtoworlds[..., :3, :4])
+            dataset["points3D_xyz"] = points
+        else:
+            transform = np.eye(4)
+        self.transform = transform
+
+        # size of the scene measured by cameras
+        camera_locations = dataset["cameras"].poses[:, :3, 3]
+        scene_center = np.mean(camera_locations, axis=0)
+        dists = np.linalg.norm(camera_locations - scene_center, axis=1)
+
+        self.scene_scale = np.max(dists)
+        self.points = dataset.get("points3D_xyz")
+        self.points_rgb = dataset.get("points3D_rgb")
+        self.dataset = dataset
+
+    def export(self):
+        return {
+            "scene_scale": self.scene_scale,
+            "num_points": len(self.points) if self.points is not None else None,
+            "transform": self.transform.tolist(),
+            "transform_base64": numpy_to_base64(self.transform),
+            "num_train_images": self.num_train_images,
+        }
+
+
+class gs_Dataset:
+    def __init__(self, parser, 
+                 split: str = "train",
+                 patch_size: Optional[int] = None,
+                 load_depths: bool = False):
+        self.parser = parser
+        self.split = split
+        self.patch_size = patch_size
+        self.load_depths = load_depths
+        self.dataset = self.preprocess_images(parser.dataset)
+
+    def __len__(self):
+        return self.parser.num_train_images
+
+    @staticmethod
+    def preprocess_images(dataset):
+        if dataset is None:
+            return dataset
+        background_color = dataset["metadata"].get("background_color", None)
+        dataset = dataset.copy()
+        dataset["images"] = [image_to_srgb(image, 
+                                           background_color=background_color, 
+                                           dtype=np.uint8) for image in dataset["images"]]
+        return dataset
+
+    def __getitem__(self, idx):
+        dataset = self.dataset
+        image = dataset["images"][idx]
+        fx, fy, cx, cy = dataset["cameras"][idx].intrinsics
+        K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+        camtoworlds = pad_poses(dataset["cameras"][idx].poses)
+        sampling_mask = None
+        if dataset.get("sampling_masks", None) is not None:
+            sampling_mask = dataset["sampling_masks"][idx]
+
+        if self.patch_size is not None:
+            # Random crop.
+            h, w = image.shape[:2]
+            x = np.random.randint(0, max(w - self.patch_size, 1))
+            y = np.random.randint(0, max(h - self.patch_size, 1))
+            image = image[y : y + self.patch_size, x : x + self.patch_size]
+            if sampling_mask is not None:
+                sampling_mask = sampling_mask[y : y + self.patch_size, x : x + self.patch_size]
+            K[0, 2] -= x
+            K[1, 2] -= y
+
+        data = {
+            "K": torch.from_numpy(K).float(),
+            "camtoworld": torch.from_numpy(camtoworlds).float(),
+            "image": torch.from_numpy(image).float(),
+            "image_id": idx,  # the index of the image in the dataset
+        }
+        if sampling_mask is not None:
+            data["sampling_mask"] = torch.from_numpy(convert_image_dtype(sampling_mask, 'float32'))
+
+        if self.load_depths:
+            # projected points to image plane to get depths
+            worldtocams = np.linalg.inv(camtoworlds)
+            points_world = dataset["points3D_xyz"][dataset["images_points3D_indides"][idx]]
+            points_cam = (worldtocams[:3, :3] @ points_world.T + worldtocams[:3, 3:4]).T
+            points_proj = (K @ points_cam.T).T
+            points = points_proj[:, :2] / points_proj[:, 2:3]  # (M, 2)
+            depths = points_cam[:, 2]  # (M,)
+            # filter out points outside the image
+            selector = (
+                (points[:, 0] >= 0)
+                & (points[:, 0] < image.shape[1])
+                & (points[:, 1] >= 0)
+                & (points[:, 1] < image.shape[0])
+                & (depths > 0)
+            )
+            points = points[selector]
+            depths = depths[selector]
+            data["points"] = torch.from_numpy(points).float()
+            data["depths"] = torch.from_numpy(depths).float()
+        return data
+
+
+# Extract code dynamically
+def _build_simple_trainer_module():
+    module_spec = importlib.util.find_spec("simple_trainer", __package__)
+    assert module_spec is not None and module_spec.origin is not None, "Failed to find simple_trainer module"
+    with open(module_spec.origin, "r") as f:
+        simple_trainer_ast = ast.parse(f.read())
+
+    # Transform simple trainer, set num_workers=0
+    class _Visitor(ast.NodeVisitor):
+        def visit(self, node):
+            if isinstance(node, ast.Call) and ast.unparse(node.func) == "torch.utils.data.DataLoader":
+                num_workers = next(x for x in node.keywords if x.arg == "num_workers")
+                num_workers.value = ast.Constant(value=0, kind=None, lineno=0, col_offset=0)
+                persistent_workers = next((x for x in node.keywords if x.arg == "persistent_workers"), None)
+                if persistent_workers is not None:
+                    node.keywords.remove(persistent_workers)
+            super().visit(node)
+    _Visitor().visit(simple_trainer_ast)
+
+    # Filter imports
+    simple_trainer_ast.body.remove(next(x for x in simple_trainer_ast.body if ast.unparse(x) == "from torch.utils.tensorboard import SummaryWriter"))
+    simple_trainer_ast.body.remove(next(x for x in simple_trainer_ast.body if ast.unparse(x) == "import viser"))
+    # simple_trainer_ast.body.remove(next(x for x in simple_trainer_ast.body if ast.unparse(x) == "import nerfview"))
+
+    # Update config
+    config_ast = next(x for x in simple_trainer_ast.body if getattr(x, "name", None) == "Config")
+    assert isinstance(config_ast, ast.ClassDef)
+    config_ast.body.insert(-1, ast.AnnAssign(
+        target=ast.Name(id="app_test_opt_steps", ctx=ast.Store(), lineno=0, col_offset=0),
+        annotation=ast.Name(id="int", ctx=ast.Load(), lineno=0, col_offset=0),
+        value=ast.Constant(value=128, kind=None, lineno=0, col_offset=0), 
+        simple=True, col_offset=0, lineno=0))
+    config_ast.body.insert(-1, ast.AnnAssign(
+        target=ast.Name(id="app_test_opt_lr", ctx=ast.Store(), lineno=0, col_offset=0),
+        annotation=ast.Name(id="float", ctx=ast.Load(), lineno=0, col_offset=0),
+        value=ast.Constant(value=0.1, kind=None, lineno=0, col_offset=0), 
+        simple=True, col_offset=0, lineno=0))
+    config_ast.body.insert(-1, ast.AnnAssign(  # type: ignore
+        target=ast.Name(id="background_color", ctx=ast.Store(), lineno=0, col_offset=0),
+        annotation=ast.parse("Optional[Tuple[float, float, float]]").body[0].value,  # type: ignore
+        value=ast.Constant(value=None, kind=None, lineno=0, col_offset=0), 
+        simple=True, col_offset=0, lineno=0))
+
+    runner_ast = next(x for x in simple_trainer_ast.body if getattr(x, "name", None) == "Runner")
+    assert isinstance(runner_ast, ast.ClassDef)
+    runner_train_ast = next(x for x in runner_ast.body if getattr(x, "name", None) == "train")
+    assert isinstance(runner_train_ast, ast.FunctionDef)
+    runner_train_ast.name = "setup_train"
+    # Training loop
+    assert isinstance(runner_train_ast.body[-1], ast.For)
+    # Train init body - we remove unused code
+    init_train_body = list(runner_train_ast.body[:-3])
+    init_train_body.pop(4)
+    init_train_body.extend(ast.parse("""
+self.trainloader=trainloader
+self.trainloader_iter=trainloader_iter
+self.schedulers=schedulers
+""").body)
+    iter_step_body = []
+    iter_step_body.extend(ast.parse("""
+trainloader_iter=self.trainloader_iter
+trainloader=self.trainloader
+schedulers=self.schedulers
+""").body)
+    iter_step_body.extend(init_train_body[:4])
+    iter_step_body.extend((runner_train_ast.body[-1].body)[1:-1])
+    iter_step_body.pop(-10)  # Remove write to tensorboard step
+    save_step = iter_step_body.pop(-9)  # Remove save() step
+    iter_step_body.pop(-2)  # Remove eval() step
+    # Remove pbar.set_description
+    iter_step_body.pop(next(i for i, step in enumerate(iter_step_body) if ast.unparse(step) == "pbar.set_description(desc)"))
+
+    # NOTE: extend gsplat to use sampling_masks
+    render_step_idx = next(i for i, step in enumerate(iter_step_body) if ast.unparse(step).startswith("(renders, alphas, info) = self.rasterize_splats"))
+    iter_step_body.insert(render_step_idx + 1, ast.parse("""if data.get("sampling_mask") is not None:
+    sampling_mask = data["sampling_mask"].to(self.device)
+    renders = renders * sampling_mask + renders.detach() * (1 - sampling_mask)
+    alphas = alphas * sampling_mask + alphas.detach() * (1 - sampling_mask)
+""").body[0])
+    
+    bkgd_blend_step_idx = next(i for i, step in enumerate(iter_step_body) if ast.unparse(step).startswith("if cfg.random_bkgd:"))
+    iter_step_body.insert(bkgd_blend_step_idx + 1, ast.parse("""if not cfg.random_bkgd and cfg.background_color is not None:
+    bkgd = torch.tensor(cfg.background_color, dtype=colors.dtype, device=self.device).view(1, 1, 1, 3)
+    colors = colors + bkgd * (1.0 - alphas)
+""").body[0])
+
+    iter_step_body.extend(ast.parse("""def _():
+    self.trainloader_iter=trainloader_iter
+    out={"loss": loss.item(), "l1loss": l1loss.item(), "ssim": ssimloss.item(), "num_gaussians": len(self.splats["means"])}
+    if cfg.depth_loss:
+        out["depthloss"] = depthloss.item()
+    return out
+""").body[0].body)  # type: ignore
+    runner_train_ast.body = init_train_body
+    runner_ast.body.append(ast.FunctionDef(lineno=0, col_offset=0,
+        name="train_iteration",
+        args=ast.arguments(  # type: ignore
+            args=[
+                ast.arg(arg="self", annotation=None, lineno=0, col_offset=0),
+                ast.arg(arg="step", annotation=None, lineno=0, col_offset=0),
+            ],
+            posonlyargs=[], kwonlyargs=[], kw_defaults=[], defaults=[],
+            kwarg=None, kwargannotation=None, return_annotation=None
+        ),
+        body=iter_step_body, decorator_list=[],
+    ))
+
+    # Save method
+    save_step_body = save_step.body[4:]  # Strip saving stats
+    save_step_body.insert(0, ast.parse("cfg=self.cfg").body[0])
+    save_step_body.insert(0, ast.parse("world_size=self.world_size").body[0])
+    # Change saving location
+    save_step_body[-1].value.args[1].values[0].value = ast.Name(id="path", ctx=ast.Load(), lineno=0, col_offset=0)
+    runner_ast.body.append(ast.FunctionDef(lineno=0, col_offset=0,
+        name="save",
+        args=ast.arguments(  # type: ignore
+            args=[
+                ast.arg(arg="self", annotation=None, lineno=0, col_offset=0),
+                ast.arg(arg="step", annotation=None, lineno=0, col_offset=0),
+                ast.arg(arg="path", annotation=None, lineno=0, col_offset=0),
+            ],
+            posonlyargs=[], kwonlyargs=[], kw_defaults=[], defaults=[],
+            kwarg=None, kwargannotation=None, return_annotation=None
+        ),
+        body=save_step_body, decorator_list=[],
+    ))
+
+    # Init method
+    init_method = next(x for x in runner_ast.body if getattr(x, "name", None) == "__init__")
+    assert isinstance(init_method, ast.FunctionDef)
+    init_method.body = init_method.body[:6] + init_method.body[14:-2]  # Remove unused code
+    init_method.args.args.append(ast.arg(arg="Parser", annotation=None, lineno=0, col_offset=0))
+    init_method.args.args.append(ast.arg(arg="Dataset", annotation=None, lineno=0, col_offset=0))
+
+    # Execute code to build module
+    module = {}
+    exec(compile(simple_trainer_ast, "<string>", "exec"), module)
+    return argparse.Namespace(**module)
+
+
+class GSplat(Method):
+    def __init__(self, *, train_dataset=None, checkpoint=None, config_overrides=None):
+        super().__init__()
+
+        self.checkpoint = checkpoint
+        self.simple_trainer = _build_simple_trainer_module()
+
+        # Build trainer
+        self.cfg = self._get_config(checkpoint, config_overrides)
+
+        # Load parser state
+        parser_state = None
+        if checkpoint is not None and os.path.exists(f"{checkpoint}/parser.json"):
+            with open(f"{checkpoint}/parser.json", "r") as f:
+                parser_state = json.load(f)
+
+        # Build runner
+        local_rank = world_rank = 0
+        world_size = 1
+        self.runner = self.simple_trainer.Runner(
+            local_rank, world_rank, world_size, self.cfg, Dataset=gs_Dataset, Parser=partial(gs_Parser, dataset=train_dataset, state=parser_state))
+        self.step = 0
+        self._loaded_step = None
+
+        # Load checkpoint if available
+        if checkpoint is not None:
+            ckpt_files = [os.path.join(checkpoint, x) for x in os.listdir(checkpoint) if x.startswith("ckpt_") and x.endswith(".pt")]
+            ckpt_files.sort(key=lambda x: int(x.split("_rank")[-1].split(".")[0]))
+            ckpts = [
+                torch.load(file, map_location=self.runner.device, weights_only=True)
+                for file in ckpt_files
+            ]
+            for k in self.runner.splats.keys():
+                self.runner.splats[k].data = torch.cat([ckpt["splats"][k] for ckpt in ckpts])
+            self.step = self._loaded_step = ckpts[0]["step"]
+
+        # Setup dataloaders if training mode
+        if train_dataset is not None:
+            self.runner.setup_train()
+
+    def _get_config(self, checkpoint, config_overrides):
+        cfg = self.simple_trainer.Config(
+            strategy=self.simple_trainer.DefaultStrategy(verbose=True),
+        )
+        cfg.data_factor = 1
+
+        if checkpoint is not None:
+            with open(f"{checkpoint}/cfg.yml", "r") as f:
+                cfg_dict = yaml.load(f, Loader=yaml.UnsafeLoader)
+            cfg.__dict__.update(cfg_dict)
+
+        # Apply config overrides
+        field_types = {k.name: k.type for k in dataclasses.fields(cfg)}
+        if "strategy" in (config_overrides or {}):
+            import gsplat.strategy as strategies  # type: ignore
+            cfg.strategy = getattr(strategies, config_overrides["strategy"])()
+        strat_types = {k.name: k.type for k in dataclasses.fields(cfg.strategy)}
+        for k, v in (config_overrides or {}).items():
+            if k == "strategy":
+                continue
+            if k.startswith("strategy."):
+                v = cast_value(strat_types[k], v)
+                setattr(cfg.strategy, k[len("strategy."):], v)
+                continue
+            v = cast_value(field_types[k], v)
+            setattr(cfg, k, v)
+
+        cfg.adjust_steps(cfg.steps_scaler)
+        cfg.steps_scaler = 1.0  # We have already adjusted the steps
+        if cfg.pose_opt:
+            warnings.warn("Pose optimization is enabled, but it will only by applied to training images. No test-time pose optimization is enabled.")
+        return cfg
+
+    @classmethod
+    def get_method_info(cls):
+        return MethodInfo(
+            method_id="",  # Will be filled in by the registry
+            required_features=frozenset(("points3D_xyz", "points3D_rgb", "color", "images_points3D_indices")),
+            supported_camera_models=frozenset(("pinhole",)),
+            supported_outputs=("color", "depth", "accumulation"),
+        )
+
+    def get_info(self):
+        return ModelInfo(
+            **self.get_method_info(),
+            hparams=flatten_hparams(self.cfg),
+            num_iterations=self.cfg.max_steps,
+            loaded_step=self._loaded_step,
+            loaded_checkpoint=self.checkpoint,
+        )
+
+    def train_iteration(self, step):
+        self.step = step
+        out = self.runner.train_iteration(step)
+        self.step = step+1
+        return out
+
+    def save(self, path):
+        os.makedirs(path, exist_ok=True)
+        with open(f"{path}/cfg.yml", "w") as f:
+            yaml.dump(vars(self.cfg), f)
+        with open(f"{path}/parser.json", "w") as f:
+            json.dump(self.runner.parser.export(), f)
+        self.runner.save(self.step, path)
+
+    @contextmanager
+    def _patch_embedding(self, embedding=None):
+        if embedding is None:
+            yield None
+            return
+        embeds_module = None
+        try:
+            was_called = False
+            if self.cfg.app_opt:
+                embeds_module = self.runner.app_module.embeds
+                class _embed(torch.nn.Module):
+                    def forward(*args, **kwargs):
+                        del args, kwargs
+                        nonlocal was_called
+                        was_called = True
+                        return embedding[None]
+                self.runner.app_module.embeds = _embed()
+            yield None
+            if self.cfg.app_opt:
+                assert was_called, "Failed to patch appearance embedding"
+        finally:
+            # Return embeds back
+            if embeds_module is not None:
+                self.runner.app_module.embeds = embeds_module
+
+    def _add_background_color(self, img, accumulation):
+        if self.cfg.background_color is None:
+            return img
+        background_color = torch.tensor(self.cfg.background_color, device=img.device, dtype=torch.float32).view(1, 1, 1, 3)
+        return img + (1.0 - accumulation) * background_color
+
+    @torch.no_grad()
+    def render(self, camera, *, options=None):
+        camera = camera.item()
+        from datasets.normalize import transform_cameras  # type: ignore
+        cfg = self.cfg
+        device = self.runner.device
+        camtoworlds_np = transform_cameras(self.runner.parser.transform, pad_poses(camera.poses[None]))
+        camtoworlds = torch.from_numpy(camtoworlds_np).float().to(device)
+        fx, fy, cx, cy = camera.intrinsics
+        Ks = torch.from_numpy(np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])).float().to(device)
+        width, height = camera.image_sizes
+
+        # Patch appearance
+        embedding_np = (options or {}).get("embedding")
+        embedding = torch.from_numpy(embedding_np).to(self.runner.device) if embedding_np is not None else None
+        outputs = (options or {}).get("outputs") or ()
+        with self._patch_embedding(embedding):
+            colors, accumulation, _ = self.runner.rasterize_splats(
+                camtoworlds=camtoworlds,
+                Ks=Ks[None],
+                width=width,
+                height=height,
+                sh_degree=cfg.sh_degree,
+                near_plane=cfg.near_plane,
+                far_plane=cfg.far_plane,
+                image_ids=object() if embedding is not None else None,
+                render_mode="RGB+ED" if "depth" in outputs else "RGB",
+            )  # [1, H, W, 3]
+            colors = self._add_background_color(colors, accumulation)
+        out = {
+            "color": torch.clamp(colors.squeeze(0)[..., :3], 0.0, 1.0).detach().cpu().numpy(),
+            "accumulation": accumulation.squeeze(0).squeeze(-1).detach().cpu().numpy(),
+        }
+        if colors.shape[-1] > 3:
+            out["depth"] = colors.squeeze(0)[..., 3].detach().cpu().numpy()
+        return out
+
+    def get_train_embedding(self, index):
+        if not self.cfg.app_opt:
+            return None
+        return self.runner.app_module.embeds.weight[index].detach().cpu().numpy()
+
+    def optimize_embedding(self, dataset: Dataset, *, embedding=None) -> OptimizeEmbeddingOutput:
+        if not self.cfg.app_opt:
+            raise NotImplementedError("Appearance optimization is not enabled, add --set app_opt=True to the command line.")
+        assert len(dataset["images"]) == 1, "Only single image optimization is supported"
+        camera = dataset["cameras"].item()
+        from datasets.normalize import transform_cameras  # type: ignore
+        cfg = self.cfg
+        device = self.runner.device
+        camtoworlds_np = transform_cameras(self.runner.parser.transform, pad_poses(camera.poses[None]))
+        camtoworlds = torch.from_numpy(camtoworlds_np).float().to(device)
+        fx, fy, cx, cy = camera.intrinsics
+        Ks = torch.from_numpy(np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])).float().to(device)
+        width, height = camera.image_sizes
+        dataset = gs_Dataset.preprocess_images(dataset)
+        pixels = torch.from_numpy(dataset["images"][0]).float().to(device)[None].div(255.)
+        # Extend gsplat, handle sampling masks
+        sampling_masks = None
+        _dataset_sampling_masks = dataset.get("sampling_masks")
+        if _dataset_sampling_masks is not None:
+            sampling_masks = torch.from_numpy(convert_image_dtype(_dataset_sampling_masks[0], np.float32)).to(device)[None]
+
+        # Patch appearance
+        if embedding is not None:
+            embedding_th = torch.from_numpy(embedding).to(self.runner.device)
+        else:
+            embedding_th = torch.zeros_like(self.runner.app_module.embeds.weight[0])
+        embedding_th = torch.nn.Parameter(embedding_th.requires_grad_())
+        optimizer = torch.optim.Adam([embedding_th], lr=self.cfg.app_test_opt_lr)
+        l1losses: List[float] = []
+        ssimlosses: List[float] = []
+        losses: List[float] = []
+        with torch.enable_grad(), self._patch_embedding(embedding_th):
+            for _ in range(self.cfg.app_test_opt_steps):
+                colors, accumulation, _ = self.runner.rasterize_splats(
+                    camtoworlds=camtoworlds,
+                    Ks=Ks[None],
+                    width=width,
+                    height=height,
+                    sh_degree=cfg.sh_degree,
+                    near_plane=cfg.near_plane,
+                    far_plane=cfg.far_plane,
+                    image_ids=object()
+                )  # [1, H, W, 3]
+                colors = self._add_background_color(colors, accumulation)
+
+                # Scale colors grad by sampling masks
+                if sampling_masks is not None:
+                    colors = colors * sampling_masks + colors.detach() * (1.0 - sampling_masks)
+
+                l1loss = F.l1_loss(colors, pixels)
+                ssimloss = 1.0 - self.runner.ssim(
+                    pixels.permute(0, 3, 1, 2), colors.permute(0, 3, 1, 2)
+                )
+                loss = l1loss * (1.0 - cfg.ssim_lambda) + ssimloss * cfg.ssim_lambda
+                l1losses.append(l1loss.item())
+                ssimlosses.append(ssimloss.item())
+                losses.append(loss.item())
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+        return {
+            "metrics": {
+                "loss": losses,
+                "l1loss": l1losses,
+                "ssim": ssimlosses,
+            },
+            "embedding": embedding_th.detach().cpu().numpy(),
+        }
diff --git a/nerfbaselines/methods/gsplat_spec.py b/nerfbaselines/methods/gsplat_spec.py
new file mode 100644
index 00000000..72b8c0c3
--- /dev/null
+++ b/nerfbaselines/methods/gsplat_spec.py
@@ -0,0 +1,86 @@
+import os
+from nerfbaselines import register
+
+
+_name = os.path.split(__file__[:-len("_spec.py")])[-1].replace("_", "-")
+
+# Scenes: Bicycle Bonsai Counter Garden Kitchen Room Stump
+# PSNRs: 25.29 32.21 29.01 27.39 31.37 31.23 26.51
+# SSIMs: 0.77 0.94 0.91 0.87 0.93 0.92 0.77
+_paper_results = {
+    "mipnerf360/bicycle": {"psnr": 25.29, "ssim": 0.77},
+    "mipnerf360/bonsai": {"psnr": 32.21, "ssim": 0.94},
+    "mipnerf360/counter": {"psnr": 29.01, "ssim": 0.91},
+    "mipnerf360/garden": {"psnr": 27.39, "ssim": 0.87},
+    "mipnerf360/kitchen": {"psnr": 31.37, "ssim": 0.93},
+    "mipnerf360/room": {"psnr": 31.23, "ssim": 0.92},
+    "mipnerf360/stump": {"psnr": 26.51, "ssim": 0.77},
+}
+
+register({
+    "id": _name,
+    "method_class": ".gsplat:GSplat",
+    "conda": {
+        "environment_name": _name,
+        "python_version": "3.10",
+        "install_script": r"""
+git clone https://github.com/nerfstudio-project/gsplat.git
+cd gsplat
+git checkout cc800d7750a891ab8684eac4ddbcf90b79d16295
+git submodule init
+git submodule update --recursive
+conda develop "$PWD/examples"
+
+# Install build dependencies
+conda install -y cuda-toolkit 'numpy<2.0.0' pytorch==2.1.2 torchvision==0.16.2 -c pytorch -c nvidia/label/cuda-11.8.0
+export LIBRARY_PATH="$CONDA_PREFIX/lib/stubs"
+if [ "$NERFBASELINES_DOCKER_BUILD" != "1" ]; then
+conda install -y gcc_linux-64=11 gxx_linux-64=11 make=4.3 cmake=3.28.3 -c conda-forge
+fi
+
+# Install dependencies
+pip install opencv-python-headless==4.10.0.84 -r examples/requirements.txt
+
+# Install and build gsplat
+pip install -e . --use-pep517 --no-build-isolation
+python -c 'from gsplat import csrc'  # Test import
+
+# Clear build dependencies
+if [ "$NERFBASELINES_DOCKER_BUILD" = "1" ]; then
+# Reduce size of the environment by removing unused files
+find "$CONDA_PREFIX" -name '*.a' -delete
+find "$CONDA_PREFIX" -type d -name 'nsight*' -exec rm -r {} +
+fi
+""",
+    },
+    "metadata": {
+        "name": "gsplat",
+        "description": """gsplat is an open-source library for CUDA accelerated rasterization of gaussians with python bindings. \
+It is inspired by the 3DGS paper, but it is faster, more memory efficient, and with a growing list of new features. In NerfBaselines, \
+the method was modified to enable appearance optimization, to support sampling masks, and to support setting background color (which is required for the Blender dataset).""",
+        "paper_title": "gsplat: An Open-Source Library for Gaussian Splatting",
+        "paper_authors": [ "Vickie Ye", "Ruilong Li", "Justin Kerr", "Matias Turkulainen", "Brent Yi", "Zhuoyang Pan", "Otto Seiskari", "Jianbo Ye", "Jeffrey Hu", "Matthew Tancik", "Angjoo Kanazawa" ],
+        "paper_link": "https://arxiv.org/pdf/2409.06765.pdf",
+        "paper_results": _paper_results,
+        "link": "https://docs.gsplat.studio/main/",
+        "licenses": [{"name": "Apache 2.0", "url": "https://raw.githubusercontent.com/nerfstudio-project/gsplat/main/LICENSE"}],
+    },
+    "presets": {
+        "blender": { 
+            "@apply": [{"dataset": "blender"}], 
+            "init_type": "random", 
+            "background_color": (1.0, 1.0, 1.0),
+            "init_extent": 0.5,
+        },
+        "phototourism": { "@apply": [{"dataset": "phototourism"}], 
+            "app_opt": True,  # Enable appearance optimization
+            "steps_scaler": 3.333334,  # 100k steps
+        },
+    },
+    "implementation_status": {
+        "mipnerf360": "reproducing",
+        "blender": "working",
+        "tanksandtemples": "working",
+        "phototourism": "working",
+    }
+})
diff --git a/nerfbaselines/methods/kplanes.py b/nerfbaselines/methods/kplanes.py
index b2f8de5a..a4c407af 100644
--- a/nerfbaselines/methods/kplanes.py
+++ b/nerfbaselines/methods/kplanes.py
@@ -23,7 +23,7 @@
 import torch.utils.data
 from functools import partial
 from nerfbaselines import (
-    Dataset, RenderOutput, OptimizeEmbeddingsOutput,
+    Dataset, RenderOutput, OptimizeEmbeddingOutput,
     Method, MethodInfo, ModelInfo, Cameras, camera_model_to_int,
 )
 from nerfbaselines.utils import convert_image_dtype
@@ -696,7 +696,7 @@ def save(self, path: str):
             with open(fname + ".sha256", "w", encoding="utf8") as f:
                 f.write(sha)
 
-    def optimize_embedding(self, dataset: Dataset, *, embedding: Optional[np.ndarray] = None) -> OptimizeEmbeddingsOutput:
+    def optimize_embedding(self, dataset: Dataset, *, embedding: Optional[np.ndarray] = None) -> OptimizeEmbeddingOutput:
         """
         Optimize embedding of a single image (passed as dataset).
 
diff --git a/nerfbaselines/methods/nerfstudio_spec.py b/nerfbaselines/methods/nerfstudio_spec.py
index d8e13500..19bf1035 100644
--- a/nerfbaselines/methods/nerfstudio_spec.py
+++ b/nerfbaselines/methods/nerfstudio_spec.py
@@ -72,6 +72,7 @@
         "paper_link": "https://arxiv.org/pdf/2302.04264.pdf",
         "paper_results": _nerfacto_paper_results,
         "link": "https://docs.nerf.studio/",
+        "licenses": [{"name": "Apache 2.0", "url": "https://raw.githubusercontent.com/nerfstudio-project/nerfstudio/main/LICENSE"}],
     },
     "presets": {
         "blender": {
diff --git a/nerfbaselines/methods/nerfw_reimpl.py b/nerfbaselines/methods/nerfw_reimpl.py
index 6a475050..3ebb583f 100644
--- a/nerfbaselines/methods/nerfw_reimpl.py
+++ b/nerfbaselines/methods/nerfw_reimpl.py
@@ -20,7 +20,7 @@
 from typing import Optional, Any
 from nerfbaselines import (
     Method, Dataset, Cameras, RenderOutput,
-    OptimizeEmbeddingsOutput, 
+    OptimizeEmbeddingOutput, 
     ModelInfo, MethodInfo, CameraModel,
 )
 from nerfbaselines.utils import invert_transform, pad_poses, convert_image_dtype
@@ -698,7 +698,7 @@ def _get_default_embedding(self):
         if self.hparams.encode_a:
             return self._model.embedding_a.weight.data.mean(0).cpu().numpy()
 
-    def optimize_embedding(self, dataset: Dataset, *, embedding: Optional[np.ndarray] = None) -> OptimizeEmbeddingsOutput:
+    def optimize_embedding(self, dataset: Dataset, *, embedding: Optional[np.ndarray] = None) -> OptimizeEmbeddingOutput:
         torch.cuda.empty_cache()
         camera = dataset["cameras"].item()  # Ensure it is a single camera
         model = self._model


Test patch:
diff --git a/nerfbaselines/cli/_test_method.py b/nerfbaselines/cli/_test_method.py
index 11247a48..c0190d85 100644
--- a/nerfbaselines/cli/_test_method.py
+++ b/nerfbaselines/cli/_test_method.py
@@ -531,6 +531,8 @@ def test_render(model2, post):
                 }))
                 match = True
                 for k, v in paper_results.items():
+                    if "note" in k:
+                        continue
                     assert k in metrics
                     tolerance = 0.2
                     if k == "ssim":
diff --git a/tests/test_render_frames.py b/tests/test_render_frames.py
index 758e893b..625eae8d 100644
--- a/tests/test_render_frames.py
+++ b/tests/test_render_frames.py
@@ -89,7 +89,7 @@ def _test_render_trajectory_command(tmp_path, out, *args):
         import json
         json.dump(_mock_trajectory(num_cams), f)
     
-    command = [x.replace("{tmp_path}", str(tmp_path)).replace("{out}", out) for x in "nerfbaselines render-trajectory --checkpoint {tmp_path}/test-checkpoint --trajectory {tmp_path}/trajectory.json --output {tmp_path}/{out}".split()]
+    command = [x.replace("{tmp_path}", str(tmp_path)).replace("{out}", out) for x in "nerfbaselines render-trajectory --checkpoint {tmp_path}/test-checkpoint --trajectory {tmp_path}/trajectory.json --output {tmp_path}/{out} --backend python".split()]
     # Patch sys.argv
     test_registry = {"test-method": {
         "id": "test-method",