[
  {
    "path": "GS_LICENSE.md",
    "content": "Gaussian-Splatting License  \n===========================  \n\n**Inria** and **the Max Planck Institut for Informatik (MPII)** hold all the ownership rights on the *Software* named **gaussian-splatting**.  \nThe *Software* is in the process of being registered with the Agence pour la Protection des  \nProgrammes (APP).  \n\nThe *Software* is still being developed by the *Licensor*.  \n\n*Licensor*'s goal is to allow the research community to use, test and evaluate  \nthe *Software*.  \n\n## 1.  Definitions  \n\n*Licensee* means any person or entity that uses the *Software* and distributes  \nits *Work*.  \n\n*Licensor* means the owners of the *Software*, i.e Inria and MPII  \n\n*Software* means the original work of authorship made available under this  \nLicense ie gaussian-splatting.  \n\n*Work* means the *Software* and any additions to or derivative works of the  \n*Software* that are made available under this License.  \n\n\n## 2.  Purpose  \nThis license is intended to define the rights granted to the *Licensee* by  \nLicensors under the *Software*.  \n\n## 3.  Rights granted  \n\nFor the above reasons Licensors have decided to distribute the *Software*.  \nLicensors grant non-exclusive rights to use the *Software* for research purposes  \nto research users (both academic and industrial), free of charge, without right  \nto sublicense.. The *Software* may be used \"non-commercially\", i.e., for research  \nand/or evaluation purposes only.  \n\nSubject to the terms and conditions of this License, you are granted a  \nnon-exclusive, royalty-free, license to reproduce, prepare derivative works of,  \npublicly display, publicly perform and distribute its *Work* and any resulting  \nderivative works in any form.  \n\n## 4.  Limitations  \n\n**4.1 Redistribution.** You may reproduce or distribute the *Work* only if (a) you do  \nso under this License, (b) you include a complete copy of this License with  \nyour distribution, and (c) you retain without modification any copyright,  \npatent, trademark, or attribution notices that are present in the *Work*.  \n\n**4.2 Derivative Works.** You may specify that additional or different terms apply  \nto the use, reproduction, and distribution of your derivative works of the *Work*  \n(\"Your Terms\") only if (a) Your Terms provide that the use limitation in  \nSection 2 applies to your derivative works, and (b) you identify the specific  \nderivative works that are subject to Your Terms. Notwithstanding Your Terms,  \nthis License (including the redistribution requirements in Section 3.1) will  \ncontinue to apply to the *Work* itself.  \n\n**4.3** Any other use without of prior consent of Licensors is prohibited. Research  \nusers explicitly acknowledge having received from Licensors all information  \nallowing to appreciate the adequacy between of the *Software* and their needs and  \nto undertake all necessary precautions for its execution and use.  \n\n**4.4** The *Software* is provided both as a compiled library file and as source  \ncode. In case of using the *Software* for a publication or other results obtained  \nthrough the use of the *Software*, users are strongly encouraged to cite the  \ncorresponding publications as explained in the documentation of the *Software*.  \n\n## 5.  Disclaimer  \n\nTHE USER CANNOT USE, EXPLOIT OR DISTRIBUTE THE *SOFTWARE* FOR COMMERCIAL PURPOSES  \nWITHOUT PRIOR AND EXPLICIT CONSENT OF LICENSORS. YOU MUST CONTACT INRIA FOR ANY  \nUNAUTHORIZED USE: stip-sophia.transfert@inria.fr . ANY SUCH ACTION WILL  \nCONSTITUTE A FORGERY. THIS *SOFTWARE* IS PROVIDED \"AS IS\" WITHOUT ANY WARRANTIES  \nOF ANY NATURE AND ANY EXPRESS OR IMPLIED WARRANTIES, WITH REGARDS TO COMMERCIAL  \nUSE, PROFESSIONNAL USE, LEGAL OR NOT, OR OTHER, OR COMMERCIALISATION OR  \nADAPTATION. UNLESS EXPLICITLY PROVIDED BY LAW, IN NO EVENT, SHALL INRIA OR THE  \nAUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR  \nCONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE  \nGOODS OR SERVICES, LOSS OF USE, DATA, OR PROFITS OR BUSINESS INTERRUPTION)  \nHOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT  \nLIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING FROM, OUT OF OR  \nIN CONNECTION WITH THE *SOFTWARE* OR THE USE OR OTHER DEALINGS IN THE *SOFTWARE*.  \n\n## 6.  Files subject to permissive licenses\nThe 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. \n\nTitle: pytorch-ssim\\\nProject code: https://github.com/Po-Hsun-Su/pytorch-ssim\\\nCopyright Evan Su, 2017\\\nLicense: https://github.com/Po-Hsun-Su/pytorch-ssim/blob/master/LICENSE.txt (MIT)"
  },
  {
    "path": "LICENSE.md",
    "content": "MIT License\n\nCopyright (c) 2025 Rafał Tobiasz Grzegorz Wilczyński\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
  },
  {
    "path": "NV_LICENSE.txt",
    "content": "Copyright (c) 2020, NVIDIA Corporation. All rights reserved.\n\n\nNvidia Source Code License (1-Way Commercial)\n\n=======================================================================\n\n1. Definitions\n\n\"Licensor\" means any person or entity that distributes its Work.\n\n\"Software\" means the original work of authorship made available under\nthis License.\n\n\"Work\" means the Software and any additions to or derivative works of\nthe Software that are made available under this License.\n\nThe terms \"reproduce,\" \"reproduction,\" \"derivative works,\" and\n\"distribution\" have the meaning as provided under U.S. copyright law;\nprovided, however, that for the purposes of this License, derivative\nworks shall not include works that remain separable from, or merely\nlink (or bind by name) to the interfaces of, the Work.\n\nWorks, including the Software, are \"made available\" under this License\nby including in or with the Work either (a) a copyright notice\nreferencing the applicability of this License to the Work, or (b) a\ncopy of this License.\n\n2. License Grants\n\n    2.1 Copyright Grant. Subject to the terms and conditions of this\n    License, each Licensor grants to you a perpetual, worldwide,\n    non-exclusive, royalty-free, copyright license to reproduce,\n    prepare derivative works of, publicly display, publicly perform,\n    sublicense and distribute its Work and any resulting derivative\n    works in any form.\n\n3. Limitations\n\n    3.1 Redistribution. You may reproduce or distribute the Work only\n    if (a) you do so under this License, (b) you include a complete\n    copy of this License with your distribution, and (c) you retain\n    without modification any copyright, patent, trademark, or\n    attribution notices that are present in the Work.\n\n    3.2 Derivative Works. You may specify that additional or different\n    terms apply to the use, reproduction, and distribution of your\n    derivative works of the Work (\"Your Terms\") only if (a) Your Terms\n    provide that the use limitation in Section 3.3 applies to your\n    derivative works, and (b) you identify the specific derivative\n    works that are subject to Your Terms. Notwithstanding Your Terms,\n    this License (including the redistribution requirements in Section\n    3.1) will continue to apply to the Work itself.\n\n    3.3 Use Limitation. The Work and any derivative works thereof only\n    may be used or intended for use non-commercially. The Work or\n    derivative works thereof may be used or intended for use by Nvidia\n    or its affiliates commercially or non-commercially. As used herein,\n    \"non-commercially\" means for research or evaluation purposes only\n    and not for any direct or indirect monetary gain.\n\n    3.4 Patent Claims. If you bring or threaten to bring a patent claim\n    against any Licensor (including any claim, cross-claim or\n    counterclaim in a lawsuit) to enforce any patents that you allege\n    are infringed by any Work, then your rights under this License from\n    such Licensor (including the grant in Section 2.1) will terminate\n    immediately.\n\n    3.5 Trademarks. This License does not grant any rights to use any\n    Licensor's or its affiliates' names, logos, or trademarks, except\n    as necessary to reproduce the notices described in this License.\n\n    3.6 Termination. If you violate any term of this License, then your\n    rights under this License (including the grant in Section 2.1) will\n    terminate immediately.\n\n4. Disclaimer of Warranty.\n\nTHE WORK IS PROVIDED \"AS IS\" WITHOUT WARRANTIES OR CONDITIONS OF ANY\nKIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF\nMERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR\nNON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER\nTHIS LICENSE.\n\n5. Limitation of Liability.\n\nEXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL\nTHEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE\nSHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,\nINDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF\nOR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK\n(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,\nLOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER\nCOMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF\nTHE POSSIBILITY OF SUCH DAMAGES.\n\n======================================================================="
  },
  {
    "path": "README.md",
    "content": "# MeshSplats\nRafał Tobiasz*, Grzegorz Wilczyński*, Marcin Mazur, Sławomir Tadeja, Przemysław Spurek\n(* indicates equal contribution)<br>\n\nThis repository contains the official authors implementation associated \nwith the paper [\"MeshSplats: Mesh-Based Rendering with Gaussian Splatting Initialization\"](https://arxiv.org/pdf/2502.07754).\n\nAbstract: *\nRecently, 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*\n\nCheck out this code if you just want to convert GS to a mesh and benefit from the advantages of both representations!\n\n<div align=\"center\">\n    <img src=\"./demo/vis_1.gif\" style=\"height:200px; width:auto;\">\n    <img src=\"./demo/vis_bicycle_2.gif\" style=\"height:200px; width:auto;\">\n</div>\n\nNote: If videos aren't visible you can find them in `demo` directory.\n\n## Run our demo in Google Colab!!!\n\nWe have prepared for you a demo (to run as quickly as possible) of the MeshSplats method.\n\n### Optimization pipeline demo\n[Here's the Colab!](https://colab.research.google.com/drive/1N6Y6TfijUw9tQ8vrnoLrAhBPlxkSzA2b?usp=sharing)\n\nFirst 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).\n\nThis data structure is as follows:\n\n```\nexamples/\n|- lego_white_background/\n|  |- checkpoints/\n|  |  |- best_model\n|  |  |- best_model.npz\n|  |- point_cloud/\n|  |  |- iteration_30000/\n|  |  |- point_cloud.ply\n|  |- pseudomeshes/\n|  |- scene_2.70_pts_8.npz\n|  |- cameras.json\n|  |- config.yaml\n|  |- colab_config.yaml\n|- lego.zip\n```\nWhere:\n- `checkpoints` contains the best model from our optimization pipeline as torch checkpoint and numpy checkpoint,\n- `point_cloud` contains the point cloud of the scene from the 30000 iteration of the GaMeS algorithm,\n- `pseudomeshes` contains the raw pseudomesh of the scene generated from the available `point_cloud.ply` file,\n- `cameras.json` contains the camera poses,\n- `config.yaml` contains the config of the experiment.  \n- `colab_config.yaml` contains the config of the experiment for the colab demo,\n- `lego.zip` contains the data of the `lego` scene from the `NeRF Synthetic` dataset.\n## Installation\n\n```bash\npip install -r requirements.txt\n```\n\nAlso this work depends on output of the following repositories:\n\n- [Gaussian Mesh Splatting](https://github.com/waczjoan/gaussian-mesh-splatting)\n- [3D Gaussian Splatting](https://github.com/graphdeco-inria/gaussian-splatting)\n- [2D Gaussian Splatting](https://github.com/hbb1/2d-gaussian-splatting)\n\nTherefore, you need to install them first.\n\n## Usage\n\nYou 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).\n\nEach script is designed to run on a single GPU.\n\nFor example, to run the experiments for the DeepBlending dataset with the GaMeS algorithm, you can use the following command:\n\n```bash\n./sh_scripts/run_games_gs-flat_sh0_db.sh\n```\n\n## Datasets\nThis repository is prepared to work with the following datasets (as you can see in the `sh_scripts` scripts):\n\n- NeRF-Synthetic\n- Tanks and Temples\n- MiP NeRF\n- DeepBlending\n\n<section class=\"section\" id=\"BibTeX\">\n  <div class=\"container is-max-desktop content\">\n    <h2 class=\"title\">BibTeX</h2>\n<h3 class=\"title\">MeshSplats</h3>\n    <pre><code>@Article{tobiasz2025meshsplats,\n        author={Rafał Tobiasz and Grzegorz Wilczyński and Marcin Mazur and Sławomir Tadeja and Przemysław Spurek},\n        year={2025},\n        eprint={2502.07754},\n        archivePrefix={arXiv},\n        primaryClass={cs.GR},\n}\n</code></pre>\n    <h3 class=\"title\">Gaussian Splatting</h3>\n    <pre><code>@Article{kerbl3Dgaussians,\n      author         = {Kerbl, Bernhard and Kopanas, Georgios and Leimk{\\\"u}hler, Thomas and Drettakis, George},\n      title          = {3D Gaussian Splatting for Real-Time Radiance Field Rendering},\n      journal        = {ACM Transactions on Graphics},\n      number         = {4},\n      volume         = {42},\n      month          = {July},\n      year           = {2023},\n      url            = {https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/}\n}</code></pre>\n</code></pre>\n    <h3 class=\"title\">NVDiffrast</h3>\n    <pre><code>@Article{laine2020modularprimitiveshighperformancedifferentiabl,\n      author         = {Kerbl, Bernhard and Kopanas, Georgios and Leimk{\\\"u}hler, Thomas and Drettakis, George},\n      title={Modular Primitives for High-Performance Differentiable Rendering}, \n      author={Samuli Laine and Janne Hellsten and Tero Karras and Yeongho Seol and Jaakko Lehtinen and Timo Aila},\n      year={2020},\n      eprint={2011.03277},\n      archivePrefix={arXiv},\n      primaryClass={cs.GR},\n}</code></pre>\n  </div>\n</section>\n"
  },
  {
    "path": "faces_notexturemap_blender.py",
    "content": "import json\nfrom math import tan, atan\nfrom pathlib import Path\nimport sys\n\nimport bpy\nimport numpy as np\n\nfrom mathutils import Matrix\n\ndef load_npy(_path):\n    _data = np.load(_path)\n    return _data\n\ndef _calculate_fov(focal, pixels):\n    return 2 * np.arctan(pixels / (2 * focal))\n\ndef setup_ambient_light():\n    if bpy.context.scene.world is None:\n        bpy.context.scene.world = bpy.data.worlds.new(\"World\")\n\n    world = bpy.context.scene.world\n    world.use_nodes = True\n    nodes = world.node_tree.nodes\n    links = world.node_tree.links\n\n    nodes.clear()\n\n    background_node = nodes.new(type='ShaderNodeBackground')\n    background_node.inputs['Color'].default_value = (1.0, 1.0, 1.0, 1.0)  # White color for ambient light\n    background_node.inputs['Strength'].default_value = 1.0  # Adjust the strength as needed\n\n    world_output_node = nodes.new(type='ShaderNodeOutputWorld')\n\n    links.new(background_node.outputs['Background'], world_output_node.inputs['Surface'])\n\n    bpy.context.view_layer.update()\n\n\ndef create_camera(name, position, rot_matrix, fovx, fovy, image_width, image_height):\n    # Create the camera object\n    bpy.ops.object.camera_add()\n    camera_obj = bpy.context.object\n    camera_obj.name = str(name)\n    camera_obj.location = position\n        \n    # Set rotation from matrix\n    rotation_matrix = Matrix(rot_matrix)\n    camera_obj.rotation_euler = rotation_matrix.to_euler()\n    \n    # Assuming the sensor width is fixed and calculating the sensor height based on the aspect ratio\n    sensor_width = 36.0  # Adjust as needed, standard full frame sensor width\n    aspect_ratio = image_width / image_height\n    sensor_height = sensor_width / aspect_ratio\n    \n    # Set camera data\n    camera = camera_obj.data\n    camera.sensor_width = sensor_width\n    camera.sensor_height = sensor_height\n    \n    # Calculate focal length from FoV using formula: focal_length = sensor_width / (2 * tan(fov / 2))\n    camera.lens = sensor_width / (2 * tan(fovx / 2))\n    \n    # Use FoVy to adjust the sensor height if needed\n    calculated_fovy = 2 * atan((sensor_height / 2) / camera.lens)\n    if calculated_fovy != fovy:\n        camera.sensor_height = 2 * camera.lens * tan(fovy / 2)\n    \n    return camera_obj\n\n\ndef gen_faces_from_texture_map(_path, _cam_json_path=None, _wanted_cam_idx=None):\n    if _wanted_cam_idx is None or _cam_json_path is None:\n        _wanted_cam_idx = []\n    elif isinstance(_wanted_cam_idx, int):\n        _wanted_cam_idx = [_wanted_cam_idx]\n    assert isinstance(_wanted_cam_idx, list)\n    \n    if _cam_json_path is not None:\n        with open(_cam_json_path, \"r\") as file:\n            _camera_data = json.load(file)\n        _camera_data = sorted(_camera_data, key=lambda x: x[\"id\"])\n        _camera_data = [{\n            \"id\": x[\"id\"],\n            \"img_name\": x[\"img_name\"],\n            \"width\": x[\"width\"],\n            \"height\": x[\"height\"],\n            \"position\": np.array(x[\"position\"]),\n            \"rotation\": np.array(x[\"rotation\"]),\n            \"fovy\": x[\"fy\"],\n            \"fovx\": x[\"fx\"]\n        } for x in _camera_data]\n        \n        if _wanted_cam_idx:\n            _new_cameras = [_camera_data[i] for i in _wanted_cam_idx]\n        else:\n            _new_cameras = _camera_data\n\n    if bpy.context.active_object.mode != 'OBJECT':\n        bpy.ops.object.mode_set(mode='OBJECT')\n    \n    bpy.ops.object.select_all(action='DESELECT')\n    bpy.ops.object.select_by_type(type='MESH')\n    bpy.ops.object.delete()\n    \n    mesh_data = load_npy(_path)\n    vertices = mesh_data[\"vertices\"]\n    faces = mesh_data[\"faces\"]\n    vert_colors = mesh_data[\"vertex_colors\"]\n    \n    mesh = bpy.data.meshes.new(name=\"MyMesh\")\n    splat_obj = bpy.data.objects.new(name=\"MyObject\", object_data=mesh)\n\n    # Link object to collection\n    bpy.context.collection.objects.link(splat_obj)\n\n    # Create mesh from vertices and faces\n    mesh.from_pydata(vertices.tolist(), [], faces.tolist())\n    mesh.update()\n\n    # Deselect all objects\n    bpy.ops.object.select_all(action='DESELECT')\n\n    # Select the newly created object\n    splat_obj.select_set(True)\n\n    # Set the newly created object as the active object\n    bpy.context.view_layer.objects.active = splat_obj\n\n    if not splat_obj.data.vertex_colors:\n        splat_obj.data.vertex_colors.new()\n\n    vertex_color_layer = splat_obj.data.vertex_colors.active\n\n    for poly in splat_obj.data.polygons:\n        for loop_index in range(poly.loop_start, poly.loop_start + poly.loop_total):\n            loop = splat_obj.data.loops[loop_index]\n            vertex_index = loop.vertex_index\n            vertex_color = vert_colors[vertex_index].tolist()\n            \n            # Assign color and opacity (RGBA) to each vertex\n            vertex_color_layer.data[loop_index].color = vertex_color\n\n    # Create a new material\n    material = bpy.data.materials.new(name=\"VertexColorMaterial\")\n    material.use_nodes = True\n\n    # Enable transparency\n    material.blend_method = 'HASHED'\n    material.shadow_method = \"NONE\"\n    material.use_backface_culling = False\n    material.alpha_threshold = 0.01\n\n    # Get the material's node tree\n    nodes = material.node_tree.nodes\n    links = material.node_tree.links\n\n    # Clear all nodes\n    for node in nodes:\n        nodes.remove(node)\n\n    # Add a Vertex Color node\n    vertex_color_node = nodes.new(type='ShaderNodeVertexColor')\n    vertex_color_node.location = (-300, 300)\n    vertex_color_node.layer_name = vertex_color_layer.name\n\n    # Add a Principled BSDF node\n    bsdf_node = nodes.new(type='ShaderNodeBsdfPrincipled')\n    bsdf_node.location = (100, 300)\n    # Set metallic and specular to 0 for flat shading\n    bsdf_node.inputs['Metallic'].default_value = 0.0\n    bsdf_node.inputs['Specular'].default_value = 0.0\n    bsdf_node.inputs['Roughness'].default_value = 1.0\n\n    # Connect the vertex color to the BSDF node\n    links.new(vertex_color_node.outputs['Color'], bsdf_node.inputs['Base Color'])\n    links.new(vertex_color_node.outputs['Alpha'], bsdf_node.inputs['Alpha'])\n\n    # Add an Output node\n    output_node = nodes.new(type='ShaderNodeOutputMaterial')\n    output_node.location = (300, 300)\n\n    # Link the BSDF node to the output node\n    links.new(bsdf_node.outputs['BSDF'], output_node.inputs['Surface'])\n\n    # Assign the material to the cube\n    if splat_obj.type == 'MESH':\n        if len(splat_obj.data.materials) == 0:\n            splat_obj.data.materials.append(material)\n        else:\n            splat_obj.data.materials[0] = material\n\n    # Update the mesh\n    splat_obj.data.update()\n\n    # Set the render engine to Eevee\n    bpy.context.scene.render.engine = 'BLENDER_EEVEE'\n    \n    # Configure Eevee settings for better quality\n    bpy.context.scene.eevee.taa_render_samples = 64  # Increase samples for better quality\n    bpy.context.scene.eevee.use_soft_shadows = True\n    bpy.context.scene.eevee.use_ssr = True  # Screen Space Reflections\n    bpy.context.scene.eevee.use_ssr_refraction = True\n    bpy.context.scene.eevee.use_taa_reprojection = True\n    \n    # Enable transparency settings in Eevee\n    bpy.context.scene.eevee.use_bloom = False\n    bpy.context.scene.eevee.use_ssr_halfres = False\n    bpy.context.scene.eevee.volumetric_tile_size = '2'\n\n    setup_ambient_light()\n    \n    # Set render settings\n    bpy.context.scene.cycles.samples = 128  # Set the number of samples for rendering\n\n    _cameras = []\n    \n    adjustment = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]], dtype=np.float64)\n    for _new_cam in _new_cameras:\n        fovx = _calculate_fov(_new_cam[\"fovx\"], _new_cam[\"width\"])\n        fovy = _calculate_fov(_new_cam[\"fovy\"], _new_cam[\"height\"])\n        _new_cam_obj = create_camera(\n            name=_new_cam[\"img_name\"],\n            position=_new_cam[\"position\"],\n            rot_matrix=_new_cam[\"rotation\"] @ adjustment,\n            fovx=fovx,\n            fovy=fovy,\n            image_width=_new_cam[\"width\"],\n            image_height=_new_cam[\"height\"]\n        )\n        _cameras.append(_new_cam_obj)\n    \n    scale = Path(_path).stem.split(\"_\")[1]\n    out_dir = Path(Path(_path).parent, \"images\", scale)\n    out_dir.mkdir(exist_ok=True, parents=True)\n    \n    # setup output img size\n    bpy.context.scene.render.resolution_x = _new_cameras[0][\"width\"]\n    bpy.context.scene.render.resolution_y = _new_cameras[0][\"height\"]\n    \n    for _cam_obj in _cameras:\n        output_path = str(Path(out_dir, f\"{_cam_obj.name}.png\"))\n        # Set camera as active camera\n        bpy.context.scene.camera = _cam_obj\n\n        # Set output path for rendered image\n        bpy.context.scene.render.filepath = output_path\n\n        # Render the image\n        bpy.ops.render.render(write_still=True)\n        print(f\"Rendered image saved to {output_path}\")\n\nargs = sys.argv\n\nargs_start_index = args.index(\"--\") + 1\nscript_args = args[args_start_index:]\n\nassert \"--npz_path\" in script_args\nassert \"--cam_path\" in script_args\n\n_npz_path = script_args[script_args.index(\"--npz_path\") + 1]\n_cam_path = script_args[script_args.index(\"--cam_path\") + 1]\n\ngen_faces_from_texture_map(\n    _npz_path,\n    _cam_path,\n)\n"
  },
  {
    "path": "generate_pseudomesh.py",
    "content": "import argparse\nfrom pathlib import Path\n\nimport numpy as np\n\nimport gs_utils\n\ndef main(ply_path: Path, algorithm: str, scale_muls: np.ndarray, no_of_points: int):\n    xyz, feat_dc, feat_res, opac, scales, rots_quaternion = gs_utils.load_ply(\n        ply_path, 0\n    )\n    if algorithm == \"games\":\n        games_ckpt_path = Path(ply_path.parent, \"model_params.pt\")\n        try:\n            games_data = gs_utils.load_games_pt(games_ckpt_path)\n        except FileNotFoundError:\n            games_data = None\n\n    if algorithm == \"2dgs\":\n        pseudomeshes = gs_utils.generate_pseudomesh_2dgs(\n            xyz=xyz,\n            features_dc=feat_dc,\n            opacities=opac,\n            scales=scales,\n            rots=rots_quaternion,\n            scale_muls=scale_muls,\n            no_of_points=no_of_points\n        )\n    elif algorithm == \"surfels\":\n        pseudomeshes = gs_utils.generate_pseudomesh_surfels(\n            xyz=xyz,\n            features_dc=feat_dc,\n            opacities=opac,\n            scales=scales,\n            rots=rots_quaternion,\n            scale_muls=scale_muls,\n            no_of_points=no_of_points\n        )\n    elif algorithm == \"games\":\n        pseudomeshes = gs_utils.generate_pseudomesh_games(\n            ckpt_data=games_data,\n            xyz=xyz,\n            features_dc=feat_dc,\n            opacities=opac,\n            scales=scales,\n            rots=rots_quaternion,\n            scale_muls=scale_muls,\n            no_of_points=no_of_points\n        )\n    elif algorithm == \"sugar2d\":\n        pseudomeshes = gs_utils.generate_pseudomesh_sugar_2d(\n            xyz=xyz,\n            features_dc=feat_dc,\n            opacities=opac,\n            scales=scales,\n            rots=rots_quaternion,\n            scale_muls=scale_muls,\n            no_of_points=no_of_points\n        )\n    elif algorithm == \"sugar3d\":\n        pseudomeshes = gs_utils.generate_pseudomesh_sugar_3d(\n            xyz=xyz,\n            features_dc=feat_dc,\n            opacities=opac,\n            scales=scales,\n            rots=rots_quaternion,\n            scale_muls=scale_muls,\n            no_of_points=no_of_points\n        )\n    elif algorithm == \"3dgs\":\n        pseudomeshes = gs_utils.generate_3dgs_pseudomesh(\n            xyz=xyz,\n            features_dc=feat_dc,\n            opacities=opac,\n            scales=scales,\n            rots=rots_quaternion,\n            scale_muls=scale_muls,\n            no_of_points=no_of_points\n        )\n    else:\n        raise NotImplementedError(f\"Algorithm {algorithm} not yet implemented\")\n    \n    # save data\n    if \"sugar\" in algorithm:\n        main_out_dir = ply_path.parent\n    else:\n        main_out_dir = ply_path.parent.parent.parent\n    out_dir = Path(main_out_dir, \"pseudomeshes\")\n    out_dir.mkdir(exist_ok=True, parents=True)\n    \n    for scale_val, pseudomesh_data in pseudomeshes.items():\n        scale_val_str = \"{:.2f}\".format(scale_val)\n        out_path = Path(out_dir, f\"scale_{scale_val_str}_pts_{str(no_of_points)}.npz\")\n        np.savez(\n            out_path,\n            **pseudomesh_data\n        )\n        print(f\"Saved {out_path}\")\n    \ndef read_args():\n    parser = argparse.ArgumentParser(description=\"This script accepts ply file and outputs pseudomesh of it\")\n    parser.add_argument(\n        \"--ply\",\n        required=True,\n        type=str,\n        help=\"Path to the ply file\"\n    )\n    parser.add_argument(\n        \"--algorithm\",\n        default=\"3dgs\",\n        type=str,\n        choices=[\"2dgs\", \"games\", \"surfels\", \"sugar2d\", \"sugar3d\", \"3dgs\"]\n    )\n    parser.add_argument(\n        \"--scale_min\",\n        default=2.3,\n        type=float\n    )\n    parser.add_argument(\n        \"--scale_max\",\n        default=2.4,\n        type=float\n    )\n    parser.add_argument(\n        \"--scale_step\",\n        default=0.2,\n        type=float\n    )\n    parser.add_argument(\n        \"--no_points\",\n        default=8,\n        type=int\n    )\n    args = parser.parse_args()\n    \n    return {\n        \"ply_path\": Path(args.ply),\n        \"algorithm\": args.algorithm,\n        \"scale_muls\": np.arange(args.scale_min, args.scale_max, args.scale_step),\n        \"no_of_points\": args.no_points\n    }\n\nif __name__ == \"__main__\":\n    main(**read_args())\n"
  },
  {
    "path": "gs_utils.py",
    "content": "import numpy as np\nfrom plyfile import PlyData\nfrom scipy.special import expit as sigmoid\nimport torch\n\ntry:\n    import sys\n    sys.path.append(\"/path/to/GaMeS_repo\")\n    import games\nexcept Exception as e:\n    print(e)\n\nC0 = 0.28209479177387814\n\n\ndef load_games_pt(path):\n    params = torch.load(path)\n    alpha = params['_alpha']\n    scale = params['_scale']\n    vertices, triangles, faces = None, None, None\n    if 'vertices' in params:\n        vertices = params['vertices']\n    if 'triangles' in params:\n        triangles = params['triangles']\n    if 'faces' in params:\n        faces = params['faces']\n    return {\n        \"alpha\": alpha,\n        \"scale\": scale,\n        \"vertices\": vertices,\n        \"triangles\": triangles,\n        \"faces\": faces\n    }\n\n\ndef get_games_scales_and_rots(data, eps=1e-8):\n    def dot(v, u):\n        return (v * u).sum(dim=-1, keepdim=True)\n    \n    def proj(v, u):\n        \"\"\"\n        projection of vector v onto subspace spanned by u\n\n        vector u is assumed to be already normalized\n        \"\"\"\n        coef = dot(v, u)\n        return coef * u\n    \n    triangles = data[\"triangles\"]\n    normals = torch.linalg.cross(\n        triangles[:, 1] - triangles[:, 0],\n        triangles[:, 2] - triangles[:, 0],\n        dim=1\n    )\n    v0 = normals / (torch.linalg.vector_norm(normals, dim=-1, keepdim=True) + eps)\n    means = torch.mean(triangles, dim=1)\n    v1 = triangles[:, 1] - means\n    v1_norm = torch.linalg.vector_norm(v1, dim=-1, keepdim=True) + eps\n    v1 = v1 / v1_norm\n    v2_init = triangles[:, 2] - means\n    v2 = v2_init - proj(v2_init, v0) - proj(v2_init, v1)  # Gram-Schmidt\n    v2 = v2 / (torch.linalg.vector_norm(v2, dim=-1, keepdim=True) + eps)\n\n    s1 = v1_norm / 2.\n    s2 = dot(v2_init, v2) / 2.\n    \n    scales = torch.concat([s1, s2], dim=1).unsqueeze(dim=1)\n    scales = scales.broadcast_to((*data[\"alpha\"].shape[:2], 2))\n    scales = torch.nn.functional.relu(data[\"scale\"] * scales.flatten(start_dim=0, end_dim=1)) + eps\n    # scales [N, 2]\n    \n    rots = torch.stack((v1, v2), dim=1).unsqueeze(dim=1)\n    rots = rots.broadcast_to((*data[\"alpha\"].shape[:2], 2, 3)).flatten(start_dim=0, end_dim=1)\n    # rots [N, 2, 3]\n    \n    return scales, rots\n\n\ndef load_ply(path, max_sh_degree):\n    plydata = PlyData.read(path)\n\n    xyz = np.stack((np.asarray(plydata.elements[0][\"x\"]),\n                    np.asarray(plydata.elements[0][\"y\"]),\n                    np.asarray(plydata.elements[0][\"z\"])),  axis=1)\n    opacities = np.asarray(plydata.elements[0][\"opacity\"])[..., np.newaxis]\n\n    features_dc = np.zeros((xyz.shape[0], 3, 1))\n    features_dc[:, 0, 0] = np.asarray(plydata.elements[0][\"f_dc_0\"])\n    features_dc[:, 1, 0] = np.asarray(plydata.elements[0][\"f_dc_1\"])\n    features_dc[:, 2, 0] = np.asarray(plydata.elements[0][\"f_dc_2\"])\n\n    extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith(\"f_rest_\")]\n    extra_f_names = sorted(extra_f_names, key = lambda x: int(x.split('_')[-1]))\n    assert len(extra_f_names)==3*(max_sh_degree + 1) ** 2 - 3\n    features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))\n    for idx, attr_name in enumerate(extra_f_names):\n        features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])\n    # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)\n    features_extra = features_extra.reshape((features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1))\n\n    scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith(\"scale_\")]\n    scale_names = sorted(scale_names, key = lambda x: int(x.split('_')[-1]))\n    scales = np.zeros((xyz.shape[0], len(scale_names)))\n    for idx, attr_name in enumerate(scale_names):\n        scales[:, idx] = np.asarray(plydata.elements[0][attr_name])\n\n    rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith(\"rot\")]\n    rot_names = sorted(rot_names, key = lambda x: int(x.split('_')[-1]))\n    rots = np.zeros((xyz.shape[0], len(rot_names)))\n    for idx, attr_name in enumerate(rot_names):\n        rots[:, idx] = np.asarray(plydata.elements[0][attr_name])\n\n    features_dc = np.transpose(features_dc, [0, 2, 1])  # TODO check if this is ok\n    features_rest = np.transpose(features_extra, [0, 2, 1])  # TODO check if this is ok\n\n    return xyz, features_dc, features_rest, opacities, scales, rots\n\n\ndef normalize_rots(mat):\n    return mat / np.linalg.norm(mat, axis=1, keepdims=True)\n\ndef build_euler_rotation(r):\n    norm = np.sqrt(r[:,0]*r[:,0] + r[:,1]*r[:,1] + r[:,2]*r[:,2] + r[:,3]*r[:,3])\n\n    q = r / norm[:, None]\n\n    R = np.zeros((q.shape[0], 3, 3), dtype=np.float64)\n\n    r = q[:, 0]\n    x = q[:, 1]\n    y = q[:, 2]\n    z = q[:, 3]\n\n    R[:, 0, 0] = 1 - 2 * (y*y + z*z)\n    R[:, 0, 1] = 2 * (x*y - r*z)\n    R[:, 0, 2] = 2 * (x*z + r*y)\n    R[:, 1, 0] = 2 * (x*y + r*z)\n    R[:, 1, 1] = 1 - 2 * (x*x + z*z)\n    R[:, 1, 2] = 2 * (y*z - r*x)\n    R[:, 2, 0] = 2 * (x*z - r*y)\n    R[:, 2, 1] = 2 * (y*z + r*x)\n    R[:, 2, 2] = 1 - 2 * (x*x + y*y)\n    return R\n\n\ndef get_scaling(scales: np.ndarray) -> np.ndarray:\n    return np.exp(scales)\n\ndef get_opacity(opacities: np.ndarray) -> np.ndarray:\n    return sigmoid(opacities)\n\n\ndef 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):\n    if ckpt_data is not None:\n        scale, rotation = get_games_scales_and_rots(ckpt_data)\n        scale = scale.cpu().detach().numpy()\n        rotation = rotation.cpu().detach().numpy()\n    else:\n        scale = get_scaling(scales)[:, :2]\n        rotation = np.transpose(build_euler_rotation(rots)[:, :, 1:], [0, 2, 1])\n    init_colors = get_rgb_colors(features_dc)\n    init_opacities = get_opacity(opacities)\n    origin = xyz\n    \n    output = gen_2d_pseudomesh(scale_muls, no_of_points, init_colors, init_opacities, scale, rotation, origin)\n    return output\n\n\ndef 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):\n    init_colors = get_rgb_colors(features_dc)\n    init_opacities = get_opacity(opacities)\n    scale = get_scaling(scales)[:, :2]\n    rotation = np.transpose(build_euler_rotation(rots)[:, :, :2], [0, 2, 1])\n    origin = xyz\n    \n    output = gen_2d_pseudomesh(scale_muls, no_of_points, init_colors, init_opacities, scale, rotation, origin)\n    return output\n\n\ndef 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):\n    init_colors = get_rgb_colors(features_dc)\n    init_opacities = get_opacity(opacities)\n    scale = get_scaling(scales)\n    rotation = np.transpose(build_euler_rotation(rots)[:, :, :2], [0, 2, 1])\n    origin = xyz\n    \n    output = gen_2d_pseudomesh(scale_muls, no_of_points, init_colors, init_opacities, scale, rotation, origin)\n    return output\n\ndef 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):\n    init_colors = get_rgb_colors(features_dc)\n    init_opacities = get_opacity(opacities)\n    scale = get_scaling(scales)\n    rotation = np.transpose(build_euler_rotation(rots), [0, 2, 1])\n    origin = xyz\n    \n    scale = scale[:, 1:]\n    rotation = rotation[:, 1:, :]\n    \n    output = gen_2d_pseudomesh(scale_muls, no_of_points, init_colors, init_opacities, scale, rotation, origin)\n    return output\n\ndef 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):\n    assert no_of_points == 8\n    # handle background generation\n    init_colors = get_rgb_colors(features_dc)\n    init_opacities = get_opacity(opacities)\n    scale = get_scaling(scales)\n    rotation = np.transpose(build_euler_rotation(rots), [0, 2, 1])\n    origin = xyz\n    \n    scales_1 = scale[:, [0, 1]]\n    rots_1 = rotation[:, [0, 1], :]\n    \n    scales_2 = scale[:, [1, 2]]\n    rots_2 = rotation[:, [1, 2], :]\n    \n    scales_3 = scale[:, [0, 2]]\n    rots_3 = rotation[:, [0, 2], :]\n    \n    output = {}\n    \n    for scale_mul in scale_muls:\n        new_points_1 = _get_vertices(origin, scales_1, rots_1, scale_mul, no_of_points)\n        new_points_2 = _get_vertices(origin, scales_2, rots_2, scale_mul, no_of_points)\n        new_points_3 = _get_vertices(origin, scales_3, rots_3, scale_mul, no_of_points)\n        \n        # remove redundant verts\n        new_points_2 = new_points_2[:, [1, 2, 3, 5, 6, 7], :]\n        new_points_3 = new_points_3[:, [1, 3, 5, 7], :]\n        \n        vertices = np.vstack([origin, *[new_points_1[:, _idx, :] for _idx in range(new_points_1.shape[1])]])\n        vertices = np.vstack([vertices, *[new_points_2[:, _idx, :] for _idx in range(new_points_2.shape[1])]])\n        vertices = np.vstack([vertices, *[new_points_3[:, _idx, :] for _idx in range(new_points_3.shape[1])]])\n        \n        origin_idxs = np.arange(0, origin.shape[0]).reshape(-1, 1)\n        indexes_1 = np.hstack([\n            np.arange(_idx * origin.shape[0], (_idx + 1) * origin.shape[0]).reshape(-1, 1) for _idx in range(1, 9)\n        ])\n        indexes_2 = np.hstack([\n            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]\n        ])\n        indexes_3 = np.hstack([\n            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]\n        ])\n        \n        # create faces\n        all_faces = []\n        \n        # faces 1\n        for face_idx in range(no_of_points):\n            last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0\n            all_faces.append(\n                np.hstack([origin_idxs, indexes_1[:, face_idx].reshape(-1, 1), indexes_1[:, last_idx].reshape(-1, 1)])\n            )\n        \n        # faces 2\n        for face_idx in range(no_of_points):\n            last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0\n            all_faces.append(\n                np.hstack([origin_idxs, indexes_2[:, face_idx].reshape(-1, 1), indexes_2[:, last_idx].reshape(-1, 1)])\n            )\n\n        # faces 3\n        for face_idx in range(no_of_points):\n            last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0\n            all_faces.append(\n                np.hstack([origin_idxs, indexes_3[:, face_idx].reshape(-1, 1), indexes_3[:, last_idx].reshape(-1, 1)])\n            )\n\n        all_faces = np.vstack(all_faces)\n        \n        vertex_colors = np.vstack([init_colors] * 19)\n        face_colors = np.vstack([init_colors] * 24)\n        \n        # vertex_opacities = np.vstack([init_opacities] * (no_of_points + 1))\n        vertex_opacities = np.vstack([init_opacities, *[init_opacities * opac_mul] * 18])\n        # face_opacities = np.vstack([init_opacities] * no_of_points)\n        face_opacities = np.vstack([init_opacities] * 24)\n        \n        vertex_rgba = np.hstack([vertex_colors, vertex_opacities])\n        face_rgba = np.hstack([face_colors, face_opacities])\n        output[scale_mul] = dict(\n            vertices=vertices,\n            vertex_colors=vertex_rgba,\n            faces=all_faces,\n            face_colors=face_rgba\n        )\n    \n    return output\n\n\ndef _get_vertices(origin, scales, rots, scale_mul, no_of_points):\n    scales = scale_mul * np.repeat(np.repeat(scales[:, None, :, None], no_of_points, 1), 3, -1)\n    rots = np.repeat(rots[:, None, :, :], no_of_points, 1)\n    scaled_rots = scales * rots\n    \n    angle = np.linspace(-np.pi, np.pi, no_of_points + 1)[:-1]\n    sins = np.repeat(np.repeat((np.sin(angle).reshape(-1, 1)[None, :, :]), origin.shape[0], 0), 3, -1)\n    coss = np.repeat(np.repeat((np.cos(angle).reshape(-1, 1)[None, :, :]), origin.shape[0], 0), 3, -1)\n    \n    _origin = np.repeat(origin[:, None, :], no_of_points, 1)\n    new_points = _origin + coss * scaled_rots[:, :, 0, :] + sins * scaled_rots[:, :, 1, :]\n    return new_points\n\ndef gen_2d_pseudomesh(scale_muls, no_of_points, init_colors, init_opacities, scale, rotation, origin, opac_mul=.2):\n    angle = np.linspace(-np.pi, np.pi, no_of_points + 1)[:-1]\n    sins = np.repeat(np.repeat((np.sin(angle).reshape(-1, 1)[None, :, :]), origin.shape[0], 0), 3, -1)\n    coss = np.repeat(np.repeat((np.cos(angle).reshape(-1, 1)[None, :, :]), origin.shape[0], 0), 3, -1)\n\n    _origin = np.repeat(origin[:, None, :], no_of_points, 1)\n    \n    output = {}\n    for scale_mul in scale_muls:\n        scales = scale_mul * np.repeat(np.repeat(scale[:, None, :, None], no_of_points, 1), 3, -1)\n        rots = np.repeat(rotation[:, None, :, :], no_of_points, 1)\n        scaled_rots = scales * rots\n        \n        new_points = _origin + coss * scaled_rots[:, :, 0, :] + sins * scaled_rots[:, :, 1, :]\n        \n        vertices = np.vstack([origin, *[new_points[:, _idx, :] for _idx in range(no_of_points)]])\n        origin_idxs = np.arange(0, origin.shape[0]).reshape(-1, 1)\n        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)])\n\n        # create faces\n        all_faces = []\n        for face_idx in range(no_of_points):\n            last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0\n            all_faces.append(\n                np.hstack([origin_idxs, indexes[:, face_idx].reshape(-1, 1), indexes[:, last_idx].reshape(-1, 1)])\n            )\n        all_faces = np.vstack(all_faces)\n        \n        vertex_colors = np.vstack([init_colors] * (no_of_points + 1))\n        face_colors = np.vstack([init_colors] * no_of_points)\n        \n        vertex_opacities = np.vstack([init_opacities, *[init_opacities * opac_mul] * no_of_points])\n        face_opacities = np.vstack([init_opacities] * no_of_points)\n        \n        vertex_rgba = np.hstack([vertex_colors, vertex_opacities])\n        face_rgba = np.hstack([face_colors, face_opacities])\n        \n        output[scale_mul] = dict(\n            vertices=vertices, \n            vertex_colors=vertex_rgba, \n            faces=all_faces, \n            face_colors=face_rgba\n        )\n        \n    return output\n    \n\ndef get_rgb_colors(color_features):\n    colors = np.transpose(color_features, [0, 2, 1])\n    result = C0 * colors[..., 0] + 0.5\n    result = result.clip(min=0., max=1.)\n    return result\n\n\ndef 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):\n    assert no_of_points == 8\n    \n    init_colors = get_rgb_colors(features_dc)\n    init_opacities = get_opacity(opacities)\n    scale = get_scaling(scales)\n    rotation = np.transpose(build_euler_rotation(rots), [0, 2, 1])\n    origin = xyz\n    \n    scales_1 = scale[:, [0, 1]]\n    rots_1 = rotation[:, [0, 1], :]\n    \n    scales_2 = scale[:, [1, 2]]\n    rots_2 = rotation[:, [1, 2], :]\n    \n    scales_3 = scale[:, [0, 2]]\n    rots_3 = rotation[:, [0, 2], :]\n    \n    output = {}\n    for scale_mul in scale_muls:\n        new_points_1 = _get_vertices(origin, scales_1, rots_1, scale_mul, no_of_points)\n        new_points_2 = _get_vertices(origin, scales_2, rots_2, scale_mul, no_of_points)\n        new_points_3 = _get_vertices(origin, scales_3, rots_3, scale_mul, no_of_points)\n        \n        # remove redundant verts\n        new_points_2 = new_points_2[:, [1, 2, 3, 5, 6, 7], :]\n        new_points_3 = new_points_3[:, [1, 3, 5, 7], :]\n        \n        vertices = np.vstack([origin, *[new_points_1[:, _idx, :] for _idx in range(new_points_1.shape[1])]])\n        vertices = np.vstack([vertices, *[new_points_2[:, _idx, :] for _idx in range(new_points_2.shape[1])]])\n        vertices = np.vstack([vertices, *[new_points_3[:, _idx, :] for _idx in range(new_points_3.shape[1])]])\n        \n        origin_idxs = np.arange(0, origin.shape[0]).reshape(-1, 1)\n        indexes_1 = np.hstack([\n            np.arange(_idx * origin.shape[0], (_idx + 1) * origin.shape[0]).reshape(-1, 1) for _idx in range(1, 9)\n        ])\n        indexes_2 = np.hstack([\n            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]\n        ])\n        indexes_3 = np.hstack([\n            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]\n        ])\n        \n        # create faces\n        all_faces = []\n        \n        # faces 1\n        for face_idx in range(no_of_points):\n            last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0\n            all_faces.append(\n                np.hstack([origin_idxs, indexes_1[:, face_idx].reshape(-1, 1), indexes_1[:, last_idx].reshape(-1, 1)])\n            )\n        \n        # faces 2\n        for face_idx in range(no_of_points):\n            last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0\n            all_faces.append(\n                np.hstack([origin_idxs, indexes_2[:, face_idx].reshape(-1, 1), indexes_2[:, last_idx].reshape(-1, 1)])\n            )\n\n        # faces 3\n        for face_idx in range(no_of_points):\n            last_idx = face_idx + 1 if face_idx + 1 < no_of_points else 0\n            all_faces.append(\n                np.hstack([origin_idxs, indexes_3[:, face_idx].reshape(-1, 1), indexes_3[:, last_idx].reshape(-1, 1)])\n            )\n\n        all_faces = np.vstack(all_faces)\n        \n        vertex_colors = np.vstack([init_colors] * 19)\n        face_colors = np.vstack([init_colors] * 24)\n        \n        vertex_opacities = np.vstack([init_opacities, *[init_opacities * 0.2] * 18])\n        face_opacities = np.vstack([init_opacities] * 24)\n        \n        vertex_rgba = np.hstack([vertex_colors, vertex_opacities])\n        face_rgba = np.hstack([face_colors, face_opacities])\n        \n        output[scale_mul] = dict(\n            vertices=vertices,\n            vertex_colors=vertex_rgba,\n            faces=all_faces,\n            face_colors=face_rgba\n        )\n    \n    return output\n"
  },
  {
    "path": "mesh_optim/2dgs_experiments_run.py",
    "content": "import os\nfrom pathlib import Path\nimport subprocess\nimport yaml\n\n\nPYTHON_PATH = Path(\"/home/olaf/miniconda3/envs/splats310/bin/python\")\nPSEUDOMESH_PY = Path(\"/home/grzegos/projects/phd/gs_raytracing/generate_pseudomesh.py\")\nPSEUDO_CWD_PATH = PSEUDOMESH_PY.parent\nPSEUDOMESH_OPTIM_PY = Path(\"/home/grzegos/projects/phd/gs_raytracing/mesh_optim/optimize_pseudomesh.py\")\nOPTIM_CWD_PATH = PSEUDOMESH_OPTIM_PY.parent\n\nDSET_MAIN_DIR = Path(\"/home/grzegos/datasets/games_set\")\nTWODGS_EXP_DIR = Path(\"/home/grzegos/projects/phd/2dgs_nerf_output\")\nOUT_EXP_DIR = Path(\"/home/grzegos/projects/phd/2dgs_nvdiff_exps\")\n\nTMP_CONFIG_PATH = Path(\"/home/grzegos/projects/phd/gs_raytracing/mesh_optim/temp_config.yaml\")\n\ndef main():\n    with open(\"/home/grzegos/projects/phd/gs_raytracing/template_config.yaml\", \"r\") as file:\n        template_cfg = yaml.load(file, Loader=yaml.FullLoader)\n    twodgs_exp_dirs = list(TWODGS_EXP_DIR.glob(\"*\"))\n    \n    env = os.environ.copy()  # Copy current environment\n    env['CUDA_VISIBLE_DEVICES'] = '1'\n    \n    for exp_dir in twodgs_exp_dirs:\n        dset_path = Path(DSET_MAIN_DIR, exp_dir.stem)\n        exp_name = exp_dir.stem\n        \n        # create pseudomesh\n        ply_path = Path(exp_dir, \"point_cloud\", \"iteration_10000\", \"point_cloud.ply\")\n        subprocess.run(\n            [\n                str(PYTHON_PATH),\n                str(PSEUDOMESH_PY),\n                \"--ply\", str(ply_path),\n                \"--algorithm\", \"2dgs\",\n                \"--scale_min\", \"2.3\",\n                \"--scale_max\", \"2.4\",\n                \"--scale_step\", \"0.2\",\n                \"--no_points\", \"8\"\n            ],\n            cwd=str(PSEUDO_CWD_PATH)\n        )\n        \n        # run optim\n        pseudomesh_path = Path(exp_dir, \"pseudomeshes\", \"scale_2.30_pts_8.npz\")\n        # change config\n        template_cfg[\"experiment_name\"] = exp_name\n        template_cfg[\"experiments_dir\"] = str(OUT_EXP_DIR)\n        template_cfg[\"pseudomesh_path\"] = str(pseudomesh_path)\n        template_cfg[\"dataset\"][\"dataset_path\"] = str(dset_path)\n        # save cfg\n        with open(TMP_CONFIG_PATH, 'w') as outfile:\n            yaml.dump(template_cfg, outfile, default_flow_style=False)\n        # pass cfg path to optimizer\n        subprocess.run(\n            [\n                str(PYTHON_PATH),\n                str(PSEUDOMESH_OPTIM_PY),\n                \"--cfg_path\", str(TMP_CONFIG_PATH)\n            ],\n            cwd=str(OPTIM_CWD_PATH),\n            env=env\n        )\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "mesh_optim/cam_utils.py",
    "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact  george.drettakis@inria.fr\n#\n\n\nimport collections\nimport struct\n\nimport numpy as np\n\n\nCameraModel = collections.namedtuple(\n    \"CameraModel\", [\"model_id\", \"model_name\", \"num_params\"])\nCamera = collections.namedtuple(\n    \"Camera\", [\"id\", \"model\", \"width\", \"height\", \"params\"])\nBaseImage = collections.namedtuple(\n    \"Image\", [\"id\", \"qvec\", \"tvec\", \"camera_id\", \"name\", \"xys\", \"point3D_ids\"])\nPoint3D = collections.namedtuple(\n    \"Point3D\", [\"id\", \"xyz\", \"rgb\", \"error\", \"image_ids\", \"point2D_idxs\"])\nCAMERA_MODELS = {\n    CameraModel(model_id=0, model_name=\"SIMPLE_PINHOLE\", num_params=3),\n    CameraModel(model_id=1, model_name=\"PINHOLE\", num_params=4),\n    CameraModel(model_id=2, model_name=\"SIMPLE_RADIAL\", num_params=4),\n    CameraModel(model_id=3, model_name=\"RADIAL\", num_params=5),\n    CameraModel(model_id=4, model_name=\"OPENCV\", num_params=8),\n    CameraModel(model_id=5, model_name=\"OPENCV_FISHEYE\", num_params=8),\n    CameraModel(model_id=6, model_name=\"FULL_OPENCV\", num_params=12),\n    CameraModel(model_id=7, model_name=\"FOV\", num_params=5),\n    CameraModel(model_id=8, model_name=\"SIMPLE_RADIAL_FISHEYE\", num_params=4),\n    CameraModel(model_id=9, model_name=\"RADIAL_FISHEYE\", num_params=5),\n    CameraModel(model_id=10, model_name=\"THIN_PRISM_FISHEYE\", num_params=12)\n}\nCAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model)\n                         for camera_model in CAMERA_MODELS])\nCAMERA_MODEL_NAMES = dict([(camera_model.model_name, camera_model)\n                           for camera_model in CAMERA_MODELS])\n\ndef qvec2rotmat(qvec):\n    return np.array([\n        [1 - 2 * qvec[2]**2 - 2 * qvec[3]**2,\n         2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],\n         2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2]],\n        [2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],\n         1 - 2 * qvec[1]**2 - 2 * qvec[3]**2,\n         2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]],\n        [2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],\n         2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],\n         1 - 2 * qvec[1]**2 - 2 * qvec[2]**2]])\n\n\nclass Image(BaseImage):\n    def qvec2rotmat(self):\n        return qvec2rotmat(self.qvec)\n\n\ndef read_next_bytes(fid, num_bytes, format_char_sequence, endian_character=\"<\"):\n    \"\"\"Read and unpack the next bytes from a binary file.\n    :param fid:\n    :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.\n    :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.\n    :param endian_character: Any of {@, =, <, >, !}\n    :return: Tuple of read and unpacked values.\n    \"\"\"\n    data = fid.read(num_bytes)\n    return struct.unpack(endian_character + format_char_sequence, data)\n\n\ndef read_extrinsics_binary(path_to_model_file):\n    \"\"\"\n    see: src/base/reconstruction.cc\n        void Reconstruction::ReadImagesBinary(const std::string& path)\n        void Reconstruction::WriteImagesBinary(const std::string& path)\n    \"\"\"\n    images = {}\n    with open(path_to_model_file, \"rb\") as fid:\n        num_reg_images = read_next_bytes(fid, 8, \"Q\")[0]\n        for _ in range(num_reg_images):\n            binary_image_properties = read_next_bytes(\n                fid, num_bytes=64, format_char_sequence=\"idddddddi\")\n            image_id = binary_image_properties[0]\n            qvec = np.array(binary_image_properties[1:5])\n            tvec = np.array(binary_image_properties[5:8])\n            camera_id = binary_image_properties[8]\n            image_name = \"\"\n            current_char = read_next_bytes(fid, 1, \"c\")[0]\n            while current_char != b\"\\x00\":   # look for the ASCII 0 entry\n                image_name += current_char.decode(\"utf-8\")\n                current_char = read_next_bytes(fid, 1, \"c\")[0]\n            num_points2D = read_next_bytes(fid, num_bytes=8,\n                                           format_char_sequence=\"Q\")[0]\n            x_y_id_s = read_next_bytes(fid, num_bytes=24*num_points2D,\n                                       format_char_sequence=\"ddq\"*num_points2D)\n            xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])),\n                                   tuple(map(float, x_y_id_s[1::3]))])\n            point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))\n            images[image_id] = Image(\n                id=image_id, qvec=qvec, tvec=tvec,\n                camera_id=camera_id, name=image_name,\n                xys=xys, point3D_ids=point3D_ids)\n    return images\n\n\ndef read_intrinsics_binary(path_to_model_file):\n    \"\"\"\n    see: src/base/reconstruction.cc\n        void Reconstruction::WriteCamerasBinary(const std::string& path)\n        void Reconstruction::ReadCamerasBinary(const std::string& path)\n    \"\"\"\n    cameras = {}\n    with open(path_to_model_file, \"rb\") as fid:\n        num_cameras = read_next_bytes(fid, 8, \"Q\")[0]\n        for _ in range(num_cameras):\n            camera_properties = read_next_bytes(\n                fid, num_bytes=24, format_char_sequence=\"iiQQ\")\n            camera_id = camera_properties[0]\n            model_id = camera_properties[1]\n            model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name\n            width = camera_properties[2]\n            height = camera_properties[3]\n            num_params = CAMERA_MODEL_IDS[model_id].num_params\n            params = read_next_bytes(fid, num_bytes=8*num_params,\n                                     format_char_sequence=\"d\"*num_params)\n            cameras[camera_id] = Camera(id=camera_id,\n                                        model=model_name,\n                                        width=width,\n                                        height=height,\n                                        params=np.array(params))\n        assert len(cameras) == num_cameras\n    return cameras\n"
  },
  {
    "path": "mesh_optim/data.py",
    "content": "import json\nfrom pathlib import Path\nfrom PIL import Image\n\nimport torch\nfrom torch.utils.data import Dataset\nimport numpy as np\n\nfrom cam_utils import qvec2rotmat, read_extrinsics_binary, read_intrinsics_binary\n\n\ndef create_blend_proj_mats(camera_angle_x, img_shape, transf_mat, far, near):\n    focal = img_shape[0] / (2 * np.tan(camera_angle_x / 2))\n    transf_matrix = torch.tensor(transf_mat, device=\"cpu\", dtype=torch.float32)\n    view_matrix = torch.inverse(transf_matrix)\n\n    proj_matrix = torch.zeros(4, 4, device=\"cpu\", dtype=torch.float32)\n    proj_matrix[0, 0] = 2 * focal / img_shape[0]\n    proj_matrix[1, 1] = -2 * focal / img_shape[1]\n    proj_matrix[2, 2] = -(far + near) / (far - near)\n    proj_matrix[2, 3] = -2.0 * far * near / (far - near)\n    proj_matrix[3, 2] = -1.0\n    \n    mvp_matrix = proj_matrix @ view_matrix\n    \n    return {\n        \"transf_matrix\": transf_matrix,\n        \"view_matrix\": view_matrix,\n        \"proj_matrix\": proj_matrix,\n        \"mvp_matrix\": mvp_matrix\n    }\n\n\ndef create_colmap_proj_mats(focal_x, focal_y, img_shape, transf_mat, far, near):  # TODO test it bruh\n    transf_matrix = torch.tensor(transf_mat, device=\"cpu\", dtype=torch.float32)\n    view_matrix = torch.inverse(transf_matrix)\n    \n    proj_matrix = torch.zeros(4, 4, device=\"cpu\", dtype=torch.float32)\n    proj_matrix[0, 0] = 2 * focal_x / img_shape[0]\n    proj_matrix[1, 1] = -2 * focal_y / img_shape[1]\n    proj_matrix[2, 2] = -(far + near) / (far - near)\n    proj_matrix[2, 3] = -2.0 * far * near / (far - near)\n    proj_matrix[3, 2] = -1.0\n    \n    mvp_matrix = proj_matrix @ view_matrix\n\n    return {\n        \"transf_matrix\": transf_matrix,\n        \"view_matrix\": view_matrix,\n        \"proj_matrix\": proj_matrix,\n        \"mvp_matrix\": mvp_matrix\n    }\n    \n\nclass ImageCamDataset(Dataset):\n    def __init__(self, dataset_path, near, far, imgs_in_ram=True, res=1, test=False):\n        self.data_dir = Path(dataset_path)\n        self.near = near\n        self.far = far\n        self._imgs_in_ram = imgs_in_ram\n        self._res = res\n        self.test = test\n        self.mode = None\n        self.shape = None\n        # print(self.test)\n        if Path(dataset_path, \"transforms_train.json\").exists():\n            self.mode = \"blender\"\n        else:\n            self.mode = \"colmap\"\n        self.load_cameras()\n        # print(len(self))\n        \n    \n    def _load_blender_cameras(self):\n        if self.test:\n            self._cam_info_path = Path(self.data_dir, \"transforms_test.json\")\n        else:\n            self._cam_info_path = Path(self.data_dir, \"transforms_train.json\")\n        # Load camera parameters from JSON file\n        with open(self._cam_info_path, 'r') as f:\n            self.camera_data = json.load(f)\n        \n        self.image_files = []\n        _cam_angle_x = self.camera_data[\"camera_angle_x\"]\n        for _idx, _cam in enumerate(self.camera_data[\"frames\"]):\n            img_path = Path(self.data_dir, _cam[\"file_path\"]).with_suffix(\".png\")\n            img = None\n            if not _idx or self._imgs_in_ram:\n                img, width, height = self.load_image(img_path, get_shape=True)\n                self.shape = [width, height]\n            trans_mat = np.array(_cam[\"transform_matrix\"], dtype=np.float32)\n            proj_mats = create_blend_proj_mats(\n                _cam_angle_x,\n                self.shape,\n                trans_mat,\n                self.far,\n                self.near\n            )\n            \n            self.image_files.append({\n                \"name\": img_path.name,\n                \"img_path\": str(img_path),\n                \"img\": img,\n                **proj_mats\n            })\n    \n    def _load_colmap_cameras(self):\n        print(self._res)\n        try:\n            cameras_extrinsic_file = Path(self.data_dir, \"sparse/0\", \"images.bin\")\n            cameras_intrinsic_file = Path(self.data_dir, \"sparse/0\", \"cameras.bin\")\n            cam_extrinsics = read_extrinsics_binary(cameras_extrinsic_file)\n            cam_intrinsics = read_intrinsics_binary(cameras_intrinsic_file)\n        except:\n            raise NotImplementedError(\".txt Colmap structure not supported\")\n        \n        image_files = []\n        for _idx, key in enumerate(cam_extrinsics):\n            extr = cam_extrinsics[key]\n            intr = cam_intrinsics[extr.camera_id]\n            height = intr.height\n            width = intr.width\n            img_path = Path(self.data_dir, \"images\", extr.name)\n            \n            img = None\n            if not _idx or self._imgs_in_ram:\n                img, width, height = self.load_image(img_path, self._res, get_shape=True)\n                self.shape = [width, height]\n            elif self._res != 1:\n                height = round(height / self._res)\n                width = round(width / self._res)\n\n            R = qvec2rotmat(extr.qvec)\n            T = np.array(extr.tvec)\n            transf_mat = np.eye(4, dtype=np.float32)\n            transf_mat[:3, :3] = R\n            transf_mat[:3, -1] = T\n            \n            transf_mat = np.linalg.inv(transf_mat)\n            transf_mat[:3, 1:3] *= -1\n\n            if intr.model==\"SIMPLE_PINHOLE\":\n                focal_length_x = intr.params[0]\n                focal_length_y = intr.params[0]\n            elif intr.model==\"PINHOLE\":\n                focal_length_x = intr.params[0]\n                focal_length_y = intr.params[1]\n            else:\n                assert False, \"Colmap camera model not handled: only undistorted datasets (PINHOLE or SIMPLE_PINHOLE cameras) supported!\"\n            \n            # focal_x = focal_length_x / (2 * self._res)  # TODO FIX FOR TANDT\n            focal_x = focal_length_x / self._res\n            # focal_y = focal_length_y / (2 * self._res)\n            focal_y = focal_length_y / self._res\n            \n            proj_mats = create_colmap_proj_mats(\n                focal_x,\n                focal_y,\n                self.shape,\n                transf_mat,\n                self.far,\n                self.near\n            )\n            \n            image_files.append({\n                \"name\": extr.name,\n                \"img_path\": str(img_path),\n                \"img\": img,\n                \"focal_x\": focal_x,\n                \"focal_y\": focal_y,\n                **proj_mats\n            })\n        sorted_imgs = sorted(image_files, key=lambda x: x[\"name\"].split(\".\")[0])\n        if self.test:\n            cond = lambda x: x % 8 == 0\n        else:\n            cond = lambda x: x % 8 != 0\n        self.image_files = [c for idx, c in enumerate(sorted_imgs) if cond(idx)]\n                \n    \n    def load_cameras(self):\n        if self.mode == \"blender\":\n            self._load_blender_cameras()\n        else: \n            self._load_colmap_cameras()\n    \n    @staticmethod\n    def load_image(path, res=None, get_shape=False):\n        _img = Image.open(path)\n        if res is not None and res != 1:\n            _img.thumbnail(\n                [\n                    round(_img.width / res),\n                    round(_img.height / res)\n                ],\n                Image.Resampling.LANCZOS\n            )\n        img = np.array(_img).astype(np.float32) / 255\n        img = torch.from_numpy(img)\n        if get_shape:\n            return img, _img.width, _img.height\n        return img\n        \n    def __len__(self):\n        return len(self.image_files)\n    \n    def __getitem__(self, idx):\n        \"\"\"\n        Get a sample from the dataset.\n        \n        Returns:\n            dict: Contains image and camera parameters\n        \"\"\"\n        if self.image_files[idx][\"img\"] is None:\n            img = self.load_image(self.image_files[idx][\"img_path\"], self._res)\n            self.image_files[idx][\"img\"] = img\n        self.image_files[idx][\"img\"] = self.image_files[idx][\"img\"]\n        return self.image_files[idx]\n\n# # Example usage:\n# if __name__ == \"__main__\":\n    \n#     # Create dataset instance\n#     dataset = ImageCamDataset(\n#         dataset_path=\"/path/to/dataset\",\n#         near=0.01,\n#         far=100.,\n#         imgs_in_ram=False\n#     )\n    \n#     # Get a sample\n#     sample = dataset[0]\n#     print(f\"Image shape: {sample['image'].shape}\")"
  },
  {
    "path": "mesh_optim/diff_render.py",
    "content": "import json\nfrom pathlib import Path\n\nimport torch\nimport numpy as np\nimport nvdiffrast.torch as dr\nimport matplotlib.pyplot as plt\n\n\ndef create_obj(device):\n    vertices = torch.tensor([\n        [-1, 1, -1],\n        [1, 1, -1],\n        [1, -1, -1],\n        [-1, -1, -1],\n        [-1, 1, 1],\n        [1, 1, 1],\n        [1, -1, 1],\n        [-1, -1, 1],\n    ], device=device, dtype=torch.float32, requires_grad=True)\n\n    _verts_add = torch.tensor([\n        [-1, 1, -1],\n        [1, 1, -1],\n        [1, -1, -1],\n        [-1, -1, -1],\n        [-1, 1, 1],\n        [1, 1, 1],\n        [1, -1, 1],\n        [-1, -1, 1],\n    ], device=device, dtype=torch.float32, requires_grad=True)\n    z_change = torch.zeros_like(vertices)\n    z_change[:, -1] = -1.1\n    \n    vertices = torch.cat(\n        [\n            vertices,\n            _verts_add * 0.2 + z_change\n        ],\n        dim=0\n    )\n\n    # vertices = _verts_add * 0.2 + z_change\n    \n    colors = torch.tensor([\n        [1, 0, 0, 0.6],\n        [0, 1, 0, 0.6],\n        [1, 0, 0, 0.6],\n        [0, 0, 1, 0.6],\n        [0, 0, 1, 1.],\n        [1, 0, 0, 1.],\n        [0, 0, 1, 1.],\n        [0, 1, 0, 1.],\n    ], device=device, dtype=torch.float32, requires_grad=True)\n    \n    _colors_add = torch.cat(\n        [\n            torch.zeros_like(colors[:, :3], requires_grad=True),\n            torch.ones_like(colors[:, 3:], requires_grad=True),\n        ],\n        dim=-1\n    )\n    colors = torch.cat(\n        [\n            colors,\n            _colors_add\n        ],\n        dim=0\n    )\n    # colors = _colors_add\n    \n    # Define faces (triangles)\n    faces = torch.tensor([\n        [0, 1, 4],\n        [1, 5, 4],\n        [1, 2, 5],\n        [2, 6, 5],\n        [2, 3, 6],\n        [3, 7, 6],\n        [3, 0, 7],\n        [0, 4, 7],\n        [0, 1, 3],\n        [1, 2, 3],\n        [4, 5, 7],\n        [5, 6, 7],\n    ], device=device, dtype=torch.int32).contiguous()\n    _new_faces = faces + 8\n    faces = torch.cat(\n        [\n            faces,\n            _new_faces\n        ],\n        dim=0\n    )\n    \n    return vertices, colors, faces\n\n\ndef load_pseudomesh(_path):\n    mesh_data = np.load(_path)\n    vertices = mesh_data[\"vertices\"]\n    faces = mesh_data[\"faces\"]\n    vert_colors = mesh_data[\"vertex_colors\"]\n    face_colors = mesh_data[\"face_colors\"]\n    return vertices, faces, vert_colors, face_colors\n\n\ndef create_proj_mats(camera_angle_x, img_size, transf_mat, far, near):\n    focal = img_size / (2 * np.tan(camera_angle_x / 2))\n    transf_matrix = torch.tensor(transf_mat, device=\"cpu\", dtype=torch.float32)\n    view_matrix = torch.inverse(transf_matrix)\n\n    proj_matrix = torch.zeros(4, 4, device=\"cpu\", dtype=torch.float32)\n    proj_matrix[0, 0] = 2 * focal / img_size\n    proj_matrix[1, 1] = 2 * focal / img_size\n    proj_matrix[2, 2] = -(far + near) / (far - near)\n    proj_matrix[2, 3] = -2.0 * far * near / (far - near)\n    proj_matrix[3, 2] = -1.0\n    \n    mvp_matrix = proj_matrix @ view_matrix\n    \n    return {\n        \"transf_matrix\": transf_matrix,\n        \"view_matrix\": view_matrix,\n        \"proj_matrix\": proj_matrix,\n        \"mvp_matrix\": mvp_matrix\n    }\n\n\ndef create_camera_mats(cam_data, img_size=512, near=0.01, far=100., device=\"cuda\"):\n    height, width = img_size, img_size\n    camera_angle_x = cam_data[\"camera_angle_x\"]\n    focal = width / (2 * np.tan(camera_angle_x / 2))\n    \n    out = {}\n    for cam in cam_data[\"frames\"]:\n        transf_matrix = cam[\"transform_matrix\"]\n        \n        transf_matrix = torch.tensor(transf_matrix, device=device, dtype=torch.float32)\n        \n        # view matrix - transforms vertices from world space \n        # vertex coords to camera/view space\n        \n        view_matrix = torch.inverse(transf_matrix)\n        \n        # Convert focal length to OpenGL-style projection matrix\n        proj_matrix = torch.zeros(4, 4, device=device, dtype=torch.float32)\n        proj_matrix[0, 0] = 2 * focal / width\n        proj_matrix[1, 1] = 2 * focal / height\n        proj_matrix[2, 2] = -(far + near) / (far - near)\n        proj_matrix[2, 3] = -2.0 * far * near / (far - near)\n        proj_matrix[3, 2] = -1.0\n        \n        # mvp (model-view-projection) matrix - combines all transformations into one matrix\n        # basically u apply it to verts and you have them on an image lol\n        mvp_matrix = proj_matrix @ view_matrix\n        \n        out[cam[\"file_path\"]] = {\n            \"transf_matrix\": transf_matrix,\n            \"view_matrix\": view_matrix,\n            \"proj_matrix\": proj_matrix,\n            \"mvp_matrix\": mvp_matrix\n        }\n        _out = create_proj_mats(\n            camera_angle_x, \n            img_size,\n            cam[\"transform_matrix\"],\n            far,\n            near\n        )\n        out[cam[\"file_path\"]] = {k: v.cuda() for k, v in _out.items()}\n    return out\n\n\ndef render(glctx, verts, colors, faces, cam_data, img_size, device):\n    # world vertices to image positions\n    pos_clip = torch.matmul(\n        torch.cat(\n            [\n                verts, \n                torch.ones_like(verts[:, :1])    \n            ], dim=1\n        ), \n        cam_data[\"mvp_matrix\"].t()\n    )[None, ...]\n    \n    with torch.no_grad():\n        # Calculate face centers in clip space\n        face_verts = pos_clip[0, faces.long()]\n        face_centers = face_verts.mean(dim=1)\n        # Sort faces by Z coordinate (depth) in descending order (back to front)\n        face_z = face_centers[..., 2]\n        face_order = torch.argsort(face_z, descending=True)\n        faces = faces[face_order]\n    \n    rast_out, rast_db = dr.rasterize(\n        glctx, \n        pos_clip, \n        faces, \n        (img_size, img_size),\n        grad_db=True  # Enable gradients for all attributes\n    )\n    \n    # Interpolate colors\n    color_interp = dr.interpolate(colors[None, ...], rast_out, faces, rast_db=rast_db,diff_attrs='all')[0]\n    \n    # Background color\n    background = torch.ones_like(color_interp)\n    \n    output = color_interp\n    # Apply built-in transparency blending\n    output = dr.antialias(\n        color_interp,\n        rast_out,\n        pos_clip,\n        faces,\n    )  # Remove batch dimension\n    \n    # Blend with background\n    alpha = output[..., 3:4]\n    rgb = output[..., :3]\n    final_color = rgb + background[..., :3] * (1 - alpha)\n    \n    output = torch.cat([final_color, alpha], dim=-1)\n    \n    output = torch.clamp(output, 0., 1.)\n    return output\n\n\n# def render_with_depth_peeling(glctx, verts, colors, faces, cam_data, img_size, num_layers=4, device=\"cuda\"):\n#     \"\"\"\n#     Render with nvdiffrast's DepthPeeler for order-independent transparency\n#     \"\"\"\n#     # Transform vertices to clip space\n#     pos_clip = torch.matmul(\n#         torch.cat([verts, torch.ones_like(verts[:, :1])], dim=1),\n#         cam_data[\"mvp_matrix\"].t()\n#     )[None, ...]\n    \n#     # Initialize DepthPeeler\n#     background = torch.ones((1, img_size, img_size, 4), device=device)\n#     final_color = torch.zeros_like(background)\n    \n#     # Iterate through layers\n#     import time\n#     start_time = time.time()\n#     with dr.DepthPeeler(glctx, pos_clip, faces, (img_size, img_size)) as peeler:\n#         for layer_idx in range(num_layers):\n#             # Get rasterization output for current layer\n#             rast_out, rast_db = peeler.rasterize_next_layer()\n            \n#             # Skip if no geometry in this layer\n#             if rast_out is None:\n#                 break\n                \n#             # Interpolate colors\n#             color_layer = dr.interpolate(\n#                 colors[None, ...],\n#                 rast_out,\n#                 faces,\n#                 rast_db=rast_db,\n#                 diff_attrs='all'\n#             )[0]\n            \n#             # # Antialias the colors\n#             # color_aa = dr.antialias(\n#             #     color_layer,\n#             #     pos_clip,\n#             #     faces,\n#             #     rast_out,\n#             #     rast_db=rast_db\n#             # )\n            \n#             # Front-to-back blending\n#             alpha = color_layer[..., 3:4]\n#             rgb = color_layer[..., :3]\n#             final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([rgb * alpha, alpha], dim=-1)\n#     print(f\"Final time: {time.time() - start_time}s\")\n#     # Blend with background\n#     alpha = final_color[..., 3:4]\n#     rgb = final_color[..., :3]\n#     final_color = rgb + background[..., :3] * (1 - alpha)\n    \n#     loss = final_color.mean()\n#     loss.backward()\n    \n#     return torch.clamp(final_color, 0., 1.)\n\n\ndef to_torch(*args, device=\"cpu\"):\n    return [torch.from_numpy(x).float().to(device) for x in args]\n    \n\ndef test_render():\n    cam_info_path = Path(\"/home/grzegos/projects/phd/gs_raytracing/data/cam_test/transforms_train.json\")\n    with open(cam_info_path, \"r\") as file:\n        cam_data = json.load(file)\n    \n    glctx = dr.RasterizeGLContext()\n    img_size = 512\n    near = 1e-2\n    far = 1e2\n    device = \"cuda:0\"\n    \n    verts, colors, faces = create_obj(device)\n    \n    cam_matrices = create_camera_mats(\n        cam_data=cam_data,\n        img_size=img_size,\n        near=near,\n        far=far,\n        device=device\n    )\n    \n    for cam_name, cam_mats in cam_matrices.items():\n        out_img = render_with_depth_peeling(\n            glctx=glctx,\n            verts=verts,\n            colors=colors,\n            faces=faces,\n            cam_data=cam_mats,\n            img_size=img_size,\n            num_layers=2,\n            device=device\n        )\n        \n        plt.imsave(f\"02_{cam_name.split('/')[-1]}\", out_img.detach().cpu().numpy()[0])\n    \n\ndef render_pseudomesh():\n    path_to_dataset = Path(\"/home/grzegos/datasets/games_set/hotdog\")\n    path_to_camera_data = Path(path_to_dataset, \"transforms_train.json\")\n    path_to_pseudomesh = Path(\"/home/grzegos/projects/phd/games_nerf_output/hotdog_test_fff_sh0/pseudomeshes/scale_2.30_pts_8.npz\")\n    \n    with open(path_to_camera_data, \"r\") as file:\n        cam_data = json.load(file)\n    \n    glctx = dr.RasterizeGLContext()\n    img_size = 800\n    near = 1e-2\n    far = 1e2\n    device = \"cuda:0\"\n    num_layers = 100\n        \n    cam_matrices = create_camera_mats(\n        cam_data=cam_data,\n        img_size=img_size,\n        near=near,\n        far=far,\n        device=device\n    )\n    \n    vertices, faces, vert_colors, face_colors = load_pseudomesh(path_to_pseudomesh)\n    \n    vertices, faces, vert_colors = to_torch(vertices, faces, vert_colors, device=device)\n    vert_colors.requires_grad = True\n    faces = faces.int()\n    \n    for cam_name, cam_mats in cam_matrices.items():\n        out_img = render_with_depth_peeling(\n            glctx=glctx,\n            verts=vertices,\n            vert_colors=vert_colors,\n            faces=faces,\n            cam_data=cam_mats,\n            img_size=img_size,\n            num_layers=num_layers,\n            device=device\n        )\n    \n        plt.imsave(f\"{str(num_layers)}_{cam_name.split('/')[-1]}.png\", out_img.detach().cpu().numpy()[0])\ndef render_with_depth_peeling(glctx, verts, vert_colors, faces, cam_data, img_size, num_layers=4, device=\"cuda\"):\n    \"\"\"\n    Render with nvdiffrast's DepthPeeler for order-independent transparency.\n    Optimized for performance, especially regarding the vert_colors tensor.\n    \"\"\"\n    # Transform vertices to clip space\n    pos_clip = torch.matmul(\n        torch.cat([verts, torch.ones_like(verts[:, :1])], dim=1),\n        cam_data[\"mvp_matrix\"].t()\n    )[None, ...]\n    \n    # Initialize DepthPeeler\n    background = torch.ones((1, img_size, img_size, 4), device=device)\n    final_color = torch.zeros_like(background)\n    \n    # Iterate through layers\n    with dr.DepthPeeler(glctx, pos_clip, faces, (img_size, img_size)) as peeler:\n        for layer_idx in range(num_layers):\n            # Get rasterization output for current layer\n            rast_out, rast_db = peeler.rasterize_next_layer()\n            \n            # Skip if no geometry in this layer\n            if rast_out is None:\n                break\n                \n            # Interpolate colors\n            # Use `interp1d` for more efficient color interpolation\n            color_layer = dr.interpolate(\n                vert_colors[None, ...],\n                rast_out,\n                faces,\n                rast_db=rast_db,\n                diff_attrs='all'\n            )[0]\n            \n            # Front-to-back blending\n            alpha = color_layer[..., 3:4]\n            rgb = color_layer[..., :3]\n            final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([rgb * alpha, alpha], dim=-1)\n    \n    # Blend with background\n    alpha = final_color[..., 3:4]\n    rgb = final_color[..., :3]\n    final_color = rgb + background[..., :3] * (1 - alpha)\n    \n    loss = final_color.mean()\n    loss.backward()\n    \n    return torch.clamp(final_color, 0., 1.)\n\nif __name__ == \"__main__\":\n    # main()\n    render_pseudomesh()\n    # Regular render\n    # save_render(angle=np.pi/6)\n    \n    # Or run optimization\n    # optimize_colors()\n"
  },
  {
    "path": "mesh_optim/generate_multi_views_circle.py",
    "content": "import argparse\nimport math\nfrom pathlib import Path\nimport time\n\nimport torch\nimport torchvision\nimport imageio.v2 as imageio\n\nfrom data import ImageCamDataset\nfrom optimize_pseudomesh import _load_best_model\nfrom render_pseudomesh import _load_config\n\ndef get_depth_map(model, mvp_mat, width, height, cam_pos, num_layers):\n    optim_depth_map, optim_depth_map_alpha = model.get_depth_map(mvp_mat.unsqueeze(0), width, height, cam_pos, num_layers)\n    optim_depth_map = (optim_depth_map - optim_depth_map.min()) / (optim_depth_map.max() - optim_depth_map.min())\n    optim_depth_map_alpha = torch.clamp(optim_depth_map_alpha, 0.0, 1.0).squeeze(-1)\n    \n    optim_depth_map = torch.cat(\n        [\n            optim_depth_map[..., 0],\n            torch.zeros_like(optim_depth_map[..., 0]),\n            1 - optim_depth_map[..., 0],\n        ],\n        dim=0\n    )\n    background = torch.tensor([1, 0, 0], dtype=torch.float32, device=optim_depth_map.device)\n    optim_depth_map = optim_depth_map * optim_depth_map_alpha + background[:3, None, None] * (1 - optim_depth_map_alpha)\n    return optim_depth_map\n\ndef get_gray_map(model, mvp_mat, width, height, color_verts, num_layers):\n    _color_verts = torch.cat([color_verts, model.pseudomesh.vertex_colors[..., -1:]], dim=1)\n    rgb, alpha = model.get_gray_map(mvp_mat.unsqueeze(0), width, height, num_layers, _color_verts)\n    return rgb.squeeze(0).permute(2, 0, 1)#, alpha.squeeze(0)\n\ndef get_normal_map(model, mvp_mat, width, height, num_layers):\n    rgb, alpha = model.get_normal_map(mvp_mat.unsqueeze(0), width, height, num_layers)\n    rgb = torch.clamp(rgb, 0.0, 1.0)\n    return rgb.squeeze(0).permute(2, 0, 1)#, alpha.squeeze(0)\n\ndef get_all_maps(model, mvp_mat, width, height, cam_pos, num_layers, color_verts=None):\n    results = model.render_all_maps(mvp_mat.unsqueeze(0), width, height, cam_pos, num_layers, color_verts)\n    color = results[\"color\"][0].squeeze(0).permute(2, 0, 1).clamp(0.0, 1.0)\n    normal = results[\"normal\"][0].squeeze(0).permute(2, 0, 1).clamp(0.0, 1.0)\n    depth = results[\"depth\"][0].squeeze(0).permute(2, 0, 1)\n    depth_alpha = results[\"depth\"][1].squeeze(0)\n    gray = results[\"gray\"][0].squeeze(0).permute(2, 0, 1).clamp(0.0, 1.0)\n    \n    depth = (depth - depth.min()) / (depth.max() - depth.min())\n    depth_alpha = torch.clamp(depth_alpha, 0.0, 1.0).squeeze(-1).unsqueeze(0)\n    \n    depth = torch.cat(\n        [\n            depth[0:1, ...],\n            torch.zeros_like(depth[0:1, ...]),\n            1 - depth[0:1, ...],\n        ],\n        dim=0\n    )\n    background = torch.tensor([1, 0, 0], dtype=torch.float32, device=depth.device)\n    depth = depth * depth_alpha + background[:3, None, None] * (1 - depth_alpha)\n    \n    return color, normal, depth, gray\n\ndef normalize(v, eps=1e-6):\n    \"\"\"Return the normalized vector v.\"\"\"\n    norm = torch.norm(v)\n    return v if norm < eps else v / norm\n\ndef look_at(eye, target, up):\n    # Compute the forward vector (from the camera toward the target)\n    forward = normalize(target - eye)\n    \n    # Compute the right vector\n    right = normalize(torch.cross(forward, up))\n    \n    # Recompute the orthogonal up vector\n    true_up = torch.cross(right, forward)\n    # true_up = up\n    \n    # Build rotation matrix.\n    # Note: We use -forward so that the camera looks along its -z axis.\n    R = torch.stack([right, true_up, -forward], dim=1)  # Shape: (3,3)\n    \n    # Build the full 4x4 transformation matrix\n    T = torch.eye(4)\n    T[:3, :3] = R\n    T[:3, 3] = eye\n    return T\n\ndef generate_camera_matrices(target, radius, num_views=20, \n                             up=torch.tensor([0, -1, 0], dtype=torch.float32),\n                             tilt_angle=0.0):\n    target = target.float()\n    cam_mats = []\n    \n    # Precompute the rotation matrix for tilting about the x-axis\n    c = math.cos(tilt_angle)\n    s = math.sin(tilt_angle)\n    R_tilt = torch.tensor([\n        [1,  0,  0],\n        [0,  c, -s],\n        [0,  s,  c]\n    ], dtype=torch.float32)\n    \n    for i in range(num_views):\n        angle = 2 * math.pi * i / num_views\n        \n        # Compute a point on the circle (initially in the XZ-plane)\n        circle_point = torch.tensor([\n            radius * math.cos(angle),\n            0.0,\n            radius * math.sin(angle)\n        ], dtype=torch.float32)\n        \n        # Apply the tilt rotation to the circle point\n        tilted_point = torch.matmul(R_tilt, circle_point)\n        \n        # Compute the camera position by adding the target (center of the circle)\n        cam_pos = target + tilted_point\n        \n        # Create the camera-to-world transformation matrix for this camera\n        cam_mat = look_at(cam_pos, target, up)\n        cam_mats.append(cam_mat)\n    \n    return torch.stack(cam_mats)\n\ndef calculate_mvp(focal_x, focal_y, img_shape, view_mat, far, near):\n    proj_matrix = torch.zeros(4, 4, device=\"cpu\", dtype=torch.float32)\n    proj_matrix[0, 0] = 2 * focal_x / img_shape[0]\n    proj_matrix[1, 1] = -2 * focal_y / img_shape[1]\n    proj_matrix[2, 2] = -(far + near) / (far - near)\n    proj_matrix[2, 3] = -2.0 * far * near / (far - near)\n    proj_matrix[3, 2] = -1.0\n    \n    mvp_matrix = proj_matrix @ view_mat\n    return mvp_matrix\n\n\ndef create_gif(image_paths, output_path, fps):\n    \"\"\"Create a GIF from a list of image paths\"\"\"\n    images = []\n    # Sort paths numerically based on the frame number in the filename\n    sorted_paths = sorted(image_paths, key=lambda x: int(x.stem.split('_')[0]))\n    for path in sorted_paths:\n        images.append(imageio.imread(path))\n    imageio.mimsave(output_path, images, fps=fps)\n\ndef calculate_alphas(img1, img2, left_int, dist_interval):\n    dist = (img1[\"view_matrix\"][:3, 3] - img2[\"view_matrix\"][:3, 3]).pow(2).sum().sqrt()\n    alphas = []\n    left_dist = dist\n    curr_pos = -left_int\n    while dist_interval < left_dist:\n        curr_pos += dist_interval\n        alphas.append(curr_pos / dist)\n        left_dist = dist - curr_pos\n    left_int = dist - curr_pos\n    return torch.tensor(alphas, dtype=torch.float32), left_int\n\ndef main(cfg_path, no_sec, fps):\n    config = _load_config(cfg_path)\n    output_dir = Path(config[\"experiment_dir\"]) / \"multi_views\"\n    output_dir.mkdir(parents=True, exist_ok=True)\n    \n    dset = ImageCamDataset(\n        **{\n            **config[\"dataset\"],\n        }\n    )\n    optim_pseudomesh = _load_best_model(config)\n    \n    circle_center = torch.cat([x[\"transf_matrix\"][:3, 3].unsqueeze(0) for x in dset.image_files], dim=0).mean(dim=0)\n    radius = 4.\n    all_frames = no_sec * fps\n    \n    tilt = math.radians(20)\n    up = torch.tensor([0, -0.78, -0.22], dtype=torch.float32)\n    up = up / torch.norm(up)\n    camera_matrices = generate_camera_matrices(circle_center, radius, all_frames, up=up, tilt_angle=tilt)\n    \n    _range = 0.3\n    color_verts = torch.rand(\n        optim_pseudomesh.pseudomesh.vertices.shape[0],\n        device=optim_pseudomesh.pseudomesh.vertices.device,\n    )[:, None].repeat(1, 3) * _range + 0.5 - _range / 2.\n    \n    _iter = 0\n    focal_x, focal_y = dset.image_files[0][\"focal_x\"], dset.image_files[0][\"focal_y\"]\n    width, height = dset.image_files[0][\"img\"].shape[1], dset.image_files[0][\"img\"].shape[0]\n    print(\"Generate views\")\n    for cam_mat in camera_matrices:\n        start_time = time.time()\n        view_mat = torch.inverse(cam_mat)\n        mvp_mat = calculate_mvp(focal_x, focal_y, [width, height], view_mat, dset.far, dset.near).to(config[\"device\"])\n        with torch.no_grad():\n            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)\n\n        optim_img_path = output_dir / f\"{_iter:05d}_optim.png\"\n        depth_img_path = output_dir / f\"{_iter:05d}_optim-depth.png\"\n        gray_img_path = output_dir / f\"{_iter:05d}_optim-gray.png\"\n        normal_img_path = output_dir / f\"{_iter:05d}_optim-normal.png\"\n        torchvision.utils.save_image(optim_img.squeeze().permute(2, 0, 1), optim_img_path)\n        torchvision.utils.save_image(optim_depth_map, depth_img_path)\n        torchvision.utils.save_image(optim_gray_map, gray_img_path)\n        torchvision.utils.save_image(optim_normal_map, normal_img_path)\n        print(f\"Saved {_iter}, took: {round(time.time() - start_time, 3)} s\")\n        _iter += 1\n\n    # Create GIFs after generating all images\n    optim_images = list(output_dir.glob(\"*_optim.png\"))\n    gray_images = list(output_dir.glob(\"*_optim-gray.png\"))\n    depth_images = list(output_dir.glob(\"*_optim-depth.png\"))\n    normal_images = list(output_dir.glob(\"*_optim-normal.png\"))\n    \n    create_gif(optim_images, output_dir / \"optim.gif\", fps)\n    create_gif(gray_images, output_dir / \"gray.gif\", fps)\n    create_gif(depth_images, output_dir / \"depth.gif\", fps)\n    create_gif(normal_images, output_dir / \"normal.gif\", fps)\n\n\nif __name__ == \"__main__\":\n    _parser = argparse.ArgumentParser(\n        description=\"Process a single string argument.\"\n    )\n    _parser.add_argument(\n        \"--cfg_path\", \"-cp\",\n        type=str, \n        help=\"Path to config\",\n        default=\"/home/grzegos/experiments/mip_games_gs-flat_sh0_bicycle/config.yaml\"\n    )\n    _parser.add_argument(\n        \"--sec\",\n        type=int,\n        help=\"How long the video should be\",\n        default=10\n    )\n    _parser.add_argument(\n        \"--fps\",\n        type=int,\n        help=\"Frames per second\",\n        default=30\n    )\n\n    _args = _parser.parse_args()\n    main(_args.cfg_path, _args.sec, _args.fps)"
  },
  {
    "path": "mesh_optim/log_config.yaml",
    "content": "version: 1\nformatters:\n  simple:\n    format: '%(asctime)s - %(levelname)s - %(module)s - %(funcName)s - %(message)s'\nhandlers:\n  console:\n    class: logging.StreamHandler\n    formatter: simple\n    level: INFO\n    stream: ext://sys.stdout\n  file_handler:\n    class: logging.FileHandler\n    formatter: simple\n    level: DEBUG\n    # change filename based on the output training directory\n    filename: /path/to/logs/logs.log\n    mode: a\nloggers:\n  my_module:\n    level: DEBUG\n    handlers: [console, file_handler]\n    propagate: no\nroot:\n  level: DEBUG\n  handlers: [console, file_handler]\n"
  },
  {
    "path": "mesh_optim/lpipsPyTorch/__init__.py",
    "content": "import torch\n\nfrom .modules.lpips import LPIPS\n\n\ndef lpips(x: torch.Tensor,\n          y: torch.Tensor,\n          net_type: str = 'alex',\n          version: str = '0.1'):\n    r\"\"\"Function that measures\n    Learned Perceptual Image Patch Similarity (LPIPS).\n\n    Arguments:\n        x, y (torch.Tensor): the input tensors to compare.\n        net_type (str): the network type to compare the features: \n                        'alex' | 'squeeze' | 'vgg'. Default: 'alex'.\n        version (str): the version of LPIPS. Default: 0.1.\n    \"\"\"\n    device = x.device\n    criterion = LPIPS(net_type, version).to(device)\n    return criterion(x, y)\n\ndef get_lpips_model(device, net_type: str = 'alex', version: str = '0.1'):\n    return LPIPS(net_type, version).to(device)\n"
  },
  {
    "path": "mesh_optim/lpipsPyTorch/modules/lpips.py",
    "content": "import torch\nimport torch.nn as nn\n\nfrom .networks import get_network, LinLayers\nfrom .utils import get_state_dict\n\n\nclass LPIPS(nn.Module):\n    r\"\"\"Creates a criterion that measures\n    Learned Perceptual Image Patch Similarity (LPIPS).\n\n    Arguments:\n        net_type (str): the network type to compare the features: \n                        'alex' | 'squeeze' | 'vgg'. Default: 'alex'.\n        version (str): the version of LPIPS. Default: 0.1.\n    \"\"\"\n    def __init__(self, net_type: str = 'alex', version: str = '0.1'):\n\n        assert version in ['0.1'], 'v0.1 is only supported now'\n\n        super(LPIPS, self).__init__()\n\n        # pretrained network\n        self.net = get_network(net_type)\n\n        # linear layers\n        self.lin = LinLayers(self.net.n_channels_list)\n        self.lin.load_state_dict(get_state_dict(net_type, version))\n\n    def forward(self, x: torch.Tensor, y: torch.Tensor):\n        feat_x, feat_y = self.net(x), self.net(y)\n\n        diff = [(fx - fy) ** 2 for fx, fy in zip(feat_x, feat_y)]\n        res = [l(d).mean((2, 3), True) for d, l in zip(diff, self.lin)]\n\n        return torch.sum(torch.cat(res, 0), 0, True)\n"
  },
  {
    "path": "mesh_optim/lpipsPyTorch/modules/networks.py",
    "content": "from typing import Sequence\n\nfrom itertools import chain\n\nimport torch\nimport torch.nn as nn\nfrom torchvision import models\n\nfrom .utils import normalize_activation\n\n\ndef get_network(net_type: str):\n    if net_type == 'alex':\n        return AlexNet()\n    elif net_type == 'squeeze':\n        return SqueezeNet()\n    elif net_type == 'vgg':\n        return VGG16()\n    else:\n        raise NotImplementedError('choose net_type from [alex, squeeze, vgg].')\n\n\nclass LinLayers(nn.ModuleList):\n    def __init__(self, n_channels_list: Sequence[int]):\n        super(LinLayers, self).__init__([\n            nn.Sequential(\n                nn.Identity(),\n                nn.Conv2d(nc, 1, 1, 1, 0, bias=False)\n            ) for nc in n_channels_list\n        ])\n\n        for param in self.parameters():\n            param.requires_grad = False\n\n\nclass BaseNet(nn.Module):\n    def __init__(self):\n        super(BaseNet, self).__init__()\n\n        # register buffer\n        self.register_buffer(\n            'mean', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])\n        self.register_buffer(\n            'std', torch.Tensor([.458, .448, .450])[None, :, None, None])\n\n    def set_requires_grad(self, state: bool):\n        for param in chain(self.parameters(), self.buffers()):\n            param.requires_grad = state\n\n    def z_score(self, x: torch.Tensor):\n        return (x - self.mean) / self.std\n\n    def forward(self, x: torch.Tensor):\n        x = self.z_score(x)\n\n        output = []\n        for i, (_, layer) in enumerate(self.layers._modules.items(), 1):\n            x = layer(x)\n            if i in self.target_layers:\n                output.append(normalize_activation(x))\n            if len(output) == len(self.target_layers):\n                break\n        return output\n\n\nclass SqueezeNet(BaseNet):\n    def __init__(self):\n        super(SqueezeNet, self).__init__()\n\n        self.layers = models.squeezenet1_1(True).features\n        self.target_layers = [2, 5, 8, 10, 11, 12, 13]\n        self.n_channels_list = [64, 128, 256, 384, 384, 512, 512]\n\n        self.set_requires_grad(False)\n\n\nclass AlexNet(BaseNet):\n    def __init__(self):\n        super(AlexNet, self).__init__()\n\n        self.layers = models.alexnet(True).features\n        self.target_layers = [2, 5, 8, 10, 12]\n        self.n_channels_list = [64, 192, 384, 256, 256]\n\n        self.set_requires_grad(False)\n\n\nclass VGG16(BaseNet):\n    def __init__(self):\n        super(VGG16, self).__init__()\n\n        self.layers = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).features\n        self.target_layers = [4, 9, 16, 23, 30]\n        self.n_channels_list = [64, 128, 256, 512, 512]\n\n        self.set_requires_grad(False)\n"
  },
  {
    "path": "mesh_optim/lpipsPyTorch/modules/utils.py",
    "content": "from collections import OrderedDict\n\nimport torch\n\n\ndef normalize_activation(x, eps=1e-10):\n    norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True))\n    return x / (norm_factor + eps)\n\n\ndef get_state_dict(net_type: str = 'alex', version: str = '0.1'):\n    # build url\n    url = 'https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/' \\\n        + f'master/lpips/weights/v{version}/{net_type}.pth'\n\n    # download\n    old_state_dict = torch.hub.load_state_dict_from_url(\n        url, progress=True,\n        map_location=None if torch.cuda.is_available() else torch.device('cpu')\n    )\n\n    # rename keys\n    new_state_dict = OrderedDict()\n    for key, val in old_state_dict.items():\n        new_key = key\n        new_key = new_key.replace('lin', '')\n        new_key = new_key.replace('model.', '')\n        new_state_dict[new_key] = val\n\n    return new_state_dict\n"
  },
  {
    "path": "mesh_optim/mesh_utils.py",
    "content": "import torch\nfrom torch import nn\n\n\n\ndef hard_prune_mesh(vertices, faces, vertex_color, mask):\n    vertices = vertices[~mask]\n    vertex_color = vertex_color[~mask]\n    \n    face_mask = torch.any(mask[faces], dim=1)\n    faces = faces[~face_mask]\n    \n    vertices_to_new_idx = torch.zeros(mask.shape[0], dtype=torch.long, device=vertices.device)\n    vertices_to_new_idx[~mask] = torch.arange(vertices.shape[0], device=vertices.device)\n    faces = vertices_to_new_idx[faces]\n    return vertices, faces, vertex_color, mask\n\n\ndef soft_prune_mesh(vertices, faces, vertex_color, mask):\n    face_mask = torch.all(mask[faces], dim=1)\n    faces = faces[~face_mask]\n    valid_vertices = faces.unique().long()\n\n    valid_verts_mask = torch.zeros_like(vertices[:, 0]).bool()\n    valid_verts_mask[valid_vertices] = True\n\n    vertices_to_new_idx = torch.zeros(valid_verts_mask.shape[0], dtype=torch.long, device=vertices.device)\n    vertices_to_new_idx[valid_verts_mask] = torch.arange(valid_verts_mask.sum(), device=vertices.device)\n\n    faces = vertices_to_new_idx[faces.long()]\n    \n    vertices = vertices[valid_verts_mask]\n    vertex_color = vertex_color[valid_verts_mask]\n    \n    return vertices, faces, vertex_color, ~valid_verts_mask\n\n\ndef prune_mesh(vertices, faces, vertex_color, mask, mode):\n    assert mode in [\"soft\", \"hard\"]\n    method = hard_prune_mesh if mode == \"hard\" else soft_prune_mesh\n    vertices, faces, vertex_color, mask = method(\n        vertices,\n        faces,\n        vertex_color,\n        mask\n    )\n    return vertices, faces.int(), vertex_color, mask\n\ndef prune_optimizer(optimizer, mask):\n    optimizable_tensors = {}\n    for group in optimizer.param_groups:\n        stored_state = optimizer.state.get(group['params'][0], None)\n        if stored_state is not None:\n            stored_state[\"exp_avg\"] = stored_state[\"exp_avg\"][~mask]\n            stored_state[\"exp_avg_sq\"] = stored_state[\"exp_avg_sq\"][~mask]\n\n            del optimizer.state[group['params'][0]]\n            group[\"params\"][0] = nn.Parameter((group[\"params\"][0][~mask].requires_grad_(True)))\n            optimizer.state[group['params'][0]] = stored_state\n\n            optimizable_tensors[group[\"name\"]] = group[\"params\"][0]\n        else:\n            group[\"params\"][0] = nn.Parameter(group[\"params\"][0][~mask].requires_grad_(True))\n            optimizable_tensors[group[\"name\"]] = group[\"params\"][0]\n    return optimizable_tensors\n"
  },
  {
    "path": "mesh_optim/metrics.py",
    "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact  george.drettakis@inria.fr\n#\nfrom argparse import ArgumentParser\nimport json\nimport os\nfrom pathlib import Path\nfrom PIL import Image\n\nimport numpy as np\nimport torch\nimport torchvision.transforms.functional as tf\nfrom tqdm import tqdm\n\nfrom lpipsPyTorch import lpips\nfrom ml_utils import psnr, ssim\n\n\ndef PILtoTorch(pil_image):\n    resized_image = torch.from_numpy(np.array(pil_image)) / 255.0\n    if len(resized_image.shape) == 3:\n        return resized_image.permute(2, 0, 1)\n    else:\n        return resized_image.unsqueeze(dim=-1).permute(2, 0, 1)\n\n\ndef readImages(renders_dir, gt_dir, mesh_renders_dir):\n    gs_renders = []\n    mesh_renders = []\n    gts = []\n    image_names = []\n    for fname in os.listdir(renders_dir):\n        mesh_render = Image.open(mesh_renders_dir / fname)\n        gs_render = Image.open(renders_dir / fname)\n        gt_path = gt_dir / fname\n        try:\n            gt = Image.open(gt_path)\n        except FileNotFoundError:\n            try:\n                gt = Image.open(gt_path.with_suffix(\".jpg\"))\n            except FileNotFoundError:\n                gt = Image.open(gt_path.with_suffix(\".JPG\"))\n        if gt.width != gs_render.width or gt.height != gs_render.height:\n            gt = gt.resize((int(gs_render.width), int(gs_render.height)))\n        if mesh_render.width != gs_render.width or mesh_render.height != gs_render.height:\n            mesh_render = mesh_render.resize((int(gs_render.width), int(gs_render.height)))\n        gs_renders.append(PILtoTorch(gs_render).cuda())\n        mesh_renders.append(PILtoTorch(mesh_render).cuda())\n        gt_render = PILtoTorch(gt).cuda()\n        if gt_render.shape[0] != 3:\n            gt_render = gt_render[:3, ...]\n            gt_mask = gt_render[3:4, ...]\n            gt_render *= gt_mask.to(gt_render.device)\n        gts.append(gt_render)\n        image_names.append(fname)\n    return gs_renders, mesh_renders, gts, image_names\n\ndef evaluate(model_path, dset_path, algorithm):\n    if Path(dset_path, \"transforms_train.json\").exists():\n        dset_type = \"nerf\"\n    else:\n        dset_type = \"colmap\"\n    \n    full_dict = {}\n    per_view_dict = {}\n    full_dict_polytopeonly = {}\n    per_view_dict_polytopeonly = {}\n\n    full_dict[model_path] = {}\n    per_view_dict[model_path] = {}\n    full_dict_polytopeonly[model_path] = {}\n    per_view_dict_polytopeonly[model_path] = {}\n\n    test_dir = Path(model_path) / \"test\"\n\n    for method in os.listdir(test_dir):\n        print(\"Method:\", method)\n\n        full_dict[model_path][method] = {}\n        per_view_dict[model_path][method] = {}\n        full_dict_polytopeonly[model_path][method] = {}\n        per_view_dict_polytopeonly[model_path][method] = {}\n\n        method_dir = test_dir / method\n        gt_dir = Path(dset_path, \"test\") if dset_type == \"nerf\" else Path(dset_path, \"images\")\n        \n        renders_dir_name = \"renders\"\n        if algorithm == \"games\":\n            renders_dir_name += \"_gs_flat\"\n        \n        gs_renders_dir = method_dir / renders_dir_name\n        mesh_renders_dir = method_dir / \"pseudomesh_renders\"\n        gs_renders, mesh_renders, gts, image_names = readImages(gs_renders_dir, gt_dir, mesh_renders_dir)\n\n        gs_ssims = []\n        gs_psnrs = []\n        gs_lpipss = []\n\n        mesh_ssims = []\n        mesh_psnrs = []\n        mesh_lpipss = []\n\n        for idx in tqdm(range(len(gs_renders)), desc=\"Metric evaluation progress\"):\n            gs_ssims.append(ssim(gs_renders[idx], gts[idx]))\n            gs_psnrs.append(psnr(gs_renders[idx], gts[idx]))\n            gs_lpipss.append(lpips(gs_renders[idx], gts[idx], net_type='vgg'))\n            mesh_ssims.append(ssim(mesh_renders[idx], gts[idx]))\n            mesh_psnrs.append(psnr(mesh_renders[idx], gts[idx]))\n            mesh_lpipss.append(lpips(mesh_renders[idx], gts[idx], net_type='vgg'))\n\n        full_dict[model_path][method].update({\n            \"GS_SSIM\": torch.tensor(gs_ssims).mean().item(),\n            \"GS_PSNR\": torch.tensor(gs_psnrs).mean().item(),\n            \"GS_LPIPS\": torch.tensor(gs_lpipss).mean().item(),\n            \"MESH_SSIM\": torch.tensor(mesh_ssims).mean().item(),\n            \"MESH_PSNR\": torch.tensor(mesh_psnrs).mean().item(),\n            \"MESH_LPIPS\": torch.tensor(mesh_lpipss).mean().item(),\n        })\n        per_view_dict[model_path][method].update({\n            \"GS_SSIM\": {name: ssim for ssim, name in zip(torch.tensor(gs_ssims).tolist(), image_names)},\n            \"GS_PSNR\": {name: psnr for psnr, name in zip(torch.tensor(gs_psnrs).tolist(), image_names)},\n            \"GS_LPIPS\": {name: lp for lp, name in zip(torch.tensor(gs_lpipss).tolist(), image_names)},\n            \"MESH_SSIM\": {name: ssim for ssim, name in zip(torch.tensor(mesh_ssims).tolist(), image_names)},\n            \"MESH_PSNR\": {name: psnr for psnr, name in zip(torch.tensor(mesh_psnrs).tolist(), image_names)},\n            \"MESH_LPIPS\": {name: lp for lp, name in zip(torch.tensor(mesh_lpipss).tolist(), image_names)},\n        })\n\n    with open(model_path + \"/results.json\", 'w') as fp:\n        json.dump(full_dict[model_path], fp, indent=True)\n    with open(model_path + \"/per_view.json\", 'w') as fp:\n        json.dump(per_view_dict[model_path], fp, indent=True)\n\nif __name__ == \"__main__\":\n    device = torch.device(\"cuda:0\")\n    torch.cuda.set_device(device)\n\n    # Set up command line argument parser\n    parser = ArgumentParser(description=\"Training script parameters\")\n    parser.add_argument('--model_path', '-m', required=True, type=str)\n    parser.add_argument('--set_path', '-s', required=True, type=str)\n    parser.add_argument('--white_background', action=\"store_true\")\n    parser.add_argument(\n        \"--method\",\n        type=str,\n        default=\"3dgs\",\n        choices=[\"3dgs\", \"games\", \"sugar\", \"2dgs\"]\n    )\n    args = parser.parse_args()\n    evaluate(args.model_path, args.set_path, args.method, args.white_background)\n"
  },
  {
    "path": "mesh_optim/ml_utils.py",
    "content": "import math\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\n\nfrom lpipsPyTorch import get_lpips_model\n\ndef _exp_schedule(init_val, gamma, timestep):\n    return init_val * gamma ** timestep\n\n\ndef _step_schedule(curr_val, steps, timestep):\n    return steps.get(timestep, curr_val)\n\n\ndef dp_schedule(_type, curr_val, init_val, epoch, params):\n    assert _type in [\"exp\", \"step\"]\n    if _type == \"exp\":\n        out_val = _exp_schedule(init_val, float(params[\"gamma\"]), epoch)\n    elif _type == \"step\":\n        out_val = _step_schedule(curr_val, params[\"steps\"], epoch)\n    return max(1, int(out_val))\n\n\nclass ColorExponentialLR(torch.optim.lr_scheduler.ExponentialLR):\n    def __init__(self, optimizer, gamma, param_group_index=0, last_epoch=-1, verbose=False):\n        self.param_group_index = param_group_index\n        super().__init__(optimizer, gamma, last_epoch, verbose)\n    \n    def get_lr(self):        \n        # Only update the learning rate for the specified parameter group\n        lrs = []\n        for i, _ in enumerate(self.optimizer.param_groups):\n            if i == self.param_group_index:\n                lrs.append(self.base_lrs[i] * self.gamma ** self.last_epoch)\n            else:\n                lrs.append(self.optimizer.param_groups[i]['lr'])\n        return lrs\n\n\ndef _calc_bs_mul(method: str, init_bs: float, curr_bs: float) -> float:\n    assert method in [\"constant\", \"linear\", \"sqrt\"]\n    if method == \"linear\":\n        return curr_bs / init_bs\n    elif method == \"sqrt\":\n        return np.sqrt(curr_bs / init_bs)\n    else:\n        return 1.\n\n\ndef get_optimizer(config: dict, params: dict) -> torch.optim.Optimizer:\n    optimizer_config = config[\"optimizer\"]\n    init_bs = 1\n    \n    lr_mul = _calc_bs_mul(\n        method=optimizer_config[\"batch_scal_method\"],\n        init_bs=init_bs,\n        curr_bs=config[\"dloader\"][\"batch_size\"]\n    )\n    \n    _params = [{\"params\": val, \"name\": key, \"lr\": lr_mul * float(optimizer_config[\"lrs\"][key])} for key, val in params.items()]\n    if optimizer_config[\"name\"] == \"adam\":\n        return torch.optim.Adam(\n            _params,\n            lr=0.\n        )\n    else:\n        raise NotImplementedError(f\"{optimizer_config['name']} not yet implemented\")\n\n\ndef get_scheduler(optimizer: torch.optim.Optimizer, params: dict) -> ColorExponentialLR:\n    if not params[\"use\"]:\n        return None\n    return ColorExponentialLR(\n        optimizer,\n        params[\"gamma\"],\n        0\n    )\n\n\nclass Losser:\n    def __init__(self, loss_cfg: dict) -> None:\n        self.loss_cfg = loss_cfg\n        self.pips = None\n        if any([\"pips\" in x for x in self.loss_cfg.keys()]):\n            self.pips = get_lpips_model(\n                \"cuda\",\n                \"vgg\"\n            )\n            \n        self._method_lookup = {\n            \"ssim\": {\n                \"method\": self.ssim_wrap,\n                \"params\": [\"img_pred\", \"img_gt\"]\n            },\n            \"img_l1\": {\n                \"method\": self.img_l1_wrap,\n                \"params\": [\"img_pred\", \"img_gt\"]\n            },\n            \"pips\": {\n                \"method\": self.pips_wrap,\n                \"params\": [\"img_pred\", \"img_gt\"]\n            },\n            \"psnr\": {\n                \"method\": self.psnr_wrap,\n                \"params\": [\"img_pred\", \"img_gt\"]\n            },\n            \"dice\": {\n                \"method\": self.dice_wrap,\n                \"params\": [\"mask_pred\", \"mask_gt\"]\n            },\n            \"delta_xyz\": {\n                \"method\": self.delta_wrap,\n                \"params\": [\"delta_xyz\"]\n            },\n            \"delta_rots\": {\n                \"method\": self.delta_wrap,\n                \"params\": [\"delta_rots\"]\n            },\n            \"delta_scales\": {\n                \"method\": self.delta_wrap,\n                \"params\": [\"delta_scales\"]\n            },\n            \"mse\": {\n                \"method\": self.mse_wrap,\n                \"params\": [\"img_pred\", \"img_gt\"]\n            },\n            \"body_displ\": {\n                \"method\": self.vector_norm_mean,\n                \"params\": [\"body_displ_pred\", \"body_displ_gt\"]\n            },\n            \"pose_displ\": {\n                \"method\": self.vector_norm_mean,\n                \"params\": [\"pose_displ_pred\", \"pose_displ_gt\"]\n            },\n            \"final_splat\": {\n                \"method\": self.vector_norm_mean,\n                \"params\": [\"final_splats_pred\", \"final_splats_gt\"]\n            }\n        }\n\n    def vector_norm_mean(self, loss_data: dict, params: list) -> torch.Tensor:\n        return torch.mean(torch.norm(\n            loss_data[params[0]] - loss_data[params[1]],\n            dim=1,\n            keepdim=True\n        ))\n    \n    def img_l1_wrap(self, loss_data: dict, params: list) -> torch.Tensor:\n        return l1_loss(loss_data[params[0]], loss_data[params[1]])\n\n    def ssim_wrap(self, loss_data: dict, params: list) -> torch.Tensor:\n        return 1. - ssim(loss_data[params[0]], loss_data[params[1]])\n\n    def mse_wrap(self, loss_data: dict, params: list) -> torch.Tensor:\n        return mse_loss(loss_data[params[0]], loss_data[params[1]])\n\n    def pips_wrap(self, loss_data: dict, params: list) -> torch.Tensor:\n        inp_data_1, inp_data_2 = loss_data[params[0]], loss_data[params[1]]\n        if inp_data_1.shape[1] != 3:\n            inp_data_1 = inp_data_1.permute(0, 3, 1, 2)\n        if inp_data_2.shape[1] != 3:\n            inp_data_2 = inp_data_2.permute(0, 3, 1, 2)\n        return self.pips(inp_data_1, inp_data_2)\n\n    def dice_wrap(self, loss_data: dict, params: list) -> torch.Tensor:\n        return dice_loss(loss_data[params[0]], loss_data[params[1]])\n\n    def delta_wrap(self, loss_data: dict, params: list) -> torch.Tensor:\n        return norm_loss(loss_data[params[0]], p=2)\n    \n    def psnr_wrap(self, loss_data: dict, params: list) -> torch.Tensor:\n        pred_img = loss_data[params[0]]\n        gt_img = loss_data[params[1]]\n        return psnr(pred_img, gt_img)\n\n    def __call__(self, loss_data: dict) -> dict:\n        out_losses = {}\n        for loss_name in self.loss_cfg.keys():\n            loss_method_name = loss_name.split(\"-\")[0]\n            loss_method = self._method_lookup[loss_method_name][\"method\"]\n            params = self._method_lookup[loss_method_name][\"params\"]\n            out_losses[loss_name] = loss_method(loss_data, params)\n        out_losses[\"loss\"] = sum([out_losses[key] * float(weight) for key, weight in self.loss_cfg.items() if \"debug\" not in key])\n        return out_losses\n\n\ndef psnr(pred_img: torch.Tensor, gt_img: torch.Tensor, reduction: str = \"mean\") -> torch.Tensor:\n    assert pred_img.shape[0] == gt_img.shape[0]\n    mse = (((pred_img - gt_img)) ** 2).view(pred_img.shape[0], -1).mean(1, keepdim=True).mean()\n    return 20 * torch.log10(1.0 / torch.sqrt(mse))\n\n\ndef dice_loss(mask_pred: torch.Tensor, mask_gt: torch.Tensor, smooth: float = 1e-6) -> torch.Tensor:\n    \"\"\"\n    masks both are tensors which have values 0 and 1 (0 for bg, 1 for mask)\n    \"\"\"\n    intersection = torch.sum(mask_pred * mask_gt)\n    cardinality = torch.sum(mask_pred) + torch.sum(mask_gt)\n    dice = (2 * intersection + smooth) / (cardinality + smooth)\n    return 1 - dice\n\n\ndef norm_loss(data: torch.Tensor, p: int = 1, dim: int = 1) -> torch.Tensor:\n    return torch.norm(data, p=p, dim=dim).mean()\n\n\ndef l1_loss(network_output: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:\n    return torch.abs((network_output - gt)).mean()\n\n\ndef mse_loss(network_output: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:\n    return torch.sum((network_output - gt) ** 2) / (np.prod(list(network_output.shape)))\n\n\ndef gaussian(window_size: int, sigma: float) -> torch.Tensor:\n    gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])\n    return gauss / gauss.sum()\n\n\ndef create_window(window_size: int, channel: int) -> torch.Tensor:\n    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)\n    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)\n    window = torch.autograd.Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())\n    return window\n\n\ndef ssim(img1: torch.Tensor, img2: torch.Tensor, window_size: int = 11, size_average: bool = True) -> torch.Tensor:\n    channel = img1.size(-3)\n    window = create_window(window_size, channel)\n\n    if img1.is_cuda:\n        window = window.cuda(img1.get_device())\n    window = window.type_as(img1)\n\n    return _ssim(img1, img2, window, window_size, channel, size_average)\n\n\ndef _ssim(img1: torch.Tensor, img2: torch.Tensor, window: int, window_size: int, channel: int, size_average: bool = True) -> torch.Tensor:\n    mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)\n    mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)\n\n    mu1_sq = mu1.pow(2)\n    mu2_sq = mu2.pow(2)\n    mu1_mu2 = mu1 * mu2\n\n    sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq\n    sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq\n    sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2\n\n    C1 = 0.01 ** 2\n    C2 = 0.03 ** 2\n\n    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))\n\n    if size_average:\n        return ssim_map.mean()\n    else:\n        return ssim_map.mean(1).mean(1).mean(1)"
  },
  {
    "path": "mesh_optim/models.py",
    "content": "import numpy as np\nimport nvdiffrast.torch as dr\nimport torch\nimport torch.nn as nn\n\n\nclass Pseudomesh(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        self.vert_requires_grad = \"vertices\" in config[\"model\"][\"optimizable_params\"]\n        self.vert_col_requires_grad = \"vertices\" in config[\"model\"][\"optimizable_params\"]\n        \n        self.register_buffer('faces', torch.tensor(0).int())\n        self.vertices = nn.Parameter(torch.tensor(0., requires_grad=self.vert_requires_grad))\n        self.vertex_colors = nn.Parameter(torch.tensor(0., requires_grad=self.vert_col_requires_grad))\n        self.register_buffer('grad_acc', torch.tensor(0).float())\n        self.register_buffer('grad_denom', torch.tensor(0).int())\n\nclass PseudomeshRenderer():\n    def __init__(self, config):\n        self.pseudomesh = Pseudomesh(config)\n        self.glctx = dr.RasterizeGLContext()\n    \n    @classmethod\n    def create_model(cls, config):\n        mesh_data = np.load(config[\"pseudomesh_path\"])\n        obj = cls(config)\n        obj.set_values(mesh_data[\"vertices\"], mesh_data[\"faces\"], mesh_data[\"vertex_colors\"])\n        return obj\n    \n    @property\n    def vert_requires_grad(self):\n        return self.pseudomesh.vert_requires_grad\n    \n    @property\n    def vert_col_requires_grad(self):\n        return self.pseudomesh.vert_col_requires_grad\n    \n    def set_values(self, vertices, faces, vertex_colors, grad_acc=None):\n        if isinstance(vertices, np.ndarray) and not isinstance(vertices, torch.nn.parameter.Parameter):\n            vertices = torch.tensor(vertices, dtype=torch.float32, requires_grad=self.vert_requires_grad)\n        if isinstance(faces, np.ndarray):\n            faces = torch.tensor(faces, dtype=torch.int32)\n        if isinstance(vertex_colors, np.ndarray) and not isinstance(vertex_colors, torch.nn.parameter.Parameter):\n            vertex_colors = torch.tensor(vertex_colors, dtype=torch.float32, requires_grad=self.vert_col_requires_grad)\n        self.pseudomesh.faces.data = faces\n        \n        if isinstance(vertices, torch.nn.parameter.Parameter):\n            self.pseudomesh.vertices = vertices\n        else:\n            self.pseudomesh.vertices = nn.Parameter(vertices)\n        \n        if isinstance(vertex_colors, torch.nn.parameter.Parameter):\n            self.pseudomesh.vertex_colors = vertex_colors\n        else:\n            self.pseudomesh.vertex_colors = nn.Parameter(vertex_colors)\n            \n        if grad_acc is None:\n            grad_acc = torch.zeros_like(self.pseudomesh.vertex_colors[:,0])\n        self.pseudomesh.grad_acc.data = grad_acc\n    \n    def acc_grad(self):\n        grad = torch.norm(self.pseudomesh.vertex_colors.grad, dim=-1)\n        self.pseudomesh.grad_acc.data = self.pseudomesh.grad_acc.data + grad\n        self.pseudomesh.grad_denom.data = self.pseudomesh.grad_denom.data + 1\n    \n    def reset_acc_grad(self):\n        self.pseudomesh.grad_acc.data = torch.zeros_like(self.pseudomesh.vertex_colors[:,0])\n        self.pseudomesh.grad_denom.data = torch.tensor(0).int()\n    \n    def to(self, device):\n        self.pseudomesh = self.pseudomesh.to(device)\n        return self\n    \n    def __call__(self, mvp_mat, width, height, num_layers):\n        pos_clip =         torch.matmul(\n            torch.cat(\n                [\n                    self.pseudomesh.vertices,\n                    torch.ones_like(self.pseudomesh.vertices[:, :1])\n                ],\n                dim=1\n            ),\n            mvp_mat.permute(0, 2, 1)\n        )\n        \n        final_color = torch.zeros((mvp_mat.shape[0], height, width, 4), device=self.pseudomesh.vertices.device, dtype=torch.float32)\n        \n        with dr.DepthPeeler(self.glctx, pos_clip, self.pseudomesh.faces, (height, width)) as peeler:\n            for layer_idx in range(num_layers):\n                rast_out, rast_db = peeler.rasterize_next_layer()\n                \n                if rast_out is None:\n                    break\n                \n                color_layer = dr.interpolate(\n                    self.pseudomesh.vertex_colors[None, ...],\n                    rast_out,\n                    self.pseudomesh.faces,\n                    rast_db=rast_db,\n                    diff_attrs='all'\n                )[0]\n                \n                alpha = color_layer[..., 3:4]\n                rgb = color_layer[..., :3]\n                final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([rgb * alpha, alpha], dim=-1)\n        \n        alpha = final_color[..., 3:4]\n        rgb = final_color[..., :3]\n        return rgb, alpha\n\n    def get_depth_map(self, mvp_mat, width, height, cam_pos, num_layers):\n        vertices = self.pseudomesh.vertices  # shape: (N, 3)\n        vertex_dist = torch.norm(vertices - cam_pos, dim=1, keepdim=True)  # shape: (N, 1)\n        \n        dmin = vertex_dist.min()\n        dmax = vertex_dist.max()\n        vertex_dist_norm = (vertex_dist - dmin) / (dmax - dmin + 1e-8)\n\n        vertex_colors = torch.cat([\n            vertex_dist_norm, \n            vertex_dist_norm, \n            vertex_dist_norm, \n            self.pseudomesh.vertex_colors[..., -1:]\n        ], dim=1)  # shape: (N, 4)\n        \n        vertices_homo = torch.cat([vertices, torch.ones_like(vertices[:, :1])], dim=1)\n        pos_clip = torch.matmul(vertices_homo, mvp_mat.permute(0, 2, 1))\n        \n        final_color = torch.zeros((mvp_mat.shape[0], height, width, 4), device=self.pseudomesh.vertices.device, dtype=torch.float32)\n        \n        with dr.DepthPeeler(self.glctx, pos_clip, self.pseudomesh.faces, (height, width)) as peeler:\n            for layer_idx in range(num_layers):\n                rast_out, rast_db = peeler.rasterize_next_layer()\n                \n                if rast_out is None:\n                    break\n                \n                color_layer = dr.interpolate(\n                    vertex_colors[None, ...],\n                    rast_out,\n                    self.pseudomesh.faces,\n                    rast_db=rast_db,\n                    diff_attrs='all'\n                )[0]\n                \n                alpha = color_layer[..., 3:4]\n                rgb = color_layer[..., :3]\n                final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([rgb * alpha, alpha], dim=-1)\n\n        alpha = final_color[..., 3:4]\n        rgb = final_color[..., :3]\n        return rgb, alpha\n    \n    def get_gray_map(self, mvp_mat, width, height, num_layers, color_verts):\n        pos_clip = torch.matmul(\n            torch.cat(\n                [\n                    self.pseudomesh.vertices,\n                    torch.ones_like(self.pseudomesh.vertices[:, :1])\n                ],\n                dim=1\n            ),\n            mvp_mat.permute(0, 2, 1)\n        )\n        \n        final_color = torch.zeros((mvp_mat.shape[0], height, width, 4), device=self.pseudomesh.vertices.device, dtype=torch.float32)\n        \n        with dr.DepthPeeler(self.glctx, pos_clip, self.pseudomesh.faces, (height, width)) as peeler:\n            for layer_idx in range(num_layers):\n                rast_out, rast_db = peeler.rasterize_next_layer()\n                \n                if rast_out is None:\n                    break\n                \n                color_layer = dr.interpolate(\n                    color_verts[None, ...],\n                    rast_out,\n                    self.pseudomesh.faces,\n                    rast_db=rast_db,\n                    diff_attrs='all'\n                )[0]\n                alpha = color_layer[..., 3:4]\n                rgb = color_layer[..., :3]\n                final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([rgb * alpha, alpha], dim=-1)\n        \n        alpha = final_color[..., 3:4]\n        rgb = final_color[..., :3]\n        return rgb, alpha\n\n    def get_normal_map(self, mvp_mat, width, height, num_layers):\n        vertices = self.pseudomesh.vertices\n        faces = self.pseudomesh.faces\n        \n        # Get face vertices\n        v0 = vertices[faces[:, 0]]\n        v1 = vertices[faces[:, 1]]\n        v2 = vertices[faces[:, 2]]\n        \n        # Calculate face normals\n        face_normals = torch.cross(v1 - v0, v2 - v0)\n        face_normals = face_normals / (torch.norm(face_normals, dim=1, keepdim=True) + 1e-8)\n        \n        # Initialize vertex normals\n        vertex_normals = torch.zeros_like(vertices)\n        \n        # Accumulate face normals to vertices\n        for i in range(3):\n            vertex_indices = faces[:, i]\n            vertex_normals.index_add_(0, vertex_indices, face_normals)\n        \n        # Normalize vertex normals\n        vertex_normals = vertex_normals / (torch.norm(vertex_normals, dim=1, keepdim=True) + 1e-8)\n        \n        # Rest of the rendering process remains the same\n        pos_clip = torch.matmul(\n            torch.cat(\n                [\n                    vertices,\n                    torch.ones_like(vertices[:, :1])\n                ],\n                dim=1\n            ),\n            mvp_mat.permute(0, 2, 1)\n        )\n        \n        vertex_attrs = torch.cat([\n            vertex_normals,\n            self.pseudomesh.vertex_colors[..., -1:]\n        ], dim=1)\n        \n        final_color = torch.zeros((mvp_mat.shape[0], height, width, 4), \n                                device=vertices.device, \n                                dtype=torch.float32)\n        \n        with dr.DepthPeeler(self.glctx, pos_clip, faces, (height, width)) as peeler:\n            for layer_idx in range(num_layers):\n                rast_out, rast_db = peeler.rasterize_next_layer()\n                \n                if rast_out is None:\n                    break\n                \n                color_layer = dr.interpolate(\n                    vertex_attrs[None, ...],\n                    rast_out,\n                    faces,\n                    rast_db=rast_db,\n                    diff_attrs='all'\n                )[0]\n                \n                alpha = color_layer[..., 3:4]\n                normals = color_layer[..., :3]\n                normals = normals / (torch.norm(normals, dim=-1, keepdim=True) + 1e-8)\n                final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([normals * alpha, alpha], dim=-1)\n        \n        alpha = final_color[..., 3:4]\n        rgb = final_color[..., :3]\n        return rgb, alpha\n\n    def render_all_maps(self, mvp_mat, width, height, cam_pos, num_layers, gray_verts=None):\n        vertices = self.pseudomesh.vertices\n        faces = self.pseudomesh.faces\n        \n        # Calculate vertex normals\n        v0 = vertices[faces[:, 0]]\n        v1 = vertices[faces[:, 1]]\n        v2 = vertices[faces[:, 2]]\n        \n        face_normals = torch.cross(v1 - v0, v2 - v0)\n        face_normals = face_normals / (torch.norm(face_normals, dim=1, keepdim=True) + 1e-8)\n        \n        vertex_normals = torch.zeros_like(vertices)\n        for i in range(3):\n            vertex_indices = faces[:, i]\n            vertex_normals.index_add_(0, vertex_indices, face_normals)\n        vertex_normals = vertex_normals / (torch.norm(vertex_normals, dim=1, keepdim=True) + 1e-8)\n        \n        # Calculate depth values\n        vertex_dist = torch.norm(vertices - cam_pos, dim=1, keepdim=True)\n        dmin, dmax = vertex_dist.min(), vertex_dist.max()\n        vertex_dist_norm = (vertex_dist - dmin) / (dmax - dmin + 1e-8)\n        depth_colors = torch.cat([\n            vertex_dist_norm.repeat(1, 3),\n            self.pseudomesh.vertex_colors[..., -1:]\n        ], dim=1)\n        \n        # Prepare normal colors\n        normal_colors = torch.cat([\n            vertex_normals,\n            self.pseudomesh.vertex_colors[..., -1:]\n        ], dim=1)\n        \n        # Initialize output tensors\n        device = vertices.device\n        shape = (mvp_mat.shape[0], height, width, 4)\n        final_color = torch.zeros(shape, device=device, dtype=torch.float32)\n        final_normal = torch.zeros(shape, device=device, dtype=torch.float32)\n        final_depth = torch.zeros(shape, device=device, dtype=torch.float32)\n        final_gray = torch.zeros(shape, device=device, dtype=torch.float32) if gray_verts is not None else None\n        \n        # Transform vertices to clip space\n        pos_clip = torch.matmul(\n            torch.cat([vertices, torch.ones_like(vertices[:, :1])], dim=1),\n            mvp_mat.permute(0, 2, 1)\n        )\n        \n        with dr.DepthPeeler(self.glctx, pos_clip, faces, (height, width)) as peeler:\n            for _ in range(num_layers):\n                rast_out, rast_db = peeler.rasterize_next_layer()\n                \n                if rast_out is None:\n                    break\n                \n                # Interpolate all attributes at once\n                color_layer = dr.interpolate(\n                    self.pseudomesh.vertex_colors[None, ...],\n                    rast_out, faces, rast_db=rast_db, diff_attrs='all'\n                )[0]\n                \n                normal_layer = dr.interpolate(\n                    normal_colors[None, ...],\n                    rast_out, faces, rast_db=rast_db, diff_attrs='all'\n                )[0]\n                \n                depth_layer = dr.interpolate(\n                    depth_colors[None, ...],\n                    rast_out, faces, rast_db=rast_db, diff_attrs='all'\n                )[0]\n                \n                # Process color map\n                alpha = color_layer[..., 3:4]\n                rgb = color_layer[..., :3]\n                final_color = final_color + (1.0 - final_color[..., 3:4]) * torch.cat([rgb * alpha, alpha], dim=-1)\n                \n                # Process normal map\n                normals = normal_layer[..., :3]\n                normals = normals / (torch.norm(normals, dim=-1, keepdim=True) + 1e-8)\n                final_normal = final_normal + (1.0 - final_normal[..., 3:4]) * torch.cat([normals * alpha, alpha], dim=-1)\n                \n                # Process depth map\n                depth_rgb = depth_layer[..., :3]\n                final_depth = final_depth + (1.0 - final_depth[..., 3:4]) * torch.cat([depth_rgb * alpha, alpha], dim=-1)\n                \n                # Process gray map if gray_verts provided\n                if gray_verts is not None:\n                    gray_layer = dr.interpolate(\n                        gray_verts[None, ...],\n                        rast_out, faces, rast_db=rast_db, diff_attrs='all'\n                    )[0]\n                    gray_rgb = gray_layer[..., :3]\n                    final_gray = final_gray + (1.0 - final_gray[..., 3:4]) * torch.cat([gray_rgb * alpha, alpha], dim=-1)\n        \n        results = {\n            'color': (final_color[..., :3], final_color[..., 3:4]),\n            'normal': (final_normal[..., :3], final_normal[..., 3:4]),\n            'depth': (final_depth[..., :3], final_depth[..., 3:4])\n        }\n        \n        if gray_verts is not None:\n            results['gray'] = (final_gray[..., :3], final_gray[..., 3:4])\n        \n        return results\n"
  },
  {
    "path": "mesh_optim/optimize_pseudo_config.yaml",
    "content": "experiment_name: exp_name\nexperiments_dir: /path/to/experiments_dir\nmode: depth\n\npseudomesh_path: /path/to/pseudomesh.npz\ndevice: cuda\nwhite_background: False\n\ndataset:\n  dataset_path: /dset/path\n  near: 0.01\n  far: 100.\n  imgs_in_ram: True\n  res: 1\n\ndloader:\n  batch_size: 4\n  shuffle: True\n  num_workers: 10\n\ntraining:\n  epochs: 50\n\nmodel:\n  optimizable_params: \n    - vertices\n    - vertex_colors\n  optim_epoch_start:\n    vertices: 0\n\nalpha_pruning:\n  type: alpha  # possible modes: alpha, gradient and both\n  start_epoch: 10\n  epoch_step: 10\n  alpha_eps: 1e-4\n  grad_eps: 1e-7\n  mode: soft  # soft or hard\n\noptimizer:\n  name: adam\n  batch_scal_method: \"linear\"  # [linear, constant, sqrt]\n  lrs:\n    vertices: 1e-6\n    vertex_colors: 5e-3\n\nlr_scheduler:\n  use: True\n  gamma: 0.999\n\nrenderer:\n  depth_steps: 50\n  dp_scheduler:\n    perform: False\n    type: step  # step of exp\n    init_depth_steps: 50\n    # params:  # for exp\n    #   gamma: 0.998\n    params:   # for step\n      steps:\n        0: 50\n        60: 1\n\nloss:  # name of the loss and it's weight\n  ssim: .4\n  img_l1: .6\n  psnr-debug: 0.\n  # pips: 1e1\n  # dice: 1.\n  # delta_xyz: 1e-2\n  # delta_scales: 1e-2\n  # delta_rots: 5e-2\n\ntest_loss:\n  ssim: 1.\n  img_l1: 1.\n  psnr: 1.\n  pips: 1.\n\nwandb:\n  use: True\n  key_path: /path/to/wandb_key.json\n  project_name: project_name\n  entity: entity_name\n  imgs_per_epoch: 4\n"
  },
  {
    "path": "mesh_optim/optimize_pseudomesh.py",
    "content": "import argparse\nfrom collections import defaultdict\nimport json\nimport logging\nimport logging.config\nfrom pathlib import Path\nimport shutil\nfrom time import time\nimport yaml\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\nimport torchvision\nfrom tqdm import tqdm\nimport wandb\n\nfrom data import ImageCamDataset\nfrom mesh_utils import prune_mesh, prune_optimizer\nfrom ml_utils import dp_schedule, get_optimizer, Losser, get_scheduler\nfrom models import PseudomeshRenderer\n\nLOGGER = None\n\ndef _load_config(path):\n    with open(path, \"r\") as file:\n        data = yaml.load(file, Loader=yaml.FullLoader)\n    return data\n\ndef setup_logger(config):\n    with open(\"./log_config.yaml\", \"r\") as file:\n        _config = yaml.load(file, Loader=yaml.FullLoader)\n    log_filepath = Path(\n        config[\"experiments_dir\"], config[\"experiment_name\"], \"logs.log\"\n    )\n    log_filepath.parent.mkdir(exist_ok=True, parents=True)\n    _config[\"handlers\"][\"file_handler\"][\"filename\"] = str(log_filepath)\n    logging.config.dictConfig(_config)\n    \n    global LOGGER\n    LOGGER = logging.getLogger(__name__)\n\n\ndef 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):\n    if not test:\n        optimizer.zero_grad()\n    \n    # get bakcgorund\n    bg_color = [1, 1, 1] if config[\"white_background\"] else [0, 0, 0]\n    background = torch.tensor(bg_color, dtype=torch.float32, device=\"cuda\")\n    \n    mvp_mat = data[\"mvp_matrix\"]\n    gt_image = data[\"img\"]\n    \n    if test:\n        start_time = time()\n    rgb, alpha = model(\n        mvp_mat,\n        gt_image.shape[2],  # width\n        gt_image.shape[1],  # height\n        config[\"renderer\"][\"depth_steps\"]\n    )\n    if test:\n        end_time = time() - start_time\n    \n    # handle background\n    final_color = rgb * alpha + background[..., :3] * (1 - alpha)\n    if gt_image.shape[-1] == 4:\n        gt_image = gt_image[... ,:3] * gt_image[..., 3:4] + background * (1 - gt_image[..., 3:4])\n    else:\n        gt_image = gt_image[..., :3]\n    \n    # calculate losses\n    data_to_loss = {\n        \"img_pred\": final_color,\n        \"img_gt\": gt_image,\n    }\n    losses = loss_obj(data_to_loss)\n    if test:\n        losses[\"render_time\"] = torch.tensor(end_time)\n    \n    # update net\n    if not test:\n        losses[\"loss\"].backward()\n        model.acc_grad()\n        optimizer.step()\n        if scheduler is not None:\n            scheduler.step()\n    return losses, data_to_loss\n\n\ndef wandb_log(losses: dict, step: int, data_to_loss: dict = None, mode: str = \"train\", epoch: int = None) -> None:\n    assert mode in [\"train\", \"val\"]\n    if losses is not None:\n        log_losses = {f\"{mode}_{_k}\": _v.item() if isinstance(_v, torch.Tensor) else _v for _k, _v in losses.items()}\n        wandb.log(\n            data=log_losses,\n            step=step\n        )\n    if data_to_loss is not None:\n        grid_data = []\n        name = \"\"\n        for data_name, data in data_to_loss.items():\n            if \"delta\" in data_name:\n                continue\n            if len(data.shape) == 4:\n                data = data.squeeze(0)\n            if data.shape[0] == 1:\n                data = data.repeat(3, 1, 1)\n            grid_data.append(data)\n            name += \"\" if not name else \"   \"\n            name += data_name\n        grid_data = torch.cat([x.permute(0, 3, 1, 2) for x in grid_data], dim=0)\n        grid_img = torchvision.utils.make_grid(grid_data, nrow=grid_data.shape[0] // 2)\n        grid_img = torch.clip(grid_img, 0., 1.)\n        grid_img = F.interpolate(\n            grid_img.unsqueeze(0), \n            size=(grid_img.shape[1]//4, grid_img.shape[2]//4), \n            mode='bilinear', \n            align_corners=False\n        ).squeeze(0).permute(1, 2, 0).detach().cpu().numpy() * 255\n        grid_img = grid_img.astype(np.uint8)\n        img_grid = wandb.Image(grid_img, mode=\"RGB\")\n        wandb.log(\n            data={name: img_grid},\n            step=step\n        )\n        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}\n        wandb.log(\n            data=deltas,\n            step=step\n        )\n\n\ndef _collect_imgs_for_logs(data_to_loss:dict, container: dict, _iter: int, step: int) -> None:\n    if not _iter % step:\n        for _img_name, _img in data_to_loss.items():\n            if container.get(_img_name, None) is None:\n                container[_img_name] = _img\n            else:\n                container[_img_name] = torch.cat(\n                    [\n                        container[_img_name],\n                        _img\n                    ],\n                    dim=0\n                )\n\n\ndef _save_ckpt(ckpt_dir: Path, epoch: int, model: PseudomeshRenderer, optimizer: torch.optim.Optimizer) -> None:\n    out_path = Path(ckpt_dir, \"best_model\")\n    LOGGER.info(f\"Save checkpoint: {out_path}\")\n    \n    ckpt = {\n        \"epoch\": epoch,\n        \"model\": model.pseudomesh.state_dict(),\n        \"optimizer\": optimizer.state_dict()\n    }\n    torch.save(ckpt, out_path)\n\n\ndef batch_to_device(data, device):\n    data[\"img\"] = data[\"img\"].detach().to(device)\n    data[\"mvp_matrix\"] = data[\"mvp_matrix\"].detach().to(device)\n\n\ndef pruning(model, optimizer, alpha_eps, grad_eps, mode, _type):\n    init_shape = model.pseudomesh.vertices.shape[0]\n    if _type in [\"alpha\", \"both\"]:\n        opacity_mask = (model.pseudomesh.vertex_colors[:, -1:] < eval(alpha_eps)).reshape(-1)\n    else: \n        opacity_mask = torch.zeros_like(model.pseudomesh.vertex_colors[:, -1:]).reshape(-1).bool()\n    if _type in [\"gradient\", \"both\"]:\n        grad_mask = (model.pseudomesh.grad_acc / model.pseudomesh.grad_denom) < eval(grad_eps)\n    else:\n        grad_mask = torch.zeros_like(model.pseudomesh.vertex_colors[:, -1:]).reshape(-1).bool()\n    grad_opac_mask = torch.logical_or(opacity_mask, grad_mask)\n    new_verts, new_faces, new_vert_colors, mask = prune_mesh(\n        model.pseudomesh.vertices,\n        model.pseudomesh.faces,\n        model.pseudomesh.vertex_colors,\n        grad_opac_mask,\n        mode\n    )\n    LOGGER.info(f\"Pruned {init_shape - new_verts.shape[0]} vertices. Left vertices: {new_verts.shape[0]}, left faces: {new_faces.shape[0]}\")\n    optimized_params = prune_optimizer(optimizer, mask)\n    \n    new_verts = optimized_params.get(\"vertices\", new_verts)\n    new_vert_colors = optimized_params.get(\"vertex_colors\", new_vert_colors)\n    model.set_values(new_verts, new_faces, new_vert_colors)\n    model.reset_acc_grad()\n\n\ndef train(pseudomesh: PseudomeshRenderer, dloader: DataLoader, config: dict) -> None:\n    # create optimizer\n    params = {\n        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\"]\n    }\n    optimizer = get_optimizer(config, params)\n    scheduler = get_scheduler(optimizer, config[\"lr_scheduler\"])\n    # create loss function\n    loss_obj = Losser(config[\"loss\"])\n    \n    _iter = 0\n    best_loss = float(\"inf\")\n    start_time = time()\n    best_time = time()\n    \n    LOGGER.info(\"Start training\")\n    for epoch in range(config[\"training\"][\"epochs\"]):\n        LOGGER.info(f\"Epoch {epoch}\")\n        # train\n        _batch_iter = 0\n        _img_iter = 0\n        batch_loss = 0\n        train_data_to_loss_epoch = {}\n        for data in tqdm(dloader):\n            batch_to_device(data, config[\"device\"])\n            losses, data_to_loss = iter_pass(\n                pseudomesh, \n                data, \n                config, \n                loss_obj, \n                optimizer,\n                scheduler\n            )\n            losses[\"depth_steps\"] = config[\"renderer\"][\"depth_steps\"]\n            losses[\"color_lr\"] = optimizer.param_groups[0]['lr']\n            if config[\"wandb\"][\"use\"]:\n                wandb_log(\n                    losses=losses,\n                    step=_iter,\n                    data_to_loss=None,\n                    mode=\"train\",\n                    epoch=epoch\n                )\n                _collect_imgs_for_logs(\n                    data_to_loss,\n                    train_data_to_loss_epoch,\n                    _img_iter,\n                    len(dloader.dataset) // config[\"wandb\"][\"imgs_per_epoch\"]\n                )\n            batch_loss += losses[\"loss\"].item()\n            _img_iter += data[\"img\"].shape[0]\n            _batch_iter += 1\n            _iter += data[\"img\"].shape[0]\n        if config[\"wandb\"][\"use\"]:\n            wandb_log(\n                losses=None,\n                step=_iter,\n                data_to_loss=train_data_to_loss_epoch,\n                mode=\"train\",\n                epoch=epoch\n            )\n        batch_loss /= _batch_iter\n        if batch_loss < best_loss:\n            best_time = time()\n            _save_ckpt(\n                ckpt_dir=config[\"ckpt_dir\"],\n                epoch=epoch,\n                model=pseudomesh,\n                optimizer=optimizer\n            )\n            best_loss = batch_loss\n\n        alpha_prun_cond = epoch >= config[\"alpha_pruning\"][\"start_epoch\"]\n        alpha_prun_cond = alpha_prun_cond and not (\n            (epoch - config[\"alpha_pruning\"][\"start_epoch\"]) % config[\"alpha_pruning\"][\"epoch_step\"]\n        )\n        if alpha_prun_cond:\n            LOGGER.info(\"Perform pruning\")\n            pruning(\n                pseudomesh, \n                optimizer, \n                config[\"alpha_pruning\"][\"alpha_eps\"], \n                config[\"alpha_pruning\"][\"grad_eps\"], \n                config[\"alpha_pruning\"][\"mode\"],\n                config[\"alpha_pruning\"][\"type\"]\n            )\n        # check if i should update optimizer params\n        for param_name, epoch_start in config[\"model\"][\"optim_epoch_start\"].items():\n            if epoch_start == epoch:\n                LOGGER.info(f\"Add {param_name} to optimizer\")\n                optimizer.add_param_group({\n                    \"name\": param_name,\n                    \"params\": getattr(pseudomesh.pseudomesh, param_name),\n                    \"lr\": float(config[\"optimizer\"][\"lrs\"][param_name])\n                })\n\n        # do the depth peeling scheduling\n        if config[\"renderer\"][\"dp_scheduler\"][\"perform\"]:  # if scheduling should be applied\n            new_dp_val = dp_schedule(\n                config[\"renderer\"][\"dp_scheduler\"][\"type\"],\n                config[\"renderer\"][\"depth_steps\"],\n                config[\"renderer\"][\"dp_scheduler\"][\"init_depth_steps\"],\n                epoch,\n                config[\"renderer\"][\"dp_scheduler\"][\"params\"]\n            )\n            if config[\"renderer\"][\"depth_steps\"] != new_dp_val:\n                config[\"renderer\"][\"depth_steps\"] = new_dp_val\n                LOGGER.info(f\"Changed depth steps to: {config['renderer']['depth_steps']}\")\n\n    return best_time - start_time\n\ndef test(pseudomesh: PseudomeshRenderer, dloader: DataLoader, config: dict) -> dict:\n    LOGGER.info(\"Start test\")\n    \n    loss_obj = Losser(config[\"test_loss\"])\n    metrics = defaultdict(lambda: 0)\n    _iter = 0\n    for data in tqdm(dloader):\n        batch_to_device(data, config[\"device\"])\n        with torch.no_grad():\n            losses, data_to_loss = iter_pass(\n                pseudomesh, \n                data, \n                config, \n                loss_obj, \n                None,\n                test=True\n            )\n        _iter += 1\n        for metric_name, metric_val in losses.items():\n            metrics[metric_name] += metric_val.item()\n        # save imgs?\n    for metric_name, metric_val in metrics.items():\n        metrics[metric_name] /= _iter\n    \n    return metrics\n\n\ndef prepare_output_dir(config: dict, cfg_path: str) -> None:\n    experiment_dir = Path(config[\"experiments_dir\"], config[\"experiment_name\"])\n    ckpt_dir = Path(experiment_dir, \"checkpoints\")\n    imgs_dir = Path(experiment_dir, \"imgs\")\n    config[\"experiment_dir\"] = experiment_dir\n    config[\"ckpt_dir\"] = ckpt_dir\n    config[\"imgs_dir\"] = imgs_dir\n    \n    ckpt_dir.mkdir(parents=True, exist_ok=True)\n    imgs_dir.mkdir(parents=True, exist_ok=True)\n\n    out_config_path = Path(config[\"experiment_dir\"], \"config.yaml\")\n    out_config = config\n    out_config[\"experiment_dir\"] = str(out_config[\"experiment_dir\"])\n    out_config[\"ckpt_dir\"] = str(out_config[\"ckpt_dir\"])\n    out_config[\"imgs_dir\"] = str(out_config[\"imgs_dir\"])\n    with open(out_config_path, \"w\") as file:\n        yaml.dump(out_config, file, default_flow_style=False)\n\n\ndef setup_wandb(config: dict) -> None:\n    LOGGER.info(\"Setup wandb\")\n    with open(config[\"wandb\"][\"key_path\"], \"r\") as file:\n        wandb_key = json.load(file)\n    wandb.login(\n        key=wandb_key,\n        relogin=True\n    )\n    wandb.init(\n        project=config[\"wandb\"][\"project_name\"],\n        entity=config[\"wandb\"][\"entity\"],\n        config=config,\n        name=config[\"experiment_name\"]\n    )\n\n\ndef _load_best_model(config: dict) -> PseudomeshRenderer:\n    pseudomesh = PseudomeshRenderer.create_model(config).to(config[\"device\"])\n    ckpt_path = Path(config[\"experiments_dir\"], config[\"experiment_name\"], \"checkpoints\", \"best_model\")\n    checkpoint = torch.load(ckpt_path)\n    pseudomesh.set_values(\n        checkpoint[\"model\"][\"vertices\"], \n        checkpoint[\"model\"][\"faces\"],                  \n        checkpoint[\"model\"][\"vertex_colors\"], \n        checkpoint[\"model\"][\"grad_acc\"]\n    )\n    pseudomesh.pseudomesh.load_state_dict(checkpoint[\"model\"])\n    return pseudomesh\n\n\ndef _save_results(results: dict, config: dict) -> None:\n    out_path = Path(config[\"experiments_dir\"], config[\"experiment_name\"], \"test_results.json\")\n    with open(out_path, \"w\") as file:\n        json.dump(results, file)\n        \n\ndef main() -> None:\n    _parser = argparse.ArgumentParser(\n        description=\"Optimize pseudomesh.\"\n    )\n    _parser.add_argument(\n        \"--cfg_path\", \"-cp\",\n        type=str, \n        help=\"Path to config\",\n        default=\"./optimize_pseudo_config.yaml\"\n    )\n    _parser.add_argument(\n        \"--exp_name\",\n        type=str, \n        help=\"Experiment name\",\n        default=\"\"\n    )\n    _parser.add_argument(\n        \"--exp_dir\",\n        type=str, \n        help=\"Experiments dir\",\n        default=\"\"\n    )\n    _parser.add_argument(\n        \"--pseudomesh_path\",\n        type=str, \n        help=\"Path to pseudomesh\",\n        default=\"\"\n    )\n    _parser.add_argument(\n        \"--dset_path\",\n        type=str, \n        help=\"Path to dset\",\n        default=\"\"\n    )\n    _parser.add_argument(\n        \"--white_background\",\n        action=\"store_true\"\n    )\n    _parser.add_argument(\n        \"--res\",\n        type=int,\n        default=-1\n    )\n\n    _args = _parser.parse_args()\n    config = _load_config(_args.cfg_path)\n    if _args.exp_name:\n        config[\"experiment_name\"] = _args.exp_name\n    if _args.pseudomesh_path:\n        config[\"pseudomesh_path\"] = _args.pseudomesh_path\n    if _args.dset_path:\n        config[\"dataset\"][\"dataset_path\"] = _args.dset_path\n    if _args.exp_dir:\n        config[\"experiments_dir\"] = _args.exp_dir\n    if _args.white_background:\n        config[\"white_background\"] = _args.white_background\n    if _args.res != -1:\n        config[\"dataset\"][\"res\"] = _args.res\n    prepare_output_dir(config, _args.cfg_path)\n    setup_logger(config)\n    if config[\"wandb\"][\"use\"]:\n        setup_wandb(config)\n    \n    # load init splat\n    pseudomesh = PseudomeshRenderer.create_model(config).to(config[\"device\"])\n    LOGGER.info(f\"Model starts with {pseudomesh.pseudomesh.vertices.shape[0]} vertices and {pseudomesh.pseudomesh.faces.shape[0]} faces.\")\n    dset = ImageCamDataset(\n        **config[\"dataset\"]\n    )\n    dloader = DataLoader(\n        dset,\n        **config[\"dloader\"]\n    )\n    \n    train_time_s = train( \n        pseudomesh=pseudomesh,\n        dloader=dloader,\n        config=config\n    )\n    \n    del pseudomesh\n    \n    # perform tests\n    torch.cuda.empty_cache()\n    \n    # load model\n    best_model = _load_best_model(config)\n    \n    # save model as npz\n    np.savez(\n        Path(config[\"experiments_dir\"], config[\"experiment_name\"], \"checkpoints\", \"best_model.npz\"), \n        vertices=best_model.pseudomesh.vertices.detach().cpu().numpy(), \n        faces=best_model.pseudomesh.faces.detach().cpu().numpy(), \n        vertex_colors=best_model.pseudomesh.vertex_colors.detach().cpu().numpy()\n    )\n    \n    final_train_losses = test(\n        pseudomesh=best_model,\n        dloader=dloader,\n        config=config\n    )\n    \n    test_dset = ImageCamDataset(\n        **{\n            \"test\": True,\n            **config[\"dataset\"],\n        }\n    )\n    test_dloader = DataLoader(\n        test_dset,\n        **config[\"dloader\"]\n    )\n    \n    final_test_losses = test(\n        pseudomesh=best_model,\n        dloader=test_dloader,\n        config=config\n    )\n    \n    out_losses = {\n        \"train_time\": train_time_s,\n        **{f\"train_{k}\": v for k, v in final_train_losses.items()},\n        **{f\"test_{k}\": v for k, v in final_test_losses.items()},\n    }\n    \n    _save_results(out_losses, config)\n\n\nif __name__ == \"__main__\":\n    main()\n\n\n\n"
  },
  {
    "path": "mesh_optim/render_pseudomesh.py",
    "content": "import argparse\nfrom pathlib import Path\nimport yaml\n\nimport torch\nfrom torch.utils.data import DataLoader\nimport torchvision\nfrom tqdm import tqdm\n\nfrom data import ImageCamDataset\nfrom optimize_pseudomesh import _load_best_model\n\n\ndef _load_config(path):\n    with open(path, \"r\") as file:\n        data = yaml.load(file, Loader=yaml.FullLoader)\n    return data\n\n\ndef render_img(model, mvp_mat, img_shape, white_background, depth_steps, device):\n    bg_color = [1, 1, 1] if white_background else [0, 0, 0]\n    background = torch.tensor(bg_color, dtype=torch.float32, device=device)\n\n    rgb, alpha = model(\n        mvp_mat,\n        img_shape[0],  # width\n        img_shape[1],  # height\n        depth_steps\n    )\n    \n    final_color = rgb * alpha + background[..., :3] * (1 - alpha)\n    return final_color\n\n\ndef main(cfg_path, method, white_background):\n    config = _load_config(cfg_path)\n    \n    method_dir_name = \"ours_30000\"\n    if method == \"2dgs\":\n        # method_dir_name = \"ours_10000\"  # for blender scenes\n        method_dir_name = \"ours_30000\"  # for real scenes\n    \n    out_imgs_dir = Path(\n        config[\"experiments_dir\"],\n        config[\"experiment_name\"],\n        \"test\",\n        method_dir_name,\n        \"pseudomesh_renders\"\n    )\n    out_imgs_dir.mkdir(parents=True, exist_ok=True)\n    \n    # setup dataloader without shuffle\n    dset = ImageCamDataset(\n        **{\n            \"test\": True,\n            **config[\"dataset\"],\n        }\n    )\n    dloader = DataLoader(\n        dset,\n        batch_size=1,\n        shuffle=False,\n        num_workers=10\n    )\n    # load model\n    pseudomesh = _load_best_model(config)\n\n    # render all images and save them in the experiment's directory\n    for data in tqdm(dloader):\n        out_path = Path(out_imgs_dir, data[\"name\"][0]).with_suffix(\".png\")\n        mvp_mat = data[\"mvp_matrix\"].detach().to(config[\"device\"])\n        \n        gt_img = data[\"img\"].squeeze()\n        height = gt_img.shape[0]\n        width = gt_img.shape[1]\n        \n        with torch.no_grad():\n            pred_img = render_img(\n                pseudomesh,\n                mvp_mat,\n                [width, height],\n                white_background,\n                config[\"renderer\"][\"depth_steps\"],\n                config[\"device\"]\n            )\n        \n        torchvision.utils.save_image(pred_img.squeeze().permute(2, 0, 1), out_path)\n\n\nif __name__ == \"__main__\":\n    _parser = argparse.ArgumentParser(\n        description=\"Process a single string argument.\"\n    )\n    _parser.add_argument(\n        \"--cfg_path\", \"-cp\",\n        type=str, \n        help=\"Path to config\",\n        required=True\n    )\n    _parser.add_argument(\n        \"--method\",\n        type=str,\n        choices=[\"3dgs\", \"games\", \"sugar\", \"2dgs\"],\n        required=True\n    )\n    _parser.add_argument(\n        '--white_background', \n        action='store_true'\n    )\n\n    _args = _parser.parse_args()\n    main(_args.cfg_path, _args.method, _args.white_background)\n"
  },
  {
    "path": "requirements.txt",
    "content": "# Core dependencies\nnumpy\ntorch>=2.0.0\ntorchvision\ntqdm\nPyYAML\nPillow\nmatplotlib\nopencv-python\n\n# Specialized libraries\nninja\nimageio\nPyOpenGL\nglfw\nxatlas\ngdown\ngit+https://github.com/NVlabs/nvdiffrast\nwandb"
  },
  {
    "path": "scripts/visualize_cameras_nerf_blender.py",
    "content": "from pathlib import Path\n\nimport bpy\nimport bpy.ops\nimport json\nfrom math import tan, atan\nfrom mathutils import Matrix\n\nimport numpy as np\n\n\nsensor_width = 36.0  # mm, typical of a full-frame camera\nsensor_height = 24.0  # mm, typical of a full-frame camera\n\n\ndef _calculate_fov(focal, pixels):\n    return 2 * np.arctan(pixels / (2 * focal))\n\n\ndef create_camera(name, position, rot_matrix, fovx, fovy, image_width, image_height):\n    # Create the camera object\n    bpy.ops.object.camera_add()\n    camera_obj = bpy.context.object\n    camera_obj.name = str(name)\n    camera_obj.location = position\n    \n    # Set rotation from matrix\n    rotation_matrix = Matrix(rot_matrix)\n    camera_obj.rotation_euler = rotation_matrix.to_euler()\n    \n    # Assuming the sensor width is fixed and calculating the sensor height based on the aspect ratio\n    sensor_width = 36.0  # Adjust as needed, standard full frame sensor width\n    aspect_ratio = image_width / image_height\n    sensor_height = sensor_width / aspect_ratio\n    \n    # Set camera data\n    camera = camera_obj.data\n    camera.sensor_width = sensor_width\n    camera.sensor_height = sensor_height\n    \n    # Calculate focal length from FoV using formula: focal_length = sensor_width / (2 * tan(fov / 2))\n    camera.lens = sensor_width / (2 * tan(fovx / 2))\n    \n    # Use FoVy to adjust the sensor height if needed\n    calculated_fovy = 2 * atan((sensor_height / 2) / camera.lens)\n    if calculated_fovy != fovy:\n        camera.sensor_height = 2 * camera.lens * tan(fovy / 2)\n    \n    return camera_obj\n\n\ndef main():\n    _cam_json_path = \"cameras.json\"\n    \n    with open(_cam_json_path, \"r\") as file:\n        _camera_data = json.load(file)\n    _camera_data = sorted(_camera_data, key=lambda x: x[\"id\"])\n    _camera_data = [{\n        \"id\": x[\"id\"],\n        \"img_name\": x[\"img_name\"],\n        \"width\": x[\"width\"],\n        \"height\": x[\"height\"],\n        \"position\": np.array(x[\"position\"]),\n        \"rotation\": np.array(x[\"rotation\"]),\n        \"fovy\": x[\"fy\"],\n        \"fovx\": x[\"fx\"]\n    } for x in _camera_data]\n    \n    adjustment = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]], dtype=np.float64)\n    \n    for _new_cam in _camera_data:\n        fovx = _calculate_fov(_new_cam[\"fovx\"], _new_cam[\"width\"])\n        fovy = _calculate_fov(_new_cam[\"fovy\"], _new_cam[\"height\"])\n        _new_cam_obj = create_camera(\n            name=_new_cam[\"id\"],\n            position=_new_cam[\"position\"],\n            rot_matrix=_new_cam[\"rotation\"] @ adjustment,\n            fovx=fovx,\n            fovy=fovy,\n            image_width=_new_cam[\"width\"],\n            image_height=_new_cam[\"height\"]\n        )\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "sh_scripts/configs/3dgs_sh0_pseudo_config.yaml",
    "content": "experiment_name: exp_name\nexperiments_dir: /path/to/experiments\nmode: depth  # possible modes:\n\npseudomesh_path: /path/to/pseudomesh.npz\ndevice: cuda\nwhite_background: True\n\ndataset:\n  dataset_path: /path/to/dataset\n  near: 0.01\n  far: 100.\n  imgs_in_ram: True\n  res: 1\n\ndloader:\n  batch_size: 4\n  shuffle: True\n  num_workers: 10\n\ntraining:\n  epochs: 15\n\nmodel:\n  optimizable_params: \n    - vertices\n    - vertex_colors\n  optim_epoch_start:\n    vertices: 0\n\nalpha_pruning:\n  type: alpha  # possible modes: alpha, gradient and both\n  start_epoch: 10\n  epoch_step: 10\n  alpha_eps: 1e-4\n  grad_eps: 1e-7\n  mode: soft  # soft or hard\n\noptimizer:\n  name: adam\n  batch_scal_method: \"linear\"  # [linear, constant, sqrt]\n  lrs:\n    vertices: 1e-6\n    vertex_colors: 5e-3\n\nlr_scheduler:\n  use: True\n  gamma: 0.999\n\nrenderer:\n  depth_steps: 50\n  dp_scheduler:\n    perform: False\n    type: step  # step of exp\n    init_depth_steps: 50\n    # params:  # for exp\n    #   gamma: 0.998\n    params:   # for step\n      steps:\n        0: 50\n        60: 1\n\nloss:  # name of the loss and it's weight\n  ssim: .4\n  img_l1: .6\n  psnr-debug: 0.\n  # pips: 1e1\n  # dice: 1.\n  # delta_xyz: 1e-2\n  # delta_scales: 1e-2\n  # delta_rots: 5e-2\n\ntest_loss:\n  ssim: 1.\n  img_l1: 1.\n  psnr: 1.\n  pips: 1.\n\nwandb:\n  use: True\n  key_path: /path/to/wandb.json\n  project_name: project_name\n  entity: entity_name\n  imgs_per_epoch: 4\n"
  },
  {
    "path": "sh_scripts/configs/db_config.yaml",
    "content": "experiment_name: exp_name\nexperiments_dir: /path/to/experiments\nmode: depth  # possible modes:\n\npseudomesh_path: /path/to/pseudomesh.npz\ndevice: cuda\nwhite_background: True\n\ndataset:\n  dataset_path: /path/to/dataset\n  near: 0.01\n  far: 100.\n  imgs_in_ram: True\n  res: 1\n\ndloader:\n  batch_size: 2\n  shuffle: True\n  num_workers: 10\n\ntraining:\n  epochs: 15\n\nmodel:\n  optimizable_params: \n    - vertices\n    - vertex_colors\n  optim_epoch_start:\n    vertices: 0\n\nalpha_pruning:\n  type: alpha  # possible modes: alpha, gradient and both\n  start_epoch: 10\n  epoch_step: 10\n  alpha_eps: 1e-4\n  grad_eps: 1e-7\n  mode: soft  # soft or hard\n\noptimizer:\n  name: adam\n  batch_scal_method: \"linear\"  # [linear, constant, sqrt]\n  lrs:\n    vertices: 1e-6\n    vertex_colors: 5e-3\n\nlr_scheduler:\n  use: True\n  gamma: 0.999\n\nrenderer:\n  depth_steps: 100\n  dp_scheduler:\n    perform: False\n    type: step  # step of exp\n    init_depth_steps: 50\n    # params:  # for exp\n    #   gamma: 0.998\n    params:   # for step\n      steps:\n        0: 50\n        60: 1\n\nloss:  # name of the loss and it's weight\n  ssim: .4\n  img_l1: .6\n  psnr-debug: 0.\n  # pips: 1e1\n  # dice: 1.\n  # delta_xyz: 1e-2\n  # delta_scales: 1e-2\n  # delta_rots: 5e-2\n\ntest_loss:\n  ssim: 1.\n  img_l1: 1.\n  psnr: 1.\n  pips: 1.\n\nwandb:\n  use: True\n  key_path: /path/to/wandb.json\n  project_name: project_name\n  entity: entity_name\n  imgs_per_epoch: 2\n"
  },
  {
    "path": "sh_scripts/configs/mip_config.yaml",
    "content": "experiment_name: exp_name\nexperiments_dir: /path/to/experiments\nmode: depth  # possible modes:\n\npseudomesh_path: /path/to/pseudomesh.npz\ndevice: cuda\nwhite_background: True\n\ndataset:\n  dataset_path: /path/to/dataset\n  near: 0.01\n  far: 100.\n  imgs_in_ram: True\n  res: 8\n\ndloader:\n  batch_size: 1\n  shuffle: True\n  num_workers: 10\n\ntraining:\n  epochs: 15\n\nmodel:\n  optimizable_params: \n    - vertices\n    - vertex_colors\n  optim_epoch_start:\n    vertices: 0\n\nalpha_pruning:\n  type: alpha  # possible modes: alpha, gradient and both\n  start_epoch: 10\n  epoch_step: 10\n  alpha_eps: 1e-4\n  grad_eps: 1e-7\n  mode: soft  # soft or hard\n\noptimizer:\n  name: adam\n  batch_scal_method: \"linear\"  # [linear, constant, sqrt]\n  lrs:\n    vertices: 1e-6\n    vertex_colors: 5e-3\n\nlr_scheduler:\n  use: True\n  gamma: 0.999\n\nrenderer:\n  depth_steps: 150\n  dp_scheduler:\n    perform: False\n    type: step  # step of exp\n    init_depth_steps: 50\n    # params:  # for exp\n    #   gamma: 0.998\n    params:   # for step\n      steps:\n        0: 50\n        60: 1\n\nloss:  # name of the loss and it's weight\n  ssim: .4\n  img_l1: .6\n  psnr-debug: 0.\n  # pips: 1e1\n  # dice: 1.\n  # delta_xyz: 1e-2\n  # delta_scales: 1e-2\n  # delta_rots: 5e-2\n\ntest_loss:\n  ssim: 1.\n  img_l1: 1.\n  psnr: 1.\n  pips: 1.\n\nwandb:\n  use: True\n  key_path: /path/to/wandb.json\n  project_name: project_name\n  entity: entity_name\n  imgs_per_epoch: 2\n"
  },
  {
    "path": "sh_scripts/configs/tandt_3dgs_config.yaml",
    "content": "experiment_name: exp_name\nexperiments_dir: /path/to/experiments\nmode: depth  # possible modes:\n\npseudomesh_path: /path/to/pseudomesh.npz\ndevice: cuda\nwhite_background: True\n\ndataset:\n  dataset_path: /path/to/dataset\n  near: 0.01\n  far: 100.\n  imgs_in_ram: True\n  res: 1\n\ndloader:\n  batch_size: 1\n  shuffle: True\n  num_workers: 10\n\ntraining:\n  epochs: 15\n\nmodel:\n  optimizable_params: \n    - vertices\n    - vertex_colors\n  optim_epoch_start:\n    vertices: 0\n\nalpha_pruning:\n  type: alpha  # possible modes: alpha, gradient and both\n  start_epoch: 10\n  epoch_step: 10\n  alpha_eps: 1e-4\n  grad_eps: 1e-7\n  mode: soft  # soft or hard\n\noptimizer:\n  name: adam\n  batch_scal_method: \"linear\"  # [linear, constant, sqrt]\n  lrs:\n    vertices: 1e-6\n    vertex_colors: 1e-2\n\nlr_scheduler:\n  use: True\n  gamma: 0.999\n\nrenderer:\n  depth_steps: 100\n  dp_scheduler:\n    perform: False\n    type: step  # step of exp\n    init_depth_steps: 50\n    # params:  # for exp\n    #   gamma: 0.998\n    params:   # for step\n      steps:\n        0: 50\n        60: 1\n\nloss:  # name of the loss and it's weight\n  ssim: .4\n  img_l1: .6\n  psnr-debug: 0.\n  # pips: 1e1\n  # dice: 1.\n  # delta_xyz: 1e-2\n  # delta_scales: 1e-2\n  # delta_rots: 5e-2\n\ntest_loss:\n  ssim: 1.\n  img_l1: 1.\n  psnr: 1.\n  pips: 1.\n\nwandb:\n  use: True\n  key_path: /path/to/wandb.json\n  project_name: project_name\n  entity: entity_name\n  imgs_per_epoch: 2\n"
  },
  {
    "path": "sh_scripts/configs/tandt_config.yaml",
    "content": "experiment_name: exp_name\nexperiments_dir: /path/to/experiments\nmode: depth  # possible modes:\n\npseudomesh_path: /path/to/pseudomesh.npz\ndevice: cuda\nwhite_background: True\n\ndataset:\n  dataset_path: /path/to/dataset\n  near: 0.01\n  far: 100.\n  imgs_in_ram: True\n  res: 1\n\ndloader:\n  batch_size: 2\n  shuffle: True\n  num_workers: 10\n\ntraining:\n  epochs: 15\n\nmodel:\n  optimizable_params: \n    - vertices\n    - vertex_colors\n  optim_epoch_start:\n    vertices: 0\n\nalpha_pruning:\n  type: alpha  # possible modes: alpha, gradient and both\n  start_epoch: 10\n  epoch_step: 10\n  alpha_eps: 1e-4\n  grad_eps: 1e-7\n  mode: soft  # soft or hard\n\noptimizer:\n  name: adam\n  batch_scal_method: \"linear\"  # [linear, constant, sqrt]\n  lrs:\n    vertices: 1e-6\n    vertex_colors: 1e-2\n\nlr_scheduler:\n  use: True\n  gamma: 0.999\n\nrenderer:\n  depth_steps: 200\n  dp_scheduler:\n    perform: False\n    type: step  # step of exp\n    init_depth_steps: 50\n    # params:  # for exp\n    #   gamma: 0.998\n    params:   # for step\n      steps:\n        0: 50\n        60: 1\n\nloss:  # name of the loss and it's weight\n  ssim: .4\n  img_l1: .6\n  psnr-debug: 0.\n  # pips: 1e1\n  # dice: 1.\n  # delta_xyz: 1e-2\n  # delta_scales: 1e-2\n  # delta_rots: 5e-2\n\ntest_loss:\n  ssim: 1.\n  img_l1: 1.\n  psnr: 1.\n  pips: 1.\n\nwandb:\n  use: True\n  key_path: /path/to/wandb.json\n  project_name: project_name\n  entity: entity_name\n  imgs_per_epoch: 2\n"
  },
  {
    "path": "sh_scripts/run_2dgs-sh0_db.sh",
    "content": "#!/bin/bash\n\n# Directory containing the elements\nDATA_DIR=\"/path/to/dataset\"\nEXP_DIR=\"/path/to/experiments\"\n\nGS_TRAIN_PATH=\"/path/to/2d-gaussian-splatting\"\nPYTHON_PATH=\"/path/to/python\"\n\nGS_PSEUDO_PATH=\"/path/to/gs_raytracing\"\nGS_OPTIM_PATH=\"/path/to/gs_raytracing/mesh_optim\"\n\nCONFIG_PATH=\"/path/to/gs_raytracing/sh_scripts/configs/db_config.yaml\"\n\nfor data_path in $DATA_DIR/*; do\n    dset_name=$(basename ${data_path})\n\n    exp_name=\"db_2dgs_sh0_${dset_name}\"\n    exp_out_dir=\"${EXP_DIR}/${exp_name}\"\n    exp_cfg_path=\"${exp_out_dir}/config.yaml\"\n    ply_path=\"${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply\"\n    pseudomesh_path=\"${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz\"\n    cd $GS_TRAIN_PATH\n    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\n    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\n    \n    cd $GS_PSEUDO_PATH\n    $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\n\n    cd $GS_OPTIM_PATH\n    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\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 2dgs --white_background\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method 2dgs\ndone"
  },
  {
    "path": "sh_scripts/run_2dgs-sh0_mip.sh",
    "content": "#!/bin/bash\n\n# Directory containing the elements\nDATA_DIR=\"/path/to/datasets/360_v2\"\nEXP_DIR=\"/path/to/experiments\"\n\nGS_TRAIN_PATH=\"/path/to/2d-gaussian-splatting\"\nPYTHON_PATH=\"/path/to/python\"\n\nGS_PSEUDO_PATH=\"/path/to/gs_raytracing\"\nGS_OPTIM_PATH=\"/path/to/gs_raytracing/mesh_optim\"\n\nCONFIG_PATH=\"/path/to/gs_raytracing/sh_scripts/configs/mip_config.yaml\"\n\nfor data_path in $DATA_DIR/*; do\n    dset_name=$(basename ${data_path})\n    # Iterate over each element in the ELEMENTS_DIR\n    exp_name=\"mip_2dgs_sh0_${dset_name}\"\n    exp_out_dir=\"${EXP_DIR}/${exp_name}\"\n    exp_cfg_path=\"${exp_out_dir}/config.yaml\"\n    ply_path=\"${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply\"\n    pseudomesh_path=\"${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz\"\n    \n    cd $GS_TRAIN_PATH\n    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\n    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\n\n\n    cd $GS_PSEUDO_PATH\n    $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\n\n    cd $GS_OPTIM_PATH\n    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\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 2dgs --white_background\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method 2dgs\ndone"
  },
  {
    "path": "sh_scripts/run_2dgs-sh0_nerf-synth.sh",
    "content": "#!/bin/bash\n\n# Directory containing the elements\nDATA_DIR=\"/path/to/datasets/games_set\"\nEXP_DIR=\"/path/to/experiments\"\n\nGS_TRAIN_PATH=\"/path/to/2d-gaussian-splatting\"\nPYTHON_PATH=\"/path/to/python\"\n\nGS_PSEUDO_PATH=\"/path/to/gs_raytracing\"\nGS_OPTIM_PATH=\"/path/to/gs_raytracing/mesh_optim\"\n\nCONFIG_PATH=\"/path/to/gs_raytracing/sh_scripts/configs/3dgs_sh0_pseudo_config.yaml\"\n\n\nfor data_path in $DATA_DIR/*; do\n    dset_name=$(basename ${data_path})\n    # Iterate over each element in the ELEMENTS_DIR\n    exp_name=\"new_nerf-synth_2dgs_sh0_${dset_name}\"\n    exp_out_dir=\"${EXP_DIR}/${exp_name}\"\n    exp_cfg_path=\"${exp_out_dir}/config.yaml\"\n    ply_path=\"${exp_out_dir}/point_cloud/iteration_10000/point_cloud.ply\"\n    pseudomesh_path=\"${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz\"\n    \n    cd $GS_TRAIN_PATH\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 10000 --sh_degree 0 --eval\n    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\n    cd $GS_PSEUDO_PATH\n    $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\n    cd $GS_OPTIM_PATH\n    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\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 2dgs\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --method 2dgs\ndone"
  },
  {
    "path": "sh_scripts/run_2dgs_sh0_tandt.sh",
    "content": "#!/bin/bash\n\n# Directory containing the elements\nDATA_DIR=\"/path/to/dataset\"\nEXP_DIR=\"/path/to/experiments\"\n\nGS_TRAIN_PATH=\"/path/to/2d-gaussian-splatting\"\nPYTHON_PATH=\"/path/to/python\"\n\nGS_PSEUDO_PATH=\"/path/to/gs_raytracing\"\nGS_OPTIM_PATH=\"/path/to/gs_raytracing/mesh_optim\"\n\nCONFIG_PATH=\"/path/to/gs_raytracing/sh_scripts/configs/tandt_config.yaml\"\n\nfor data_path in $DATA_DIR/*; do\n    dset_name=$(basename ${data_path})\n\n    # Iterate over each element in the ELEMENTS_DIR\n    exp_name=\"tandt_2dgs_sh0_${dset_name}\"\n    exp_out_dir=\"${EXP_DIR}/${exp_name}\"\n    exp_cfg_path=\"${exp_out_dir}/config.yaml\"\n    ply_path=\"${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply\"\n    pseudomesh_path=\"${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz\"\n\n    cd $GS_TRAIN_PATH\n    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\n    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\n\n    cd $GS_PSEUDO_PATH\n    $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\n\n    cd $GS_OPTIM_PATH\n    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\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 2dgs --white_background\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method 2dgs\ndone"
  },
  {
    "path": "sh_scripts/run_3dgs-sh0_db.sh",
    "content": "#!/bin/bash\n\n# Directory containing the elements\nDATA_DIR=\"/path/to/datasets/db\"\nEXP_DIR=\"/path/to/experiments\"\n\nGS_TRAIN_PY=\"/path/to/gaussian-splatting\"\nPYTHON_PATH=\"/path/to/python\"\n\nGS_PSEUDO_PATH=\"/path/to/gs_raytracing\"\nGS_OPTIM_PATH=\"/path/to/gs_raytracing/mesh_optim\"\n\nCONFIG_PATH=\"/path/to/gs_raytracing/sh_scripts/configs/db_config.yaml\"\n\nfor data_path in $DATA_DIR/*; do\n    dset_name=$(basename ${data_path})\n    exp_name=\"db_3dgs_white_sh0_${dset_name}\"\n    exp_out_dir=\"${EXP_DIR}/${exp_name}\"\n    exp_cfg_path=\"${exp_out_dir}/config.yaml\"\n    ply_path=\"${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply\"\n    pseudomesh_path=\"${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz\"\n\n    cd $GS_TRAIN_PY\n    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\n    \n    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\n\n    cd $GS_PSEUDO_PATH\n    $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\n\n    cd $GS_OPTIM_PATH\n    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\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 3dgs --white_background\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method 3dgs\ndone"
  },
  {
    "path": "sh_scripts/run_3dgs-sh0_mip.sh",
    "content": "#!/bin/bash\n\n# Directory containing the elements\nDATA_DIR=\"/path/to/datasets/360_v2\"\nEXP_DIR=\"/path/to/experiments\"\n\nGS_TRAIN_PY=\"/path/to/gaussian-splatting\"\nPYTHON_PATH=\"/path/to/python\"\n\nGS_PSEUDO_PATH=\"/path/to/gs_raytracing\"\nGS_OPTIM_PATH=\"/path/to/gs_raytracing/mesh_optim\"\n\nCONFIG_PATH=\"/path/to/gs_raytracing/sh_scripts/configs/mip_config.yaml\"\n\nfor data_path in $DATA_DIR/*; do\n    dset_name=$(basename ${data_path})\n    # Iterate over each element in the ELEMENTS_DIR\n    exp_name=\"mip_3dgs_sh0_${dset_name}\"\n    exp_out_dir=\"${EXP_DIR}/${exp_name}\"\n    exp_cfg_path=\"${exp_out_dir}/config.yaml\"\n    ply_path=\"${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply\"\n    pseudomesh_path=\"${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz\"\n\n    cd $GS_TRAIN_PY\n    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\n    \n    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\n\n    cd $GS_PSEUDO_PATH\n    $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\n\n    cd $GS_OPTIM_PATH\n    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\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 3dgs --white_background\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method 3dgs\ndone"
  },
  {
    "path": "sh_scripts/run_3dgs-sh0_nerf-synth.sh",
    "content": "#!/bin/bash\nDATA_DIR=\"/path/to/datasets/games_set\"\nEXP_DIR=\"/path/to/experiments\"\n\nGS_TRAIN_PY=\"/path/to/gaussian-splatting\"\nPYTHON_PATH=\"/path/to/python\"\n\nGS_PSEUDO_PATH=\"/path/to/gs_raytracing\"\nGS_OPTIM_PATH=\"/path/to/gs_raytracing/mesh_optim\"\n\nCONFIG_PATH=\"/path/to/gs_raytracing/sh_scripts/configs/3dgs_sh0_pseudo_config.yaml\"\n\nfor data_path in $DATA_DIR/*; do\n    dset_name=$(basename ${data_path})\n\n    # Iterate over each element in the ELEMENTS_DIR\n    exp_name=\"nerf-synth_3dgs_sh0_${dset_name}\"\n    exp_out_dir=\"${EXP_DIR}/${exp_name}\"\n    exp_cfg_path=\"${exp_out_dir}/config.yaml\"\n    ply_path=\"${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply\"\n    pseudomesh_path=\"${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz\"\n\n    cd $GS_TRAIN_PY\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH train.py --resolution 1 -s $data_path -m $exp_out_dir --iteration 30000 --sh_degree 0 --eval\n    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\n\n    cd $GS_PSEUDO_PATH\n    $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\n\n    cd $GS_OPTIM_PATH\n    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\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 3dgs\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --method 3dgs\ndone"
  },
  {
    "path": "sh_scripts/run_3dgs-sh0_tandt.sh",
    "content": "#!/bin/bash\n\n# Directory containing the elements\nDATA_DIR=\"/path/to/datasets/tandt\"\nEXP_DIR=\"/path/to/experiments\"\n\nGS_TRAIN_PY=\"/path/to/gaussian-splatting\"\nPYTHON_PATH=\"/path/to/python\"\n\nGS_PSEUDO_PATH=\"/path/to/gs_raytracing\"\nGS_OPTIM_PATH=\"/path/to/gs_raytracing/mesh_optim\"\n\nCONFIG_PATH=\"/path/to/gs_raytracing/sh_scripts/configs/tandt_3dgs_config.yaml\"\n\nfor data_path in $DATA_DIR/*; do\n    dset_name=$(basename ${data_path})\n    exp_name=\"tandt_3dgs_white_sh0_${dset_name}\"\n    exp_out_dir=\"${EXP_DIR}/${exp_name}\"\n    exp_cfg_path=\"${exp_out_dir}/config.yaml\"\n    ply_path=\"${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply\"\n    pseudomesh_path=\"${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz\"\n\n    cd $GS_TRAIN_PY\n    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\n    \n    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\n\n    cd $GS_PSEUDO_PATH\n    $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\n\n    cd $GS_OPTIM_PATH\n    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\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method 3dgs --white_background\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method 3dgs\n    break\ndone"
  },
  {
    "path": "sh_scripts/run_games_gs-flat_sh0_db.sh",
    "content": "#!/bin/bash\n\nDATA_DIR=\"/path/to/datasets/db\"\nEXP_DIR=\"/path/to/experiments\"\n\nGS_TRAIN_PATH=\"/path/to/gaussian-mesh-splatting\"\nGS_RENDER_PY=\"/path/to/gaussian-mesh-splatting/scripts/render.py\"\nPYTHON_PATH=\"/path/to/python\"\n\nGS_PSEUDO_PATH=\"/path/to/gs_raytracing\"\nGS_OPTIM_PATH=\"/path/to/gs_raytracing/mesh_optim\"\n\nCONFIG_PATH=\"/path/to/gs_raytracing/sh_scripts/configs/db_config.yaml\"\n\nexport PYTHONPATH=$GS_TRAIN_PATH:$PYTHON_PATH\n\nfor data_path in $DATA_DIR/*; do\n    dset_name=$(basename ${data_path})\n    exp_name=\"db_games_gs-flat_sh0_${dset_name}\"\n    exp_out_dir=\"${EXP_DIR}/${exp_name}\"\n    exp_cfg_path=\"${exp_out_dir}/config.yaml\"\n    ply_path=\"${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply\"\n    pseudomesh_path=\"${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz\"\n\n    cd $GS_TRAIN_PATH\n    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\n    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\n\n    cd $GS_PSEUDO_PATH\n    $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\n\n    cd $GS_OPTIM_PATH\n    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\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method games --white_background\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method games\ndone"
  },
  {
    "path": "sh_scripts/run_games_gs-flat_sh0_mip.sh",
    "content": "#!/bin/bash\n\nDATA_DIR=\"/path/to/datasets/360_v2\"\nEXP_DIR=\"/path/to/experiments\"\n\nGS_TRAIN_PATH=\"/path/to/gaussian-mesh-splatting\"\nGS_RENDER_PY=\"/path/to/gaussian-mesh-splatting/scripts/render.py\"\nPYTHON_PATH=\"/path/to/python\"\n\nGS_PSEUDO_PATH=\"/path/to/gs_raytracing\"\nGS_OPTIM_PATH=\"/path/to/gs_raytracing/mesh_optim\"\n\nCONFIG_PATH=\"/path/to/gs_raytracing/sh_scripts/configs/mip_config.yaml\"\n\nexport PYTHONPATH=$GS_TRAIN_PATH:$PYTHON_PATH\n\nfor data_path in $DATA_DIR/*; do\n    dset_name=$(basename ${data_path})\n\n    exp_name=\"mip_games_gs-flat_sh0_${dset_name}\"\n    exp_out_dir=\"${EXP_DIR}/${exp_name}\"\n    exp_cfg_path=\"${exp_out_dir}/config.yaml\"\n    ply_path=\"${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply\"\n    pseudomesh_path=\"${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz\"\n\n    cd $GS_TRAIN_PATH\n    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\n    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\n\n    cd $GS_PSEUDO_PATH\n    $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\n\n    cd $GS_OPTIM_PATH\n    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\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method games --white_background\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method games\ndone"
  },
  {
    "path": "sh_scripts/run_games_gs-flat_sh0_nerf-synth.sh",
    "content": "#!/bin/bash\n\nDATA_DIR=\"/path/to/datasets/games_set\"\nEXP_DIR=\"/path/to/experiments\"\n\nGS_TRAIN_PATH=\"/path/to/gaussian-mesh-splatting\"\nGS_RENDER_PY=\"/path/to/gaussian-mesh-splatting/scripts/render.py\"\nPYTHON_PATH=\"/path/to/python\"\n\nGS_PSEUDO_PATH=\"/path/to/gs_raytracing\"\nGS_OPTIM_PATH=\"/path/to/gs_raytracing/mesh_optim\"\n\nCONFIG_PATH=\"/path/to/gs_raytracing/sh_scripts/configs/3dgs_sh0_pseudo_config.yaml\"\n\nexport PYTHONPATH=$GS_TRAIN_PATH:$PYTHON_PATH\n\nfor data_path in $DATA_DIR/*; do\n    dset_name=$(basename ${data_path})\n\n    exp_name=\"games_gs-flat_white_sh0_${dset_name}\"\n    exp_out_dir=\"${EXP_DIR}/${exp_name}\"\n    exp_cfg_path=\"${exp_out_dir}/config.yaml\"\n    ply_path=\"${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply\"\n    pseudomesh_path=\"${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz\"\n\n    cd $GS_TRAIN_PATH\n    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\n    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\n\n    cd $GS_PSEUDO_PATH\n    $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\n\n    cd $GS_OPTIM_PATH\n    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\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method games\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir  --method games\ndone\n"
  },
  {
    "path": "sh_scripts/run_games_gs-flat_sh0_tandt.sh",
    "content": "#!/bin/bash\n\nDATA_DIR=\"/path/to/datasets/tandt\"\nEXP_DIR=\"/path/to/experiments\"\n\nGS_TRAIN_PATH=\"/path/to/gaussian-mesh-splatting\"\nGS_RENDER_PY=\"/path/to/gaussian-mesh-splatting/scripts/render.py\"\nPYTHON_PATH=\"/path/to/python\"\n\nGS_PSEUDO_PATH=\"/path/to/gs_raytracing\"\nGS_OPTIM_PATH=\"/path/to/gs_raytracing/mesh_optim\"\n\nCONFIG_PATH=\"/path/to/gs_raytracing/sh_scripts/configs/tandt_config.yaml\"\n\nexport PYTHONPATH=$GS_TRAIN_PATH:$PYTHON_PATH\nfor data_path in $DATA_DIR/*; do\n    dset_name=$(basename ${data_path})\n\n    exp_name=\"tandt_games_gs-flat_sh0_${dset_name}\"\n    exp_out_dir=\"${EXP_DIR}/${exp_name}\"\n    exp_cfg_path=\"${exp_out_dir}/config.yaml\"\n    ply_path=\"${exp_out_dir}/point_cloud/iteration_30000/point_cloud.ply\"\n    pseudomesh_path=\"${exp_out_dir}/pseudomeshes/scale_2.70_pts_8.npz\"\n\n    cd $GS_TRAIN_PATH\n    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\n    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\n\n    cd $GS_PSEUDO_PATH\n    $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\n\n    cd $GS_OPTIM_PATH\n    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\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH render_pseudomesh.py --cfg_path $exp_cfg_path --method games --white_background\n    CUDA_VISIBLE_DEVICES=0 $PYTHON_PATH metrics.py -s $data_path -m $exp_out_dir --white_background --method games\ndone"
  }
]