Repository: gwilczynski95/meshsplats Branch: main Commit: 6cf6c2399ed0 Files: 43 Total size: 173.7 KB Directory structure: gitextract_j8tpcnjm/ ├── GS_LICENSE.md ├── LICENSE.md ├── NV_LICENSE.txt ├── README.md ├── faces_notexturemap_blender.py ├── generate_pseudomesh.py ├── gs_utils.py ├── mesh_optim/ │ ├── 2dgs_experiments_run.py │ ├── cam_utils.py │ ├── data.py │ ├── diff_render.py │ ├── generate_multi_views_circle.py │ ├── log_config.yaml │ ├── lpipsPyTorch/ │ │ ├── __init__.py │ │ └── modules/ │ │ ├── lpips.py │ │ ├── networks.py │ │ └── utils.py │ ├── mesh_utils.py │ ├── metrics.py │ ├── ml_utils.py │ ├── models.py │ ├── optimize_pseudo_config.yaml │ ├── optimize_pseudomesh.py │ └── render_pseudomesh.py ├── requirements.txt ├── scripts/ │ └── visualize_cameras_nerf_blender.py └── sh_scripts/ ├── configs/ │ ├── 3dgs_sh0_pseudo_config.yaml │ ├── db_config.yaml │ ├── mip_config.yaml │ ├── tandt_3dgs_config.yaml │ └── tandt_config.yaml ├── run_2dgs-sh0_db.sh ├── run_2dgs-sh0_mip.sh ├── run_2dgs-sh0_nerf-synth.sh ├── run_2dgs_sh0_tandt.sh ├── run_3dgs-sh0_db.sh ├── run_3dgs-sh0_mip.sh ├── run_3dgs-sh0_nerf-synth.sh ├── run_3dgs-sh0_tandt.sh ├── run_games_gs-flat_sh0_db.sh ├── run_games_gs-flat_sh0_mip.sh ├── run_games_gs-flat_sh0_nerf-synth.sh └── run_games_gs-flat_sh0_tandt.sh ================================================ FILE CONTENTS ================================================ ================================================ FILE: GS_LICENSE.md ================================================ 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) ================================================ FILE: LICENSE.md ================================================ MIT License Copyright (c) 2025 Rafał Tobiasz Grzegorz Wilczyński Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ================================================ FILE: NV_LICENSE.txt ================================================ Copyright (c) 2020, NVIDIA Corporation. All rights reserved. Nvidia Source Code License (1-Way Commercial) ======================================================================= 1. Definitions "Licensor" means any person or entity that distributes its Work. "Software" means the original work of authorship made available under this License. "Work" means the Software and any additions to or derivative works of the Software that are made available under this License. The terms "reproduce," "reproduction," "derivative works," and "distribution" have the meaning as provided under U.S. copyright law; provided, however, that for the purposes of this License, derivative works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work. Works, including the Software, are "made available" under this License by including in or with the Work either (a) a copyright notice referencing the applicability of this License to the Work, or (b) a copy of this License. 2. License Grants 2.1 Copyright Grant. Subject to the terms and conditions of this License, each Licensor grants to you a perpetual, worldwide, non-exclusive, royalty-free, copyright license to reproduce, prepare derivative works of, publicly display, publicly perform, sublicense and distribute its Work and any resulting derivative works in any form. 3. Limitations 3.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. 3.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 3.3 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. 3.3 Use Limitation. The Work and any derivative works thereof only may be used or intended for use non-commercially. The Work or derivative works thereof may be used or intended for use by Nvidia or its affiliates commercially or non-commercially. As used herein, "non-commercially" means for research or evaluation purposes only and not for any direct or indirect monetary gain. 3.4 Patent Claims. If you bring or threaten to bring a patent claim against any Licensor (including any claim, cross-claim or counterclaim in a lawsuit) to enforce any patents that you allege are infringed by any Work, then your rights under this License from such Licensor (including the grant in Section 2.1) will terminate immediately. 3.5 Trademarks. This License does not grant any rights to use any Licensor's or its affiliates' names, logos, or trademarks, except as necessary to reproduce the notices described in this License. 3.6 Termination. If you violate any term of this License, then your rights under this License (including the grant in Section 2.1) will terminate immediately. 4. Disclaimer of Warranty. THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER THIS LICENSE. 5. Limitation of Liability. EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK (INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. ======================================================================= ================================================ FILE: README.md ================================================ # MeshSplats Rafał Tobiasz*, Grzegorz Wilczyński*, Marcin Mazur, Sławomir Tadeja, Przemysław Spurek (* indicates equal contribution)
This repository contains the official authors implementation associated with the paper ["MeshSplats: Mesh-Based Rendering with Gaussian Splatting Initialization"](https://arxiv.org/pdf/2502.07754). Abstract: * Recently, a range of neural network-based methods for image rendering have been introduced. Gaussian Splatting (GS) is a recent and pivotal technique in 3D computer graphics. GS-based algorithms almost always bypass classical methods such as ray tracing, which offers numerous inherent advantages for rendering. For example, ray tracing is able to handle incoherent rays for advanced lighting effects, including shadows and reflections. To address this limitation, we introduce MeshSplats, a method which converts GS to a mesh-like format. Following the completion of training, MeshSplats transforms Gaussian elements into mesh faces, enabling rendering using ray tracing methods with all their associated benefits. Our model can be utilized immediately following transformation, yielding a mesh of slightly reduced quality without additional training. Furthermore, we can enhance the reconstruction quality through the application of a dedicated optimization algorithm that operates on mesh faces rather than Gaussian components. The efficacy of our method is substantiated by experimental results, underscoring its extensive applications in computer graphics and image processing* Check out this code if you just want to convert GS to a mesh and benefit from the advantages of both representations!
Note: If videos aren't visible you can find them in `demo` directory. ## Run our demo in Google Colab!!! We have prepared for you a demo (to run as quickly as possible) of the MeshSplats method. ### Optimization pipeline demo [Here's the Colab!](https://colab.research.google.com/drive/1N6Y6TfijUw9tQ8vrnoLrAhBPlxkSzA2b?usp=sharing) First of all, all of the data you need is [here](https://drive.google.com/drive/folders/1KjREkxiKRY-_7iDgV9q-5ItTn9rn4vPu?usp=sharing). Download it as it contains full data of one of the experiments we performed (`lego` scene from the `NeRF Synthetic` dataset with white background). This data structure is as follows: ``` examples/ |- lego_white_background/ | |- checkpoints/ | | |- best_model | | |- best_model.npz | |- point_cloud/ | | |- iteration_30000/ | | |- point_cloud.ply | |- pseudomeshes/ | |- scene_2.70_pts_8.npz | |- cameras.json | |- config.yaml | |- colab_config.yaml |- lego.zip ``` Where: - `checkpoints` contains the best model from our optimization pipeline as torch checkpoint and numpy checkpoint, - `point_cloud` contains the point cloud of the scene from the 30000 iteration of the GaMeS algorithm, - `pseudomeshes` contains the raw pseudomesh of the scene generated from the available `point_cloud.ply` file, - `cameras.json` contains the camera poses, - `config.yaml` contains the config of the experiment. - `colab_config.yaml` contains the config of the experiment for the colab demo, - `lego.zip` contains the data of the `lego` scene from the `NeRF Synthetic` dataset. ## Installation ```bash pip install -r requirements.txt ``` Also this work depends on output of the following repositories: - [Gaussian Mesh Splatting](https://github.com/waczjoan/gaussian-mesh-splatting) - [3D Gaussian Splatting](https://github.com/graphdeco-inria/gaussian-splatting) - [2D Gaussian Splatting](https://github.com/hbb1/2d-gaussian-splatting) Therefore, you need to install them first. ## Usage You can find the scripts for running the experiments in the `sh_scripts` folder. Remember to change all paths to the correct ones in the scripts. We provided configs for the experiments in the `sh_scripts/configs` folder (once again remember to change the paths). Each script is designed to run on a single GPU. For example, to run the experiments for the DeepBlending dataset with the GaMeS algorithm, you can use the following command: ```bash ./sh_scripts/run_games_gs-flat_sh0_db.sh ``` ## Datasets This repository is prepared to work with the following datasets (as you can see in the `sh_scripts` scripts): - NeRF-Synthetic - Tanks and Temples - MiP NeRF - DeepBlending

BibTeX

MeshSplats

@Article{tobiasz2025meshsplats,
        author={Rafał Tobiasz and Grzegorz Wilczyński and Marcin Mazur and Sławomir Tadeja and Przemysław Spurek},
        year={2025},
        eprint={2502.07754},
        archivePrefix={arXiv},
        primaryClass={cs.GR},
}

Gaussian Splatting

@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/}
}

NVDiffrast

@Article{laine2020modularprimitiveshighperformancedifferentiabl,
      author         = {Kerbl, Bernhard and Kopanas, Georgios and Leimk{\"u}hler, Thomas and Drettakis, George},
      title={Modular Primitives for High-Performance Differentiable Rendering}, 
      author={Samuli Laine and Janne Hellsten and Tero Karras and Yeongho Seol and Jaakko Lehtinen and Timo Aila},
      year={2020},
      eprint={2011.03277},
      archivePrefix={arXiv},
      primaryClass={cs.GR},
}
================================================ FILE: faces_notexturemap_blender.py ================================================ import json from math import tan, atan from pathlib import Path import sys import bpy import numpy as np from mathutils import Matrix def load_npy(_path): _data = np.load(_path) return _data def _calculate_fov(focal, pixels): return 2 * np.arctan(pixels / (2 * focal)) def setup_ambient_light(): if bpy.context.scene.world is None: bpy.context.scene.world = bpy.data.worlds.new("World") world = bpy.context.scene.world world.use_nodes = True nodes = world.node_tree.nodes links = world.node_tree.links nodes.clear() background_node = nodes.new(type='ShaderNodeBackground') background_node.inputs['Color'].default_value = (1.0, 1.0, 1.0, 1.0) # White color for ambient light background_node.inputs['Strength'].default_value = 1.0 # Adjust the strength as needed world_output_node = nodes.new(type='ShaderNodeOutputWorld') links.new(background_node.outputs['Background'], world_output_node.inputs['Surface']) bpy.context.view_layer.update() def create_camera(name, position, rot_matrix, fovx, fovy, image_width, image_height): # Create the camera object bpy.ops.object.camera_add() camera_obj = bpy.context.object camera_obj.name = str(name) camera_obj.location = position # Set rotation from matrix rotation_matrix = Matrix(rot_matrix) camera_obj.rotation_euler = rotation_matrix.to_euler() # Assuming the sensor width is fixed and calculating the sensor height based on the aspect ratio sensor_width = 36.0 # Adjust as needed, standard full frame sensor width aspect_ratio = image_width / image_height sensor_height = sensor_width / aspect_ratio # Set camera data camera = camera_obj.data camera.sensor_width = sensor_width camera.sensor_height = sensor_height # Calculate focal length from FoV using formula: focal_length = sensor_width / (2 * tan(fov / 2)) camera.lens = sensor_width / (2 * tan(fovx / 2)) # Use FoVy to adjust the sensor height if needed calculated_fovy = 2 * atan((sensor_height / 2) / camera.lens) if calculated_fovy != fovy: camera.sensor_height = 2 * camera.lens * tan(fovy / 2) return camera_obj def gen_faces_from_texture_map(_path, _cam_json_path=None, _wanted_cam_idx=None): if _wanted_cam_idx is None or _cam_json_path is None: _wanted_cam_idx = [] elif isinstance(_wanted_cam_idx, int): _wanted_cam_idx = [_wanted_cam_idx] assert isinstance(_wanted_cam_idx, list) if _cam_json_path is not None: with open(_cam_json_path, "r") as file: _camera_data = json.load(file) _camera_data = sorted(_camera_data, key=lambda x: x["id"]) _camera_data = [{ "id": x["id"], "img_name": x["img_name"], "width": x["width"], "height": x["height"], "position": np.array(x["position"]), "rotation": np.array(x["rotation"]), "fovy": x["fy"], "fovx": x["fx"] } for x in _camera_data] if _wanted_cam_idx: _new_cameras = [_camera_data[i] for i in _wanted_cam_idx] else: _new_cameras = _camera_data if bpy.context.active_object.mode != 'OBJECT': bpy.ops.object.mode_set(mode='OBJECT') bpy.ops.object.select_all(action='DESELECT') bpy.ops.object.select_by_type(type='MESH') bpy.ops.object.delete() mesh_data = load_npy(_path) vertices = mesh_data["vertices"] faces = mesh_data["faces"] vert_colors = mesh_data["vertex_colors"] mesh = bpy.data.meshes.new(name="MyMesh") splat_obj = bpy.data.objects.new(name="MyObject", object_data=mesh) # Link object to collection bpy.context.collection.objects.link(splat_obj) # Create mesh from vertices and faces mesh.from_pydata(vertices.tolist(), [], faces.tolist()) mesh.update() # Deselect all objects bpy.ops.object.select_all(action='DESELECT') # Select the newly created object splat_obj.select_set(True) # Set the newly created object as the active object bpy.context.view_layer.objects.active = splat_obj if not splat_obj.data.vertex_colors: splat_obj.data.vertex_colors.new() vertex_color_layer = splat_obj.data.vertex_colors.active for poly in splat_obj.data.polygons: for loop_index in range(poly.loop_start, poly.loop_start + poly.loop_total): loop = splat_obj.data.loops[loop_index] vertex_index = loop.vertex_index vertex_color = vert_colors[vertex_index].tolist() # Assign color and opacity (RGBA) to each vertex vertex_color_layer.data[loop_index].color = vertex_color # Create a new material material = bpy.data.materials.new(name="VertexColorMaterial") material.use_nodes = True # Enable transparency material.blend_method = 'HASHED' material.shadow_method = "NONE" material.use_backface_culling = False material.alpha_threshold = 0.01 # Get the material's node tree nodes = material.node_tree.nodes links = material.node_tree.links # Clear all nodes for node in nodes: nodes.remove(node) # Add a Vertex Color node vertex_color_node = nodes.new(type='ShaderNodeVertexColor') vertex_color_node.location = (-300, 300) vertex_color_node.layer_name = vertex_color_layer.name # Add a Principled BSDF node bsdf_node = nodes.new(type='ShaderNodeBsdfPrincipled') bsdf_node.location = (100, 300) # Set metallic and specular to 0 for flat shading bsdf_node.inputs['Metallic'].default_value = 0.0 bsdf_node.inputs['Specular'].default_value = 0.0 bsdf_node.inputs['Roughness'].default_value = 1.0 # Connect the vertex color to the BSDF node links.new(vertex_color_node.outputs['Color'], bsdf_node.inputs['Base Color']) links.new(vertex_color_node.outputs['Alpha'], bsdf_node.inputs['Alpha']) # Add an Output node output_node = nodes.new(type='ShaderNodeOutputMaterial') output_node.location = (300, 300) # Link the BSDF node to the output node links.new(bsdf_node.outputs['BSDF'], output_node.inputs['Surface']) # Assign the material to the cube if splat_obj.type == 'MESH': if len(splat_obj.data.materials) == 0: splat_obj.data.materials.append(material) else: splat_obj.data.materials[0] = material # Update the mesh splat_obj.data.update() # Set the render engine to Eevee bpy.context.scene.render.engine = 'BLENDER_EEVEE' # Configure Eevee settings for better quality bpy.context.scene.eevee.taa_render_samples = 64 # Increase samples for better quality bpy.context.scene.eevee.use_soft_shadows = True bpy.context.scene.eevee.use_ssr = True # Screen Space Reflections bpy.context.scene.eevee.use_ssr_refraction = True bpy.context.scene.eevee.use_taa_reprojection = True # Enable transparency settings in Eevee bpy.context.scene.eevee.use_bloom = False bpy.context.scene.eevee.use_ssr_halfres = False bpy.context.scene.eevee.volumetric_tile_size = '2' setup_ambient_light() # Set render settings bpy.context.scene.cycles.samples = 128 # Set the number of samples for rendering _cameras = [] adjustment = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]], dtype=np.float64) for _new_cam in _new_cameras: fovx = _calculate_fov(_new_cam["fovx"], _new_cam["width"]) fovy = _calculate_fov(_new_cam["fovy"], _new_cam["height"]) _new_cam_obj = create_camera( name=_new_cam["img_name"], position=_new_cam["position"], rot_matrix=_new_cam["rotation"] @ adjustment, fovx=fovx, fovy=fovy, image_width=_new_cam["width"], image_height=_new_cam["height"] ) _cameras.append(_new_cam_obj) scale = Path(_path).stem.split("_")[1] out_dir = Path(Path(_path).parent, "images", scale) out_dir.mkdir(exist_ok=True, parents=True) # setup output img size bpy.context.scene.render.resolution_x = _new_cameras[0]["width"] bpy.context.scene.render.resolution_y = _new_cameras[0]["height"] for _cam_obj in _cameras: output_path = str(Path(out_dir, f"{_cam_obj.name}.png")) # Set camera as active camera bpy.context.scene.camera = _cam_obj # Set output path for rendered image bpy.context.scene.render.filepath = output_path # Render the image bpy.ops.render.render(write_still=True) print(f"Rendered image saved to {output_path}") args = sys.argv args_start_index = args.index("--") + 1 script_args = args[args_start_index:] assert "--npz_path" in script_args assert "--cam_path" in script_args _npz_path = script_args[script_args.index("--npz_path") + 1] _cam_path = script_args[script_args.index("--cam_path") + 1] gen_faces_from_texture_map( _npz_path, _cam_path, ) ================================================ FILE: generate_pseudomesh.py ================================================ import argparse from pathlib import Path import numpy as np import gs_utils def main(ply_path: Path, algorithm: str, scale_muls: np.ndarray, no_of_points: int): xyz, feat_dc, feat_res, opac, scales, rots_quaternion = gs_utils.load_ply( ply_path, 0 ) if algorithm == "games": games_ckpt_path = Path(ply_path.parent, "model_params.pt") try: games_data = gs_utils.load_games_pt(games_ckpt_path) except FileNotFoundError: games_data = None if algorithm == "2dgs": pseudomeshes = gs_utils.generate_pseudomesh_2dgs( xyz=xyz, features_dc=feat_dc, opacities=opac, scales=scales, rots=rots_quaternion, scale_muls=scale_muls, no_of_points=no_of_points ) elif algorithm == "surfels": pseudomeshes = gs_utils.generate_pseudomesh_surfels( xyz=xyz, features_dc=feat_dc, opacities=opac, scales=scales, rots=rots_quaternion, scale_muls=scale_muls, no_of_points=no_of_points ) elif algorithm == "games": pseudomeshes = gs_utils.generate_pseudomesh_games( ckpt_data=games_data, xyz=xyz, features_dc=feat_dc, opacities=opac, scales=scales, rots=rots_quaternion, scale_muls=scale_muls, no_of_points=no_of_points ) elif algorithm == "sugar2d": pseudomeshes = gs_utils.generate_pseudomesh_sugar_2d( xyz=xyz, features_dc=feat_dc, opacities=opac, scales=scales, rots=rots_quaternion, scale_muls=scale_muls, no_of_points=no_of_points ) elif algorithm == "sugar3d": pseudomeshes = gs_utils.generate_pseudomesh_sugar_3d( xyz=xyz, features_dc=feat_dc, opacities=opac, scales=scales, rots=rots_quaternion, scale_muls=scale_muls, no_of_points=no_of_points ) elif algorithm == "3dgs": pseudomeshes = gs_utils.generate_3dgs_pseudomesh( xyz=xyz, features_dc=feat_dc, opacities=opac, scales=scales, rots=rots_quaternion, scale_muls=scale_muls, no_of_points=no_of_points ) else: raise NotImplementedError(f"Algorithm {algorithm} not yet implemented") # save data if "sugar" in algorithm: main_out_dir = ply_path.parent else: main_out_dir = ply_path.parent.parent.parent out_dir = Path(main_out_dir, "pseudomeshes") out_dir.mkdir(exist_ok=True, parents=True) for scale_val, pseudomesh_data in pseudomeshes.items(): scale_val_str = "{:.2f}".format(scale_val) out_path = Path(out_dir, f"scale_{scale_val_str}_pts_{str(no_of_points)}.npz") np.savez( out_path, **pseudomesh_data ) print(f"Saved {out_path}") def read_args(): parser = argparse.ArgumentParser(description="This script accepts ply file and outputs pseudomesh of it") parser.add_argument( "--ply", required=True, type=str, help="Path to the ply file" ) parser.add_argument( "--algorithm", default="3dgs", type=str, choices=["2dgs", "games", "surfels", "sugar2d", "sugar3d", "3dgs"] ) parser.add_argument( "--scale_min", default=2.3, type=float ) parser.add_argument( "--scale_max", default=2.4, type=float ) parser.add_argument( "--scale_step", default=0.2, type=float ) parser.add_argument( "--no_points", default=8, type=int ) args = parser.parse_args() return { "ply_path": Path(args.ply), "algorithm": args.algorithm, "scale_muls": np.arange(args.scale_min, args.scale_max, args.scale_step), "no_of_points": args.no_points } if __name__ == "__main__": main(**read_args()) ================================================ FILE: gs_utils.py ================================================ import numpy as np from plyfile import PlyData from scipy.special import expit as sigmoid import torch try: import sys sys.path.append("/path/to/GaMeS_repo") import games except Exception as e: print(e) C0 = 0.28209479177387814 def load_games_pt(path): params = torch.load(path) alpha = params['_alpha'] scale = params['_scale'] vertices, triangles, faces = None, None, None if 'vertices' in params: vertices = params['vertices'] if 'triangles' in params: triangles = params['triangles'] if 'faces' in params: faces = params['faces'] return { "alpha": alpha, "scale": scale, "vertices": vertices, "triangles": triangles, "faces": faces } def get_games_scales_and_rots(data, eps=1e-8): def dot(v, u): return (v * u).sum(dim=-1, keepdim=True) def proj(v, u): """ projection of vector v onto subspace spanned by u vector u is assumed to be already normalized """ coef = dot(v, u) return coef * u triangles = data["triangles"] normals = torch.linalg.cross( triangles[:, 1] - triangles[:, 0], triangles[:, 2] - triangles[:, 0], dim=1 ) v0 = normals / (torch.linalg.vector_norm(normals, dim=-1, keepdim=True) + eps) means = torch.mean(triangles, dim=1) v1 = triangles[:, 1] - means v1_norm = torch.linalg.vector_norm(v1, dim=-1, keepdim=True) + eps v1 = v1 / v1_norm v2_init = triangles[:, 2] - means v2 = v2_init - proj(v2_init, v0) - proj(v2_init, v1) # Gram-Schmidt v2 = v2 / (torch.linalg.vector_norm(v2, dim=-1, keepdim=True) + eps) s1 = v1_norm / 2. s2 = dot(v2_init, v2) / 2. scales = torch.concat([s1, s2], dim=1).unsqueeze(dim=1) scales = scales.broadcast_to((*data["alpha"].shape[:2], 2)) scales = torch.nn.functional.relu(data["scale"] * scales.flatten(start_dim=0, end_dim=1)) + eps # scales [N, 2] rots = torch.stack((v1, v2), dim=1).unsqueeze(dim=1) rots = rots.broadcast_to((*data["alpha"].shape[:2], 2, 3)).flatten(start_dim=0, end_dim=1) # rots [N, 2, 3] return scales, rots def load_ply(path, max_sh_degree): plydata = PlyData.read(path) xyz = np.stack((np.asarray(plydata.elements[0]["x"]), np.asarray(plydata.elements[0]["y"]), np.asarray(plydata.elements[0]["z"])), axis=1) opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis] features_dc = np.zeros((xyz.shape[0], 3, 1)) features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"]) features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"]) features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"]) extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")] extra_f_names = sorted(extra_f_names, key = lambda x: int(x.split('_')[-1])) assert len(extra_f_names)==3*(max_sh_degree + 1) ** 2 - 3 features_extra = np.zeros((xyz.shape[0], len(extra_f_names))) for idx, attr_name in enumerate(extra_f_names): features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name]) # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC) features_extra = features_extra.reshape((features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1)) scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")] scale_names = sorted(scale_names, key = lambda x: int(x.split('_')[-1])) scales = np.zeros((xyz.shape[0], len(scale_names))) for idx, attr_name in enumerate(scale_names): scales[:, idx] = np.asarray(plydata.elements[0][attr_name]) rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")] rot_names = sorted(rot_names, key = lambda x: int(x.split('_')[-1])) rots = np.zeros((xyz.shape[0], len(rot_names))) for idx, attr_name in enumerate(rot_names): rots[:, idx] = np.asarray(plydata.elements[0][attr_name]) features_dc = np.transpose(features_dc, [0, 2, 1]) # TODO check if this is ok features_rest = np.transpose(features_extra, [0, 2, 1]) # TODO check if this is ok return xyz, features_dc, features_rest, opacities, scales, rots def normalize_rots(mat): return mat / np.linalg.norm(mat, axis=1, keepdims=True) def build_euler_rotation(r): norm = np.sqrt(r[:,0]*r[:,0] + r[:,1]*r[:,1] + r[:,2]*r[:,2] + r[:,3]*r[:,3]) q = r / norm[:, None] R = np.zeros((q.shape[0], 3, 3), dtype=np.float64) r = q[:, 0] x = q[:, 1] y = q[:, 2] z = q[:, 3] R[:, 0, 0] = 1 - 2 * (y*y + z*z) R[:, 0, 1] = 2 * (x*y - r*z) R[:, 0, 2] = 2 * (x*z + r*y) R[:, 1, 0] = 2 * (x*y + r*z) R[:, 1, 1] = 1 - 2 * (x*x + z*z) R[:, 1, 2] = 2 * (y*z - r*x) R[:, 2, 0] = 2 * (x*z - r*y) R[:, 2, 1] = 2 * (y*z + r*x) R[:, 2, 2] = 1 - 2 * (x*x + y*y) return R def get_scaling(scales: np.ndarray) -> np.ndarray: return np.exp(scales) def get_opacity(opacities: np.ndarray) -> np.ndarray: return sigmoid(opacities) def generate_pseudomesh_games(ckpt_data: dict, xyz: np.ndarray, features_dc: np.ndarray, opacities: np.ndarray, scales: np.ndarray, rots: np.ndarray, scale_muls: np.ndarray, no_of_points: int = 4): if ckpt_data is not None: scale, rotation = get_games_scales_and_rots(ckpt_data) scale = scale.cpu().detach().numpy() rotation = rotation.cpu().detach().numpy() else: scale = get_scaling(scales)[:, :2] rotation = np.transpose(build_euler_rotation(rots)[:, :, 1:], [0, 2, 1]) init_colors = get_rgb_colors(features_dc) init_opacities = get_opacity(opacities) origin = xyz output = gen_2d_pseudomesh(scale_muls, no_of_points, init_colors, init_opacities, scale, rotation, origin) return output def generate_pseudomesh_surfels(xyz: np.ndarray, features_dc: np.ndarray, opacities: np.ndarray, scales: np.ndarray, rots: np.ndarray, scale_muls: np.ndarray, no_of_points: int = 4): init_colors = get_rgb_colors(features_dc) init_opacities = get_opacity(opacities) scale = get_scaling(scales)[:, :2] rotation = np.transpose(build_euler_rotation(rots)[:, :, :2], [0, 2, 1]) origin = xyz output = gen_2d_pseudomesh(scale_muls, no_of_points, init_colors, init_opacities, scale, rotation, origin) return output def generate_pseudomesh_2dgs(xyz: np.ndarray, features_dc: np.ndarray, opacities: np.ndarray, scales: np.ndarray, rots: np.ndarray, scale_muls: np.ndarray, no_of_points: int = 4): init_colors = get_rgb_colors(features_dc) init_opacities = get_opacity(opacities) scale = get_scaling(scales) rotation = np.transpose(build_euler_rotation(rots)[:, :, :2], [0, 2, 1]) origin = xyz output = gen_2d_pseudomesh(scale_muls, no_of_points, init_colors, init_opacities, scale, rotation, origin) return output def generate_pseudomesh_sugar_2d(xyz: np.ndarray, features_dc: np.ndarray, opacities: np.ndarray, scales: np.ndarray, rots: np.ndarray, scale_muls: np.ndarray, no_of_points: int = 8): init_colors = get_rgb_colors(features_dc) init_opacities = get_opacity(opacities) scale = get_scaling(scales) rotation = np.transpose(build_euler_rotation(rots), [0, 2, 1]) origin = xyz scale = scale[:, 1:] rotation = rotation[:, 1:, :] output = gen_2d_pseudomesh(scale_muls, no_of_points, init_colors, init_opacities, scale, rotation, origin) return output def generate_pseudomesh_sugar_3d(xyz: np.ndarray, features_dc: np.ndarray, opacities: np.ndarray, scales: np.ndarray, rots: np.ndarray, scale_muls: np.ndarray, no_of_points: int = 8, opac_mul: float = .2): assert no_of_points == 8 # handle background generation init_colors = get_rgb_colors(features_dc) init_opacities = get_opacity(opacities) scale = get_scaling(scales) rotation = np.transpose(build_euler_rotation(rots), [0, 2, 1]) origin = xyz scales_1 = scale[:, [0, 1]] rots_1 = rotation[:, [0, 1], :] scales_2 = scale[:, [1, 2]] rots_2 = rotation[:, [1, 2], :] scales_3 = scale[:, [0, 2]] rots_3 = rotation[:, [0, 2], :] output = {} for scale_mul in scale_muls: new_points_1 = _get_vertices(origin, scales_1, rots_1, scale_mul, no_of_points) new_points_2 = _get_vertices(origin, scales_2, rots_2, scale_mul, no_of_points) new_points_3 = _get_vertices(origin, scales_3, rots_3, scale_mul, no_of_points) # remove redundant verts new_points_2 = new_points_2[:, [1, 2, 3, 5, 6, 7], :] new_points_3 = new_points_3[:, [1, 3, 5, 7], :] vertices = np.vstack([origin, *[new_points_1[:, _idx, :] for _idx in range(new_points_1.shape[1])]]) vertices = np.vstack([vertices, *[new_points_2[:, _idx, :] for _idx in range(new_points_2.shape[1])]]) vertices = np.vstack([vertices, *[new_points_3[:, _idx, :] for _idx in range(new_points_3.shape[1])]]) origin_idxs = np.arange(0, origin.shape[0]).reshape(-1, 1) indexes_1 = np.hstack([ np.arange(_idx * origin.shape[0], (_idx + 1) * origin.shape[0]).reshape(-1, 1) for _idx in range(1, 9) ]) indexes_2 = np.hstack([ np.arange(_idx * origin.shape[0], (_idx + 1) * origin.shape[0]).reshape(-1, 1) for _idx in [3, 9, 10, 11, 7, 12, 13, 14] ]) indexes_3 = np.hstack([ np.arange(_idx * origin.shape[0], (_idx + 1) * origin.shape[0]).reshape(-1, 1) for _idx in [1, 15, 10, 16, 5, 17, 13, 18] ]) # create faces all_faces = [] # faces 1 for face_idx in range(no_of_points): last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0 all_faces.append( np.hstack([origin_idxs, indexes_1[:, face_idx].reshape(-1, 1), indexes_1[:, last_idx].reshape(-1, 1)]) ) # faces 2 for face_idx in range(no_of_points): last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0 all_faces.append( np.hstack([origin_idxs, indexes_2[:, face_idx].reshape(-1, 1), indexes_2[:, last_idx].reshape(-1, 1)]) ) # faces 3 for face_idx in range(no_of_points): last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0 all_faces.append( np.hstack([origin_idxs, indexes_3[:, face_idx].reshape(-1, 1), indexes_3[:, last_idx].reshape(-1, 1)]) ) all_faces = np.vstack(all_faces) vertex_colors = np.vstack([init_colors] * 19) face_colors = np.vstack([init_colors] * 24) # vertex_opacities = np.vstack([init_opacities] * (no_of_points + 1)) vertex_opacities = np.vstack([init_opacities, *[init_opacities * opac_mul] * 18]) # face_opacities = np.vstack([init_opacities] * no_of_points) face_opacities = np.vstack([init_opacities] * 24) vertex_rgba = np.hstack([vertex_colors, vertex_opacities]) face_rgba = np.hstack([face_colors, face_opacities]) output[scale_mul] = dict( vertices=vertices, vertex_colors=vertex_rgba, faces=all_faces, face_colors=face_rgba ) return output def _get_vertices(origin, scales, rots, scale_mul, no_of_points): scales = scale_mul * np.repeat(np.repeat(scales[:, None, :, None], no_of_points, 1), 3, -1) rots = np.repeat(rots[:, None, :, :], no_of_points, 1) scaled_rots = scales * rots angle = np.linspace(-np.pi, np.pi, no_of_points + 1)[:-1] sins = np.repeat(np.repeat((np.sin(angle).reshape(-1, 1)[None, :, :]), origin.shape[0], 0), 3, -1) coss = np.repeat(np.repeat((np.cos(angle).reshape(-1, 1)[None, :, :]), origin.shape[0], 0), 3, -1) _origin = np.repeat(origin[:, None, :], no_of_points, 1) new_points = _origin + coss * scaled_rots[:, :, 0, :] + sins * scaled_rots[:, :, 1, :] return new_points def gen_2d_pseudomesh(scale_muls, no_of_points, init_colors, init_opacities, scale, rotation, origin, opac_mul=.2): angle = np.linspace(-np.pi, np.pi, no_of_points + 1)[:-1] sins = np.repeat(np.repeat((np.sin(angle).reshape(-1, 1)[None, :, :]), origin.shape[0], 0), 3, -1) coss = np.repeat(np.repeat((np.cos(angle).reshape(-1, 1)[None, :, :]), origin.shape[0], 0), 3, -1) _origin = np.repeat(origin[:, None, :], no_of_points, 1) output = {} for scale_mul in scale_muls: scales = scale_mul * np.repeat(np.repeat(scale[:, None, :, None], no_of_points, 1), 3, -1) rots = np.repeat(rotation[:, None, :, :], no_of_points, 1) scaled_rots = scales * rots new_points = _origin + coss * scaled_rots[:, :, 0, :] + sins * scaled_rots[:, :, 1, :] vertices = np.vstack([origin, *[new_points[:, _idx, :] for _idx in range(no_of_points)]]) origin_idxs = np.arange(0, origin.shape[0]).reshape(-1, 1) indexes = np.hstack([np.arange(_idx * origin.shape[0], (_idx + 1) * origin.shape[0]).reshape(-1, 1) for _idx in range(1, no_of_points + 1)]) # create faces all_faces = [] for face_idx in range(no_of_points): last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0 all_faces.append( np.hstack([origin_idxs, indexes[:, face_idx].reshape(-1, 1), indexes[:, last_idx].reshape(-1, 1)]) ) all_faces = np.vstack(all_faces) vertex_colors = np.vstack([init_colors] * (no_of_points + 1)) face_colors = np.vstack([init_colors] * no_of_points) vertex_opacities = np.vstack([init_opacities, *[init_opacities * opac_mul] * no_of_points]) face_opacities = np.vstack([init_opacities] * no_of_points) vertex_rgba = np.hstack([vertex_colors, vertex_opacities]) face_rgba = np.hstack([face_colors, face_opacities]) output[scale_mul] = dict( vertices=vertices, vertex_colors=vertex_rgba, faces=all_faces, face_colors=face_rgba ) return output def get_rgb_colors(color_features): colors = np.transpose(color_features, [0, 2, 1]) result = C0 * colors[..., 0] + 0.5 result = result.clip(min=0., max=1.) return result def generate_3dgs_pseudomesh(xyz: np.ndarray, features_dc: np.ndarray, opacities: np.ndarray, scales: np.ndarray, rots: np.ndarray, scale_muls: np.ndarray, no_of_points: int = 8): assert no_of_points == 8 init_colors = get_rgb_colors(features_dc) init_opacities = get_opacity(opacities) scale = get_scaling(scales) rotation = np.transpose(build_euler_rotation(rots), [0, 2, 1]) origin = xyz scales_1 = scale[:, [0, 1]] rots_1 = rotation[:, [0, 1], :] scales_2 = scale[:, [1, 2]] rots_2 = rotation[:, [1, 2], :] scales_3 = scale[:, [0, 2]] rots_3 = rotation[:, [0, 2], :] output = {} for scale_mul in scale_muls: new_points_1 = _get_vertices(origin, scales_1, rots_1, scale_mul, no_of_points) new_points_2 = _get_vertices(origin, scales_2, rots_2, scale_mul, no_of_points) new_points_3 = _get_vertices(origin, scales_3, rots_3, scale_mul, no_of_points) # remove redundant verts new_points_2 = new_points_2[:, [1, 2, 3, 5, 6, 7], :] new_points_3 = new_points_3[:, [1, 3, 5, 7], :] vertices = np.vstack([origin, *[new_points_1[:, _idx, :] for _idx in range(new_points_1.shape[1])]]) vertices = np.vstack([vertices, *[new_points_2[:, _idx, :] for _idx in range(new_points_2.shape[1])]]) vertices = np.vstack([vertices, *[new_points_3[:, _idx, :] for _idx in range(new_points_3.shape[1])]]) origin_idxs = np.arange(0, origin.shape[0]).reshape(-1, 1) indexes_1 = np.hstack([ np.arange(_idx * origin.shape[0], (_idx + 1) * origin.shape[0]).reshape(-1, 1) for _idx in range(1, 9) ]) indexes_2 = np.hstack([ np.arange(_idx * origin.shape[0], (_idx + 1) * origin.shape[0]).reshape(-1, 1) for _idx in [3, 9, 10, 11, 7, 12, 13, 14] ]) indexes_3 = np.hstack([ np.arange(_idx * origin.shape[0], (_idx + 1) * origin.shape[0]).reshape(-1, 1) for _idx in [1, 15, 10, 16, 5, 17, 13, 18] ]) # create faces all_faces = [] # faces 1 for face_idx in range(no_of_points): last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0 all_faces.append( np.hstack([origin_idxs, indexes_1[:, face_idx].reshape(-1, 1), indexes_1[:, last_idx].reshape(-1, 1)]) ) # faces 2 for face_idx in range(no_of_points): last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0 all_faces.append( np.hstack([origin_idxs, indexes_2[:, face_idx].reshape(-1, 1), indexes_2[:, last_idx].reshape(-1, 1)]) ) # faces 3 for face_idx in range(no_of_points): last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0 all_faces.append( np.hstack([origin_idxs, indexes_3[:, face_idx].reshape(-1, 1), indexes_3[:, last_idx].reshape(-1, 1)]) ) all_faces = np.vstack(all_faces) vertex_colors = np.vstack([init_colors] * 19) face_colors = np.vstack([init_colors] * 24) vertex_opacities = np.vstack([init_opacities, *[init_opacities * 0.2] * 18]) face_opacities = np.vstack([init_opacities] * 24) vertex_rgba = np.hstack([vertex_colors, vertex_opacities]) face_rgba = np.hstack([face_colors, face_opacities]) output[scale_mul] = dict( vertices=vertices, vertex_colors=vertex_rgba, faces=all_faces, face_colors=face_rgba ) return output ================================================ FILE: mesh_optim/2dgs_experiments_run.py ================================================ import os from pathlib import Path import subprocess import yaml PYTHON_PATH = Path("/home/olaf/miniconda3/envs/splats310/bin/python") PSEUDOMESH_PY = Path("/home/grzegos/projects/phd/gs_raytracing/generate_pseudomesh.py") PSEUDO_CWD_PATH = PSEUDOMESH_PY.parent PSEUDOMESH_OPTIM_PY = Path("/home/grzegos/projects/phd/gs_raytracing/mesh_optim/optimize_pseudomesh.py") OPTIM_CWD_PATH = PSEUDOMESH_OPTIM_PY.parent DSET_MAIN_DIR = Path("/home/grzegos/datasets/games_set") TWODGS_EXP_DIR = Path("/home/grzegos/projects/phd/2dgs_nerf_output") OUT_EXP_DIR = Path("/home/grzegos/projects/phd/2dgs_nvdiff_exps") TMP_CONFIG_PATH = Path("/home/grzegos/projects/phd/gs_raytracing/mesh_optim/temp_config.yaml") def main(): with open("/home/grzegos/projects/phd/gs_raytracing/template_config.yaml", "r") as file: template_cfg = yaml.load(file, Loader=yaml.FullLoader) twodgs_exp_dirs = list(TWODGS_EXP_DIR.glob("*")) env = os.environ.copy() # Copy current environment env['CUDA_VISIBLE_DEVICES'] = '1' for exp_dir in twodgs_exp_dirs: dset_path = Path(DSET_MAIN_DIR, exp_dir.stem) exp_name = exp_dir.stem # create pseudomesh ply_path = Path(exp_dir, "point_cloud", "iteration_10000", "point_cloud.ply") subprocess.run( [ str(PYTHON_PATH), str(PSEUDOMESH_PY), "--ply", str(ply_path), "--algorithm", "2dgs", "--scale_min", "2.3", "--scale_max", "2.4", "--scale_step", "0.2", "--no_points", "8" ], cwd=str(PSEUDO_CWD_PATH) ) # run optim pseudomesh_path = Path(exp_dir, "pseudomeshes", "scale_2.30_pts_8.npz") # change config template_cfg["experiment_name"] = exp_name template_cfg["experiments_dir"] = str(OUT_EXP_DIR) template_cfg["pseudomesh_path"] = str(pseudomesh_path) template_cfg["dataset"]["dataset_path"] = str(dset_path) # save cfg with open(TMP_CONFIG_PATH, 'w') as outfile: yaml.dump(template_cfg, outfile, default_flow_style=False) # pass cfg path to optimizer subprocess.run( [ str(PYTHON_PATH), str(PSEUDOMESH_OPTIM_PY), "--cfg_path", str(TMP_CONFIG_PATH) ], cwd=str(OPTIM_CWD_PATH), env=env ) if __name__ == "__main__": main() ================================================ FILE: mesh_optim/cam_utils.py ================================================ # # Copyright (C) 2023, Inria # GRAPHDECO research group, https://team.inria.fr/graphdeco # 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 george.drettakis@inria.fr # import collections import struct import numpy as np CameraModel = collections.namedtuple( "CameraModel", ["model_id", "model_name", "num_params"]) Camera = collections.namedtuple( "Camera", ["id", "model", "width", "height", "params"]) BaseImage = collections.namedtuple( "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"]) Point3D = collections.namedtuple( "Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"]) CAMERA_MODELS = { CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3), CameraModel(model_id=1, model_name="PINHOLE", num_params=4), CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4), CameraModel(model_id=3, model_name="RADIAL", num_params=5), CameraModel(model_id=4, model_name="OPENCV", num_params=8), CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8), CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12), CameraModel(model_id=7, model_name="FOV", num_params=5), CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4), CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5), CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12) } CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS]) CAMERA_MODEL_NAMES = dict([(camera_model.model_name, camera_model) for camera_model in CAMERA_MODELS]) def qvec2rotmat(qvec): return np.array([ [1 - 2 * qvec[2]**2 - 2 * qvec[3]**2, 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2]], [2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], 1 - 2 * qvec[1]**2 - 2 * qvec[3]**2, 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]], [2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], 1 - 2 * qvec[1]**2 - 2 * qvec[2]**2]]) class Image(BaseImage): def qvec2rotmat(self): return qvec2rotmat(self.qvec) def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"): """Read and unpack the next bytes from a binary file. :param fid: :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc. :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. :param endian_character: Any of {@, =, <, >, !} :return: Tuple of read and unpacked values. """ data = fid.read(num_bytes) return struct.unpack(endian_character + format_char_sequence, data) def read_extrinsics_binary(path_to_model_file): """ see: src/base/reconstruction.cc void Reconstruction::ReadImagesBinary(const std::string& path) void Reconstruction::WriteImagesBinary(const std::string& path) """ images = {} with open(path_to_model_file, "rb") as fid: num_reg_images = read_next_bytes(fid, 8, "Q")[0] for _ in range(num_reg_images): binary_image_properties = read_next_bytes( fid, num_bytes=64, format_char_sequence="idddddddi") image_id = binary_image_properties[0] qvec = np.array(binary_image_properties[1:5]) tvec = np.array(binary_image_properties[5:8]) camera_id = binary_image_properties[8] image_name = "" current_char = read_next_bytes(fid, 1, "c")[0] while current_char != b"\x00": # look for the ASCII 0 entry image_name += current_char.decode("utf-8") current_char = read_next_bytes(fid, 1, "c")[0] num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0] x_y_id_s = read_next_bytes(fid, num_bytes=24*num_points2D, format_char_sequence="ddq"*num_points2D) xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))]) point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3]))) images[image_id] = Image( id=image_id, qvec=qvec, tvec=tvec, camera_id=camera_id, name=image_name, xys=xys, point3D_ids=point3D_ids) return images def read_intrinsics_binary(path_to_model_file): """ see: src/base/reconstruction.cc void Reconstruction::WriteCamerasBinary(const std::string& path) void Reconstruction::ReadCamerasBinary(const std::string& path) """ cameras = {} with open(path_to_model_file, "rb") as fid: num_cameras = read_next_bytes(fid, 8, "Q")[0] for _ in range(num_cameras): camera_properties = read_next_bytes( fid, num_bytes=24, format_char_sequence="iiQQ") camera_id = camera_properties[0] model_id = camera_properties[1] model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name width = camera_properties[2] height = camera_properties[3] num_params = CAMERA_MODEL_IDS[model_id].num_params params = read_next_bytes(fid, num_bytes=8*num_params, format_char_sequence="d"*num_params) cameras[camera_id] = Camera(id=camera_id, model=model_name, width=width, height=height, params=np.array(params)) assert len(cameras) == num_cameras return cameras ================================================ FILE: mesh_optim/data.py ================================================ import json from pathlib import Path from PIL import Image import torch from torch.utils.data import Dataset import numpy as np from cam_utils import qvec2rotmat, read_extrinsics_binary, read_intrinsics_binary def create_blend_proj_mats(camera_angle_x, img_shape, transf_mat, far, near): focal = img_shape[0] / (2 * np.tan(camera_angle_x / 2)) transf_matrix = torch.tensor(transf_mat, device="cpu", dtype=torch.float32) view_matrix = torch.inverse(transf_matrix) proj_matrix = torch.zeros(4, 4, device="cpu", dtype=torch.float32) proj_matrix[0, 0] = 2 * focal / img_shape[0] proj_matrix[1, 1] = -2 * focal / img_shape[1] proj_matrix[2, 2] = -(far + near) / (far - near) proj_matrix[2, 3] = -2.0 * far * near / (far - near) proj_matrix[3, 2] = -1.0 mvp_matrix = proj_matrix @ view_matrix return { "transf_matrix": transf_matrix, "view_matrix": view_matrix, "proj_matrix": proj_matrix, "mvp_matrix": mvp_matrix } def create_colmap_proj_mats(focal_x, focal_y, img_shape, transf_mat, far, near): # TODO test it bruh transf_matrix = torch.tensor(transf_mat, device="cpu", dtype=torch.float32) view_matrix = torch.inverse(transf_matrix) proj_matrix = torch.zeros(4, 4, device="cpu", dtype=torch.float32) proj_matrix[0, 0] = 2 * focal_x / img_shape[0] proj_matrix[1, 1] = -2 * focal_y / img_shape[1] proj_matrix[2, 2] = -(far + near) / (far - near) proj_matrix[2, 3] = -2.0 * far * near / (far - near) proj_matrix[3, 2] = -1.0 mvp_matrix = proj_matrix @ view_matrix return { "transf_matrix": transf_matrix, "view_matrix": view_matrix, "proj_matrix": proj_matrix, "mvp_matrix": mvp_matrix } class ImageCamDataset(Dataset): def __init__(self, dataset_path, near, far, imgs_in_ram=True, res=1, test=False): self.data_dir = Path(dataset_path) self.near = near self.far = far self._imgs_in_ram = imgs_in_ram self._res = res self.test = test self.mode = None self.shape = None # print(self.test) if Path(dataset_path, "transforms_train.json").exists(): self.mode = "blender" else: self.mode = "colmap" self.load_cameras() # print(len(self)) def _load_blender_cameras(self): if self.test: self._cam_info_path = Path(self.data_dir, "transforms_test.json") else: self._cam_info_path = Path(self.data_dir, "transforms_train.json") # Load camera parameters from JSON file with open(self._cam_info_path, 'r') as f: self.camera_data = json.load(f) self.image_files = [] _cam_angle_x = self.camera_data["camera_angle_x"] for _idx, _cam in enumerate(self.camera_data["frames"]): img_path = Path(self.data_dir, _cam["file_path"]).with_suffix(".png") img = None if not _idx or self._imgs_in_ram: img, width, height = self.load_image(img_path, get_shape=True) self.shape = [width, height] trans_mat = np.array(_cam["transform_matrix"], dtype=np.float32) proj_mats = create_blend_proj_mats( _cam_angle_x, self.shape, trans_mat, self.far, self.near ) self.image_files.append({ "name": img_path.name, "img_path": str(img_path), "img": img, **proj_mats }) def _load_colmap_cameras(self): print(self._res) try: cameras_extrinsic_file = Path(self.data_dir, "sparse/0", "images.bin") cameras_intrinsic_file = Path(self.data_dir, "sparse/0", "cameras.bin") cam_extrinsics = read_extrinsics_binary(cameras_extrinsic_file) cam_intrinsics = read_intrinsics_binary(cameras_intrinsic_file) except: raise NotImplementedError(".txt Colmap structure not supported") image_files = [] for _idx, key in enumerate(cam_extrinsics): extr = cam_extrinsics[key] intr = cam_intrinsics[extr.camera_id] height = intr.height width = intr.width img_path = Path(self.data_dir, "images", extr.name) img = None if not _idx or self._imgs_in_ram: img, width, height = self.load_image(img_path, self._res, get_shape=True) self.shape = [width, height] elif self._res != 1: height = round(height / self._res) width = round(width / self._res) R = qvec2rotmat(extr.qvec) T = np.array(extr.tvec) transf_mat = np.eye(4, dtype=np.float32) transf_mat[:3, :3] = R transf_mat[:3, -1] = T transf_mat = np.linalg.inv(transf_mat) transf_mat[:3, 1:3] *= -1 if intr.model=="SIMPLE_PINHOLE": focal_length_x = intr.params[0] focal_length_y = intr.params[0] elif intr.model=="PINHOLE": focal_length_x = intr.params[0] focal_length_y = intr.params[1] else: assert False, "Colmap camera model not handled: only undistorted datasets (PINHOLE or SIMPLE_PINHOLE cameras) supported!" # focal_x = focal_length_x / (2 * self._res) # TODO FIX FOR TANDT focal_x = focal_length_x / self._res # focal_y = focal_length_y / (2 * self._res) focal_y = focal_length_y / self._res proj_mats = create_colmap_proj_mats( focal_x, focal_y, self.shape, transf_mat, self.far, self.near ) image_files.append({ "name": extr.name, "img_path": str(img_path), "img": img, "focal_x": focal_x, "focal_y": focal_y, **proj_mats }) sorted_imgs = sorted(image_files, key=lambda x: x["name"].split(".")[0]) if self.test: cond = lambda x: x % 8 == 0 else: cond = lambda x: x % 8 != 0 self.image_files = [c for idx, c in enumerate(sorted_imgs) if cond(idx)] def load_cameras(self): if self.mode == "blender": self._load_blender_cameras() else: self._load_colmap_cameras() @staticmethod def load_image(path, res=None, get_shape=False): _img = Image.open(path) if res is not None and res != 1: _img.thumbnail( [ round(_img.width / res), round(_img.height / res) ], Image.Resampling.LANCZOS ) img = np.array(_img).astype(np.float32) / 255 img = torch.from_numpy(img) if get_shape: return img, _img.width, _img.height return img def __len__(self): return len(self.image_files) def __getitem__(self, idx): """ Get a sample from the dataset. Returns: dict: Contains image and camera parameters """ if self.image_files[idx]["img"] is None: img = self.load_image(self.image_files[idx]["img_path"], self._res) self.image_files[idx]["img"] = img self.image_files[idx]["img"] = self.image_files[idx]["img"] return self.image_files[idx] # # Example usage: # if __name__ == "__main__": # # Create dataset instance # dataset = ImageCamDataset( # dataset_path="/path/to/dataset", # near=0.01, # far=100., # imgs_in_ram=False # ) # # Get a sample # sample = dataset[0] # print(f"Image shape: {sample['image'].shape}") ================================================ FILE: mesh_optim/diff_render.py ================================================ import json from pathlib import Path import torch import numpy as np import nvdiffrast.torch as dr import matplotlib.pyplot as plt def create_obj(device): vertices = torch.tensor([ [-1, 1, -1], [1, 1, -1], [1, -1, -1], [-1, -1, -1], [-1, 1, 1], [1, 1, 1], [1, -1, 1], [-1, -1, 1], ], device=device, dtype=torch.float32, requires_grad=True) _verts_add = torch.tensor([ [-1, 1, -1], [1, 1, -1], [1, -1, -1], [-1, -1, -1], [-1, 1, 1], [1, 1, 1], [1, -1, 1], [-1, -1, 1], ], device=device, dtype=torch.float32, requires_grad=True) z_change = torch.zeros_like(vertices) z_change[:, -1] = -1.1 vertices = torch.cat( [ vertices, _verts_add * 0.2 + z_change ], dim=0 ) # vertices = _verts_add * 0.2 + z_change colors = torch.tensor([ [1, 0, 0, 0.6], [0, 1, 0, 0.6], [1, 0, 0, 0.6], [0, 0, 1, 0.6], [0, 0, 1, 1.], [1, 0, 0, 1.], [0, 0, 1, 1.], [0, 1, 0, 1.], ], device=device, dtype=torch.float32, requires_grad=True) _colors_add = torch.cat( [ torch.zeros_like(colors[:, :3], requires_grad=True), torch.ones_like(colors[:, 3:], requires_grad=True), ], dim=-1 ) colors = torch.cat( [ colors, _colors_add ], dim=0 ) # colors = _colors_add # Define faces (triangles) faces = torch.tensor([ [0, 1, 4], [1, 5, 4], [1, 2, 5], [2, 6, 5], [2, 3, 6], [3, 7, 6], [3, 0, 7], [0, 4, 7], [0, 1, 3], [1, 2, 3], [4, 5, 7], [5, 6, 7], ], device=device, dtype=torch.int32).contiguous() _new_faces = faces + 8 faces = torch.cat( [ faces, _new_faces ], dim=0 ) return vertices, colors, faces def load_pseudomesh(_path): mesh_data = np.load(_path) vertices = mesh_data["vertices"] faces = mesh_data["faces"] vert_colors = mesh_data["vertex_colors"] face_colors = mesh_data["face_colors"] return vertices, faces, vert_colors, face_colors def create_proj_mats(camera_angle_x, img_size, transf_mat, far, near): focal = img_size / (2 * np.tan(camera_angle_x / 2)) transf_matrix = torch.tensor(transf_mat, device="cpu", dtype=torch.float32) view_matrix = torch.inverse(transf_matrix) proj_matrix = torch.zeros(4, 4, device="cpu", dtype=torch.float32) proj_matrix[0, 0] = 2 * focal / img_size proj_matrix[1, 1] = 2 * focal / img_size proj_matrix[2, 2] = -(far + near) / (far - near) proj_matrix[2, 3] = -2.0 * far * near / (far - near) proj_matrix[3, 2] = -1.0 mvp_matrix = proj_matrix @ view_matrix return { "transf_matrix": transf_matrix, "view_matrix": view_matrix, "proj_matrix": proj_matrix, "mvp_matrix": mvp_matrix } def create_camera_mats(cam_data, img_size=512, near=0.01, far=100., device="cuda"): height, width = img_size, img_size camera_angle_x = cam_data["camera_angle_x"] focal = width / (2 * np.tan(camera_angle_x / 2)) out = {} for cam in cam_data["frames"]: transf_matrix = cam["transform_matrix"] transf_matrix = torch.tensor(transf_matrix, device=device, dtype=torch.float32) # view matrix - transforms vertices from world space # vertex coords to camera/view space view_matrix = torch.inverse(transf_matrix) # Convert focal length to OpenGL-style projection matrix proj_matrix = torch.zeros(4, 4, device=device, dtype=torch.float32) proj_matrix[0, 0] = 2 * focal / width proj_matrix[1, 1] = 2 * focal / height proj_matrix[2, 2] = -(far + near) / (far - near) proj_matrix[2, 3] = -2.0 * far * near / (far - near) proj_matrix[3, 2] = -1.0 # mvp (model-view-projection) matrix - combines all transformations into one matrix # basically u apply it to verts and you have them on an image lol mvp_matrix = proj_matrix @ view_matrix out[cam["file_path"]] = { "transf_matrix": transf_matrix, "view_matrix": view_matrix, "proj_matrix": proj_matrix, "mvp_matrix": mvp_matrix } _out = create_proj_mats( camera_angle_x, img_size, cam["transform_matrix"], far, near ) out[cam["file_path"]] = {k: v.cuda() for k, v in _out.items()} return out def render(glctx, verts, colors, faces, cam_data, img_size, device): # world vertices to image positions pos_clip = torch.matmul( torch.cat( [ verts, torch.ones_like(verts[:, :1]) ], dim=1 ), cam_data["mvp_matrix"].t() )[None, ...] with torch.no_grad(): # Calculate face centers in clip space face_verts = pos_clip[0, faces.long()] face_centers = face_verts.mean(dim=1) # Sort faces by Z coordinate (depth) in descending order (back to front) face_z = face_centers[..., 2] face_order = torch.argsort(face_z, descending=True) faces = faces[face_order] rast_out, rast_db = dr.rasterize( glctx, pos_clip, faces, (img_size, img_size), grad_db=True # Enable gradients for all attributes ) # Interpolate colors color_interp = dr.interpolate(colors[None, ...], rast_out, faces, rast_db=rast_db,diff_attrs='all')[0] # Background color background = torch.ones_like(color_interp) output = color_interp # Apply built-in transparency blending output = dr.antialias( color_interp, rast_out, pos_clip, faces, ) # Remove batch dimension # Blend with background alpha = output[..., 3:4] rgb = output[..., :3] final_color = rgb + background[..., :3] * (1 - alpha) output = torch.cat([final_color, alpha], dim=-1) output = torch.clamp(output, 0., 1.) return output # def render_with_depth_peeling(glctx, verts, colors, faces, cam_data, img_size, num_layers=4, device="cuda"): # """ # Render with nvdiffrast's DepthPeeler for order-independent transparency # """ # # Transform vertices to clip space # pos_clip = torch.matmul( # torch.cat([verts, torch.ones_like(verts[:, :1])], dim=1), # cam_data["mvp_matrix"].t() # )[None, ...] # # Initialize DepthPeeler # background = torch.ones((1, img_size, img_size, 4), device=device) # final_color = torch.zeros_like(background) # # Iterate through layers # import time # start_time = time.time() # with dr.DepthPeeler(glctx, pos_clip, faces, (img_size, img_size)) as peeler: # for layer_idx in range(num_layers): # # Get rasterization output for current layer # rast_out, rast_db = peeler.rasterize_next_layer() # # Skip if no geometry in this layer # if rast_out is None: # break # # Interpolate colors # color_layer = dr.interpolate( # colors[None, ...], # rast_out, # faces, # rast_db=rast_db, # diff_attrs='all' # )[0] # # # Antialias the colors # # color_aa = dr.antialias( # # color_layer, # # pos_clip, # # faces, # # rast_out, # # rast_db=rast_db # # ) # # Front-to-back blending # alpha = color_layer[..., 3:4] # rgb = color_layer[..., :3] # final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([rgb * alpha, alpha], dim=-1) # print(f"Final time: {time.time() - start_time}s") # # Blend with background # alpha = final_color[..., 3:4] # rgb = final_color[..., :3] # final_color = rgb + background[..., :3] * (1 - alpha) # loss = final_color.mean() # loss.backward() # return torch.clamp(final_color, 0., 1.) def to_torch(*args, device="cpu"): return [torch.from_numpy(x).float().to(device) for x in args] def test_render(): cam_info_path = Path("/home/grzegos/projects/phd/gs_raytracing/data/cam_test/transforms_train.json") with open(cam_info_path, "r") as file: cam_data = json.load(file) glctx = dr.RasterizeGLContext() img_size = 512 near = 1e-2 far = 1e2 device = "cuda:0" verts, colors, faces = create_obj(device) cam_matrices = create_camera_mats( cam_data=cam_data, img_size=img_size, near=near, far=far, device=device ) for cam_name, cam_mats in cam_matrices.items(): out_img = render_with_depth_peeling( glctx=glctx, verts=verts, colors=colors, faces=faces, cam_data=cam_mats, img_size=img_size, num_layers=2, device=device ) plt.imsave(f"02_{cam_name.split('/')[-1]}", out_img.detach().cpu().numpy()[0]) def render_pseudomesh(): path_to_dataset = Path("/home/grzegos/datasets/games_set/hotdog") path_to_camera_data = Path(path_to_dataset, "transforms_train.json") path_to_pseudomesh = Path("/home/grzegos/projects/phd/games_nerf_output/hotdog_test_fff_sh0/pseudomeshes/scale_2.30_pts_8.npz") with open(path_to_camera_data, "r") as file: cam_data = json.load(file) glctx = dr.RasterizeGLContext() img_size = 800 near = 1e-2 far = 1e2 device = "cuda:0" num_layers = 100 cam_matrices = create_camera_mats( cam_data=cam_data, img_size=img_size, near=near, far=far, device=device ) vertices, faces, vert_colors, face_colors = load_pseudomesh(path_to_pseudomesh) vertices, faces, vert_colors = to_torch(vertices, faces, vert_colors, device=device) vert_colors.requires_grad = True faces = faces.int() for cam_name, cam_mats in cam_matrices.items(): out_img = render_with_depth_peeling( glctx=glctx, verts=vertices, vert_colors=vert_colors, faces=faces, cam_data=cam_mats, img_size=img_size, num_layers=num_layers, device=device ) plt.imsave(f"{str(num_layers)}_{cam_name.split('/')[-1]}.png", out_img.detach().cpu().numpy()[0]) def render_with_depth_peeling(glctx, verts, vert_colors, faces, cam_data, img_size, num_layers=4, device="cuda"): """ Render with nvdiffrast's DepthPeeler for order-independent transparency. Optimized for performance, especially regarding the vert_colors tensor. """ # Transform vertices to clip space pos_clip = torch.matmul( torch.cat([verts, torch.ones_like(verts[:, :1])], dim=1), cam_data["mvp_matrix"].t() )[None, ...] # Initialize DepthPeeler background = torch.ones((1, img_size, img_size, 4), device=device) final_color = torch.zeros_like(background) # Iterate through layers with dr.DepthPeeler(glctx, pos_clip, faces, (img_size, img_size)) as peeler: for layer_idx in range(num_layers): # Get rasterization output for current layer rast_out, rast_db = peeler.rasterize_next_layer() # Skip if no geometry in this layer if rast_out is None: break # Interpolate colors # Use `interp1d` for more efficient color interpolation color_layer = dr.interpolate( vert_colors[None, ...], rast_out, faces, rast_db=rast_db, diff_attrs='all' )[0] # Front-to-back blending alpha = color_layer[..., 3:4] rgb = color_layer[..., :3] final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([rgb * alpha, alpha], dim=-1) # Blend with background alpha = final_color[..., 3:4] rgb = final_color[..., :3] final_color = rgb + background[..., :3] * (1 - alpha) loss = final_color.mean() loss.backward() return torch.clamp(final_color, 0., 1.) if __name__ == "__main__": # main() render_pseudomesh() # Regular render # save_render(angle=np.pi/6) # Or run optimization # optimize_colors() ================================================ FILE: mesh_optim/generate_multi_views_circle.py ================================================ import argparse import math from pathlib import Path import time import torch import torchvision import imageio.v2 as imageio from data import ImageCamDataset from optimize_pseudomesh import _load_best_model from render_pseudomesh import _load_config def get_depth_map(model, mvp_mat, width, height, cam_pos, num_layers): optim_depth_map, optim_depth_map_alpha = model.get_depth_map(mvp_mat.unsqueeze(0), width, height, cam_pos, num_layers) optim_depth_map = (optim_depth_map - optim_depth_map.min()) / (optim_depth_map.max() - optim_depth_map.min()) optim_depth_map_alpha = torch.clamp(optim_depth_map_alpha, 0.0, 1.0).squeeze(-1) optim_depth_map = torch.cat( [ optim_depth_map[..., 0], torch.zeros_like(optim_depth_map[..., 0]), 1 - optim_depth_map[..., 0], ], dim=0 ) background = torch.tensor([1, 0, 0], dtype=torch.float32, device=optim_depth_map.device) optim_depth_map = optim_depth_map * optim_depth_map_alpha + background[:3, None, None] * (1 - optim_depth_map_alpha) return optim_depth_map def get_gray_map(model, mvp_mat, width, height, color_verts, num_layers): _color_verts = torch.cat([color_verts, model.pseudomesh.vertex_colors[..., -1:]], dim=1) rgb, alpha = model.get_gray_map(mvp_mat.unsqueeze(0), width, height, num_layers, _color_verts) return rgb.squeeze(0).permute(2, 0, 1)#, alpha.squeeze(0) def get_normal_map(model, mvp_mat, width, height, num_layers): rgb, alpha = model.get_normal_map(mvp_mat.unsqueeze(0), width, height, num_layers) rgb = torch.clamp(rgb, 0.0, 1.0) return rgb.squeeze(0).permute(2, 0, 1)#, alpha.squeeze(0) def get_all_maps(model, mvp_mat, width, height, cam_pos, num_layers, color_verts=None): results = model.render_all_maps(mvp_mat.unsqueeze(0), width, height, cam_pos, num_layers, color_verts) color = results["color"][0].squeeze(0).permute(2, 0, 1).clamp(0.0, 1.0) normal = results["normal"][0].squeeze(0).permute(2, 0, 1).clamp(0.0, 1.0) depth = results["depth"][0].squeeze(0).permute(2, 0, 1) depth_alpha = results["depth"][1].squeeze(0) gray = results["gray"][0].squeeze(0).permute(2, 0, 1).clamp(0.0, 1.0) depth = (depth - depth.min()) / (depth.max() - depth.min()) depth_alpha = torch.clamp(depth_alpha, 0.0, 1.0).squeeze(-1).unsqueeze(0) depth = torch.cat( [ depth[0:1, ...], torch.zeros_like(depth[0:1, ...]), 1 - depth[0:1, ...], ], dim=0 ) background = torch.tensor([1, 0, 0], dtype=torch.float32, device=depth.device) depth = depth * depth_alpha + background[:3, None, None] * (1 - depth_alpha) return color, normal, depth, gray def normalize(v, eps=1e-6): """Return the normalized vector v.""" norm = torch.norm(v) return v if norm < eps else v / norm def look_at(eye, target, up): # Compute the forward vector (from the camera toward the target) forward = normalize(target - eye) # Compute the right vector right = normalize(torch.cross(forward, up)) # Recompute the orthogonal up vector true_up = torch.cross(right, forward) # true_up = up # Build rotation matrix. # Note: We use -forward so that the camera looks along its -z axis. R = torch.stack([right, true_up, -forward], dim=1) # Shape: (3,3) # Build the full 4x4 transformation matrix T = torch.eye(4) T[:3, :3] = R T[:3, 3] = eye return T def generate_camera_matrices(target, radius, num_views=20, up=torch.tensor([0, -1, 0], dtype=torch.float32), tilt_angle=0.0): target = target.float() cam_mats = [] # Precompute the rotation matrix for tilting about the x-axis c = math.cos(tilt_angle) s = math.sin(tilt_angle) R_tilt = torch.tensor([ [1, 0, 0], [0, c, -s], [0, s, c] ], dtype=torch.float32) for i in range(num_views): angle = 2 * math.pi * i / num_views # Compute a point on the circle (initially in the XZ-plane) circle_point = torch.tensor([ radius * math.cos(angle), 0.0, radius * math.sin(angle) ], dtype=torch.float32) # Apply the tilt rotation to the circle point tilted_point = torch.matmul(R_tilt, circle_point) # Compute the camera position by adding the target (center of the circle) cam_pos = target + tilted_point # Create the camera-to-world transformation matrix for this camera cam_mat = look_at(cam_pos, target, up) cam_mats.append(cam_mat) return torch.stack(cam_mats) def calculate_mvp(focal_x, focal_y, img_shape, view_mat, far, near): proj_matrix = torch.zeros(4, 4, device="cpu", dtype=torch.float32) proj_matrix[0, 0] = 2 * focal_x / img_shape[0] proj_matrix[1, 1] = -2 * focal_y / img_shape[1] proj_matrix[2, 2] = -(far + near) / (far - near) proj_matrix[2, 3] = -2.0 * far * near / (far - near) proj_matrix[3, 2] = -1.0 mvp_matrix = proj_matrix @ view_mat return mvp_matrix def create_gif(image_paths, output_path, fps): """Create a GIF from a list of image paths""" images = [] # Sort paths numerically based on the frame number in the filename sorted_paths = sorted(image_paths, key=lambda x: int(x.stem.split('_')[0])) for path in sorted_paths: images.append(imageio.imread(path)) imageio.mimsave(output_path, images, fps=fps) def calculate_alphas(img1, img2, left_int, dist_interval): dist = (img1["view_matrix"][:3, 3] - img2["view_matrix"][:3, 3]).pow(2).sum().sqrt() alphas = [] left_dist = dist curr_pos = -left_int while dist_interval < left_dist: curr_pos += dist_interval alphas.append(curr_pos / dist) left_dist = dist - curr_pos left_int = dist - curr_pos return torch.tensor(alphas, dtype=torch.float32), left_int def main(cfg_path, no_sec, fps): config = _load_config(cfg_path) output_dir = Path(config["experiment_dir"]) / "multi_views" output_dir.mkdir(parents=True, exist_ok=True) dset = ImageCamDataset( **{ **config["dataset"], } ) optim_pseudomesh = _load_best_model(config) circle_center = torch.cat([x["transf_matrix"][:3, 3].unsqueeze(0) for x in dset.image_files], dim=0).mean(dim=0) radius = 4. all_frames = no_sec * fps tilt = math.radians(20) up = torch.tensor([0, -0.78, -0.22], dtype=torch.float32) up = up / torch.norm(up) camera_matrices = generate_camera_matrices(circle_center, radius, all_frames, up=up, tilt_angle=tilt) _range = 0.3 color_verts = torch.rand( optim_pseudomesh.pseudomesh.vertices.shape[0], device=optim_pseudomesh.pseudomesh.vertices.device, )[:, None].repeat(1, 3) * _range + 0.5 - _range / 2. _iter = 0 focal_x, focal_y = dset.image_files[0]["focal_x"], dset.image_files[0]["focal_y"] width, height = dset.image_files[0]["img"].shape[1], dset.image_files[0]["img"].shape[0] print("Generate views") for cam_mat in camera_matrices: start_time = time.time() view_mat = torch.inverse(cam_mat) mvp_mat = calculate_mvp(focal_x, focal_y, [width, height], view_mat, dset.far, dset.near).to(config["device"]) with torch.no_grad(): optim_img, optim_normal_map, optim_depth_map, optim_gray_map = get_all_maps(optim_pseudomesh, mvp_mat, width, height, cam_mat[:3, 3].to(config["device"]), config["renderer"]["depth_steps"], color_verts) optim_img_path = output_dir / f"{_iter:05d}_optim.png" depth_img_path = output_dir / f"{_iter:05d}_optim-depth.png" gray_img_path = output_dir / f"{_iter:05d}_optim-gray.png" normal_img_path = output_dir / f"{_iter:05d}_optim-normal.png" torchvision.utils.save_image(optim_img.squeeze().permute(2, 0, 1), optim_img_path) torchvision.utils.save_image(optim_depth_map, depth_img_path) torchvision.utils.save_image(optim_gray_map, gray_img_path) torchvision.utils.save_image(optim_normal_map, normal_img_path) print(f"Saved {_iter}, took: {round(time.time() - start_time, 3)} s") _iter += 1 # Create GIFs after generating all images optim_images = list(output_dir.glob("*_optim.png")) gray_images = list(output_dir.glob("*_optim-gray.png")) depth_images = list(output_dir.glob("*_optim-depth.png")) normal_images = list(output_dir.glob("*_optim-normal.png")) create_gif(optim_images, output_dir / "optim.gif", fps) create_gif(gray_images, output_dir / "gray.gif", fps) create_gif(depth_images, output_dir / "depth.gif", fps) create_gif(normal_images, output_dir / "normal.gif", fps) if __name__ == "__main__": _parser = argparse.ArgumentParser( description="Process a single string argument." ) _parser.add_argument( "--cfg_path", "-cp", type=str, help="Path to config", default="/home/grzegos/experiments/mip_games_gs-flat_sh0_bicycle/config.yaml" ) _parser.add_argument( "--sec", type=int, help="How long the video should be", default=10 ) _parser.add_argument( "--fps", type=int, help="Frames per second", default=30 ) _args = _parser.parse_args() main(_args.cfg_path, _args.sec, _args.fps) ================================================ FILE: mesh_optim/log_config.yaml ================================================ version: 1 formatters: simple: format: '%(asctime)s - %(levelname)s - %(module)s - %(funcName)s - %(message)s' handlers: console: class: logging.StreamHandler formatter: simple level: INFO stream: ext://sys.stdout file_handler: class: logging.FileHandler formatter: simple level: DEBUG # change filename based on the output training directory filename: /path/to/logs/logs.log mode: a loggers: my_module: level: DEBUG handlers: [console, file_handler] propagate: no root: level: DEBUG handlers: [console, file_handler] ================================================ FILE: mesh_optim/lpipsPyTorch/__init__.py ================================================ import torch from .modules.lpips import LPIPS def lpips(x: torch.Tensor, y: torch.Tensor, net_type: str = 'alex', version: str = '0.1'): r"""Function that measures Learned Perceptual Image Patch Similarity (LPIPS). Arguments: x, y (torch.Tensor): the input tensors to compare. net_type (str): the network type to compare the features: 'alex' | 'squeeze' | 'vgg'. Default: 'alex'. version (str): the version of LPIPS. Default: 0.1. """ device = x.device criterion = LPIPS(net_type, version).to(device) return criterion(x, y) def get_lpips_model(device, net_type: str = 'alex', version: str = '0.1'): return LPIPS(net_type, version).to(device) ================================================ FILE: mesh_optim/lpipsPyTorch/modules/lpips.py ================================================ import torch import torch.nn as nn from .networks import get_network, LinLayers from .utils import get_state_dict class LPIPS(nn.Module): r"""Creates a criterion that measures Learned Perceptual Image Patch Similarity (LPIPS). Arguments: net_type (str): the network type to compare the features: 'alex' | 'squeeze' | 'vgg'. Default: 'alex'. version (str): the version of LPIPS. Default: 0.1. """ def __init__(self, net_type: str = 'alex', version: str = '0.1'): assert version in ['0.1'], 'v0.1 is only supported now' super(LPIPS, self).__init__() # pretrained network self.net = get_network(net_type) # linear layers self.lin = LinLayers(self.net.n_channels_list) self.lin.load_state_dict(get_state_dict(net_type, version)) def forward(self, x: torch.Tensor, y: torch.Tensor): feat_x, feat_y = self.net(x), self.net(y) diff = [(fx - fy) ** 2 for fx, fy in zip(feat_x, feat_y)] res = [l(d).mean((2, 3), True) for d, l in zip(diff, self.lin)] return torch.sum(torch.cat(res, 0), 0, True) ================================================ FILE: mesh_optim/lpipsPyTorch/modules/networks.py ================================================ from typing import Sequence from itertools import chain import torch import torch.nn as nn from torchvision import models from .utils import normalize_activation def get_network(net_type: str): if net_type == 'alex': return AlexNet() elif net_type == 'squeeze': return SqueezeNet() elif net_type == 'vgg': return VGG16() else: raise NotImplementedError('choose net_type from [alex, squeeze, vgg].') class LinLayers(nn.ModuleList): def __init__(self, n_channels_list: Sequence[int]): super(LinLayers, self).__init__([ nn.Sequential( nn.Identity(), nn.Conv2d(nc, 1, 1, 1, 0, bias=False) ) for nc in n_channels_list ]) for param in self.parameters(): param.requires_grad = False class BaseNet(nn.Module): def __init__(self): super(BaseNet, self).__init__() # register buffer self.register_buffer( 'mean', torch.Tensor([-.030, -.088, -.188])[None, :, None, None]) self.register_buffer( 'std', torch.Tensor([.458, .448, .450])[None, :, None, None]) def set_requires_grad(self, state: bool): for param in chain(self.parameters(), self.buffers()): param.requires_grad = state def z_score(self, x: torch.Tensor): return (x - self.mean) / self.std def forward(self, x: torch.Tensor): x = self.z_score(x) output = [] for i, (_, layer) in enumerate(self.layers._modules.items(), 1): x = layer(x) if i in self.target_layers: output.append(normalize_activation(x)) if len(output) == len(self.target_layers): break return output class SqueezeNet(BaseNet): def __init__(self): super(SqueezeNet, self).__init__() self.layers = models.squeezenet1_1(True).features self.target_layers = [2, 5, 8, 10, 11, 12, 13] self.n_channels_list = [64, 128, 256, 384, 384, 512, 512] self.set_requires_grad(False) class AlexNet(BaseNet): def __init__(self): super(AlexNet, self).__init__() self.layers = models.alexnet(True).features self.target_layers = [2, 5, 8, 10, 12] self.n_channels_list = [64, 192, 384, 256, 256] self.set_requires_grad(False) class VGG16(BaseNet): def __init__(self): super(VGG16, self).__init__() self.layers = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).features self.target_layers = [4, 9, 16, 23, 30] self.n_channels_list = [64, 128, 256, 512, 512] self.set_requires_grad(False) ================================================ FILE: mesh_optim/lpipsPyTorch/modules/utils.py ================================================ from collections import OrderedDict import torch def normalize_activation(x, eps=1e-10): norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True)) return x / (norm_factor + eps) def get_state_dict(net_type: str = 'alex', version: str = '0.1'): # build url url = 'https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/' \ + f'master/lpips/weights/v{version}/{net_type}.pth' # download old_state_dict = torch.hub.load_state_dict_from_url( url, progress=True, map_location=None if torch.cuda.is_available() else torch.device('cpu') ) # rename keys new_state_dict = OrderedDict() for key, val in old_state_dict.items(): new_key = key new_key = new_key.replace('lin', '') new_key = new_key.replace('model.', '') new_state_dict[new_key] = val return new_state_dict ================================================ FILE: mesh_optim/mesh_utils.py ================================================ import torch from torch import nn def hard_prune_mesh(vertices, faces, vertex_color, mask): vertices = vertices[~mask] vertex_color = vertex_color[~mask] face_mask = torch.any(mask[faces], dim=1) faces = faces[~face_mask] vertices_to_new_idx = torch.zeros(mask.shape[0], dtype=torch.long, device=vertices.device) vertices_to_new_idx[~mask] = torch.arange(vertices.shape[0], device=vertices.device) faces = vertices_to_new_idx[faces] return vertices, faces, vertex_color, mask def soft_prune_mesh(vertices, faces, vertex_color, mask): face_mask = torch.all(mask[faces], dim=1) faces = faces[~face_mask] valid_vertices = faces.unique().long() valid_verts_mask = torch.zeros_like(vertices[:, 0]).bool() valid_verts_mask[valid_vertices] = True vertices_to_new_idx = torch.zeros(valid_verts_mask.shape[0], dtype=torch.long, device=vertices.device) vertices_to_new_idx[valid_verts_mask] = torch.arange(valid_verts_mask.sum(), device=vertices.device) faces = vertices_to_new_idx[faces.long()] vertices = vertices[valid_verts_mask] vertex_color = vertex_color[valid_verts_mask] return vertices, faces, vertex_color, ~valid_verts_mask def prune_mesh(vertices, faces, vertex_color, mask, mode): assert mode in ["soft", "hard"] method = hard_prune_mesh if mode == "hard" else soft_prune_mesh vertices, faces, vertex_color, mask = method( vertices, faces, vertex_color, mask ) return vertices, faces.int(), vertex_color, mask def prune_optimizer(optimizer, mask): optimizable_tensors = {} for group in optimizer.param_groups: stored_state = optimizer.state.get(group['params'][0], None) if stored_state is not None: stored_state["exp_avg"] = stored_state["exp_avg"][~mask] stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][~mask] del optimizer.state[group['params'][0]] group["params"][0] = nn.Parameter((group["params"][0][~mask].requires_grad_(True))) optimizer.state[group['params'][0]] = stored_state optimizable_tensors[group["name"]] = group["params"][0] else: group["params"][0] = nn.Parameter(group["params"][0][~mask].requires_grad_(True)) optimizable_tensors[group["name"]] = group["params"][0] return optimizable_tensors ================================================ FILE: mesh_optim/metrics.py ================================================ # # Copyright (C) 2023, Inria # GRAPHDECO research group, https://team.inria.fr/graphdeco # 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 george.drettakis@inria.fr # from argparse import ArgumentParser import json import os from pathlib import Path from PIL import Image import numpy as np import torch import torchvision.transforms.functional as tf from tqdm import tqdm from lpipsPyTorch import lpips from ml_utils import psnr, ssim def PILtoTorch(pil_image): resized_image = torch.from_numpy(np.array(pil_image)) / 255.0 if len(resized_image.shape) == 3: return resized_image.permute(2, 0, 1) else: return resized_image.unsqueeze(dim=-1).permute(2, 0, 1) def readImages(renders_dir, gt_dir, mesh_renders_dir): gs_renders = [] mesh_renders = [] gts = [] image_names = [] for fname in os.listdir(renders_dir): mesh_render = Image.open(mesh_renders_dir / fname) gs_render = Image.open(renders_dir / fname) gt_path = gt_dir / fname try: gt = Image.open(gt_path) except FileNotFoundError: try: gt = Image.open(gt_path.with_suffix(".jpg")) except FileNotFoundError: gt = Image.open(gt_path.with_suffix(".JPG")) if gt.width != gs_render.width or gt.height != gs_render.height: gt = gt.resize((int(gs_render.width), int(gs_render.height))) if mesh_render.width != gs_render.width or mesh_render.height != gs_render.height: mesh_render = mesh_render.resize((int(gs_render.width), int(gs_render.height))) gs_renders.append(PILtoTorch(gs_render).cuda()) mesh_renders.append(PILtoTorch(mesh_render).cuda()) gt_render = PILtoTorch(gt).cuda() if gt_render.shape[0] != 3: gt_render = gt_render[:3, ...] gt_mask = gt_render[3:4, ...] gt_render *= gt_mask.to(gt_render.device) gts.append(gt_render) image_names.append(fname) return gs_renders, mesh_renders, gts, image_names def evaluate(model_path, dset_path, algorithm): if Path(dset_path, "transforms_train.json").exists(): dset_type = "nerf" else: dset_type = "colmap" full_dict = {} per_view_dict = {} full_dict_polytopeonly = {} per_view_dict_polytopeonly = {} full_dict[model_path] = {} per_view_dict[model_path] = {} full_dict_polytopeonly[model_path] = {} per_view_dict_polytopeonly[model_path] = {} test_dir = Path(model_path) / "test" for method in os.listdir(test_dir): print("Method:", method) full_dict[model_path][method] = {} per_view_dict[model_path][method] = {} full_dict_polytopeonly[model_path][method] = {} per_view_dict_polytopeonly[model_path][method] = {} method_dir = test_dir / method gt_dir = Path(dset_path, "test") if dset_type == "nerf" else Path(dset_path, "images") renders_dir_name = "renders" if algorithm == "games": renders_dir_name += "_gs_flat" gs_renders_dir = method_dir / renders_dir_name mesh_renders_dir = method_dir / "pseudomesh_renders" gs_renders, mesh_renders, gts, image_names = readImages(gs_renders_dir, gt_dir, mesh_renders_dir) gs_ssims = [] gs_psnrs = [] gs_lpipss = [] mesh_ssims = [] mesh_psnrs = [] mesh_lpipss = [] for idx in tqdm(range(len(gs_renders)), desc="Metric evaluation progress"): gs_ssims.append(ssim(gs_renders[idx], gts[idx])) gs_psnrs.append(psnr(gs_renders[idx], gts[idx])) gs_lpipss.append(lpips(gs_renders[idx], gts[idx], net_type='vgg')) mesh_ssims.append(ssim(mesh_renders[idx], gts[idx])) mesh_psnrs.append(psnr(mesh_renders[idx], gts[idx])) mesh_lpipss.append(lpips(mesh_renders[idx], gts[idx], net_type='vgg')) full_dict[model_path][method].update({ "GS_SSIM": torch.tensor(gs_ssims).mean().item(), "GS_PSNR": torch.tensor(gs_psnrs).mean().item(), "GS_LPIPS": torch.tensor(gs_lpipss).mean().item(), "MESH_SSIM": torch.tensor(mesh_ssims).mean().item(), "MESH_PSNR": torch.tensor(mesh_psnrs).mean().item(), "MESH_LPIPS": torch.tensor(mesh_lpipss).mean().item(), }) per_view_dict[model_path][method].update({ "GS_SSIM": {name: ssim for ssim, name in zip(torch.tensor(gs_ssims).tolist(), image_names)}, "GS_PSNR": {name: psnr for psnr, name in zip(torch.tensor(gs_psnrs).tolist(), image_names)}, "GS_LPIPS": {name: lp for lp, name in zip(torch.tensor(gs_lpipss).tolist(), image_names)}, "MESH_SSIM": {name: ssim for ssim, name in zip(torch.tensor(mesh_ssims).tolist(), image_names)}, "MESH_PSNR": {name: psnr for psnr, name in zip(torch.tensor(mesh_psnrs).tolist(), image_names)}, "MESH_LPIPS": {name: lp for lp, name in zip(torch.tensor(mesh_lpipss).tolist(), image_names)}, }) with open(model_path + "/results.json", 'w') as fp: json.dump(full_dict[model_path], fp, indent=True) with open(model_path + "/per_view.json", 'w') as fp: json.dump(per_view_dict[model_path], fp, indent=True) if __name__ == "__main__": device = torch.device("cuda:0") torch.cuda.set_device(device) # Set up command line argument parser parser = ArgumentParser(description="Training script parameters") parser.add_argument('--model_path', '-m', required=True, type=str) parser.add_argument('--set_path', '-s', required=True, type=str) parser.add_argument('--white_background', action="store_true") parser.add_argument( "--method", type=str, default="3dgs", choices=["3dgs", "games", "sugar", "2dgs"] ) args = parser.parse_args() evaluate(args.model_path, args.set_path, args.method, args.white_background) ================================================ FILE: mesh_optim/ml_utils.py ================================================ import math import numpy as np import torch import torch.nn.functional as F from lpipsPyTorch import get_lpips_model def _exp_schedule(init_val, gamma, timestep): return init_val * gamma ** timestep def _step_schedule(curr_val, steps, timestep): return steps.get(timestep, curr_val) def dp_schedule(_type, curr_val, init_val, epoch, params): assert _type in ["exp", "step"] if _type == "exp": out_val = _exp_schedule(init_val, float(params["gamma"]), epoch) elif _type == "step": out_val = _step_schedule(curr_val, params["steps"], epoch) return max(1, int(out_val)) class ColorExponentialLR(torch.optim.lr_scheduler.ExponentialLR): def __init__(self, optimizer, gamma, param_group_index=0, last_epoch=-1, verbose=False): self.param_group_index = param_group_index super().__init__(optimizer, gamma, last_epoch, verbose) def get_lr(self): # Only update the learning rate for the specified parameter group lrs = [] for i, _ in enumerate(self.optimizer.param_groups): if i == self.param_group_index: lrs.append(self.base_lrs[i] * self.gamma ** self.last_epoch) else: lrs.append(self.optimizer.param_groups[i]['lr']) return lrs def _calc_bs_mul(method: str, init_bs: float, curr_bs: float) -> float: assert method in ["constant", "linear", "sqrt"] if method == "linear": return curr_bs / init_bs elif method == "sqrt": return np.sqrt(curr_bs / init_bs) else: return 1. def get_optimizer(config: dict, params: dict) -> torch.optim.Optimizer: optimizer_config = config["optimizer"] init_bs = 1 lr_mul = _calc_bs_mul( method=optimizer_config["batch_scal_method"], init_bs=init_bs, curr_bs=config["dloader"]["batch_size"] ) _params = [{"params": val, "name": key, "lr": lr_mul * float(optimizer_config["lrs"][key])} for key, val in params.items()] if optimizer_config["name"] == "adam": return torch.optim.Adam( _params, lr=0. ) else: raise NotImplementedError(f"{optimizer_config['name']} not yet implemented") def get_scheduler(optimizer: torch.optim.Optimizer, params: dict) -> ColorExponentialLR: if not params["use"]: return None return ColorExponentialLR( optimizer, params["gamma"], 0 ) class Losser: def __init__(self, loss_cfg: dict) -> None: self.loss_cfg = loss_cfg self.pips = None if any(["pips" in x for x in self.loss_cfg.keys()]): self.pips = get_lpips_model( "cuda", "vgg" ) self._method_lookup = { "ssim": { "method": self.ssim_wrap, "params": ["img_pred", "img_gt"] }, "img_l1": { "method": self.img_l1_wrap, "params": ["img_pred", "img_gt"] }, "pips": { "method": self.pips_wrap, "params": ["img_pred", "img_gt"] }, "psnr": { "method": self.psnr_wrap, "params": ["img_pred", "img_gt"] }, "dice": { "method": self.dice_wrap, "params": ["mask_pred", "mask_gt"] }, "delta_xyz": { "method": self.delta_wrap, "params": ["delta_xyz"] }, "delta_rots": { "method": self.delta_wrap, "params": ["delta_rots"] }, "delta_scales": { "method": self.delta_wrap, "params": ["delta_scales"] }, "mse": { "method": self.mse_wrap, "params": ["img_pred", "img_gt"] }, "body_displ": { "method": self.vector_norm_mean, "params": ["body_displ_pred", "body_displ_gt"] }, "pose_displ": { "method": self.vector_norm_mean, "params": ["pose_displ_pred", "pose_displ_gt"] }, "final_splat": { "method": self.vector_norm_mean, "params": ["final_splats_pred", "final_splats_gt"] } } def vector_norm_mean(self, loss_data: dict, params: list) -> torch.Tensor: return torch.mean(torch.norm( loss_data[params[0]] - loss_data[params[1]], dim=1, keepdim=True )) def img_l1_wrap(self, loss_data: dict, params: list) -> torch.Tensor: return l1_loss(loss_data[params[0]], loss_data[params[1]]) def ssim_wrap(self, loss_data: dict, params: list) -> torch.Tensor: return 1. - ssim(loss_data[params[0]], loss_data[params[1]]) def mse_wrap(self, loss_data: dict, params: list) -> torch.Tensor: return mse_loss(loss_data[params[0]], loss_data[params[1]]) def pips_wrap(self, loss_data: dict, params: list) -> torch.Tensor: inp_data_1, inp_data_2 = loss_data[params[0]], loss_data[params[1]] if inp_data_1.shape[1] != 3: inp_data_1 = inp_data_1.permute(0, 3, 1, 2) if inp_data_2.shape[1] != 3: inp_data_2 = inp_data_2.permute(0, 3, 1, 2) return self.pips(inp_data_1, inp_data_2) def dice_wrap(self, loss_data: dict, params: list) -> torch.Tensor: return dice_loss(loss_data[params[0]], loss_data[params[1]]) def delta_wrap(self, loss_data: dict, params: list) -> torch.Tensor: return norm_loss(loss_data[params[0]], p=2) def psnr_wrap(self, loss_data: dict, params: list) -> torch.Tensor: pred_img = loss_data[params[0]] gt_img = loss_data[params[1]] return psnr(pred_img, gt_img) def __call__(self, loss_data: dict) -> dict: out_losses = {} for loss_name in self.loss_cfg.keys(): loss_method_name = loss_name.split("-")[0] loss_method = self._method_lookup[loss_method_name]["method"] params = self._method_lookup[loss_method_name]["params"] out_losses[loss_name] = loss_method(loss_data, params) out_losses["loss"] = sum([out_losses[key] * float(weight) for key, weight in self.loss_cfg.items() if "debug" not in key]) return out_losses def psnr(pred_img: torch.Tensor, gt_img: torch.Tensor, reduction: str = "mean") -> torch.Tensor: assert pred_img.shape[0] == gt_img.shape[0] mse = (((pred_img - gt_img)) ** 2).view(pred_img.shape[0], -1).mean(1, keepdim=True).mean() return 20 * torch.log10(1.0 / torch.sqrt(mse)) def dice_loss(mask_pred: torch.Tensor, mask_gt: torch.Tensor, smooth: float = 1e-6) -> torch.Tensor: """ masks both are tensors which have values 0 and 1 (0 for bg, 1 for mask) """ intersection = torch.sum(mask_pred * mask_gt) cardinality = torch.sum(mask_pred) + torch.sum(mask_gt) dice = (2 * intersection + smooth) / (cardinality + smooth) return 1 - dice def norm_loss(data: torch.Tensor, p: int = 1, dim: int = 1) -> torch.Tensor: return torch.norm(data, p=p, dim=dim).mean() def l1_loss(network_output: torch.Tensor, gt: torch.Tensor) -> torch.Tensor: return torch.abs((network_output - gt)).mean() def mse_loss(network_output: torch.Tensor, gt: torch.Tensor) -> torch.Tensor: return torch.sum((network_output - gt) ** 2) / (np.prod(list(network_output.shape))) def gaussian(window_size: int, sigma: float) -> torch.Tensor: gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) return gauss / gauss.sum() def create_window(window_size: int, channel: int) -> torch.Tensor: _1D_window = gaussian(window_size, 1.5).unsqueeze(1) _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) window = torch.autograd.Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) return window def ssim(img1: torch.Tensor, img2: torch.Tensor, window_size: int = 11, size_average: bool = True) -> torch.Tensor: channel = img1.size(-3) window = create_window(window_size, channel) if img1.is_cuda: window = window.cuda(img1.get_device()) window = window.type_as(img1) return _ssim(img1, img2, window, window_size, channel, size_average) def _ssim(img1: torch.Tensor, img2: torch.Tensor, window: int, window_size: int, channel: int, size_average: bool = True) -> torch.Tensor: mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) mu1_sq = mu1.pow(2) mu2_sq = mu2.pow(2) mu1_mu2 = mu1 * mu2 sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 C1 = 0.01 ** 2 C2 = 0.03 ** 2 ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) if size_average: return ssim_map.mean() else: return ssim_map.mean(1).mean(1).mean(1) ================================================ FILE: mesh_optim/models.py ================================================ import numpy as np import nvdiffrast.torch as dr import torch import torch.nn as nn class Pseudomesh(nn.Module): def __init__(self, config): super().__init__() self.vert_requires_grad = "vertices" in config["model"]["optimizable_params"] self.vert_col_requires_grad = "vertices" in config["model"]["optimizable_params"] self.register_buffer('faces', torch.tensor(0).int()) self.vertices = nn.Parameter(torch.tensor(0., requires_grad=self.vert_requires_grad)) self.vertex_colors = nn.Parameter(torch.tensor(0., requires_grad=self.vert_col_requires_grad)) self.register_buffer('grad_acc', torch.tensor(0).float()) self.register_buffer('grad_denom', torch.tensor(0).int()) class PseudomeshRenderer(): def __init__(self, config): self.pseudomesh = Pseudomesh(config) self.glctx = dr.RasterizeGLContext() @classmethod def create_model(cls, config): mesh_data = np.load(config["pseudomesh_path"]) obj = cls(config) obj.set_values(mesh_data["vertices"], mesh_data["faces"], mesh_data["vertex_colors"]) return obj @property def vert_requires_grad(self): return self.pseudomesh.vert_requires_grad @property def vert_col_requires_grad(self): return self.pseudomesh.vert_col_requires_grad def set_values(self, vertices, faces, vertex_colors, grad_acc=None): if isinstance(vertices, np.ndarray) and not isinstance(vertices, torch.nn.parameter.Parameter): vertices = torch.tensor(vertices, dtype=torch.float32, requires_grad=self.vert_requires_grad) if isinstance(faces, np.ndarray): faces = torch.tensor(faces, dtype=torch.int32) if isinstance(vertex_colors, np.ndarray) and not isinstance(vertex_colors, torch.nn.parameter.Parameter): vertex_colors = torch.tensor(vertex_colors, dtype=torch.float32, requires_grad=self.vert_col_requires_grad) self.pseudomesh.faces.data = faces if isinstance(vertices, torch.nn.parameter.Parameter): self.pseudomesh.vertices = vertices else: self.pseudomesh.vertices = nn.Parameter(vertices) if isinstance(vertex_colors, torch.nn.parameter.Parameter): self.pseudomesh.vertex_colors = vertex_colors else: self.pseudomesh.vertex_colors = nn.Parameter(vertex_colors) if grad_acc is None: grad_acc = torch.zeros_like(self.pseudomesh.vertex_colors[:,0]) self.pseudomesh.grad_acc.data = grad_acc def acc_grad(self): grad = torch.norm(self.pseudomesh.vertex_colors.grad, dim=-1) self.pseudomesh.grad_acc.data = self.pseudomesh.grad_acc.data + grad self.pseudomesh.grad_denom.data = self.pseudomesh.grad_denom.data + 1 def reset_acc_grad(self): self.pseudomesh.grad_acc.data = torch.zeros_like(self.pseudomesh.vertex_colors[:,0]) self.pseudomesh.grad_denom.data = torch.tensor(0).int() def to(self, device): self.pseudomesh = self.pseudomesh.to(device) return self def __call__(self, mvp_mat, width, height, num_layers): pos_clip = torch.matmul( torch.cat( [ self.pseudomesh.vertices, torch.ones_like(self.pseudomesh.vertices[:, :1]) ], dim=1 ), mvp_mat.permute(0, 2, 1) ) final_color = torch.zeros((mvp_mat.shape[0], height, width, 4), device=self.pseudomesh.vertices.device, dtype=torch.float32) with dr.DepthPeeler(self.glctx, pos_clip, self.pseudomesh.faces, (height, width)) as peeler: for layer_idx in range(num_layers): rast_out, rast_db = peeler.rasterize_next_layer() if rast_out is None: break color_layer = dr.interpolate( self.pseudomesh.vertex_colors[None, ...], rast_out, self.pseudomesh.faces, rast_db=rast_db, diff_attrs='all' )[0] alpha = color_layer[..., 3:4] rgb = color_layer[..., :3] final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([rgb * alpha, alpha], dim=-1) alpha = final_color[..., 3:4] rgb = final_color[..., :3] return rgb, alpha def get_depth_map(self, mvp_mat, width, height, cam_pos, num_layers): vertices = self.pseudomesh.vertices # shape: (N, 3) vertex_dist = torch.norm(vertices - cam_pos, dim=1, keepdim=True) # shape: (N, 1) dmin = vertex_dist.min() dmax = vertex_dist.max() vertex_dist_norm = (vertex_dist - dmin) / (dmax - dmin + 1e-8) vertex_colors = torch.cat([ vertex_dist_norm, vertex_dist_norm, vertex_dist_norm, self.pseudomesh.vertex_colors[..., -1:] ], dim=1) # shape: (N, 4) vertices_homo = torch.cat([vertices, torch.ones_like(vertices[:, :1])], dim=1) pos_clip = torch.matmul(vertices_homo, mvp_mat.permute(0, 2, 1)) final_color = torch.zeros((mvp_mat.shape[0], height, width, 4), device=self.pseudomesh.vertices.device, dtype=torch.float32) with dr.DepthPeeler(self.glctx, pos_clip, self.pseudomesh.faces, (height, width)) as peeler: for layer_idx in range(num_layers): rast_out, rast_db = peeler.rasterize_next_layer() if rast_out is None: break color_layer = dr.interpolate( vertex_colors[None, ...], rast_out, self.pseudomesh.faces, rast_db=rast_db, diff_attrs='all' )[0] alpha = color_layer[..., 3:4] rgb = color_layer[..., :3] final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([rgb * alpha, alpha], dim=-1) alpha = final_color[..., 3:4] rgb = final_color[..., :3] return rgb, alpha def get_gray_map(self, mvp_mat, width, height, num_layers, color_verts): pos_clip = torch.matmul( torch.cat( [ self.pseudomesh.vertices, torch.ones_like(self.pseudomesh.vertices[:, :1]) ], dim=1 ), mvp_mat.permute(0, 2, 1) ) final_color = torch.zeros((mvp_mat.shape[0], height, width, 4), device=self.pseudomesh.vertices.device, dtype=torch.float32) with dr.DepthPeeler(self.glctx, pos_clip, self.pseudomesh.faces, (height, width)) as peeler: for layer_idx in range(num_layers): rast_out, rast_db = peeler.rasterize_next_layer() if rast_out is None: break color_layer = dr.interpolate( color_verts[None, ...], rast_out, self.pseudomesh.faces, rast_db=rast_db, diff_attrs='all' )[0] alpha = color_layer[..., 3:4] rgb = color_layer[..., :3] final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([rgb * alpha, alpha], dim=-1) alpha = final_color[..., 3:4] rgb = final_color[..., :3] return rgb, alpha def get_normal_map(self, mvp_mat, width, height, num_layers): vertices = self.pseudomesh.vertices faces = self.pseudomesh.faces # Get face vertices v0 = vertices[faces[:, 0]] v1 = vertices[faces[:, 1]] v2 = vertices[faces[:, 2]] # Calculate face normals face_normals = torch.cross(v1 - v0, v2 - v0) face_normals = face_normals / (torch.norm(face_normals, dim=1, keepdim=True) + 1e-8) # Initialize vertex normals vertex_normals = torch.zeros_like(vertices) # Accumulate face normals to vertices for i in range(3): vertex_indices = faces[:, i] vertex_normals.index_add_(0, vertex_indices, face_normals) # Normalize vertex normals vertex_normals = vertex_normals / (torch.norm(vertex_normals, dim=1, keepdim=True) + 1e-8) # Rest of the rendering process remains the same pos_clip = torch.matmul( torch.cat( [ vertices, torch.ones_like(vertices[:, :1]) ], dim=1 ), mvp_mat.permute(0, 2, 1) ) vertex_attrs = torch.cat([ vertex_normals, self.pseudomesh.vertex_colors[..., -1:] ], dim=1) final_color = torch.zeros((mvp_mat.shape[0], height, width, 4), device=vertices.device, dtype=torch.float32) with dr.DepthPeeler(self.glctx, pos_clip, faces, (height, width)) as peeler: for layer_idx in range(num_layers): rast_out, rast_db = peeler.rasterize_next_layer() if rast_out is None: break color_layer = dr.interpolate( vertex_attrs[None, ...], rast_out, faces, rast_db=rast_db, diff_attrs='all' )[0] alpha = color_layer[..., 3:4] normals = color_layer[..., :3] normals = normals / (torch.norm(normals, dim=-1, keepdim=True) + 1e-8) final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([normals * alpha, alpha], dim=-1) alpha = final_color[..., 3:4] rgb = final_color[..., :3] return rgb, alpha def render_all_maps(self, mvp_mat, width, height, cam_pos, num_layers, gray_verts=None): vertices = self.pseudomesh.vertices faces = self.pseudomesh.faces # Calculate vertex normals v0 = vertices[faces[:, 0]] v1 = vertices[faces[:, 1]] v2 = vertices[faces[:, 2]] face_normals = torch.cross(v1 - v0, v2 - v0) face_normals = face_normals / (torch.norm(face_normals, dim=1, keepdim=True) + 1e-8) vertex_normals = torch.zeros_like(vertices) for i in range(3): vertex_indices = faces[:, i] vertex_normals.index_add_(0, vertex_indices, face_normals) vertex_normals = vertex_normals / (torch.norm(vertex_normals, dim=1, keepdim=True) + 1e-8) # Calculate depth values vertex_dist = torch.norm(vertices - cam_pos, dim=1, keepdim=True) dmin, dmax = vertex_dist.min(), vertex_dist.max() vertex_dist_norm = (vertex_dist - dmin) / (dmax - dmin + 1e-8) depth_colors = torch.cat([ vertex_dist_norm.repeat(1, 3), self.pseudomesh.vertex_colors[..., -1:] ], dim=1) # Prepare normal colors normal_colors = torch.cat([ vertex_normals, self.pseudomesh.vertex_colors[..., -1:] ], dim=1) # Initialize output tensors device = vertices.device shape = (mvp_mat.shape[0], height, width, 4) final_color = torch.zeros(shape, device=device, dtype=torch.float32) final_normal = torch.zeros(shape, device=device, dtype=torch.float32) final_depth = torch.zeros(shape, device=device, dtype=torch.float32) final_gray = torch.zeros(shape, device=device, dtype=torch.float32) if gray_verts is not None else None # Transform vertices to clip space pos_clip = torch.matmul( torch.cat([vertices, torch.ones_like(vertices[:, :1])], dim=1), mvp_mat.permute(0, 2, 1) ) with dr.DepthPeeler(self.glctx, pos_clip, faces, (height, width)) as peeler: for _ in range(num_layers): rast_out, rast_db = peeler.rasterize_next_layer() if rast_out is None: break # Interpolate all attributes at once color_layer = dr.interpolate( self.pseudomesh.vertex_colors[None, ...], rast_out, faces, rast_db=rast_db, diff_attrs='all' )[0] normal_layer = dr.interpolate( normal_colors[None, ...], rast_out, faces, rast_db=rast_db, diff_attrs='all' )[0] depth_layer = dr.interpolate( depth_colors[None, ...], rast_out, faces, rast_db=rast_db, diff_attrs='all' )[0] # Process color map alpha = color_layer[..., 3:4] rgb = color_layer[..., :3] final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([rgb * alpha, alpha], dim=-1) # Process normal map normals = normal_layer[..., :3] normals = normals / (torch.norm(normals, dim=-1, keepdim=True) + 1e-8) final_normal = final_normal + (1.0 - final_normal[..., 3:4]) * torch.cat([normals * alpha, alpha], dim=-1) # Process depth map depth_rgb = depth_layer[..., :3] final_depth = final_depth + (1.0 - final_depth[..., 3:4]) * torch.cat([depth_rgb * alpha, alpha], dim=-1) # Process gray map if gray_verts provided if gray_verts is not None: gray_layer = dr.interpolate( gray_verts[None, ...], rast_out, faces, rast_db=rast_db, diff_attrs='all' )[0] gray_rgb = gray_layer[..., :3] final_gray = final_gray + (1.0 - final_gray[..., 3:4]) * torch.cat([gray_rgb * alpha, alpha], dim=-1) results = { 'color': (final_color[..., :3], final_color[..., 3:4]), 'normal': (final_normal[..., :3], final_normal[..., 3:4]), 'depth': (final_depth[..., :3], final_depth[..., 3:4]) } if gray_verts is not None: results['gray'] = (final_gray[..., :3], final_gray[..., 3:4]) return results ================================================ FILE: mesh_optim/optimize_pseudo_config.yaml ================================================ experiment_name: exp_name experiments_dir: /path/to/experiments_dir mode: depth pseudomesh_path: /path/to/pseudomesh.npz device: cuda white_background: False dataset: dataset_path: /dset/path near: 0.01 far: 100. imgs_in_ram: True res: 1 dloader: batch_size: 4 shuffle: True num_workers: 10 training: epochs: 50 model: optimizable_params: - vertices - vertex_colors optim_epoch_start: vertices: 0 alpha_pruning: type: alpha # possible modes: alpha, gradient and both start_epoch: 10 epoch_step: 10 alpha_eps: 1e-4 grad_eps: 1e-7 mode: soft # soft or hard optimizer: name: adam batch_scal_method: "linear" # [linear, constant, sqrt] lrs: vertices: 1e-6 vertex_colors: 5e-3 lr_scheduler: use: True gamma: 0.999 renderer: depth_steps: 50 dp_scheduler: perform: False type: step # step of exp init_depth_steps: 50 # params: # for exp # gamma: 0.998 params: # for step steps: 0: 50 60: 1 loss: # name of the loss and it's weight ssim: .4 img_l1: .6 psnr-debug: 0. # pips: 1e1 # dice: 1. # delta_xyz: 1e-2 # delta_scales: 1e-2 # delta_rots: 5e-2 test_loss: ssim: 1. img_l1: 1. psnr: 1. pips: 1. wandb: use: True key_path: /path/to/wandb_key.json project_name: project_name entity: entity_name imgs_per_epoch: 4 ================================================ FILE: mesh_optim/optimize_pseudomesh.py ================================================ import argparse from collections import defaultdict import json import logging import logging.config from pathlib import Path import shutil from time import time import yaml import numpy as np import torch import torch.nn.functional as F from torch.utils.data import DataLoader import torchvision from tqdm import tqdm import wandb from data import ImageCamDataset from mesh_utils import prune_mesh, prune_optimizer from ml_utils import dp_schedule, get_optimizer, Losser, get_scheduler from models import PseudomeshRenderer LOGGER = None def _load_config(path): with open(path, "r") as file: data = yaml.load(file, Loader=yaml.FullLoader) return data def setup_logger(config): with open("./log_config.yaml", "r") as file: _config = yaml.load(file, Loader=yaml.FullLoader) log_filepath = Path( config["experiments_dir"], config["experiment_name"], "logs.log" ) log_filepath.parent.mkdir(exist_ok=True, parents=True) _config["handlers"]["file_handler"]["filename"] = str(log_filepath) logging.config.dictConfig(_config) global LOGGER LOGGER = logging.getLogger(__name__) def iter_pass(model: PseudomeshRenderer, data: dict, config: dict, loss_obj: Losser, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler = None, test: bool = False): if not test: optimizer.zero_grad() # get bakcgorund bg_color = [1, 1, 1] if config["white_background"] else [0, 0, 0] background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") mvp_mat = data["mvp_matrix"] gt_image = data["img"] if test: start_time = time() rgb, alpha = model( mvp_mat, gt_image.shape[2], # width gt_image.shape[1], # height config["renderer"]["depth_steps"] ) if test: end_time = time() - start_time # handle background final_color = rgb * alpha + background[..., :3] * (1 - alpha) if gt_image.shape[-1] == 4: gt_image = gt_image[... ,:3] * gt_image[..., 3:4] + background * (1 - gt_image[..., 3:4]) else: gt_image = gt_image[..., :3] # calculate losses data_to_loss = { "img_pred": final_color, "img_gt": gt_image, } losses = loss_obj(data_to_loss) if test: losses["render_time"] = torch.tensor(end_time) # update net if not test: losses["loss"].backward() model.acc_grad() optimizer.step() if scheduler is not None: scheduler.step() return losses, data_to_loss def wandb_log(losses: dict, step: int, data_to_loss: dict = None, mode: str = "train", epoch: int = None) -> None: assert mode in ["train", "val"] if losses is not None: log_losses = {f"{mode}_{_k}": _v.item() if isinstance(_v, torch.Tensor) else _v for _k, _v in losses.items()} wandb.log( data=log_losses, step=step ) if data_to_loss is not None: grid_data = [] name = "" for data_name, data in data_to_loss.items(): if "delta" in data_name: continue if len(data.shape) == 4: data = data.squeeze(0) if data.shape[0] == 1: data = data.repeat(3, 1, 1) grid_data.append(data) name += "" if not name else " " name += data_name grid_data = torch.cat([x.permute(0, 3, 1, 2) for x in grid_data], dim=0) grid_img = torchvision.utils.make_grid(grid_data, nrow=grid_data.shape[0] // 2) grid_img = torch.clip(grid_img, 0., 1.) grid_img = F.interpolate( grid_img.unsqueeze(0), size=(grid_img.shape[1]//4, grid_img.shape[2]//4), mode='bilinear', align_corners=False ).squeeze(0).permute(1, 2, 0).detach().cpu().numpy() * 255 grid_img = grid_img.astype(np.uint8) img_grid = wandb.Image(grid_img, mode="RGB") wandb.log( data={name: img_grid}, step=step ) deltas = {f"{mode}_{_k}": torch.norm(_v, p=2, dim=1).mean().item() for _k, _v in data_to_loss.items() if "delta_xyz" in _k} wandb.log( data=deltas, step=step ) def _collect_imgs_for_logs(data_to_loss:dict, container: dict, _iter: int, step: int) -> None: if not _iter % step: for _img_name, _img in data_to_loss.items(): if container.get(_img_name, None) is None: container[_img_name] = _img else: container[_img_name] = torch.cat( [ container[_img_name], _img ], dim=0 ) def _save_ckpt(ckpt_dir: Path, epoch: int, model: PseudomeshRenderer, optimizer: torch.optim.Optimizer) -> None: out_path = Path(ckpt_dir, "best_model") LOGGER.info(f"Save checkpoint: {out_path}") ckpt = { "epoch": epoch, "model": model.pseudomesh.state_dict(), "optimizer": optimizer.state_dict() } torch.save(ckpt, out_path) def batch_to_device(data, device): data["img"] = data["img"].detach().to(device) data["mvp_matrix"] = data["mvp_matrix"].detach().to(device) def pruning(model, optimizer, alpha_eps, grad_eps, mode, _type): init_shape = model.pseudomesh.vertices.shape[0] if _type in ["alpha", "both"]: opacity_mask = (model.pseudomesh.vertex_colors[:, -1:] < eval(alpha_eps)).reshape(-1) else: opacity_mask = torch.zeros_like(model.pseudomesh.vertex_colors[:, -1:]).reshape(-1).bool() if _type in ["gradient", "both"]: grad_mask = (model.pseudomesh.grad_acc / model.pseudomesh.grad_denom) < eval(grad_eps) else: grad_mask = torch.zeros_like(model.pseudomesh.vertex_colors[:, -1:]).reshape(-1).bool() grad_opac_mask = torch.logical_or(opacity_mask, grad_mask) new_verts, new_faces, new_vert_colors, mask = prune_mesh( model.pseudomesh.vertices, model.pseudomesh.faces, model.pseudomesh.vertex_colors, grad_opac_mask, mode ) LOGGER.info(f"Pruned {init_shape - new_verts.shape[0]} vertices. Left vertices: {new_verts.shape[0]}, left faces: {new_faces.shape[0]}") optimized_params = prune_optimizer(optimizer, mask) new_verts = optimized_params.get("vertices", new_verts) new_vert_colors = optimized_params.get("vertex_colors", new_vert_colors) model.set_values(new_verts, new_faces, new_vert_colors) model.reset_acc_grad() def train(pseudomesh: PseudomeshRenderer, dloader: DataLoader, config: dict) -> None: # create optimizer params = { param_name: getattr(pseudomesh.pseudomesh, param_name) for param_name in config["model"]["optimizable_params"] if param_name not in config["model"]["optim_epoch_start"] } optimizer = get_optimizer(config, params) scheduler = get_scheduler(optimizer, config["lr_scheduler"]) # create loss function loss_obj = Losser(config["loss"]) _iter = 0 best_loss = float("inf") start_time = time() best_time = time() LOGGER.info("Start training") for epoch in range(config["training"]["epochs"]): LOGGER.info(f"Epoch {epoch}") # train _batch_iter = 0 _img_iter = 0 batch_loss = 0 train_data_to_loss_epoch = {} for data in tqdm(dloader): batch_to_device(data, config["device"]) losses, data_to_loss = iter_pass( pseudomesh, data, config, loss_obj, optimizer, scheduler ) losses["depth_steps"] = config["renderer"]["depth_steps"] losses["color_lr"] = optimizer.param_groups[0]['lr'] if config["wandb"]["use"]: wandb_log( losses=losses, step=_iter, data_to_loss=None, mode="train", epoch=epoch ) _collect_imgs_for_logs( data_to_loss, train_data_to_loss_epoch, _img_iter, len(dloader.dataset) // config["wandb"]["imgs_per_epoch"] ) batch_loss += losses["loss"].item() _img_iter += data["img"].shape[0] _batch_iter += 1 _iter += data["img"].shape[0] if config["wandb"]["use"]: wandb_log( losses=None, step=_iter, data_to_loss=train_data_to_loss_epoch, mode="train", epoch=epoch ) batch_loss /= _batch_iter if batch_loss < best_loss: best_time = time() _save_ckpt( ckpt_dir=config["ckpt_dir"], epoch=epoch, model=pseudomesh, optimizer=optimizer ) best_loss = batch_loss alpha_prun_cond = epoch >= config["alpha_pruning"]["start_epoch"] alpha_prun_cond = alpha_prun_cond and not ( (epoch - config["alpha_pruning"]["start_epoch"]) % config["alpha_pruning"]["epoch_step"] ) if alpha_prun_cond: LOGGER.info("Perform pruning") pruning( pseudomesh, optimizer, config["alpha_pruning"]["alpha_eps"], config["alpha_pruning"]["grad_eps"], config["alpha_pruning"]["mode"], config["alpha_pruning"]["type"] ) # check if i should update optimizer params for param_name, epoch_start in config["model"]["optim_epoch_start"].items(): if epoch_start == epoch: LOGGER.info(f"Add {param_name} to optimizer") optimizer.add_param_group({ "name": param_name, "params": getattr(pseudomesh.pseudomesh, param_name), "lr": float(config["optimizer"]["lrs"][param_name]) }) # do the depth peeling scheduling if config["renderer"]["dp_scheduler"]["perform"]: # if scheduling should be applied new_dp_val = dp_schedule( config["renderer"]["dp_scheduler"]["type"], config["renderer"]["depth_steps"], config["renderer"]["dp_scheduler"]["init_depth_steps"], epoch, config["renderer"]["dp_scheduler"]["params"] ) if config["renderer"]["depth_steps"] != new_dp_val: config["renderer"]["depth_steps"] = new_dp_val LOGGER.info(f"Changed depth steps to: {config['renderer']['depth_steps']}") return best_time - start_time def test(pseudomesh: PseudomeshRenderer, dloader: DataLoader, config: dict) -> dict: LOGGER.info("Start test") loss_obj = Losser(config["test_loss"]) metrics = defaultdict(lambda: 0) _iter = 0 for data in tqdm(dloader): batch_to_device(data, config["device"]) with torch.no_grad(): losses, data_to_loss = iter_pass( pseudomesh, data, config, loss_obj, None, test=True ) _iter += 1 for metric_name, metric_val in losses.items(): metrics[metric_name] += metric_val.item() # save imgs? for metric_name, metric_val in metrics.items(): metrics[metric_name] /= _iter return metrics def prepare_output_dir(config: dict, cfg_path: str) -> None: experiment_dir = Path(config["experiments_dir"], config["experiment_name"]) ckpt_dir = Path(experiment_dir, "checkpoints") imgs_dir = Path(experiment_dir, "imgs") config["experiment_dir"] = experiment_dir config["ckpt_dir"] = ckpt_dir config["imgs_dir"] = imgs_dir ckpt_dir.mkdir(parents=True, exist_ok=True) imgs_dir.mkdir(parents=True, exist_ok=True) out_config_path = Path(config["experiment_dir"], "config.yaml") out_config = config out_config["experiment_dir"] = str(out_config["experiment_dir"]) out_config["ckpt_dir"] = str(out_config["ckpt_dir"]) out_config["imgs_dir"] = str(out_config["imgs_dir"]) with open(out_config_path, "w") as file: yaml.dump(out_config, file, default_flow_style=False) def setup_wandb(config: dict) -> None: LOGGER.info("Setup wandb") with open(config["wandb"]["key_path"], "r") as file: wandb_key = json.load(file) wandb.login( key=wandb_key, relogin=True ) wandb.init( project=config["wandb"]["project_name"], entity=config["wandb"]["entity"], config=config, name=config["experiment_name"] ) def _load_best_model(config: dict) -> PseudomeshRenderer: pseudomesh = PseudomeshRenderer.create_model(config).to(config["device"]) ckpt_path = Path(config["experiments_dir"], config["experiment_name"], "checkpoints", "best_model") checkpoint = torch.load(ckpt_path) pseudomesh.set_values( checkpoint["model"]["vertices"], checkpoint["model"]["faces"], checkpoint["model"]["vertex_colors"], checkpoint["model"]["grad_acc"] ) pseudomesh.pseudomesh.load_state_dict(checkpoint["model"]) return pseudomesh def _save_results(results: dict, config: dict) -> None: out_path = Path(config["experiments_dir"], config["experiment_name"], "test_results.json") with open(out_path, "w") as file: json.dump(results, file) def main() -> None: _parser = argparse.ArgumentParser( description="Optimize pseudomesh." ) _parser.add_argument( "--cfg_path", "-cp", type=str, help="Path to config", default="./optimize_pseudo_config.yaml" ) _parser.add_argument( "--exp_name", type=str, help="Experiment name", default="" ) _parser.add_argument( "--exp_dir", type=str, help="Experiments dir", default="" ) _parser.add_argument( "--pseudomesh_path", type=str, help="Path to pseudomesh", default="" ) _parser.add_argument( "--dset_path", type=str, help="Path to dset", default="" ) _parser.add_argument( "--white_background", action="store_true" ) _parser.add_argument( "--res", type=int, default=-1 ) _args = _parser.parse_args() config = _load_config(_args.cfg_path) if _args.exp_name: config["experiment_name"] = _args.exp_name if _args.pseudomesh_path: config["pseudomesh_path"] = _args.pseudomesh_path if _args.dset_path: config["dataset"]["dataset_path"] = _args.dset_path if _args.exp_dir: config["experiments_dir"] = _args.exp_dir if _args.white_background: config["white_background"] = _args.white_background if _args.res != -1: config["dataset"]["res"] = _args.res prepare_output_dir(config, _args.cfg_path) setup_logger(config) if config["wandb"]["use"]: setup_wandb(config) # load init splat pseudomesh = PseudomeshRenderer.create_model(config).to(config["device"]) LOGGER.info(f"Model starts with {pseudomesh.pseudomesh.vertices.shape[0]} vertices and {pseudomesh.pseudomesh.faces.shape[0]} faces.") dset = ImageCamDataset( **config["dataset"] ) dloader = DataLoader( dset, **config["dloader"] ) train_time_s = train( pseudomesh=pseudomesh, dloader=dloader, config=config ) del pseudomesh # perform tests torch.cuda.empty_cache() # load model best_model = _load_best_model(config) # save model as npz np.savez( Path(config["experiments_dir"], config["experiment_name"], "checkpoints", "best_model.npz"), vertices=best_model.pseudomesh.vertices.detach().cpu().numpy(), faces=best_model.pseudomesh.faces.detach().cpu().numpy(), vertex_colors=best_model.pseudomesh.vertex_colors.detach().cpu().numpy() ) final_train_losses = test( pseudomesh=best_model, dloader=dloader, config=config ) test_dset = ImageCamDataset( **{ "test": True, **config["dataset"], } ) test_dloader = DataLoader( test_dset, **config["dloader"] ) final_test_losses = test( pseudomesh=best_model, dloader=test_dloader, config=config ) out_losses = { "train_time": train_time_s, **{f"train_{k}": v for k, v in final_train_losses.items()}, **{f"test_{k}": v for k, v in final_test_losses.items()}, } _save_results(out_losses, config) if __name__ == "__main__": main() ================================================ FILE: mesh_optim/render_pseudomesh.py ================================================ import argparse from pathlib import Path import yaml import torch from torch.utils.data import DataLoader import torchvision from tqdm import tqdm from data import ImageCamDataset from optimize_pseudomesh import _load_best_model def _load_config(path): with open(path, "r") as file: data = yaml.load(file, Loader=yaml.FullLoader) return data def render_img(model, mvp_mat, img_shape, white_background, depth_steps, device): bg_color = [1, 1, 1] if white_background else [0, 0, 0] background = torch.tensor(bg_color, dtype=torch.float32, device=device) rgb, alpha = model( mvp_mat, img_shape[0], # width img_shape[1], # height depth_steps ) final_color = rgb * alpha + background[..., :3] * (1 - alpha) return final_color def main(cfg_path, method, white_background): config = _load_config(cfg_path) method_dir_name = "ours_30000" if method == "2dgs": # method_dir_name = "ours_10000" # for blender scenes method_dir_name = "ours_30000" # for real scenes out_imgs_dir = Path( config["experiments_dir"], config["experiment_name"], "test", method_dir_name, "pseudomesh_renders" ) out_imgs_dir.mkdir(parents=True, exist_ok=True) # setup dataloader without shuffle dset = ImageCamDataset( **{ "test": True, **config["dataset"], } ) dloader = DataLoader( dset, batch_size=1, shuffle=False, num_workers=10 ) # load model pseudomesh = _load_best_model(config) # render all images and save them in the experiment's directory for data in tqdm(dloader): out_path = Path(out_imgs_dir, data["name"][0]).with_suffix(".png") mvp_mat = data["mvp_matrix"].detach().to(config["device"]) gt_img = data["img"].squeeze() height = gt_img.shape[0] width = gt_img.shape[1] with torch.no_grad(): pred_img = render_img( pseudomesh, mvp_mat, [width, height], white_background, config["renderer"]["depth_steps"], config["device"] ) torchvision.utils.save_image(pred_img.squeeze().permute(2, 0, 1), out_path) if __name__ == "__main__": _parser = argparse.ArgumentParser( description="Process a single string argument." ) _parser.add_argument( "--cfg_path", "-cp", type=str, help="Path to config", required=True ) _parser.add_argument( "--method", type=str, choices=["3dgs", "games", "sugar", "2dgs"], required=True ) _parser.add_argument( '--white_background', action='store_true' ) _args = _parser.parse_args() main(_args.cfg_path, _args.method, _args.white_background) ================================================ FILE: requirements.txt ================================================ # Core dependencies numpy torch>=2.0.0 torchvision tqdm PyYAML Pillow matplotlib opencv-python # Specialized libraries ninja imageio PyOpenGL glfw xatlas gdown git+https://github.com/NVlabs/nvdiffrast wandb ================================================ FILE: scripts/visualize_cameras_nerf_blender.py ================================================ from pathlib import Path import bpy import bpy.ops import json from math import tan, atan from mathutils import Matrix import numpy as np sensor_width = 36.0 # mm, typical of a full-frame camera sensor_height = 24.0 # mm, typical of a full-frame camera def _calculate_fov(focal, pixels): return 2 * np.arctan(pixels / (2 * focal)) def create_camera(name, position, rot_matrix, fovx, fovy, image_width, image_height): # Create the camera object bpy.ops.object.camera_add() camera_obj = bpy.context.object camera_obj.name = str(name) camera_obj.location = position # Set rotation from matrix rotation_matrix = Matrix(rot_matrix) camera_obj.rotation_euler = rotation_matrix.to_euler() # Assuming the sensor width is fixed and calculating the sensor height based on the aspect ratio sensor_width = 36.0 # Adjust as needed, standard full frame sensor width aspect_ratio = image_width / image_height sensor_height = sensor_width / aspect_ratio # Set camera data camera = camera_obj.data camera.sensor_width = sensor_width camera.sensor_height = sensor_height # Calculate focal length from FoV using formula: focal_length = sensor_width / (2 * tan(fov / 2)) camera.lens = sensor_width / (2 * tan(fovx / 2)) # Use FoVy to adjust the sensor height if needed calculated_fovy = 2 * atan((sensor_height / 2) / camera.lens) if calculated_fovy != fovy: camera.sensor_height = 2 * camera.lens * tan(fovy / 2) return camera_obj def main(): _cam_json_path = "cameras.json" with open(_cam_json_path, "r") as file: _camera_data = json.load(file) _camera_data = sorted(_camera_data, key=lambda x: x["id"]) _camera_data = [{ "id": x["id"], "img_name": x["img_name"], "width": x["width"], "height": x["height"], "position": np.array(x["position"]), "rotation": np.array(x["rotation"]), "fovy": x["fy"], "fovx": x["fx"] } for x in _camera_data] adjustment = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]], dtype=np.float64) for _new_cam in _camera_data: fovx = _calculate_fov(_new_cam["fovx"], _new_cam["width"]) fovy = _calculate_fov(_new_cam["fovy"], _new_cam["height"]) _new_cam_obj = create_camera( name=_new_cam["id"], position=_new_cam["position"], rot_matrix=_new_cam["rotation"] @ adjustment, fovx=fovx, fovy=fovy, image_width=_new_cam["width"], image_height=_new_cam["height"] ) if __name__ == "__main__": main() ================================================ FILE: sh_scripts/configs/3dgs_sh0_pseudo_config.yaml ================================================ experiment_name: exp_name experiments_dir: /path/to/experiments mode: depth # possible modes: pseudomesh_path: /path/to/pseudomesh.npz device: cuda white_background: True dataset: dataset_path: /path/to/dataset near: 0.01 far: 100. imgs_in_ram: True res: 1 dloader: batch_size: 4 shuffle: True num_workers: 10 training: epochs: 15 model: optimizable_params: - vertices - vertex_colors optim_epoch_start: vertices: 0 alpha_pruning: type: alpha # possible modes: alpha, gradient and both start_epoch: 10 epoch_step: 10 alpha_eps: 1e-4 grad_eps: 1e-7 mode: soft # soft or hard optimizer: name: adam batch_scal_method: "linear" # [linear, constant, sqrt] lrs: vertices: 1e-6 vertex_colors: 5e-3 lr_scheduler: use: True gamma: 0.999 renderer: depth_steps: 50 dp_scheduler: perform: False type: step # step of exp init_depth_steps: 50 # params: # for exp # gamma: 0.998 params: # for step steps: 0: 50 60: 1 loss: # name of the loss and it's weight ssim: .4 img_l1: .6 psnr-debug: 0. # pips: 1e1 # dice: 1. # delta_xyz: 1e-2 # delta_scales: 1e-2 # delta_rots: 5e-2 test_loss: ssim: 1. img_l1: 1. psnr: 1. pips: 1. wandb: use: True key_path: /path/to/wandb.json project_name: project_name entity: entity_name imgs_per_epoch: 4 ================================================ FILE: sh_scripts/configs/db_config.yaml ================================================ experiment_name: exp_name experiments_dir: /path/to/experiments mode: depth # possible modes: pseudomesh_path: /path/to/pseudomesh.npz device: cuda white_background: True dataset: dataset_path: /path/to/dataset near: 0.01 far: 100. imgs_in_ram: True res: 1 dloader: batch_size: 2 shuffle: True num_workers: 10 training: epochs: 15 model: optimizable_params: - vertices - vertex_colors optim_epoch_start: vertices: 0 alpha_pruning: type: alpha # possible modes: alpha, gradient and both start_epoch: 10 epoch_step: 10 alpha_eps: 1e-4 grad_eps: 1e-7 mode: soft # soft or hard optimizer: name: adam batch_scal_method: "linear" # [linear, constant, sqrt] lrs: vertices: 1e-6 vertex_colors: 5e-3 lr_scheduler: use: True gamma: 0.999 renderer: depth_steps: 100 dp_scheduler: perform: False type: step # step of exp init_depth_steps: 50 # params: # for exp # gamma: 0.998 params: # for step steps: 0: 50 60: 1 loss: # name of the loss and it's weight ssim: .4 img_l1: .6 psnr-debug: 0. # pips: 1e1 # dice: 1. # delta_xyz: 1e-2 # delta_scales: 1e-2 # delta_rots: 5e-2 test_loss: ssim: 1. img_l1: 1. psnr: 1. pips: 1. wandb: use: True key_path: /path/to/wandb.json project_name: project_name entity: entity_name imgs_per_epoch: 2 ================================================ FILE: sh_scripts/configs/mip_config.yaml ================================================ experiment_name: exp_name experiments_dir: /path/to/experiments mode: depth # possible modes: pseudomesh_path: /path/to/pseudomesh.npz device: cuda white_background: True dataset: dataset_path: /path/to/dataset near: 0.01 far: 100. imgs_in_ram: True res: 8 dloader: batch_size: 1 shuffle: True num_workers: 10 training: epochs: 15 model: optimizable_params: - vertices - vertex_colors optim_epoch_start: vertices: 0 alpha_pruning: type: alpha # possible modes: alpha, gradient and both start_epoch: 10 epoch_step: 10 alpha_eps: 1e-4 grad_eps: 1e-7 mode: soft # soft or hard optimizer: name: adam batch_scal_method: "linear" # [linear, constant, sqrt] lrs: vertices: 1e-6 vertex_colors: 5e-3 lr_scheduler: use: True gamma: 0.999 renderer: depth_steps: 150 dp_scheduler: perform: False type: step # step of exp init_depth_steps: 50 # params: # for exp # gamma: 0.998 params: # for step steps: 0: 50 60: 1 loss: # name of the loss and it's weight ssim: .4 img_l1: .6 psnr-debug: 0. # pips: 1e1 # dice: 1. # delta_xyz: 1e-2 # delta_scales: 1e-2 # delta_rots: 5e-2 test_loss: ssim: 1. img_l1: 1. psnr: 1. pips: 1. wandb: use: True key_path: /path/to/wandb.json project_name: project_name entity: entity_name imgs_per_epoch: 2 ================================================ FILE: sh_scripts/configs/tandt_3dgs_config.yaml ================================================ experiment_name: exp_name experiments_dir: /path/to/experiments mode: depth # possible modes: pseudomesh_path: /path/to/pseudomesh.npz device: cuda white_background: True dataset: dataset_path: /path/to/dataset near: 0.01 far: 100. imgs_in_ram: True res: 1 dloader: batch_size: 1 shuffle: True num_workers: 10 training: epochs: 15 model: optimizable_params: - vertices - vertex_colors optim_epoch_start: vertices: 0 alpha_pruning: type: alpha # possible modes: alpha, gradient and both start_epoch: 10 epoch_step: 10 alpha_eps: 1e-4 grad_eps: 1e-7 mode: soft # soft or hard optimizer: name: adam batch_scal_method: "linear" # [linear, constant, sqrt] lrs: vertices: 1e-6 vertex_colors: 1e-2 lr_scheduler: use: True gamma: 0.999 renderer: depth_steps: 100 dp_scheduler: perform: False type: step # step of exp init_depth_steps: 50 # params: # for exp # gamma: 0.998 params: # for step steps: 0: 50 60: 1 loss: # name of the loss and it's weight ssim: .4 img_l1: .6 psnr-debug: 0. # pips: 1e1 # dice: 1. # delta_xyz: 1e-2 # delta_scales: 1e-2 # delta_rots: 5e-2 test_loss: ssim: 1. img_l1: 1. psnr: 1. pips: 1. wandb: use: True key_path: /path/to/wandb.json project_name: project_name entity: entity_name imgs_per_epoch: 2 ================================================ FILE: sh_scripts/configs/tandt_config.yaml ================================================ experiment_name: exp_name experiments_dir: /path/to/experiments mode: depth # possible modes: pseudomesh_path: /path/to/pseudomesh.npz device: cuda white_background: True dataset: dataset_path: /path/to/dataset near: 0.01 far: 100. imgs_in_ram: True res: 1 dloader: batch_size: 2 shuffle: True num_workers: 10 training: epochs: 15 model: optimizable_params: - vertices - vertex_colors optim_epoch_start: vertices: 0 alpha_pruning: type: alpha # possible modes: alpha, gradient and both start_epoch: 10 epoch_step: 10 alpha_eps: 1e-4 grad_eps: 1e-7 mode: soft # soft or hard optimizer: name: adam batch_scal_method: "linear" # [linear, constant, sqrt] lrs: vertices: 1e-6 vertex_colors: 1e-2 lr_scheduler: use: True gamma: 0.999 renderer: depth_steps: 200 dp_scheduler: perform: False type: step # step of exp init_depth_steps: 50 # params: # for exp # gamma: 0.998 params: # for step steps: 0: 50 60: 1 loss: # name of the loss and it's weight ssim: .4 img_l1: .6 psnr-debug: 0. # pips: 1e1 # dice: 1. # delta_xyz: 1e-2 # delta_scales: 1e-2 # delta_rots: 5e-2 test_loss: ssim: 1. img_l1: 1. psnr: 1. pips: 1. wandb: use: True key_path: /path/to/wandb.json project_name: project_name entity: entity_name imgs_per_epoch: 2 ================================================ FILE: sh_scripts/run_2dgs-sh0_db.sh ================================================ #!/bin/bash # Directory containing the elements DATA_DIR="/path/to/dataset" EXP_DIR="/path/to/experiments" GS_TRAIN_PATH="/path/to/2d-gaussian-splatting" PYTHON_PATH="/path/to/python" GS_PSEUDO_PATH="/path/to/gs_raytracing" GS_OPTIM_PATH="/path/to/gs_raytracing/mesh_optim" CONFIG_PATH="/path/to/gs_raytracing/sh_scripts/configs/db_config.yaml" for data_path in $DATA_DIR/*; do dset_name=$(basename ${data_path}) exp_name="db_2dgs_sh0_${dset_name}" exp_out_dir="${EXP_DIR}/${exp_name}" exp_cfg_path="${exp_out_dir}/config.yaml" ply_path="${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply" pseudomesh_path="${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz" cd $GS_TRAIN_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --white_background --eval CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --white_background --skip_train --eval cd $GS_PSEUDO_PATH $PYTHON_PATH generate_pseudomesh.py --ply $ply_path --algorithm 2dgs --scale_min 2.7 --scale_max 2.8 --scale_step 0.2 --no_points 8 cd $GS_OPTIM_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH optimize_pseudomesh.py --cfg_path $CONFIG_PATH --exp_name $exp_name --pseudomesh_path $pseudomesh_path --dset_path $data_path --white_background --res 2 CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 2dgs --white_background CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method 2dgs done ================================================ FILE: sh_scripts/run_2dgs-sh0_mip.sh ================================================ #!/bin/bash # Directory containing the elements DATA_DIR="/path/to/datasets/360_v2" EXP_DIR="/path/to/experiments" GS_TRAIN_PATH="/path/to/2d-gaussian-splatting" PYTHON_PATH="/path/to/python" GS_PSEUDO_PATH="/path/to/gs_raytracing" GS_OPTIM_PATH="/path/to/gs_raytracing/mesh_optim" CONFIG_PATH="/path/to/gs_raytracing/sh_scripts/configs/mip_config.yaml" for data_path in $DATA_DIR/*; do dset_name=$(basename ${data_path}) # Iterate over each element in the ELEMENTS_DIR exp_name="mip_2dgs_sh0_${dset_name}" exp_out_dir="${EXP_DIR}/${exp_name}" exp_cfg_path="${exp_out_dir}/config.yaml" ply_path="${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply" pseudomesh_path="${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz" cd $GS_TRAIN_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 4 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --white_background --eval CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render.py --resolution 4 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --white_background --skip_train --eval cd $GS_PSEUDO_PATH $PYTHON_PATH generate_pseudomesh.py --ply $ply_path --algorithm 2dgs --scale_min 2.7 --scale_max 2.8 --scale_step 0.2 --no_points 8 cd $GS_OPTIM_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH optimize_pseudomesh.py --cfg_path $CONFIG_PATH --exp_name $exp_name --pseudomesh_path $pseudomesh_path --dset_path $data_path --white_background --res 8 CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 2dgs --white_background CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method 2dgs done ================================================ FILE: sh_scripts/run_2dgs-sh0_nerf-synth.sh ================================================ #!/bin/bash # Directory containing the elements DATA_DIR="/path/to/datasets/games_set" EXP_DIR="/path/to/experiments" GS_TRAIN_PATH="/path/to/2d-gaussian-splatting" PYTHON_PATH="/path/to/python" GS_PSEUDO_PATH="/path/to/gs_raytracing" GS_OPTIM_PATH="/path/to/gs_raytracing/mesh_optim" CONFIG_PATH="/path/to/gs_raytracing/sh_scripts/configs/3dgs_sh0_pseudo_config.yaml" for data_path in $DATA_DIR/*; do dset_name=$(basename ${data_path}) # Iterate over each element in the ELEMENTS_DIR exp_name="new_nerf-synth_2dgs_sh0_${dset_name}" exp_out_dir="${EXP_DIR}/${exp_name}" exp_cfg_path="${exp_out_dir}/config.yaml" ply_path="${exp_out_dir}/point_cloud/iteration_10000/point_cloud.ply" pseudomesh_path="${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz" cd $GS_TRAIN_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 10000 --sh_degree 0 --eval CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_meshsplat.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 10000 --sh_degree 0 --skip_train --eval cd $GS_PSEUDO_PATH $PYTHON_PATH generate_pseudomesh.py --ply $ply_path --algorithm 2dgs --scale_min 2.7 --scale_max 2.8 --scale_step 0.2 --no_points 8 cd $GS_OPTIM_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH optimize_pseudomesh.py --cfg_path $CONFIG_PATH --exp_name $exp_name --pseudomesh_path $pseudomesh_path --dset_path $data_path --res 1 CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 2dgs CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --method 2dgs done ================================================ FILE: sh_scripts/run_2dgs_sh0_tandt.sh ================================================ #!/bin/bash # Directory containing the elements DATA_DIR="/path/to/dataset" EXP_DIR="/path/to/experiments" GS_TRAIN_PATH="/path/to/2d-gaussian-splatting" PYTHON_PATH="/path/to/python" GS_PSEUDO_PATH="/path/to/gs_raytracing" GS_OPTIM_PATH="/path/to/gs_raytracing/mesh_optim" CONFIG_PATH="/path/to/gs_raytracing/sh_scripts/configs/tandt_config.yaml" for data_path in $DATA_DIR/*; do dset_name=$(basename ${data_path}) # Iterate over each element in the ELEMENTS_DIR exp_name="tandt_2dgs_sh0_${dset_name}" exp_out_dir="${EXP_DIR}/${exp_name}" exp_cfg_path="${exp_out_dir}/config.yaml" ply_path="${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply" pseudomesh_path="${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz" cd $GS_TRAIN_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --white_background --eval CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --white_background --skip_train --eval cd $GS_PSEUDO_PATH $PYTHON_PATH generate_pseudomesh.py --ply $ply_path --algorithm 2dgs --scale_min 2.7 --scale_max 2.8 --scale_step 0.2 --no_points 8 cd $GS_OPTIM_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH optimize_pseudomesh.py --cfg_path $CONFIG_PATH --exp_name $exp_name --pseudomesh_path $pseudomesh_path --dset_path $data_path --white_background --res 1 CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 2dgs --white_background CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method 2dgs done ================================================ FILE: sh_scripts/run_3dgs-sh0_db.sh ================================================ #!/bin/bash # Directory containing the elements DATA_DIR="/path/to/datasets/db" EXP_DIR="/path/to/experiments" GS_TRAIN_PY="/path/to/gaussian-splatting" PYTHON_PATH="/path/to/python" GS_PSEUDO_PATH="/path/to/gs_raytracing" GS_OPTIM_PATH="/path/to/gs_raytracing/mesh_optim" CONFIG_PATH="/path/to/gs_raytracing/sh_scripts/configs/db_config.yaml" for data_path in $DATA_DIR/*; do dset_name=$(basename ${data_path}) exp_name="db_3dgs_white_sh0_${dset_name}" exp_out_dir="${EXP_DIR}/${exp_name}" exp_cfg_path="${exp_out_dir}/config.yaml" ply_path="${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply" pseudomesh_path="${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz" cd $GS_TRAIN_PY CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --white_background --eval CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --white_background --skip_train --eval cd $GS_PSEUDO_PATH $PYTHON_PATH generate_pseudomesh.py --ply $ply_path --algorithm 3dgs --scale_min 2.7 --scale_max 2.8 --scale_step 0.2 --no_points 8 cd $GS_OPTIM_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH optimize_pseudomesh.py --cfg_path $CONFIG_PATH --exp_name $exp_name --pseudomesh_path $pseudomesh_path --dset_path $data_path --white_background --res 2 CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 3dgs --white_background CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method 3dgs done ================================================ FILE: sh_scripts/run_3dgs-sh0_mip.sh ================================================ #!/bin/bash # Directory containing the elements DATA_DIR="/path/to/datasets/360_v2" EXP_DIR="/path/to/experiments" GS_TRAIN_PY="/path/to/gaussian-splatting" PYTHON_PATH="/path/to/python" GS_PSEUDO_PATH="/path/to/gs_raytracing" GS_OPTIM_PATH="/path/to/gs_raytracing/mesh_optim" CONFIG_PATH="/path/to/gs_raytracing/sh_scripts/configs/mip_config.yaml" for data_path in $DATA_DIR/*; do dset_name=$(basename ${data_path}) # Iterate over each element in the ELEMENTS_DIR exp_name="mip_3dgs_sh0_${dset_name}" exp_out_dir="${EXP_DIR}/${exp_name}" exp_cfg_path="${exp_out_dir}/config.yaml" ply_path="${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply" pseudomesh_path="${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz" cd $GS_TRAIN_PY CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 8 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --white_background --eval CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render.py --resolution 8 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --white_background --skip_train --eval cd $GS_PSEUDO_PATH $PYTHON_PATH generate_pseudomesh.py --ply $ply_path --algorithm 3dgs --scale_min 2.7 --scale_max 2.8 --scale_step 0.2 --no_points 8 cd $GS_OPTIM_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH optimize_pseudomesh.py --cfg_path $CONFIG_PATH --exp_name $exp_name --pseudomesh_path $pseudomesh_path --dset_path $data_path --white_background --res 8 CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 3dgs --white_background CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method 3dgs done ================================================ FILE: sh_scripts/run_3dgs-sh0_nerf-synth.sh ================================================ #!/bin/bash DATA_DIR="/path/to/datasets/games_set" EXP_DIR="/path/to/experiments" GS_TRAIN_PY="/path/to/gaussian-splatting" PYTHON_PATH="/path/to/python" GS_PSEUDO_PATH="/path/to/gs_raytracing" GS_OPTIM_PATH="/path/to/gs_raytracing/mesh_optim" CONFIG_PATH="/path/to/gs_raytracing/sh_scripts/configs/3dgs_sh0_pseudo_config.yaml" for data_path in $DATA_DIR/*; do dset_name=$(basename ${data_path}) # Iterate over each element in the ELEMENTS_DIR exp_name="nerf-synth_3dgs_sh0_${dset_name}" exp_out_dir="${EXP_DIR}/${exp_name}" exp_cfg_path="${exp_out_dir}/config.yaml" ply_path="${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply" pseudomesh_path="${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz" cd $GS_TRAIN_PY CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --eval CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --skip_train --eval cd $GS_PSEUDO_PATH $PYTHON_PATH generate_pseudomesh.py --ply $ply_path --algorithm 3dgs --scale_min 2.7 --scale_max 2.8 --scale_step 0.2 --no_points 8 cd $GS_OPTIM_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH optimize_pseudomesh.py --cfg_path $CONFIG_PATH --exp_name $exp_name --pseudomesh_path $pseudomesh_path --dset_path $data_path --res 1 CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 3dgs CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --method 3dgs done ================================================ FILE: sh_scripts/run_3dgs-sh0_tandt.sh ================================================ #!/bin/bash # Directory containing the elements DATA_DIR="/path/to/datasets/tandt" EXP_DIR="/path/to/experiments" GS_TRAIN_PY="/path/to/gaussian-splatting" PYTHON_PATH="/path/to/python" GS_PSEUDO_PATH="/path/to/gs_raytracing" GS_OPTIM_PATH="/path/to/gs_raytracing/mesh_optim" CONFIG_PATH="/path/to/gs_raytracing/sh_scripts/configs/tandt_3dgs_config.yaml" for data_path in $DATA_DIR/*; do dset_name=$(basename ${data_path}) exp_name="tandt_3dgs_white_sh0_${dset_name}" exp_out_dir="${EXP_DIR}/${exp_name}" exp_cfg_path="${exp_out_dir}/config.yaml" ply_path="${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply" pseudomesh_path="${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz" cd $GS_TRAIN_PY CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --white_background --eval CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --white_background --skip_train --eval cd $GS_PSEUDO_PATH $PYTHON_PATH generate_pseudomesh.py --ply $ply_path --algorithm 3dgs --scale_min 2.7 --scale_max 2.8 --scale_step 0.2 --no_points 8 cd $GS_OPTIM_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH optimize_pseudomesh.py --cfg_path $CONFIG_PATH --exp_name $exp_name --pseudomesh_path $pseudomesh_path --dset_path $data_path --white_background --res 2 CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 3dgs --white_background CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method 3dgs break done ================================================ FILE: sh_scripts/run_games_gs-flat_sh0_db.sh ================================================ #!/bin/bash DATA_DIR="/path/to/datasets/db" EXP_DIR="/path/to/experiments" GS_TRAIN_PATH="/path/to/gaussian-mesh-splatting" GS_RENDER_PY="/path/to/gaussian-mesh-splatting/scripts/render.py" PYTHON_PATH="/path/to/python" GS_PSEUDO_PATH="/path/to/gs_raytracing" GS_OPTIM_PATH="/path/to/gs_raytracing/mesh_optim" CONFIG_PATH="/path/to/gs_raytracing/sh_scripts/configs/db_config.yaml" export PYTHONPATH=$GS_TRAIN_PATH:$PYTHON_PATH for data_path in $DATA_DIR/*; do dset_name=$(basename ${data_path}) exp_name="db_games_gs-flat_sh0_${dset_name}" exp_out_dir="${EXP_DIR}/${exp_name}" exp_cfg_path="${exp_out_dir}/config.yaml" ply_path="${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply" pseudomesh_path="${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz" cd $GS_TRAIN_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 1 -s $data_path -m $exp_out_dir --gs_type gs_flat --iteration 30000 --sh_degree 0 --white_background --eval CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH $GS_RENDER_PY --resolution 1 -s $data_path -m $exp_out_dir --gs_type gs_flat --iteration 30000 --sh_degree 0 --white_background --skip_train --eval cd $GS_PSEUDO_PATH $PYTHON_PATH generate_pseudomesh.py --ply $ply_path --algorithm games --scale_min 2.7 --scale_max 2.8 --scale_step 0.2 --no_points 8 cd $GS_OPTIM_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH optimize_pseudomesh.py --cfg_path $CONFIG_PATH --exp_name $exp_name --pseudomesh_path $pseudomesh_path --dset_path $data_path --white_background --res 2 CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method games --white_background CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method games done ================================================ FILE: sh_scripts/run_games_gs-flat_sh0_mip.sh ================================================ #!/bin/bash DATA_DIR="/path/to/datasets/360_v2" EXP_DIR="/path/to/experiments" GS_TRAIN_PATH="/path/to/gaussian-mesh-splatting" GS_RENDER_PY="/path/to/gaussian-mesh-splatting/scripts/render.py" PYTHON_PATH="/path/to/python" GS_PSEUDO_PATH="/path/to/gs_raytracing" GS_OPTIM_PATH="/path/to/gs_raytracing/mesh_optim" CONFIG_PATH="/path/to/gs_raytracing/sh_scripts/configs/mip_config.yaml" export PYTHONPATH=$GS_TRAIN_PATH:$PYTHON_PATH for data_path in $DATA_DIR/*; do dset_name=$(basename ${data_path}) exp_name="mip_games_gs-flat_sh0_${dset_name}" exp_out_dir="${EXP_DIR}/${exp_name}" exp_cfg_path="${exp_out_dir}/config.yaml" ply_path="${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply" pseudomesh_path="${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz" cd $GS_TRAIN_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 4 -s $data_path -m $exp_out_dir --gs_type gs_flat --iteration 30000 --sh_degree 0 --white_background --eval CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH $GS_RENDER_PY --resolution 4 -s $data_path -m $exp_out_dir --gs_type gs_flat --iteration 30000 --sh_degree 0 --white_background --skip_train --eval cd $GS_PSEUDO_PATH $PYTHON_PATH generate_pseudomesh.py --ply $ply_path --algorithm games --scale_min 2.7 --scale_max 2.8 --scale_step 0.2 --no_points 8 cd $GS_OPTIM_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH optimize_pseudomesh.py --cfg_path $CONFIG_PATH --exp_name $exp_name --pseudomesh_path $pseudomesh_path --dset_path $data_path --white_background --res 8 CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method games --white_background CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method games done ================================================ FILE: sh_scripts/run_games_gs-flat_sh0_nerf-synth.sh ================================================ #!/bin/bash DATA_DIR="/path/to/datasets/games_set" EXP_DIR="/path/to/experiments" GS_TRAIN_PATH="/path/to/gaussian-mesh-splatting" GS_RENDER_PY="/path/to/gaussian-mesh-splatting/scripts/render.py" PYTHON_PATH="/path/to/python" GS_PSEUDO_PATH="/path/to/gs_raytracing" GS_OPTIM_PATH="/path/to/gs_raytracing/mesh_optim" CONFIG_PATH="/path/to/gs_raytracing/sh_scripts/configs/3dgs_sh0_pseudo_config.yaml" export PYTHONPATH=$GS_TRAIN_PATH:$PYTHON_PATH for data_path in $DATA_DIR/*; do dset_name=$(basename ${data_path}) exp_name="games_gs-flat_white_sh0_${dset_name}" exp_out_dir="${EXP_DIR}/${exp_name}" exp_cfg_path="${exp_out_dir}/config.yaml" ply_path="${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply" pseudomesh_path="${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz" cd $GS_TRAIN_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 1 -s $data_path -m $exp_out_dir --gs_type gs_flat --iteration 30000 --sh_degree 0 --eval CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH $GS_RENDER_PY --resolution 1 -s $data_path -m $exp_out_dir --gs_type gs_flat --iteration 30000 --sh_degree 0 --skip_train --eval cd $GS_PSEUDO_PATH $PYTHON_PATH generate_pseudomesh.py --ply $ply_path --algorithm games --scale_min 2.7 --scale_max 2.8 --scale_step 0.2 --no_points 8 cd $GS_OPTIM_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH optimize_pseudomesh.py --cfg_path $CONFIG_PATH --exp_name $exp_name --pseudomesh_path $pseudomesh_path --dset_path $data_path --res 1 CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method games CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --method games done ================================================ FILE: sh_scripts/run_games_gs-flat_sh0_tandt.sh ================================================ #!/bin/bash DATA_DIR="/path/to/datasets/tandt" EXP_DIR="/path/to/experiments" GS_TRAIN_PATH="/path/to/gaussian-mesh-splatting" GS_RENDER_PY="/path/to/gaussian-mesh-splatting/scripts/render.py" PYTHON_PATH="/path/to/python" GS_PSEUDO_PATH="/path/to/gs_raytracing" GS_OPTIM_PATH="/path/to/gs_raytracing/mesh_optim" CONFIG_PATH="/path/to/gs_raytracing/sh_scripts/configs/tandt_config.yaml" export PYTHONPATH=$GS_TRAIN_PATH:$PYTHON_PATH for data_path in $DATA_DIR/*; do dset_name=$(basename ${data_path}) exp_name="tandt_games_gs-flat_sh0_${dset_name}" exp_out_dir="${EXP_DIR}/${exp_name}" exp_cfg_path="${exp_out_dir}/config.yaml" ply_path="${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply" pseudomesh_path="${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz" cd $GS_TRAIN_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 1 -s $data_path -m $exp_out_dir --gs_type gs_flat --iteration 30000 --sh_degree 0 --white_background --eval CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH $GS_RENDER_PY --resolution 1 -s $data_path -m $exp_out_dir --gs_type gs_flat --iteration 30000 --sh_degree 0 --white_background --skip_train --eval cd $GS_PSEUDO_PATH $PYTHON_PATH generate_pseudomesh.py --ply $ply_path --algorithm games --scale_min 2.7 --scale_max 2.8 --scale_step 0.2 --no_points 8 cd $GS_OPTIM_PATH CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH optimize_pseudomesh.py --cfg_path $CONFIG_PATH --exp_name $exp_name --pseudomesh_path $pseudomesh_path --dset_path $data_path --white_background --res 1 CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method games --white_background CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method games done