Training code

.gitmodules

@@ -7,3 +7,6 @@
 [submodule "SIBR_viewers"]
 	path = SIBR_viewers
 	url = https://gitlab.inria.fr/sibr/sibr_core.git
+[submodule "submodules/PLAS"]
+	path = submodules/PLAS
+	url = https://github.com/fraunhoferhhi/PLAS.git

.vscode/launch.json

@@ -0,0 +1,35 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Train truck",
+            "type": "python",
+            "request": "launch",
+            "program": "train.py",
+            "console": "integratedTerminal",
+            "justMyCode": true,
+            "args": [
+                "--config-name",
+                "ours_q_sh_local_test",
+                "hydra.run.dir=/data/output/${now:%Y-%m-%d}/${now:%H-%M-%S}-${run.name}",
+                "dataset.source_path=/data/gaussian_splatting/tandt_db/tandt/truck",
+                "run.no_progress_bar=false",
+                // "local_window_debug_view.enabled=true",
+                // "run.save_iterations=[1200]",
+                // "optimization.iterations=1200",
+                "run.name=vs-code-debug",
+                // "dataset.sh_degree=3",
+                // "sorting.shuffle=true",
+                // "run.compress_iterations=[10,20]",
+                // "run.test_iterations=[10,20]",
+                // "run.save_iterations=[15,20]",
+                // "optimization.densification_interval=10",
+                // "optimization.iterations=20",
+                "run.test_lpips=false",
+            ]
+        },
+    ]
+}

README.md

@@ -19,10 +19,70 @@ This repository is a fork of the official authors implementation associated with
 
 The code for "Compact 3D Scene Representation via Self-Organizing Gaussian Grids" consists of multiple parts. The multi-dimensional sorting algorithm, PLAS, is available under the Apache License at [fraunhoferhhi/PLAS](https://github.com/fraunhoferhhi/PLAS).
 
-The integration of the sorting, the smoothness regularization and the compression code for training 3D scenes with the extended 3D Gaussian Splatting will become available in this repository.
+The integration of the sorting, the smoothness regularization and the compression code for training and compressing 3D scenes is available in this repository.
+
+## Cloning the Repository
+
+The repository contains submodules, thus please check it out with 
+```shell
+# SSH
+git clone git@github.com:fraunhoferhhi/Self-Organizing-Gaussians.git --recursive
+```
+or
+```shell
+# HTTPS
+git clone https://github.com/fraunhoferhhi/Self-Organizing-Gaussians.git --recursive
+```
+
+## Python Environment
+
+The code is using a few additional Python packages on top of graphdeco-inria/gaussian-splatting. We provide an extended environment.yml:
+
+Installation with [micromamba](https://mamba.readthedocs.io/en/latest/installation/micromamba-installation.html):
+
+```shell
+micromamba env create --file environment.yml --channel-priority flexible -y
+micromamba activate sogs
+```
+
+## Example training
+
+Download a dataset, e.g. [T&T](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/datasets/input/tandt_db.zip).
+
+The train.py script expects a name to a .yaml config file in the [config/](config/) folder. All parameters for the run are by default loaded from the yaml file. An example launch file can be found in .vscode/launch.json, for launching from Visual Studio Code.
+
+Example:
+
+```shell
+python train.py \
+  --config-name ours_q_sh_local_test \
+  hydra.run.dir=/data/output/${now:%Y-%m-%d}/${now:%H-%M-%S}-${run.name} \
+  dataset.source_path=/data/gaussian_splatting/tandt_db/tandt/truck \
+  run.no_progress_bar=false \
+  run.name=vs-code-debug
+```
+
+The parameter configurations can be overriden in the launch as shown (using [Hydra](https://hydra.cc/)).
+
+
+## Differences with graphdeco-inria/gaussian-splatting
+
+Code differences can be found in this diff: https://github.com/fraunhoferhhi/Self-Organizing-Gaussians/pull/1/files
+
+### Usage
+- different command-line interface for train.py (using Hydra)
+- wandb.ai used for logging
+
+### Code extensions
+- post-training quantization, compression/decompression
+- xyz log activation (gaussian_model.py)
+- grid sorting, neighbor loss (gaussian_model.py)
+- option to disable spherical harmonics
+
 
 ### Updates
 
+- 2024-06-13: Training code available
 - 2024-05-14: Improved compression scores! New results for paper v2 available on the [project website](https://fraunhoferhhi.github.io/Self-Organizing-Gaussians/)
 - 2024-05-02: Revised [paper v2](https://arxiv.org/pdf/2312.13299) on arXiv: Added compression of spherical harmonics, updated compression method with improved results (all attributes compressed with JPEG XL now), added qualitative comparison of additional scenes, moved compression explanation and comparison to main paper, added comparison with "Making Gaussian Splats smaller".
 - 2024-02-22: The code for the sorting algorithm is now available at [fraunhoferhhi/PLAS](https://github.com/fraunhoferhhi/PLAS)

arguments/__init__.py

@@ -12,6 +12,7 @@
 from argparse import ArgumentParser, Namespace
 import sys
 import os
+from omegaconf import OmegaConf
 
 class GroupParams:
     pass
@@ -110,3 +111,14 @@ def get_combined_args(parser : ArgumentParser):
         if v != None:
             merged_dict[k] = v
     return Namespace(**merged_dict)
+
+def get_hydra_training_args(model_path):
+    try:
+        cfgfilepath = os.path.join(model_path, "training_config.yaml")
+        print("Looking for config file in", cfgfilepath)
+        with open(cfgfilepath, 'r') as file:
+            print("Config file found: {}".format(cfgfilepath))
+            training_cfg = OmegaConf.load(cfgfilepath)
+            return training_cfg
+    except TypeError:
+        print("Config file not found!")

compression/codec.py

@@ -0,0 +1,81 @@
+from abc import ABC
+
+def normalize_img(img, min_val, max_val):
+
+    # min_clipped_count = (img < min_val).sum()
+    # max_clipped_count = (img > max_val).sum()
+    # print(f"Clipped {(min_clipped_count + max_clipped_count) / img.size * 100}% of values")
+
+    img = img.clip(min_val, max_val)
+    img = (img - min_val) / (max_val - min_val)
+    return img
+
+# from print_ranges.py
+min_thresholds = {
+    "_features_dc": -2,
+    "_features_rest": -1,
+    "_scaling": -13,
+    "_rotation": -1, # TODO better range
+    "_opacity": -6,
+}
+
+max_thresholds = {
+    "_features_dc": 4,
+    "_features_rest": 1,
+
+    # from print_ranges.py
+    # "_scaling": -1,
+
+    # manually overriden, because clipping large scaled gaussians to smaller scales messes up the results big time
+    "_scaling": 3,
+    "_rotation": 2,
+    "_opacity": 12,
+}
+
+class Codec(ABC):
+
+    def encode_image(self, image, out_file, **kwargs):
+        raise NotImplementedError("Subclasses should implement this!")
+
+    def decode_image(self, file_name):
+        raise NotImplementedError("Subclasses should implement this!")
+
+    def file_ending(self):
+        raise NotImplementedError("Subclasses should implement this!")
+
+    def normalize_to_thresholds(self, img, attr_name):
+
+        # normalize coordinates to 0...1
+        if attr_name == "_xyz":
+            xyz_min = img.min()
+            xyz_max = img.max()
+            return normalize_img(img, xyz_min, xyz_max), xyz_min, xyz_max
+
+        min_val = min_thresholds[attr_name]
+        max_val = max_thresholds[attr_name]
+
+        return normalize_img(img, min_val, max_val), min_val, max_val
+
+    def read_file_bytes(self, file_path):
+        with open(file_path, "rb") as f:
+            return f.read()
+
+    def write_file_bytes(self, file_path, bytes):
+        with open(file_path, "wb") as f:
+            f.write(bytes)
+
+    def encode(self, image, out_file, **kwargs):
+        self.encode_image(image, out_file, **kwargs)
+
+    def decode(self, image):
+        return self.decode_image(image)
+
+    def encode_with_normalization(self, image, attr_name, out_file, **kwargs):
+        img_norm, min_val, max_val = self.normalize_to_thresholds(image, attr_name)
+        self.encode(img_norm, out_file, **kwargs)
+        return min_val, max_val
+
+    def decode_with_normalization(self, file_name, min_val, max_val):
+        img_norm = self.decode(file_name)
+        return img_norm * (max_val - min_val) + min_val
+