support unbounded mesh

README.md

@@ -1,13 +1,88 @@
 # 2D Gaussian Splatting for Geometrically Accurate Radiance Fields
 
-[Project page](https://surfsplatting.github.io/) | [Paper](https://arxiv.org/pdf/2403.17888) | [Video](https://www.youtube.com/watch?v=oaHCtB6yiKU) | [Surfel Rasterizer (CUDA)](https://github.com/hbb1/diff-surfel-rasterization) | [Surfel Rasterizer (python)](https://colab.research.google.com/drive/1qoclD7HJ3-o0O1R8cvV3PxLhoDCMsH8W?usp=sharing) | [DTU+COLMAP (3.5GB)](https://drive.google.com/drive/folders/1SJFgt8qhQomHX55Q4xSvYE2C6-8tFll9) |<br>
+[Project page](https://surfsplatting.github.io/) | [Paper](https://arxiv.org/pdf/2403.17888) | [Video](https://www.youtube.com/watch?v=oaHCtB6yiKU) | [Surfel Rasterizer (CUDA)](https://github.com/hbb1/diff-surfel-rasterization) | [Surfel Rasterizer (Python)](https://colab.research.google.com/drive/1qoclD7HJ3-o0O1R8cvV3PxLhoDCMsH8W?usp=sharing) | [DTU+COLMAP (3.5GB)](https://drive.google.com/drive/folders/1SJFgt8qhQomHX55Q4xSvYE2C6-8tFll9) |<br>
 
 ![Teaser image](assets/teaser.jpg)
 
-This repo contains the official implementation for the paper "2D Gaussian Splatting for Geometrically Accurate Radiance Fields". Our work represents a scene with a set of 2D oriented disks (surface elements) and rasterizes the surfels with [perspective correct differentiable raseterization](https://colab.research.google.com/drive/1qoclD7HJ3-o0O1R8cvV3PxLhoDCMsH8W?usp=sharing). Our work also develops regularizations that enhance the reconstruction quality.
+This repo contains the official implementation for the paper "2D Gaussian Splatting for Geometrically Accurate Radiance Fields". Our work represents a scene with a set of 2D oriented disks (surface elements) and rasterizes the surfels with [perspective correct differentiable raseterization](https://colab.research.google.com/drive/1qoclD7HJ3-o0O1R8cvV3PxLhoDCMsH8W?usp=sharing). Our work also develops regularizations that enhance the reconstruction quality. We also devise meshing approaches for Gaussian splatting.
 
-We are in the process of finalizing the training and rasterization code (CUDA), which may take a few days (or weeks) to complete. Feel free to contact us at huangbb@@shanghaitech.edu.cn if you have any questions.
 
+## ⭐ New Features 
+- 2024/05/05: Important updates - Now our algorithm supports **unbounded mesh extraction**!
+Our key idea is to contract the space into a sphere and then perform **adaptive TSDF truncation**. 
+
+![visualization](assets/unbounded.gif)
+
+## Installation
+
+```bash
+# download
+https://github.com/hbb1/2d-gaussian-splatting.git
+
+# if you have an environment used for 3dgs, use it
+# if not, create a new environment
+conda env create --file environment.yml
+conda activate surfel_splatting
+```
+## Training
+To train a scene, simply use
+```bash
+python train.py -s <path to COLMAP or NeRF Synthetic dataset>
+```
+Commandline arguments for regularizations
+```bash
+--lambda_normal  # hyperparameter for normal consistency
+--lambda_distortion # hyperparameter for depth distortion
+--depth_ratio # 0 for mean depth and 1 for median depth, 0 works for most cases
+```
+**Tips for adjusting the parameters on your own dataset:**
+- For unbounded/large scenes, we suggest using mean depth, i.e., ``depth_trunc=0``,  for less "disk-aliasing" artefacts.
+
+## Testing
+### Bounded Mesh Extraction
+To export a mesh within a bounded volume, simply use
+```bash
+render.py -m <path to pre-trained model> -s <path to COLMAP dataset> 
+```
+Commandline arguments you should adjust accordingly for meshing for bounded TSDF fusion, use
+```bash
+--depth_ratio # 0 for mean depth and 1 for median depth
+--voxel_size # voxel size
+--depth_trunc # depth truncation
+```
+### Unbounded Mesh Extraction
+To export a mesh with an arbitrary size, we devised an unbounded TSDF fusion with space contraction and adaptive truncation.
+```bash
+render.py -m <path to pre-trained model> -s <path to COLMAP dataset> --resolution 1024
+```
+
+### Quick Examples
+Assuming you have downloaded MipNeRF360, simply use
+```bash
+python train.py -s <path to m360>/<garden> -m output/m360/garden
+# use our unbounded mesh extraction!!
+python render.py -s <path to m360>/<garden> -m output/m360/garden --unbounded --skip_test --skip_train
+# or use the bounded mesh extraction if you focus on foreground
+python render.py -s <path to m360>/<garden> -m output/m360/garden -depth_trunc 6 --voxel_size 0.008 --skip_test --skip_train
+```
+If you have downloaded the DTU dataset, you can use
+```bash
+python train.py -s <path to dtu>/<scan105> -m output/date/scan105 -r 2 --depth_ratio 1
+python render.py -r 2 --depth_ratio 1 --depth_trunc 3 --voxel_size 0.004 --skip_test --skip_train
+```
+**Custom Dataset**: We use the same COLMAP loader as 3DGS, you can prepare your data following [here](https://github.com/graphdeco-inria/gaussian-splatting?tab=readme-ov-file#processing-your-own-scenes). 
+
+## Full evaluation
+We provide two scripts to evaluate our method of novel view synthesis and geometric reconstruction.
+For novel view synthesis on MipNeRF360 (which also works for other colmap datasets), use
+```bash
+python scripts/mipnerf_eval.py -m60 <path to the MipNeRF360 dataset>
+```
+For geometry reconstruction on DTU dataset, please download the preprocessed [data](). You also need to download the ground truth [DTU point cloud](https://roboimagedata.compute.dtu.dk/?page_id=36). 
+```bash
+python scripts/dtu_eval.py --dtu <path to the preprocessed DTU dataset>   \
+     --DTU_Official <path to the official DTU dataset>
+```
 
 ## Acknowledgements
 This project is built upon [3DGS](https://github.com/graphdeco-inria/gaussian-splatting). The TSDF fusion for extracting mesh is based on [Open3D](https://github.com/isl-org/Open3D). The rendering script for MipNeRF360 is adopted from [Multinerf](https://github.com/google-research/multinerf/), while the evaluation scripts for DTU and Tanks and Temples dataset are taken from [DTUeval-python](https://github.com/jzhangbs/DTUeval-python) and [TanksAndTemples](https://github.com/isl-org/TanksAndTemples/tree/master/python_toolbox/evaluation), respectively. We thank all the authors for their great repos. 
@@ -16,10 +91,12 @@ This project is built upon [3DGS](https://github.com/graphdeco-inria/gaussian-sp
 ## Citation
 If you find our code or paper helps, please consider citing:
 ```bibtex
-@article{Huang2DGS2024,
+@inproceedings{Huang2DGS2024,
     title={2D Gaussian Splatting for Geometrically Accurate Radiance Fields},
     author={Huang, Binbin and Yu, Zehao and Chen, Anpei and Geiger, Andreas and Gao, Shenghua},
-    journal={SIGGRAPH},
-    year={2024}
+    publisher = {Association for Computing Machinery},
+    booktitle = {SIGGRAPH 2024 Conference Papers},
+    year      = {2024},
+    doi       = {10.1145/3641519.3657428}
 }
 ```

metrics.py

@@ -70,9 +70,7 @@ def evaluate(model_paths):
                 for idx in tqdm(range(len(renders)), desc="Metric evaluation progress"):
                     ssims.append(ssim(renders[idx], gts[idx]))
                     psnrs.append(psnr(renders[idx], gts[idx]))
-                    #lpipss.append(lpips(renders[idx], gts[idx], net_type='vgg'))
-                    # hack to save time
-                    lpipss.append(ssims[-1])
+                    lpipss.append(lpips(renders[idx], gts[idx], net_type='vgg'))
 
                 print("  SSIM : {:>12.7f}".format(torch.tensor(ssims).mean(), ".5"))
                 print("  PSNR : {:>12.7f}".format(torch.tensor(psnrs).mean(), ".5"))

render.py

@@ -39,6 +39,8 @@
     parser.add_argument("--voxel_size", default=0.004, type=float, help='Mesh: voxel size for TSDF')
     parser.add_argument("--depth_trunc", default=3.0, type=float, help='Mesh: Max depth range for TSDF')
     parser.add_argument("--num_cluster", default=1000, type=int, help='Mesh: number of connected clusters to export')
+    parser.add_argument("--unbounded", action="store_true", help='Mesh: using unbounded mode for meshing')
+    parser.add_argument("--mesh_res", default=1024, type=int, help='Mesh: resolution for unbounded mesh extraction')
     args = get_combined_args(parser)
     print("Rendering " + args.model_path)
 
@@ -83,14 +85,20 @@
     if not args.skip_mesh:
         print("export mesh ...")
         os.makedirs(train_dir, exist_ok=True)
-        # set the active_sh to export only diffuse color texture
+        # set the active_sh to 0 to export only diffuse texture
         gaussExtractor.gaussians.active_sh_degree = 0
         gaussExtractor.reconstruction(scene.getTrainCameras())
         # extract the mesh and save
-        mesh = gaussExtractor.extract_mesh_bounded(voxel_size=args.voxel_size, sdf_trunc=5*args.voxel_size, depth_trunc=args.depth_trunc)
-        o3d.io.write_triangle_mesh(os.path.join(train_dir, 'fuse.ply'), mesh)
+        if args.unbounded:
+            name = 'fuse_unbounded.ply'
+            mesh = gaussExtractor.extract_mesh_unbounded(resolution=args.mesh_res)
+        else:
+            name = 'fuse.ply'
+            mesh = gaussExtractor.extract_mesh_bounded(voxel_size=args.voxel_size, sdf_trunc=5*args.voxel_size, depth_trunc=args.depth_trunc)
+
+        o3d.io.write_triangle_mesh(os.path.join(train_dir, name), mesh)
+        print("mesh saved at {}".format(os.path.join(train_dir, name)))
         # post-process the mesh and save, saving the largest N clusters
         mesh_post = post_process_mesh(mesh, cluster_to_keep=args.num_cluster)
-        o3d.io.write_triangle_mesh(os.path.join(train_dir, 'fuse_post.ply'), mesh_post)
-        print("mesh saved at {}".format(os.path.join(train_dir, 'fuse.ply')))
-        print("mesh post processed saved at {}".format(os.path.join(train_dir, 'fuse_post.ply')))
+        o3d.io.write_triangle_mesh(os.path.join(train_dir, name.replace('.ply', '_post.ply')), mesh_post)
+        print("mesh post processed saved at {}".format(os.path.join(train_dir, name.replace('.ply', '_post.ply'))))

utils/loss_utils.py

@@ -34,13 +34,6 @@ def smooth_loss(disp, img):
     grad_disp_y *= torch.exp(-grad_img_y)
     return grad_disp_x.mean() + grad_disp_y.mean()
 
-# def compute_image_weight(img, scale=5.):
-#     weight = torch.zeros(1, *img.shape[1:]).to(img.device)
-#     grad_img_x = torch.mean(torch.abs(img[:, 1:-1, :-2] - img[:, 1:-1, 2:]), 0, keepdim=True) * 0.5
-#     grad_img_y = torch.mean(torch.abs(img[:, :-2, 1:-1] - img[:, 2:, 1:-1]), 0, keepdim=True) * 0.5
-#     weight[:,1:-1, 1:-1] = 0.5 * (torch.exp(-scale*grad_img_x) + torch.exp(-scale*grad_img_y))
-#     return weight
-
 def create_window(window_size, channel):
     _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
     _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)

utils/mcube_utils.py

@@ -0,0 +1,95 @@
+#
+# Copyright (C) 2024, ShanghaiTech
+# SVIP research group, https://github.com/svip-lab
+# All rights reserved.
+#
+# This software is free for non-commercial, research and evaluation use 
+# under the terms of the LICENSE.md file.
+#
+# For inquiries contact  [email protected]
+#
+
+import numpy as np
+import torch
+import trimesh
+from skimage import measure
+# modified from here https://github.com/autonomousvision/sdfstudio/blob/370902a10dbef08cb3fe4391bd3ed1e227b5c165/nerfstudio/utils/marching_cubes.py#L201
+def marching_cubes_with_contraction(
+    sdf,
+    resolution=512,
+    bounding_box_min=(-1.0, -1.0, -1.0),
+    bounding_box_max=(1.0, 1.0, 1.0),
+    return_mesh=False,
+    level=0,
+    simplify_mesh=True,
+    inv_contraction=None,
+    max_range=32.0,
+):
+    assert resolution % 512 == 0
+
+    resN = resolution
+    cropN = 512
+    level = 0
+    N = resN // cropN
+
+    grid_min = bounding_box_min
+    grid_max = bounding_box_max
+    xs = np.linspace(grid_min[0], grid_max[0], N + 1)
+    ys = np.linspace(grid_min[1], grid_max[1], N + 1)
+    zs = np.linspace(grid_min[2], grid_max[2], N + 1)
+
+    meshes = []
+    for i in range(N):
+        for j in range(N):
+            for k in range(N):
+                print(i, j, k)
+                x_min, x_max = xs[i], xs[i + 1]
+                y_min, y_max = ys[j], ys[j + 1]
+                z_min, z_max = zs[k], zs[k + 1]
+
+                x = np.linspace(x_min, x_max, cropN)
+                y = np.linspace(y_min, y_max, cropN)
+                z = np.linspace(z_min, z_max, cropN)
+
+                xx, yy, zz = np.meshgrid(x, y, z, indexing="ij")
+                points = torch.tensor(np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T, dtype=torch.float).cuda()
+
+                @torch.no_grad()
+                def evaluate(points):
+                    z = []
+                    for _, pnts in enumerate(torch.split(points, 256**3, dim=0)):
+                        z.append(sdf(pnts))
+                    z = torch.cat(z, axis=0)
+                    return z
+
+                # construct point pyramids
+                points = points.reshape(cropN, cropN, cropN, 3)
+                points = points.reshape(-1, 3)
+                pts_sdf = evaluate(points.contiguous())
+                z = pts_sdf.detach().cpu().numpy()
+                if not (np.min(z) > level or np.max(z) < level):
+                    z = z.astype(np.float32)
+                    verts, faces, normals, _ = measure.marching_cubes(
+                        volume=z.reshape(cropN, cropN, cropN),
+                        level=level,
+                        spacing=(
+                            (x_max - x_min) / (cropN - 1),
+                            (y_max - y_min) / (cropN - 1),
+                            (z_max - z_min) / (cropN - 1),
+                        ),
+                    )
+                    verts = verts + np.array([x_min, y_min, z_min])
+                    meshcrop = trimesh.Trimesh(verts, faces, normals)
+                    meshes.append(meshcrop)
+
+                print("finished one block")
+
+    combined = trimesh.util.concatenate(meshes)
+    combined.merge_vertices(digits_vertex=6)
+
+    # inverse contraction and clipping the points range
+    if inv_contraction is not None:
+        combined.vertices = inv_contraction(torch.from_numpy(combined.vertices).float().cuda()).cpu().numpy()
+        combined.vertices = np.clip(combined.vertices, -max_range, max_range)
+
+    return combined