gaussian-splatting 1.0__tar.gz

gaussian_splatting-1.0/LICENSE.md

@@ -0,0 +1,91 @@
+Gaussian-Splatting License  
+===========================  
+
+**Inria** and **the Max Planck Institut for Informatik (MPII)** hold all the ownership rights on the *Software* named **gaussian-splatting**.  
+The *Software* is in the process of being registered with the Agence pour la Protection des  
+Programmes (APP).  
+
+The *Software* is still being developed by the *Licensor*.  
+
+*Licensor*'s goal is to allow the research community to use, test and evaluate  
+the *Software*.  
+
+## 1.  Definitions  
+
+*Licensee* means any person or entity that uses the *Software* and distributes  
+its *Work*.  
+
+*Licensor* means the owners of the *Software*, i.e Inria and MPII  
+
+*Software* means the original work of authorship made available under this  
+License ie gaussian-splatting.  
+
+*Work* means the *Software* and any additions to or derivative works of the  
+*Software* that are made available under this License.  
+
+
+## 2.  Purpose  
+This license is intended to define the rights granted to the *Licensee* by  
+Licensors under the *Software*.  
+
+## 3.  Rights granted  
+
+For the above reasons Licensors have decided to distribute the *Software*.  
+Licensors grant non-exclusive rights to use the *Software* for research purposes  
+to research users (both academic and industrial), free of charge, without right  
+to sublicense.. The *Software* may be used "non-commercially", i.e., for research  
+and/or evaluation purposes only.  
+
+Subject to the terms and conditions of this License, you are granted a  
+non-exclusive, royalty-free, license to reproduce, prepare derivative works of,  
+publicly display, publicly perform and distribute its *Work* and any resulting  
+derivative works in any form.  
+
+## 4.  Limitations  
+
+**4.1 Redistribution.** You may reproduce or distribute the *Work* only if (a) you do  
+so under this License, (b) you include a complete copy of this License with  
+your distribution, and (c) you retain without modification any copyright,  
+patent, trademark, or attribution notices that are present in the *Work*.  
+
+**4.2 Derivative Works.** You may specify that additional or different terms apply  
+to the use, reproduction, and distribution of your derivative works of the *Work*  
+("Your Terms") only if (a) Your Terms provide that the use limitation in  
+Section 2 applies to your derivative works, and (b) you identify the specific  
+derivative works that are subject to Your Terms. Notwithstanding Your Terms,  
+this License (including the redistribution requirements in Section 3.1) will  
+continue to apply to the *Work* itself.  
+
+**4.3** Any other use without of prior consent of Licensors is prohibited. Research  
+users explicitly acknowledge having received from Licensors all information  
+allowing to appreciate the adequacy between of the *Software* and their needs and  
+to undertake all necessary precautions for its execution and use.  
+
+**4.4** The *Software* is provided both as a compiled library file and as source  
+code. In case of using the *Software* for a publication or other results obtained  
+through the use of the *Software*, users are strongly encouraged to cite the  
+corresponding publications as explained in the documentation of the *Software*.  
+
+## 5.  Disclaimer  
+
+THE USER CANNOT USE, EXPLOIT OR DISTRIBUTE THE *SOFTWARE* FOR COMMERCIAL PURPOSES  
+WITHOUT PRIOR AND EXPLICIT CONSENT OF LICENSORS. YOU MUST CONTACT INRIA FOR ANY  
+UNAUTHORIZED USE: stip-sophia.transfert@inria.fr . ANY SUCH ACTION WILL  
+CONSTITUTE A FORGERY. THIS *SOFTWARE* IS PROVIDED "AS IS" WITHOUT ANY WARRANTIES  
+OF ANY NATURE AND ANY EXPRESS OR IMPLIED WARRANTIES, WITH REGARDS TO COMMERCIAL  
+USE, PROFESSIONNAL USE, LEGAL OR NOT, OR OTHER, OR COMMERCIALISATION OR  
+ADAPTATION. UNLESS EXPLICITLY PROVIDED BY LAW, IN NO EVENT, SHALL INRIA OR THE  
+AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR  
+CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE  
+GOODS OR SERVICES, LOSS OF USE, DATA, OR PROFITS OR BUSINESS INTERRUPTION)  
+HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT  
+LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING FROM, OUT OF OR  
+IN CONNECTION WITH THE *SOFTWARE* OR THE USE OR OTHER DEALINGS IN THE *SOFTWARE*.  
+
+## 6.  Files subject to permissive licenses
+The contents of the file ```utils/loss_utils.py``` are based on publicly available code authored by Evan Su, which falls under the permissive MIT license. 
+
+Title: pytorch-ssim\
+Project code: https://github.com/Po-Hsun-Su/pytorch-ssim\
+Copyright Evan Su, 2017\
+License: https://github.com/Po-Hsun-Su/pytorch-ssim/blob/master/LICENSE.txt (MIT)

gaussian_splatting-1.0/PKG-INFO

@@ -0,0 +1,11 @@
+Metadata-Version: 2.1
+Name: gaussian_splatting
+Version: 1.0
+Summary: UNKNOWN
+Home-page: UNKNOWN
+License: UNKNOWN
+Platform: UNKNOWN
+License-File: LICENSE.md
+
+UNKNOWN
+

gaussian_splatting-1.0/README.md

@@ -0,0 +1,196 @@
+# packaged 3D Gaussian Splatting
+
+This repo is the **refactored python training and inference code for [3D Gaussian Splatting](https://github.com/graphdeco-inria/gaussian-splatting)**.
+Forked from commit [a2a91d9093fd791fb01f556fa717f8d9f2cfbdd7](https://github.com/graphdeco-inria/gaussian-splatting/tree/a2a91d9093fd791fb01f556fa717f8d9f2cfbdd7).
+We **refactored the original code following the standard Python package structure**, while **keeping the algorithms used in the code identical to the original version**.
+
+## Features
+
+* [x] organize the code as a standard Python package
+* [x] exposure compensation
+* [x] camera and 3DGS parameters joint training
+* [ ] depth regularization
+* [ ] integrated 2DGS (integrated [gsplat](https://github.com/nerfstudio-project/gsplat) backend)
+
+## Install
+
+### Requirements
+
+Install Pytorch and torchvision following the official guideline: [pytorch.org](https://pytorch.org/)
+
+### Local Install
+
+```shell
+git clone https://github.com/yindaheng98/gaussian-splatting --recursive
+cd gaussian-splatting
+pip install tqdm plyfile
+pip install --target . --upgrade . --no-deps
+```
+
+### Pip Install
+
+You can download wheel from [Releases](https://github.com/yindaheng98/gaussian-splatting/releases), or
+install latest from source:
+```shell
+pip install --upgrade git+https://github.com/yindaheng98/gaussian-splatting.git@master
+```
+
+## Running
+
+Download dataset [T&T+DB COLMAP (650MB)](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/datasets/input/tandt_db.zip) and extract to `./data` directory.
+
+1. Train 3DGS with densification (same with original 3DGS)
+```shell
+python -m gaussian_splatting.train -s data/truck -d output/truck -i 30000 --mode densify
+```
+
+2. Render it
+```shell
+python -m gaussian_splatting.render -s data/truck -d output/truck -i 30000 --mode densify
+```
+
+3. Joint training 3DGS and camera (load the trained 3DGS)
+```shell
+python -m gaussian_splatting.train -s data/truck -d output/truck-camera -i 30000 --mode camera -l output/truck/point_cloud/iteration_30000/point_cloud.ply
+```
+
+4. Render it with trained 3DGS
+```shell
+python -m gaussian_splatting.render -s data/truck -d output/truck-camera -i 30000 --mode camera --load_camera output/truck-camera/cameras.json
+```
+
+This repo do not contrains code for initialization.
+If you want to create your own scene, please refer to [InstantSplat](https://github.com/yindaheng98/InstantSplat) or use [convert.py](https://github.com/graphdeco-inria/gaussian-splatting/blob/main/convert.py).
+
+## Usage
+
+**See [.vscode/launch.json](.vscode/launch.json) for more example.**
+
+**See [gaussian_splatting.train](gaussian_splatting/train.py) and [gaussian_splatting.render](gaussian_splatting/render.py) for full options.**
+
+### Gaussian models
+
+`GaussianModel` is the basic 3DGS model.
+```python
+from gaussian_splatting import GaussianModel
+gaussians = GaussianModel(sh_degree).to(device)
+```
+
+If you want cameras-3DGS joint training, use `CameraTrainableGaussianModel`, the rendering process is different.
+```python
+from gaussian_splatting import CameraTrainableGaussianModel
+gaussians = CameraTrainableGaussianModel(sh_degree).to(device)
+```
+
+save and load params:
+```python
+gaussians.save_ply("output/truck/point_cloud/iteration_30000/point_cloud.ply")
+gaussians.load_ply("output/truck/point_cloud/iteration_30000/point_cloud.ply")
+```
+
+init 3DGS with sparse point cloud extracted by colmap:
+```python
+from gaussian_splatting.dataset.colmap import colmap_init
+colmap_init(gaussians, "data/truck")
+```
+
+### Dataset
+
+Basic colmap dataset:
+```python
+from gaussian_splatting.dataset.colmap import ColmapCameraDataset, colmap_init
+dataset = ColmapCameraDataset("data/truck")
+```
+
+save to JSON and load JSON dataset:
+```python
+dataset.save_cameras("output/truck/cameras.json")
+from gaussian_splatting import JSONCameraDataset
+dataset = JSONCameraDataset("output/truck/cameras.json")
+```
+
+Dataset with trainable cameras:
+```python
+from gaussian_splatting import TrainableCameraDataset
+dataset = TrainableCameraDataset("data/truck") # init cameras from colmap
+dataset = TrainableCameraDataset.from_json("output/truck/cameras.json") # init cameras from saved json
+```
+
+### Inference
+
+```python
+for camera in dataset:
+  out = gaussians(camera)
+  image = out["render"]
+  ... # compute loss, save image or others
+```
+
+### Training
+
+`BaseTrainer` only optimize the 3DGS parameters, without densification or joint training with cameras.
+```python
+from gaussian_splatting.trainer import BaseTrainer
+trainer = BaseTrainer(
+    gaussians,
+    spatial_lr_scale=dataset.scene_extent(),
+    ... # see gaussian_splatting/trainer/trainer.py for full options
+)
+```
+
+`DensificationTrainer` optimize the 3DGS parameters and densify it.
+```python
+from gaussian_splatting.trainer import DensificationTrainer
+trainer = DensificationTrainer(
+    gaussians,
+    scene_extent=dataset.scene_extent(),
+    ... # see gaussian_splatting/trainer/densifier.py for full options
+)
+```
+
+`CameraTrainer` jointly optimize the 3DGS parameters and cameras, without densification
+```python
+from gaussian_splatting.trainer import CameraTrainer
+trainer = CameraTrainer(
+    gaussians,
+    scene_extent=dataset.scene_extent(),
+    dataset=dataset,
+    ... # see gaussian_splatting/trainer/camera_trainable.py for full options
+)
+```
+
+Train it:
+```python
+for camera in dataset:
+    loss, out = trainer.step(camera)
+```
+
+# 3D Gaussian Splatting for Real-Time Radiance Field Rendering
+Bernhard Kerbl*, Georgios Kopanas*, Thomas Leimkühler, George Drettakis (* indicates equal contribution)<br>
+| [Webpage](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/) | [Full Paper](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/3d_gaussian_splatting_high.pdf) | [Video](https://youtu.be/T_kXY43VZnk) | [Other GRAPHDECO Publications](http://www-sop.inria.fr/reves/publis/gdindex.php) | [FUNGRAPH project page](https://fungraph.inria.fr) |<br>
+| [T&T+DB COLMAP (650MB)](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/datasets/input/tandt_db.zip) | [Pre-trained Models (14 GB)](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/datasets/pretrained/models.zip) | [Viewers for Windows (60MB)](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/binaries/viewers.zip) | [Evaluation Images (7 GB)](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/evaluation/images.zip) |<br>
+![Teaser image](assets/teaser.png)
+
+This repository contains the official authors implementation associated with the paper "3D Gaussian Splatting for Real-Time Radiance Field Rendering", which can be found [here](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/). We further provide the reference images used to create the error metrics reported in the paper, as well as recently created, pre-trained models. 
+
+<a href="https://www.inria.fr/"><img height="100" src="assets/logo_inria.png"> </a>
+<a href="https://univ-cotedazur.eu/"><img height="100" src="assets/logo_uca.png"> </a>
+<a href="https://www.mpi-inf.mpg.de"><img height="100" src="assets/logo_mpi.png"> </a> 
+<a href="https://team.inria.fr/graphdeco/"> <img style="width:100%;" src="assets/logo_graphdeco.png"></a>
+
+Abstract: *Radiance Field methods have recently revolutionized novel-view synthesis of scenes captured with multiple photos or videos. However, achieving high visual quality still requires neural networks that are costly to train and render, while recent faster methods inevitably trade off speed for quality. For unbounded and complete scenes (rather than isolated objects) and 1080p resolution rendering, no current method can achieve real-time display rates. We introduce three key elements that allow us to achieve state-of-the-art visual quality while maintaining competitive training times and importantly allow high-quality real-time (≥ 30 fps) novel-view synthesis at 1080p resolution. First, starting from sparse points produced during camera calibration, we represent the scene with 3D Gaussians that preserve desirable properties of continuous volumetric radiance fields for scene optimization while avoiding unnecessary computation in empty space; Second, we perform interleaved optimization/density control of the 3D Gaussians, notably optimizing anisotropic covariance to achieve an accurate representation of the scene; Third, we develop a fast visibility-aware rendering algorithm that supports anisotropic splatting and both accelerates training and allows realtime rendering. We demonstrate state-of-the-art visual quality and real-time rendering on several established datasets.*
+
+<section class="section" id="BibTeX">
+  <div class="container is-max-desktop content">
+    <h2 class="title">BibTeX</h2>
+    <pre><code>@Article{kerbl3Dgaussians,
+      author       = {Kerbl, Bernhard and Kopanas, Georgios and Leimk{\"u}hler, Thomas and Drettakis, George},
+      title        = {3D Gaussian Splatting for Real-Time Radiance Field Rendering},
+      journal      = {ACM Transactions on Graphics},
+      number       = {4},
+      volume       = {42},
+      month        = {July},
+      year         = {2023},
+      url          = {https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/}
+}</code></pre>
+  </div>
+</section>

gaussian_splatting-1.0/gaussian_splatting/__init__.py

@@ -0,0 +1,3 @@
+from .camera import Camera
+from .gaussian_model import GaussianModel
+from .camera_trainable import CameraTrainableGaussianModel

gaussian_splatting-1.0/gaussian_splatting/camera.py

@@ -0,0 +1,106 @@
+import os
+from typing import NamedTuple, Callable
+import torch
+from PIL import Image
+from .utils import fov2focal, focal2fov, getProjectionMatrix, getWorld2View2, PILtoTorch, matrix_to_quaternion
+
+
+class Camera(NamedTuple):
+    image_height: int
+    image_width: int
+    FoVx: float
+    FoVy: float
+    R: torch.Tensor
+    T: torch.Tensor
+    world_view_transform: torch.Tensor
+    projection_matrix: torch.Tensor
+    full_proj_transform: torch.Tensor
+    camera_center: torch.Tensor
+    quaternion: torch.Tensor
+    ground_truth_image_path: str
+    ground_truth_image: torch.Tensor = None
+    postprocess: Callable[['Camera', torch.Tensor], torch.Tensor] = lambda camera, x: x
+    bg_color: torch.Tensor = torch.tensor([0., 0., 0.])
+    custom_data: dict = {}
+
+
+def camera2dict(camera: Camera, id):
+    Rt = torch.zeros((4, 4))
+    Rt[:3, :3] = camera.R
+    Rt[:3, 3] = camera.T
+    Rt[3, 3] = 1.0
+
+    W2C = torch.linalg.inv(Rt)
+    pos = W2C[:3, 3]
+    rot = W2C[:3, :3]
+    serializable_array_2d = [x.tolist() for x in rot]
+    camera_entry = {
+        'id': id,
+        'width': camera.image_width,
+        'height': camera.image_height,
+        'position': pos.tolist(),
+        'rotation': serializable_array_2d,
+        'fy': fov2focal(camera.FoVx, camera.image_height),
+        'fx': fov2focal(camera.FoVy, camera.image_width),
+        'ground_truth_image_path': camera.ground_truth_image_path.replace("\\", "/"),
+        "img_name": os.path.basename(camera.ground_truth_image_path),
+    }
+    return camera_entry
+
+
+def build_camera(
+        image_height: int, image_width: int,
+        FoVx: float, FoVy: float,
+        R: torch.Tensor, T: torch.Tensor,
+        image_path: str = None, device="cuda"
+):
+    zfar = 100.0
+    znear = 0.01
+    trans = torch.zeros(3)
+    scale = 1.0
+    world_view_transform = getWorld2View2(R, T, trans, scale).to(device).transpose(0, 1)
+    projection_matrix = getProjectionMatrix(znear=znear, zfar=zfar, fovX=FoVx, fovY=FoVy).to(device).transpose(0, 1)
+    full_proj_transform = (world_view_transform.unsqueeze(0).bmm(projection_matrix.unsqueeze(0))).squeeze(0)
+    camera_center = world_view_transform.inverse()[3, :3]
+    quaternion = matrix_to_quaternion(R)
+    gt_image = None
+    if image_path is not None:
+        pil_image = Image.open(image_path)
+        torch_image = PILtoTorch(pil_image)
+        gt_image = torch_image[:3, ...].clamp(0.0, 1.0).to(device)
+        image_height = gt_image.shape[1]
+        image_width = gt_image.shape[2]
+    return Camera(
+        # image_height=colmap_camera.image_height, # colmap_camera.image_height is read from cameras.bin, maybe dfferent from the actual image size
+        # image_width=colmap_camera.image_width, # colmap_camera.image_width is read from cameras.bin, maybe dfferent from the actual image size
+        image_height=image_height, image_width=image_width,
+        FoVx=FoVx, FoVy=FoVy,
+        R=R.to(device), T=T.to(device),
+        world_view_transform=world_view_transform,
+        projection_matrix=projection_matrix,
+        full_proj_transform=full_proj_transform,
+        camera_center=camera_center,
+        quaternion=quaternion.to(device),
+        ground_truth_image_path=image_path,
+        ground_truth_image=gt_image
+    )
+
+
+def dict2camera(camera_dict, device="cuda"):
+    W2C = torch.zeros((4, 4))
+    W2C[:3, 3] = torch.tensor(camera_dict['position'])
+    W2C[:3, :3] = torch.tensor(camera_dict['rotation'])
+    W2C[3, 3] = 1.0
+    Rt = torch.linalg.inv(W2C)
+    T = Rt[:3, 3]
+    R = Rt[:3, :3]
+    return build_camera(
+        image_width=camera_dict['width'],
+        image_height=camera_dict['height'],
+        FoVx=focal2fov(camera_dict['fy'], camera_dict['height']),
+        FoVy=focal2fov(camera_dict['fx'], camera_dict['width']),
+        R=R,
+        T=T,
+        image_path=camera_dict['ground_truth_image_path'],
+        device=device
+    )

gaussian_splatting-1.0/gaussian_splatting/camera_trainable.py

@@ -0,0 +1,95 @@
+import torch
+import math
+from .gaussian_model import GaussianModel, Camera
+from .utils import normalize_quaternion, quaternion_to_matrix, quaternion_raw_multiply
+
+from gaussian_splatting.diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
+from gaussian_splatting.simple_knn._C import distCUDA2
+
+
+class CameraTrainableGaussianModel(GaussianModel):
+
+    def forward(self, viewpoint_camera: Camera):
+        # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
+        screenspace_points = torch.zeros_like(self.get_xyz, dtype=self.get_xyz.dtype, requires_grad=True, device=self._xyz.device) + 0
+        try:
+            screenspace_points.retain_grad()
+        except:
+            pass
+
+        # Set up rasterization configuration
+        tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
+        tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
+
+        # Set camera pose as identity. Then, we will transform the Gaussians around camera_pose
+        w2c = torch.eye(4, device=self._xyz.device)
+        projmatrix = (
+            w2c.unsqueeze(0).bmm(viewpoint_camera.projection_matrix.unsqueeze(0))
+        ).squeeze(0)
+        campos = w2c.inverse()[3, :3]
+        raster_settings = GaussianRasterizationSettings(
+            image_height=int(viewpoint_camera.image_height),
+            image_width=int(viewpoint_camera.image_width),
+            tanfovx=tanfovx,
+            tanfovy=tanfovy,
+            bg=viewpoint_camera.bg_color.to(self._xyz.device),
+            scale_modifier=self.scale_modifier,
+            # viewmatrix=viewpoint_camera.world_view_transform,
+            # projmatrix=viewpoint_camera.full_proj_transform,
+            viewmatrix=w2c,
+            projmatrix=projmatrix,
+            sh_degree=self.active_sh_degree,
+            # campos=viewpoint_camera.camera_center,
+            campos=campos,
+            prefiltered=False,
+            debug=self.debug,
+            antialiasing=self.antialiasing
+        )
+
+        rasterizer = GaussianRasterizer(raster_settings=raster_settings)
+
+        # means3D = pc.get_xyz
+        rel_w2c = torch.eye(4, device=self._xyz.device)
+        quaternion = viewpoint_camera.quaternion
+        rel_w2c[:3, :3] = quaternion_to_matrix(normalize_quaternion(quaternion.unsqueeze(0))).squeeze(0)
+        rel_w2c[:3, 3] = viewpoint_camera.T
+        # Transform mean and rot of Gaussians to camera frame
+        gaussians_xyz = self._xyz.clone()
+        gaussians_rot = self._rotation.clone()
+
+        xyz_ones = torch.ones(gaussians_xyz.shape[0], 1).cuda().float()
+        xyz_homo = torch.cat((gaussians_xyz, xyz_ones), dim=1)
+        gaussians_xyz_trans = (rel_w2c @ xyz_homo.T).T[:, :3]
+        gaussians_rot_trans = quaternion_raw_multiply(quaternion, gaussians_rot)
+        means3D = gaussians_xyz_trans
+        means2D = screenspace_points
+        opacity = self.get_opacity
+
+        scales = self.get_scaling
+        rotations = gaussians_rot_trans  # pc.get_rotation
+
+        shs = self.get_features
+
+        # Rasterize visible Gaussians to image, obtain their radii (on screen).
+        rendered_image, radii, depth_image = rasterizer(
+            means3D=means3D,
+            means2D=means2D,
+            shs=shs,
+            colors_precomp=None,
+            opacities=opacity,
+            scales=scales,
+            rotations=rotations,
+            cov3D_precomp=None)
+        rendered_image = viewpoint_camera.postprocess(viewpoint_camera, rendered_image)
+
+        # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
+        # They will be excluded from value updates used in the splitting criteria.
+        rendered_image = rendered_image.clamp(0, 1)
+        out = {
+            "render": rendered_image,
+            "viewspace_points": screenspace_points,
+            "visibility_filter": (radii > 0).nonzero(),
+            "radii": radii,
+            "depth": depth_image
+        }
+        return out

gaussian_splatting-1.0/gaussian_splatting/dataset/__init__.py

@@ -0,0 +1,2 @@
+from .dataset import CameraDataset, JSONCameraDataset
+from .camera_trainable import TrainableCameraDataset

gaussian_splatting-1.0/gaussian_splatting/dataset/camera_trainable.py

@@ -0,0 +1,70 @@
+import json
+from typing import List
+from gaussian_splatting import Camera
+import torch
+import torch.nn as nn
+
+from gaussian_splatting.camera import camera2dict
+from gaussian_splatting.utils import quaternion_to_matrix
+from .dataset import CameraDataset, JSONCameraDataset
+
+
+def exposure_postprocess(camera: Camera, x: torch.Tensor):
+    exposure = camera.custom_data['exposures']
+    return torch.matmul(x.permute(1, 2, 0), exposure[:3, :3]).permute(2, 0, 1) + exposure[:3, 3, None, None]
+
+
+class TrainableCameraDataset(CameraDataset):
+
+    def __init__(self, cameras: List[Camera], exposures: List[torch.Tensor] = []):
+        super().__init__()
+        self.cameras = cameras
+        self.quaternions = nn.Parameter(torch.stack([camera.quaternion for camera in cameras]))
+        self.Ts = nn.Parameter(torch.stack([camera.T for camera in cameras]))
+        self.exposures = nn.Parameter(torch.stack([torch.eye(3, 4, device=camera.T.device) for camera in cameras]))
+        if len(exposures) > 0:
+            assert len(exposures) == len(cameras), "Number of exposures must match number of cameras"
+            with torch.no_grad():
+                for idx, exposure in enumerate(exposures):
+                    self.exposures[idx, ...] = exposure.to(self.exposures.device)
+
+    def __len__(self):
+        return len(self.cameras)
+
+    def __getitem__(self, idx) -> Camera:
+        return Camera(**{
+            **self.cameras[idx]._asdict(),
+            'quaternion': self.quaternions[idx, ...],
+            'T': self.Ts[idx, ...],
+            'postprocess': exposure_postprocess,
+            'custom_data': {
+                **self.cameras[idx].custom_data,
+                'exposures': self.exposures[idx, ...]
+            }
+        })
+
+    def to(self, device):
+        self.quaternions.to(device)
+        self.Ts.to(device)
+        self.exposures.to(device)
+        return self
+
+    def save_cameras(self, path):
+        cameras = []
+        for idx, camera in enumerate(self):
+            cameras.append({
+                **camera2dict(Camera(**{
+                    **camera._asdict(),
+                    'R': quaternion_to_matrix(self.quaternions[idx, ...]),
+                    'T': self.Ts[idx, ...],
+                }), idx),
+                "exposure": self.exposures[idx, ...].detach().tolist(),
+            })
+        with open(path, 'w') as f:
+            json.dump(cameras, f, indent=2)
+
+    @classmethod
+    def from_json(cls, path):
+        cameras = JSONCameraDataset(path)
+        exposures = [(torch.tensor(camera['exposure']) if 'exposure' in camera else torch.eye(3, 4)) for camera in cameras.json_cameras]
+        return cls(cameras, exposures)

gaussian_splatting-1.0/gaussian_splatting/dataset/colmap/__init__.py

@@ -0,0 +1,5 @@
+from .dataset import ColmapCameraDataset, ColmapTrainableCameraDataset
+from .params_init import colmap_init
+from .dataset import ColmapCamera, read_colmap_cameras
+from .params_init import read_colmap_points3D, read_points3D_binary, read_points3D_text
+from .dataset import read_cameras_binary, read_cameras_text, parse_colmap_camera