Full Code of HLinChen/VCR-GauS for AI

main aa715d19bfac cached
99 files
474.8 KB
130.4k tokens
492 symbols
1 requests
Download .txt
Showing preview only (503K chars total). Download the full file or copy to clipboard to get everything.
Repository: HLinChen/VCR-GauS
Branch: main
Commit: aa715d19bfac
Files: 99
Total size: 474.8 KB

Directory structure:
gitextract_oa0m8dnw/

├── .gitmodules
├── LICENSE.md
├── README.md
├── arguments/
│   └── __init__.py
├── bash_scripts/
│   ├── 0_train.sh
│   ├── 1_preprocess_tnt.sh
│   ├── 2_extract_normal_dsine.sh
│   ├── 3_extract_mask.sh
│   ├── 4_extract_normal_geow.sh
│   ├── convert.sh
│   └── install.sh
├── configs/
│   ├── 360_v2/
│   │   └── base.yaml
│   ├── config.py
│   ├── config_base.yaml
│   ├── dtu/
│   │   ├── base.yaml
│   │   └── dtu_scan24.yaml
│   ├── reconstruct.yaml
│   ├── scannetpp/
│   │   └── base.yaml
│   └── tnt/
│       ├── Barn.yaml
│       ├── Caterpillar.yaml
│       ├── Courthouse.yaml
│       ├── Ignatius.yaml
│       ├── Meetingroom.yaml
│       ├── Truck.yaml
│       └── base.yaml
├── environment.yml
├── evaluation/
│   ├── crop_mesh.py
│   ├── eval_dtu/
│   │   ├── eval.py
│   │   ├── evaluate_single_scene.py
│   │   └── render_utils.py
│   ├── eval_tnt.py
│   ├── full_eval.py
│   ├── lpipsPyTorch/
│   │   ├── __init__.py
│   │   └── modules/
│   │       ├── lpips.py
│   │       ├── networks.py
│   │       └── utils.py
│   ├── metrics.py
│   ├── render.py
│   └── tnt_eval/
│       ├── README.md
│       ├── config.py
│       ├── evaluation.py
│       ├── plot.py
│       ├── registration.py
│       ├── requirements.txt
│       ├── run.py
│       ├── trajectory_io.py
│       └── util.py
├── gaussian_renderer/
│   ├── __init__.py
│   └── network_gui.py
├── process_data/
│   ├── convert.py
│   ├── convert_360_to_json.py
│   ├── convert_data_to_json.py
│   ├── convert_dtu_to_json.py
│   ├── convert_tnt_to_json.py
│   ├── extract_mask.py
│   ├── extract_normal.py
│   ├── extract_normal_geo.py
│   ├── visualize_colmap.ipynb
│   └── visualize_transforms.ipynb
├── pyproject.toml
├── python_scripts/
│   ├── run_base.py
│   ├── run_dtu.py
│   ├── run_mipnerf360.py
│   ├── run_tnt.py
│   ├── show_360.py
│   ├── show_dtu.py
│   └── show_tnt.py
├── requirements.txt
├── scene/
│   ├── __init__.py
│   ├── appearance_network.py
│   ├── cameras.py
│   ├── colmap_loader.py
│   ├── dataset_readers.py
│   └── gaussian_model.py
├── tools/
│   ├── __init__.py
│   ├── camera.py
│   ├── camera_utils.py
│   ├── crop_mesh.py
│   ├── denoise_pcd.py
│   ├── depth2mesh.py
│   ├── distributed.py
│   ├── general_utils.py
│   ├── graphics_utils.py
│   ├── image_utils.py
│   ├── loss_utils.py
│   ├── math_utils.py
│   ├── mcube_utils.py
│   ├── mesh_utils.py
│   ├── normal_utils.py
│   ├── prune.py
│   ├── render_utils.py
│   ├── semantic_id.py
│   ├── sh_utils.py
│   ├── system_utils.py
│   ├── termcolor.py
│   ├── visualization.py
│   └── visualize.py
├── train.py
└── trainer.py

================================================
FILE CONTENTS
================================================

================================================
FILE: .gitmodules
================================================
[submodule "submodules/simple-knn"]
	path = submodules/simple-knn
	url = https://gitlab.inria.fr/bkerbl/simple-knn.git
[submodule "submodules/diff-gaussian-rasterization"]
	path = submodules/diff-gaussian-rasterization
	url = https://github.com/HLinChen/diff-gaussian-rasterization
[submodule "SIBR_viewers"]
	path = SIBR_viewers
	url = https://gitlab.inria.fr/sibr/sibr_core.git
[submodule "submodules/colmap"]
	path = submodules/colmap
	url = https://github.com/colmap/colmap.git


================================================
FILE: LICENSE.md
================================================
Gaussian-Splatting License  
===========================  

**Inria** and **the Max Planck Institut for Informatik (MPII)** hold all the ownership rights on the *Software* named **gaussian-splatting**.  
The *Software* is in the process of being registered with the Agence pour la Protection des  
Programmes (APP).  

The *Software* is still being developed by the *Licensor*.  

*Licensor*'s goal is to allow the research community to use, test and evaluate  
the *Software*.  

## 1.  Definitions  

*Licensee* means any person or entity that uses the *Software* and distributes  
its *Work*.  

*Licensor* means the owners of the *Software*, i.e Inria and MPII  

*Software* means the original work of authorship made available under this  
License ie gaussian-splatting.  

*Work* means the *Software* and any additions to or derivative works of the  
*Software* that are made available under this License.  


## 2.  Purpose  
This license is intended to define the rights granted to the *Licensee* by  
Licensors under the *Software*.  

## 3.  Rights granted  

For the above reasons Licensors have decided to distribute the *Software*.  
Licensors grant non-exclusive rights to use the *Software* for research purposes  
to research users (both academic and industrial), free of charge, without right  
to sublicense.. The *Software* may be used "non-commercially", i.e., for research  
and/or evaluation purposes only.  

Subject to the terms and conditions of this License, you are granted a  
non-exclusive, royalty-free, license to reproduce, prepare derivative works of,  
publicly display, publicly perform and distribute its *Work* and any resulting  
derivative works in any form.  

## 4.  Limitations  

**4.1 Redistribution.** You may reproduce or distribute the *Work* only if (a) you do  
so under this License, (b) you include a complete copy of this License with  
your distribution, and (c) you retain without modification any copyright,  
patent, trademark, or attribution notices that are present in the *Work*.  

**4.2 Derivative Works.** You may specify that additional or different terms apply  
to the use, reproduction, and distribution of your derivative works of the *Work*  
("Your Terms") only if (a) Your Terms provide that the use limitation in  
Section 2 applies to your derivative works, and (b) you identify the specific  
derivative works that are subject to Your Terms. Notwithstanding Your Terms,  
this License (including the redistribution requirements in Section 3.1) will  
continue to apply to the *Work* itself.  

**4.3** Any other use without of prior consent of Licensors is prohibited. Research  
users explicitly acknowledge having received from Licensors all information  
allowing to appreciate the adequacy between of the *Software* and their needs and  
to undertake all necessary precautions for its execution and use.  

**4.4** The *Software* is provided both as a compiled library file and as source  
code. In case of using the *Software* for a publication or other results obtained  
through the use of the *Software*, users are strongly encouraged to cite the  
corresponding publications as explained in the documentation of the *Software*.  

## 5.  Disclaimer  

THE USER CANNOT USE, EXPLOIT OR DISTRIBUTE THE *SOFTWARE* FOR COMMERCIAL PURPOSES  
WITHOUT PRIOR AND EXPLICIT CONSENT OF LICENSORS. YOU MUST CONTACT INRIA FOR ANY  
UNAUTHORIZED USE: stip-sophia.transfert@inria.fr . ANY SUCH ACTION WILL  
CONSTITUTE A FORGERY. THIS *SOFTWARE* IS PROVIDED "AS IS" WITHOUT ANY WARRANTIES  
OF ANY NATURE AND ANY EXPRESS OR IMPLIED WARRANTIES, WITH REGARDS TO COMMERCIAL  
USE, PROFESSIONNAL USE, LEGAL OR NOT, OR OTHER, OR COMMERCIALISATION OR  
ADAPTATION. UNLESS EXPLICITLY PROVIDED BY LAW, IN NO EVENT, SHALL INRIA OR THE  
AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR  
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE  
GOODS OR SERVICES, LOSS OF USE, DATA, OR PROFITS OR BUSINESS INTERRUPTION)  
HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT  
LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING FROM, OUT OF OR  
IN CONNECTION WITH THE *SOFTWARE* OR THE USE OR OTHER DEALINGS IN THE *SOFTWARE*.  


================================================
FILE: README.md
================================================
<p align="center">

  <h1 align="center">VCR-GauS: View Consistent Depth-Normal Regularizer for Gaussian Surface Reconstruction</h1>
  <p align="center">
    <a href="https://hlinchen.github.io/">Hanlin Chen</a>,
    <a href="https://weify627.github.io/">Fangyin Wei</a>,
    <a href="https://chaneyddtt.github.io/">Chen Li</a>,
    <a href="https://tianxinhuang.github.io/">Tianxin Huang</a>,
    <a href="https://scholar.google.com/citations?user=vv1uLeUAAAAJ&hl=en">Yunsong Wang</a>,
    <a href="https://www.comp.nus.edu.sg/~leegh/">Gim Hee Lee</a>

  </p>

  <h2 align="center">NeurIPS 2024</h2>

  <h3 align="center"><a href="https://arxiv.org/pdf/2406.05774">arXiv</a> | <a href="https://hlinchen.github.io/projects/VCR-GauS/">Project Page</a>  </h3>
  <div align="center"></div>
</p>


<p align="center">
  <a href="">
    <img src="./media/VCR-GauS.jpg" alt="Logo" width="95%">
  </a>
</p>

<p align="left">
VCR-GauS formulates a novel multi-view D-Normal regularizer that enables full optimization of the Gaussian geometric parameters to achieve better surface reconstruction. We further design a confidence term to weigh our D-Normal regularizer to mitigate inconsistencies of normal predictions across multiple views.</p>
<br>

# Updates

* **[2024.10.31]**: We uploaded a new version to arXiv, adding theoretical proofs and visualization results for the D-Normal Regularizer.
* **[2024.09.24]**: VCR-GauS is accepted to NeurIPS 2024.

# Installation
Clone the repository and create an anaconda environment using
```
git clone https://github.com/HLinChen/VCR-GauS.git --recursive
cd VCR-GauS
git pull --recurse-submodules

env=vcr
conda create -n $env -y python=3.10
conda activate $env
pip install -e ".[train]"
# you can specify your own cuda path
export CUDA_HOME=/usr/local/cuda-11.8
pip install -r requirements.txt
```
We also uploaded a built anaconda environment [here](https://huggingface.co/hanlin-chen/VCR-GauS/resolve/main/vcr.zip?download=true); you can download it and unzip and put it in your_anaconda_path/envs/ .

For eval TNT with the official scripts, you need to build a new environment with open3d==0.10:
```
env=f1eval
conda create -n $env -y python=3.8
conda activate $env
pip install -e ".[f1eval]"
```

For extract normal maps based on [DSINE](https://baegwangbin.github.io/DSINE/), you need to build a new environment:
```
conda create --name dsine python=3.10
conda activate dsine

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
python -m pip install geffnet
```


Similar to Gaussian Splatting, we also use colmap to process data and you can follow [COLMAP website](https://colmap.github.io/) to install it.


# Dataset

<!-- Please download the Mip-NeRF 360 dataset from the [official webiste](https://jonbarron.info/mipnerf360/), the preprocessed DTU dataset from [2DGS](https://surfsplatting.github.io/), the proprocessed Tanks and Temples dataset from [here](https://huggingface.co/datasets/ZehaoYu/gaussian-opacity-fields/tree/main). You need to download the ground truth point clouds from the [DTU dataset](https://roboimagedata.compute.dtu.dk/?page_id=36) and save to `dtu_eval/Offical_DTU_Dataset` to evaluate the geometry reconstruction. For the [Tanks and Temples](https://www.tanksandtemples.org/download/) dataset, you need to download the ground truth point clouds, alignments and cropfiles and save to `eval_tnt/TrainingSet`, such as `eval_tnt/TrainingSet/Caterpillar/Caterpillar.ply`. -->


## Tanks and Temples dataset
You can download the proprocessed Tanks and Temples dataset from [here](https://huggingface.co/hanlin-chen/VCR-GauS/resolve/main/tnt.zip?download=true). Or proprocess it by your self:
Download the data from [Tanks and Temples](https://tanksandtemples.org/download/) website.
You will also need to download additional [COLMAP/camera/alignment](https://drive.google.com/file/d/1jAr3IDvhVmmYeDWi0D_JfgiHcl70rzVE/view?resourcekey=) and the images of each scene.  
The file structure should look like (you need to move the downloaded images to folder `images_raw`):
```
tanks_and_temples
├─ Barn
│  ├─ Barn_COLMAP_SfM.log   (camera poses)
│  ├─ Barn.json             (cropfiles)
│  ├─ Barn.ply              (ground-truth point cloud)
│  ├─ Barn_trans.txt        (colmap-to-ground-truth transformation)
│  └─ images_raw            (raw input images downloaded from Tanks and Temples website)
│     ├─ 000001.png
│     ├─ 000002.png
│     ...
├─ Caterpillar
│  ├─ ...
...
```
#### 1. Colmap and bounding box json
Run the following command to generate json and colmap files:
```bash
# Modify --tnt_path to be the Tanks and Temples root directory.
sh bash_scripts/1_preprocess_tnt.sh
```

#### 2. Normal maps
You need to download the [code](https://github.com/baegwangbin/DSINE) and [model weight](https://drive.google.com/drive/folders/1t3LMJIIrSnCGwOEf53Cyg0lkSXd3M4Hm) of DSINE first. Then, modify **CODE_PATH** to be the DSINE root directory, **CKPT** to be the DSINE model path, **DATADIR** to be the TNT root directory in the bash script.
Run the following command to generate normal maps:

```bash
sh bash_scripts/2_extract_normal_dsine.sh
```

#### 3. Semantic masks (optional)

If you don't want to use the semantic masks, you can set **optim.loss_weight.semantic=0** and skip the mask generation.

You need to download the [code](https://github.com/IDEA-Research/Grounded-Segment-Anything) and model of Grounded-SAM first. Then, install the environment based on 'Install without Docker' in the [webside](https://github.com/IDEA-Research/Grounded-Segment-Anything). Next, modify **GSAM_PATH** to be the GSAM root directory, **DATADIR** to be the TNT root directory in the bash script. Run the following command to generate semantic masks:

```bash
sh bash_scripts/3_extract_mask.sh
```

## Other datasets
Please download the Mip-NeRF 360 dataset from the official [webiste](https://jonbarron.info/mipnerf360/), the preprocessed DTU dataset from [2DGS](https://drive.google.com/drive/folders/1SJFgt8qhQomHX55Q4xSvYE2C6-8tFll9). And extract normal maps with DSINE following the above scripts. You can also use [GeoWizard](https://github.com/fuxiao0719/GeoWizard) to extract normal maps by following the script: 'bash_scripts/4_extract_normal_geow.sh', and please install the corresponding environment and download the code as well as model weights first.

# Training and Evaluation
## From the scratch:
```
# you might need to update the data path in the script accordingly

# Tanks and Temples dataset
python python_scripts/run_tnt.py

# Mip-NeRF 360 dataset
python python_scripts/run_mipnerf360.py
```

## Only eval the metrics
We have uploaded the extracted meshes, you can download and eval them by yourselves ([TNT](https://huggingface.co/hanlin-chen/VCR-GauS/resolve/main/tnt_mesh.zip?download=true) and [DTU](https://huggingface.co/Chiller3/VCR-GauS/resolve/main/dtu_mesh.zip?download=true)). You might need to update the **mesh and data path** in the script accordingly. And set **do_train** and **do_extract_mesh** to be False.

```
# Tanks and Temples dataset
python python_scripts/run_tnt.py

# DTU dataset
python python_scripts/run_dtu.py
```

## Additional regularizations:
We also incorporate some regularizations, like depth distortion loss and normal consistency loss, following [2DGS](https://surfsplatting.github.io/) and [GOF](https://niujinshuchong.github.io/gaussian-opacity-fields/). You can play with it by:
- normal consistency loss: setting optim.loss_weight.consistent_normal > 0;
- depth distortion loss:
  1. set optim.loss_weight.depth_var > 0
  2. set NUM_DIST = 1 in submodules/diff-gaussian-rasterization/cuda_rasterizer/config.h, and reinstall diff-gaussian-rasterization


# Custom Dataset
We use the same data format from 3DGS, please follow [here](https://github.com/graphdeco-inria/gaussian-splatting?tab=readme-ov-file#processing-your-own-scenes) to prepare the your dataset. Then you can train your model and extract a mesh.
```
# Generate bounding box
python process_data/convert_data_to_json.py \
        --scene_type outdoor \
        --data_dir /your/data/path

# Extract normal maps
# Use DSINE:
python -W ignore process_data/extract_normal.py \
    --dsine_path /your/dsine/code/path \
    --ckpt /your/ckpt/path \
    --img_path /your/data/path/images \
    --intrins_path /your/data/path/ \
    --output_path /your/data/path/normals

# Or use GeoWizard
python process_data/extract_normal_geo.py \
  --code_path ${CODE_PATH} \
  --input_dir /your/data/path/images/ \
  --output_dir /your/data/path/ \
  --ensemble_size 3 \
  --denoise_steps 10 \
  --seed 0 \
  --domain ${DOMAIN_TYPE} # outdoor indoor object

# training
# --model.resolution=2 for using downsampled images with factor 2
# --model.use_decoupled_appearance=True to enable decoupled appearance modeling if your images has changing lighting conditions
python train.py \
  --config=configs/reconstruct.yaml \
  --logdir=/your/log/path/ \
  --model.source_path=/your/data/path/ \
  --model.data_device=cpu \
  --model.resolution=2 \
  --wandb \
  --wandb_name vcr-gaus"

# extract the mesh after training
python tools/depth2mesh.py \
  --voxel_size 5e-3 \
  --max_depth 8 \
  --clean \
  --cfg_path /your/gaussian/path/config.yaml"
```

# Acknowledgements
This project is built upon [3DGS](https://github.com/graphdeco-inria/gaussian-splatting). Evaluation scripts for DTU and Tanks and Temples dataset are taken from [DTUeval-python](https://github.com/jzhangbs/DTUeval-python) and [TanksAndTemples](https://github.com/isl-org/TanksAndTemples/tree/master/python_toolbox/evaluation) respectively. We also utilize the normal estimation [DSINE](https://github.com/baegwangbin/DSINE) as well as [GeoWizard](https://fuxiao0719.github.io/projects/geowizard/), and semantic segmentation [SAM](https://github.com/facebookresearch/segment-anything) and [Grounded-SAM](https://github.com/IDEA-Research/Grounded-Segment-Anything?tab=readme-ov-file#install-without-docker). In addition, we use the pruning method in [LightGaussin](https://lightgaussian.github.io/). We thank all the authors for their great work and repos. 


# Citation
If you find our code or paper useful, please cite
```bibtex
@article{chen2024vcr,
  author    = {Chen, Hanlin and Wei, Fangyin and Li, Chen and Huang, Tianxin and Wang, Yunsong and Lee, Gim Hee},
  title     = {VCR-GauS: View Consistent Depth-Normal Regularizer for Gaussian Surface Reconstruction},
  journal   = {arXiv preprint arXiv:2406.05774},
  year      = {2024},
}
```

If you find the flatten 3D Gaussian useful, please kindly cite
```bibtex
@article{chen2023neusg,
  title={Neusg: Neural implicit surface reconstruction with 3d gaussian splatting guidance},
  author={Chen, Hanlin and Li, Chen and Lee, Gim Hee},
  journal={arXiv preprint arXiv:2312.00846},
  year={2023}
}
```


================================================
FILE: arguments/__init__.py
================================================
#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

from argparse import ArgumentParser, Namespace
import sys
import os

class GroupParams:
    pass

class ParamGroup:
    def __init__(self, parser: ArgumentParser, name : str, fill_none = False):
        group = parser.add_argument_group(name)
        for key, value in vars(self).items():
            shorthand = False
            if key.startswith("_"):
                shorthand = True
                key = key[1:]
            t = type(value)
            value = value if not fill_none else None 
            if shorthand:
                if t == bool:
                    group.add_argument("--" + key, ("-" + key[0:1]), default=value, action="store_true")
                else:
                    group.add_argument("--" + key, ("-" + key[0:1]), default=value, type=t)
            else:
                if t == bool:
                    group.add_argument("--" + key, default=value, action="store_true")
                else:
                    group.add_argument("--" + key, default=value, type=t)

    def extract(self, args):
        group = GroupParams()
        for arg in vars(args).items():
            if arg[0] in vars(self) or ("_" + arg[0]) in vars(self):
                setattr(group, arg[0], arg[1])
        return group

class ModelParams(ParamGroup): 
    def __init__(self, parser, sentinel=False):
        self.sh_degree = 3
        self._source_path = ""
        self._model_path = ""
        self._images = "images"
        self._resolution = -1
        self._white_background = False
        self.data_device = "cuda"
        self.eval = False
        super().__init__(parser, "Loading Parameters", sentinel)

    def extract(self, args):
        g = super().extract(args)
        g.source_path = os.path.abspath(g.source_path)
        return g

class PipelineParams(ParamGroup):
    def __init__(self, parser):
        self.convert_SHs_python = False
        self.compute_cov3D_python = False
        self.debug = False
        super().__init__(parser, "Pipeline Parameters")

class OptimizationParams(ParamGroup):
    def __init__(self, parser):
        self.iterations = 30_000
        self.position_lr_init = 0.00016
        self.position_lr_final = 0.0000016
        self.position_lr_delay_mult = 0.01
        self.position_lr_max_steps = 30_000
        self.feature_lr = 0.0025
        self.opacity_lr = 0.05
        self.scaling_lr = 0.005
        self.rotation_lr = 0.001
        self.percent_dense = 0.01
        self.lambda_dssim = 0.2
        self.densification_interval = 100
        self.opacity_reset_interval = 3000
        self.densify_from_iter = 500
        self.densify_until_iter = 15_000
        self.densify_grad_threshold = 0.0002
        self.random_background = False
        super().__init__(parser, "Optimization Parameters")

def get_combined_args(parser : ArgumentParser):
    cmdlne_string = sys.argv[1:]
    cfgfile_string = "Namespace()"
    args_cmdline = parser.parse_args(cmdlne_string)

    try:
        cfgfilepath = os.path.join(args_cmdline.model_path, "cfg_args")
        print("Looking for config file in", cfgfilepath)
        with open(cfgfilepath) as cfg_file:
            print("Config file found: {}".format(cfgfilepath))
            cfgfile_string = cfg_file.read()
    except TypeError:
        print("Config file not found at")
        pass
    args_cfgfile = eval(cfgfile_string)

    merged_dict = vars(args_cfgfile).copy()
    for k,v in vars(args_cmdline).items():
        if v != None:
            merged_dict[k] = v
    return Namespace(**merged_dict)


================================================
FILE: bash_scripts/0_train.sh
================================================
GPU=0
export CUDA_VISIBLE_DEVICES=${GPU}
ls


DATASET=tnt
SCENE=Barn
NAME=${SCENE}

PROJECT=vcr_gaus

TRIAL_NAME=vcr_gaus

CFG=configs/${DATASET}/${SCENE}.yaml

DIR=/your/log/path/${PROJECT}/${DATASET}/${NAME}/${TRIAL_NAME}

python train.py \
    --config=${CFG} \
    --port=-1 \
    --logdir=${DIR} \
    --model.source_path=/your/data/path/${DATASET}/${SCENE}/ \
    --model.resolution=1 \
    --model.data_device=cpu \
    --wandb \
    --wandb_name ${PROJECT}


================================================
FILE: bash_scripts/1_preprocess_tnt.sh
================================================
echo "Compute intrinsics, undistort images and generate json files. This may take a while"
python process_data/convert_tnt_to_json.py \
        --tnt_path /your/data/path \
        --run_colmap \
        --export_json 

================================================
FILE: bash_scripts/2_extract_normal_dsine.sh
================================================
export CUDA_VISIBLE_DEVICES=0

DOMAIN_TYPE=indoor
DATADIR=/your/data/path

CODE_PATH=/your/dsine/code/path
CKPT=/your/dsine/code/path/checkpoints/dsine.pt

for SCENE in Barn Caterpillar Courthouse Ignatius Meetingroom Truck;
do
    SCENE_PATH=${DATADIR}/${SCENE}
    # dsine
    python -W ignore process_data/extract_normal.py \
            --dsine_path ${CODE_PATH} \
            --ckpt ${CKPT} \
            --img_path ${SCENE_PATH}/images \
            --intrins_path ${SCENE_PATH}/ \
            --output_path ${SCENE_PATH}/normals
done

================================================
FILE: bash_scripts/3_extract_mask.sh
================================================
export CUDA_VISIBLE_DEVICES=0

DATADIR=/your/data/path
GSAM_PATH=~/code/gsam
CKPT_PATH=${GSAM_PATH}

for SCENE in Barn Caterpillar Courthouse Ignatius Meetingroom Truck;
do
    SCENE_PATH=${DATADIR}/${SCENE}
    # meething room scene_tye: indoor, others: outdoor
        if [ ${SCENE} = "Meetingroom" ]; then
            SCENE_TYPE="indoor"
        else
            SCENE_TYPE="outdoor"
        fi
    python -W ignore process_data/extract_mask.py \
            --config ${GSAM_PATH}/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py \
            --grounded_checkpoint ${CKPT_PATH}/groundingdino_swint_ogc.pth \
            --sam_hq_checkpoint ${CKPT_PATH}/sam_hq_vit_h.pth \
            --gsam_path ${GSAM_PATH} \
            --use_sam_hq \
            --input_image ${SCENE_PATH}/images/ \
            --output_dir ${SCENE_PATH}/masks \
            --box_threshold 0.5 \
            --text_threshold 0.2 \
            --scene ${SCENE} \
            --scene_type ${SCENE_TYPE} \
            --device "cuda"
done


================================================
FILE: bash_scripts/4_extract_normal_geow.sh
================================================
export CUDA_VISIBLE_DEVICES=0

# DOMAIN_TYPE=outdoor
# DOMAIN_TYPE=indoor
DOMAIN_TYPE=object
DATADIR=/your/data/path/DTU_mask

CODE_PATH=/your/geowizard/path


for SCENE in scan106  scan114  scan122  scan37  scan55  scan65  scan83 scan105    scan110  scan118  scan24   scan40  scan63  scan69  scan97;
do
    SCENE_PATH=${DATADIR}/${SCENE}
    python process_data/extract_normal_geo.py \
        --code_path ${CODE_PATH} \
        --input_dir ${SCENE_PATH}/images/ \
        --output_dir ${SCENE_PATH}/ \
        --ensemble_size 3 \
        --denoise_steps 10 \
        --seed 0 \
        --domain ${DOMAIN_TYPE}
done

================================================
FILE: bash_scripts/convert.sh
================================================
SCENE=Truck
DATA_ROOT=/your/data/path/${SCENE}

python convert.py -s $DATA_ROOT # [--resize] #If not resizing, ImageMagick is not needed




================================================
FILE: bash_scripts/install.sh
================================================
env=vcr
conda create -n $env -y python=3.10
conda activate $env
pip install -e ".[train]"
export CUDA_HOME=/usr/local/cuda-11.2
pip install -r requirements.txt

================================================
FILE: configs/360_v2/base.yaml
================================================
_parent_: configs/reconstruct.yaml

model:
    eval: True
    llffhold: 8
    split: False

optim:
    mask_depth_thr: 1
    densify_large:
        percent_dense: 5e-2
        sample_cams:
            random: False
            num: 100
    loss_weight:
        semantic: 0
        l1_scale: 1

================================================
FILE: configs/config.py
================================================
'''
-----------------------------------------------------------------------------
Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.

NVIDIA CORPORATION and its licensors retain all intellectual property
and proprietary rights in and to this software, related documentation
and any modifications thereto. Any use, reproduction, disclosure or
distribution of this software and related documentation without an express
license agreement from NVIDIA CORPORATION is strictly prohibited.
-----------------------------------------------------------------------------
'''

import collections
import functools
import os
import re

import yaml
from tools.distributed import master_only_print as print
from tools.termcolor import cyan, green, yellow

DEBUG = False
USE_JIT = False


class AttrDict(dict):
    """Dict as attribute trick."""

    def __init__(self, *args, **kwargs):
        super(AttrDict, self).__init__(*args, **kwargs)
        self.__dict__ = self
        for key, value in self.__dict__.items():
            if isinstance(value, dict):
                self.__dict__[key] = AttrDict(value)
            elif isinstance(value, (list, tuple)):
                if value and isinstance(value[0], dict):
                    self.__dict__[key] = [AttrDict(item) for item in value]
                else:
                    self.__dict__[key] = value

    def yaml(self):
        """Convert object to yaml dict and return."""
        yaml_dict = {}
        for key, value in self.__dict__.items():
            if isinstance(value, AttrDict):
                yaml_dict[key] = value.yaml()
            elif isinstance(value, list):
                if value and isinstance(value[0], AttrDict):
                    new_l = []
                    for item in value:
                        new_l.append(item.yaml())
                    yaml_dict[key] = new_l
                else:
                    yaml_dict[key] = value
            else:
                yaml_dict[key] = value
        return yaml_dict

    def __repr__(self):
        """Print all variables."""
        ret_str = []
        for key, value in self.__dict__.items():
            if isinstance(value, AttrDict):
                ret_str.append('{}:'.format(key))
                child_ret_str = value.__repr__().split('\n')
                for item in child_ret_str:
                    ret_str.append('    ' + item)
            elif isinstance(value, list):
                if value and isinstance(value[0], AttrDict):
                    ret_str.append('{}:'.format(key))
                    for item in value:
                        # Treat as AttrDict above.
                        child_ret_str = item.__repr__().split('\n')
                        for item in child_ret_str:
                            ret_str.append('    ' + item)
                else:
                    ret_str.append('{}: {}'.format(key, value))
            else:
                ret_str.append('{}: {}'.format(key, value))
        return '\n'.join(ret_str)


class Config(AttrDict):
    r"""Configuration class. This should include every human specifiable
    hyperparameter values for your training."""

    def __init__(self, filename=None, verbose=False):
        super(Config, self).__init__()
        self.source_filename = filename

        # Load the base configuration file.
        base_filename = os.path.join(
            os.path.dirname(__file__), './config_base.yaml'
        )
        cfg_base = self.load_config(base_filename)
        recursive_update(self, cfg_base)

        # Update with given configurations.
        cfg_dict = self.load_config(filename)
        recursive_update(self, cfg_dict)

        if verbose:
            print(' imaginaire config '.center(80, '-'))
            print(self.__repr__())
            print(''.center(80, '-'))

    def load_config(self, filename):
        # Update with given configurations.
        assert os.path.exists(filename), f'File {filename} not exist.'
        yaml_loader = yaml.SafeLoader
        yaml_loader.add_implicit_resolver(
            u'tag:yaml.org,2002:float',
            re.compile(u'''^(?:
             [-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)?
            |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
            |\\.[0-9_]+(?:[eE][-+][0-9]+)?
            |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]*
            |[-+]?\\.(?:inf|Inf|INF)
            |\\.(?:nan|NaN|NAN))$''', re.X),
            list(u'-+0123456789.'))
        try:
            with open(filename) as file:
                cfg_dict = yaml.load(file, Loader=yaml_loader)
                cfg_dict = AttrDict(cfg_dict)
        except EnvironmentError:
            print(f'Please check the file with name of "{filename}"')
        # Inherit configurations from parent
        parent_key = "_parent_"
        if parent_key in cfg_dict:
            parent_filename = cfg_dict.pop(parent_key)
            cfg_parent = self.load_config(parent_filename)
            recursive_update(cfg_parent, cfg_dict)
            cfg_dict = cfg_parent
        return cfg_dict

    def print_config(self, level=0):
        """Recursively print the configuration (with termcolor)."""
        for key, value in sorted(self.items()):
            if isinstance(value, (dict, Config)):
                print("   " * level + cyan("* ") + green(key) + ":")
                Config.print_config(value, level + 1)
            else:
                print("   " * level + cyan("* ") + green(key) + ":", yellow(value))

    def save_config(self, logdir):
        """Save the final configuration to a yaml file."""
        cfg_fname = f"{logdir}/config.yaml"
        with open(cfg_fname, "w") as file:
            yaml.safe_dump(self.yaml(), file, default_flow_style=False, indent=4)

def rsetattr(obj, attr, val):
    """Recursively find object and set value"""
    pre, _, post = attr.rpartition('.')
    return setattr(rgetattr(obj, pre) if pre else obj, post, val)


def rgetattr(obj, attr, *args):
    """Recursively find object and return value"""

    def _getattr(obj, attr):
        r"""Get attribute."""
        return getattr(obj, attr, *args)

    return functools.reduce(_getattr, [obj] + attr.split('.'))


def recursive_update(d, u):
    """Recursively update AttrDict d with AttrDict u"""
    for key, value in u.items():
        if isinstance(value, collections.abc.Mapping):
            d.__dict__[key] = recursive_update(d.get(key, AttrDict({})), value)
        elif isinstance(value, (list, tuple)):
            if value and isinstance(value[0], dict):
                d.__dict__[key] = [AttrDict(item) for item in value]
            else:
                d.__dict__[key] = value
        else:
            d.__dict__[key] = value
    return d


def recursive_update_strict(d, u, stack=[]):
    """Recursively update AttrDict d with AttrDict u with strict matching"""
    for key, value in u.items():
        if key not in d:
            key_full = ".".join(stack + [key])
            raise KeyError(f"The input key '{key_full}; does not exist in the config files.")
        if isinstance(value, collections.abc.Mapping):
            d.__dict__[key] = recursive_update_strict(d.get(key, AttrDict({})), value, stack + [key])
        elif isinstance(value, (list, tuple)):
            if value and isinstance(value[0], dict):
                d.__dict__[key] = [AttrDict(item) for item in value]
            else:
                d.__dict__[key] = value
        else:
            d.__dict__[key] = value
    return d


def parse_cmdline_arguments(args):
    """
    Parse arguments from command line.
    Syntax: --key1.key2.key3=value --> value
            --key1.key2.key3=      --> None
            --key1.key2.key3       --> True
            --key1.key2.key3!      --> False
    """
    cfg_cmd = {}
    for arg in args:
        assert arg.startswith("--")
        if "=" not in arg[2:]:
            key_str, value = (arg[2:-1], "false") if arg[-1] == "!" else (arg[2:], "true")
        else:
            key_str, value = arg[2:].split("=")
        keys_sub = key_str.split(".")
        cfg_sub = cfg_cmd
        for k in keys_sub[:-1]:
            cfg_sub.setdefault(k, {})
            cfg_sub = cfg_sub[k]
        assert keys_sub[-1] not in cfg_sub, keys_sub[-1]
        cfg_sub[keys_sub[-1]] = yaml.safe_load(value)
    return cfg_cmd


================================================
FILE: configs/config_base.yaml
================================================
logdir: "/your/log/path/debug/"
ip: 127.0.0.1
port: -1
detect_anomaly: False
silent: 0
seed: 0

model:
    sh_degree: 3
    source_path: "/your/data/path/tnt/Barn/"
    model_path: "/your/log/path/"
    images: "images"
    resolution: -1
    white_background: False
    data_device: "cuda"
    eval: False
    llffhold: 1
    init_ply: "sparse/points3D.ply"
    max_init_points:
    split: False
    sphere: False
    load_depth: False
    load_normal: False
    load_mask: False
    normal_folder: 'normals'
    depth_folder: 'depths'
    use_decoupled_appearance: False
    ch_sem_feat: 0
    num_cls: 0
    max_mem: 22
    load_mask: False
    use_decoupled_appearance: False
    use_decoupled_dnormal: False
    ratio: 0
    mesh:
        voxel_size: 3e-3
    depth_type: 'traditional'

optim:
    iterations: 30000
    position_lr_init: 0.00016
    position_lr_final: 0.0000016
    position_lr_delay_mult: 0.01
    position_lr_max_steps: 30000
    feature_lr: 0.0025
    sdf_lr: 0.001
    weight_decay: 1e-2
    opacity_lr: 0.05
    scaling_lr: 0.005
    rotation_lr: 0.001
    appearance_embeddings_lr: 0.001
    appearance_network_lr: 0.001
    cls_lr: 5e-4
    percent_dense: 0.01
    densification_interval: 100
    opacity_reset_interval: 3000
    densify_from_iter: 500
    densify_until_iter: 15000
    densify_grad_threshold: 0.0005
    random_background: False
    rand_pts: 20000
    edge_thr: 0
    mask_depth_thr: 0
    loss_weight:
        l1: 0.8
        ssim: 0.2
        distortion: 0.
        semantic: 0
        mono_depth: 0
        mono_normal: 0
        depth_normal: 0
    prune:
        iterations: []
        percent: 0.5
        decay: 0.6
        v_pow: 0.1

pipline:
    convert_SHs_python: False
    compute_cov3D_python: False
    debug: False
    
data:
    name: dummy

train:
    test_iterations: [7000, 30000]
    save_iterations: [7000, 30000]
    checkpoint_iterations: [30000]
    save_splat: False
    start_checkpoint: 
    debug_from: -1



================================================
FILE: configs/dtu/base.yaml
================================================
_parent_: configs/reconstruct.yaml

model:
    use_decoupled_appearance: False
    use_decoupled_dnormal: False
    normal_folder: 'normal_npz_indoor'
    eval: False

optim:
    exp_t: 0.01
    mask_depth_thr: 0
    loss_weight:
        l1_scale: 0.5
    consistent_normal_from_iter: 15000
    close_depth_from_iter: 15000
    densify_large:
        percent_dense: 1e-2
        sample_cams:
            random: False
            num: 30
    loss_weight:
        semantic: 0
        depth_normal: 0
        mono_normal: 0.01
        consistent_normal: 0.05
        distortion: 1000
        depth_var: 0
    random_background: False
        

================================================
FILE: configs/dtu/dtu_scan24.yaml
================================================
_parent_: configs/dtu/base.yaml


================================================
FILE: configs/reconstruct.yaml
================================================
_parent_: configs/config_base.yaml


model:
    load_mask: False
    use_decoupled_appearance: False
    use_decoupled_dnormal: False
    ch_sem_feat: 2
    num_cls: 2
    depth_type: 'intersection'
optim:
    mask_depth_thr: 0.8
    edge_thr: 0
    exp_t: 0.01
    cos_thr: -1
    close_depth_from_iter: 0
    normal_from_iter: 0
    dnormal_from_iter: 0
    consistent_normal_from_iter: 0
    curv_from_iter: 0
    loss_weight:
        l1: 0.8
        ssim: 0.2
        l1_scale: 1
        entropy: 0
        depth_var: 0.
        mono_depth: 0
        mono_normal: 0.01
        depth_normal: 0.01
        consistent_normal: 0
    prune:
        iterations: [15000, 25000]
        percent: 0.5
        decay: 0.6
        v_pow: 0.1
    densify_large:
        percent_dense: 2e-3
        interval: 1
        sample_cams:
            random: True
            num: 200
            up: True
            around: True
            look_mode: 'target'
    random_background: True


train:
    checkpoint_iterations: []
    save_mesh: False
    save_iterations: [30000]

================================================
FILE: configs/scannetpp/base.yaml
================================================
_parent_: configs/reconstruct.yaml

model:
    split: True
    eval: True
    use_decoupled_appearance: False
    use_decoupled_dnormal: False
    mesh:
        voxel_size: 1.5e-2

optim:
    mask_depth_thr: 0
    curv_from_iter: 15000
    densify_large:
        percent_dense: 1e-2
        sample_cams:
            random: False
    loss_weight:
        semantic: 0
        curv: 0.05

================================================
FILE: configs/tnt/Barn.yaml
================================================
_parent_: configs/tnt/base.yaml


================================================
FILE: configs/tnt/Caterpillar.yaml
================================================
_parent_: configs/tnt/base.yaml


================================================
FILE: configs/tnt/Courthouse.yaml
================================================
_parent_: configs/tnt/base.yaml


================================================
FILE: configs/tnt/Ignatius.yaml
================================================
_parent_: configs/tnt/base.yaml


================================================
FILE: configs/tnt/Meetingroom.yaml
================================================
_parent_: configs/tnt/base.yaml

optim:
    exp_t: 1e-3
    mask_depth_thr: 0
    densify_large:
        percent_dense: 5e-3
        sample_cams:
            random: False
    loss_weight:
        semantic: 0
model:
    num_cls: 3
    use_decoupled_appearance: False

================================================
FILE: configs/tnt/Truck.yaml
================================================
_parent_: configs/tnt/base.yaml


================================================
FILE: configs/tnt/base.yaml
================================================
_parent_: configs/reconstruct.yaml

model:
    use_decoupled_appearance: True
    use_decoupled_dnormal: False
    eval: False
    llffhold: 5

optim:
    exp_t: 5e-3
    loss_weight:
        depth_normal: 0.015
        semantic: 0.005
        l1_scale: 1

================================================
FILE: environment.yml
================================================
name: fast_render
channels:
  - pytorch
  - nvidia
  - conda-forge
  - defaults
dependencies:
  - python=3.10
  - pytorch==2.0.1
  - torchvision==0.15.2
  - torchaudio==2.0.2
  - pytorch-cuda=11.8
  - pip:
    - open3d
    - plyfile
    - ninja
    - GPUtil
    - opencv-python
    - lpips
    - trimesh
    - pymeshlab
    - termcolor
    - wandb
    - imageio
    - scikit-image
    - torchmetrics
    - mediapy
    - "git+https://github.com/facebookresearch/pytorch3d.git"
    - submodules/diff-gaussian-rasterization
    - submodules/simple-knn

================================================
FILE: evaluation/crop_mesh.py
================================================
import os
import json
import plyfile
import argparse
# import open3d as o3d
import numpy as np
# from tqdm import tqdm
import trimesh
from sklearn.cluster import DBSCAN


def align_gt_with_cam(pts, trans):
    trans_inv = np.linalg.inv(trans)
    pts_aligned = pts @ trans_inv[:3, :3].transpose(-1, -2) + trans_inv[:3, -1]
    return pts_aligned


def main(args):
    assert os.path.exists(args.ply_path), f"PLY file {args.ply_path} does not exist."
    gt_trans = np.loadtxt(args.align_path)
    
    mesh_rec = trimesh.load(args.ply_path, process=False)
    mesh_gt = trimesh.load(args.gt_path, process=False)
    
    mesh_gt.vertices = align_gt_with_cam(mesh_gt.vertices, gt_trans)
    
    to_align, _ = trimesh.bounds.oriented_bounds(mesh_gt)
    mesh_gt.vertices = (to_align[:3, :3] @ mesh_gt.vertices.T + to_align[:3, 3:]).T
    mesh_rec.vertices = (to_align[:3, :3] @ mesh_rec.vertices.T + to_align[:3, 3:]).T
    
    min_points = mesh_gt.vertices.min(axis=0)
    max_points = mesh_gt.vertices.max(axis=0)

    mask_min = (mesh_rec.vertices - min_points[None]) > 0
    mask_max = (mesh_rec.vertices - max_points[None]) < 0

    mask = np.concatenate((mask_min, mask_max), axis=1).all(axis=1)
    face_mask = mask[mesh_rec.faces].all(axis=1)

    mesh_rec.update_vertices(mask)
    mesh_rec.update_faces(face_mask)
    
    mesh_rec.vertices = (to_align[:3, :3].T @ mesh_rec.vertices.T - to_align[:3, :3].T @ to_align[:3, 3:]).T
    mesh_gt.vertices = (to_align[:3, :3].T @ mesh_gt.vertices.T - to_align[:3, :3].T @ to_align[:3, 3:]).T
    
    # save mesh_rec and mesh_rec in args.out_path
    mesh_rec.export(args.out_path)
    
    # downsample mesh_gt
    
    idx = np.random.choice(np.arange(len(mesh_gt.vertices)), 5000000)
    mesh_gt.vertices = mesh_gt.vertices[idx]
    mesh_gt.colors = mesh_gt.colors[idx]
    
    mesh_gt.export(args.gt_path.replace('.ply', '_trans.ply'))
    
    
    return



if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--gt_path",
        type=str,
        default='/your/path//Barn_GT.ply',
        help="path to a dataset/scene directory containing X.json, X.ply, ...",
    )
    parser.add_argument(
        "--align_path",
        type=str,
        default='/your/path//Barn_trans.txt',
        help="path to a dataset/scene directory containing X.json, X.ply, ...",
    )
    parser.add_argument(
        "--ply_path",
        type=str,
        default='/your/path//Barn_lowres.ply',
        help="path to reconstruction ply file",
    )
    parser.add_argument(
        "--scene",
        type=str,
        default='Barn',
        help="path to reconstruction ply file",
    )
    parser.add_argument(
        "--out_path",
        type=str,
        default='/your/path//Barn_lowres_crop.ply',
        help=
        "output directory, default: an evaluation directory is created in the directory of the ply file",
    )
    args = parser.parse_args()
    
    main(args)

================================================
FILE: evaluation/eval_dtu/eval.py
================================================
# adapted from https://github.com/jzhangbs/DTUeval-python
import numpy as np
import open3d as o3d
import sklearn.neighbors as skln
from tqdm import tqdm
from scipy.io import loadmat
import multiprocessing as mp
import argparse

def sample_single_tri(input_):
    n1, n2, v1, v2, tri_vert = input_
    c = np.mgrid[:n1+1, :n2+1]
    c += 0.5
    c[0] /= max(n1, 1e-7)
    c[1] /= max(n2, 1e-7)
    c = np.transpose(c, (1,2,0))
    k = c[c.sum(axis=-1) < 1]  # m2
    q = v1 * k[:,:1] + v2 * k[:,1:] + tri_vert
    return q

def write_vis_pcd(file, points, colors):
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    pcd.colors = o3d.utility.Vector3dVector(colors)
    o3d.io.write_point_cloud(file, pcd)

if __name__ == '__main__':
    mp.freeze_support()

    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default='data_in.ply')
    parser.add_argument('--scan', type=int, default=1)
    parser.add_argument('--mode', type=str, default='mesh', choices=['mesh', 'pcd'])
    parser.add_argument('--dataset_dir', type=str, default='.')
    parser.add_argument('--vis_out_dir', type=str, default='.')
    parser.add_argument('--downsample_density', type=float, default=0.2)
    parser.add_argument('--patch_size', type=float, default=60)
    parser.add_argument('--max_dist', type=float, default=20)
    parser.add_argument('--visualize_threshold', type=float, default=10)
    args = parser.parse_args()

    thresh = args.downsample_density
    if args.mode == 'mesh':
        pbar = tqdm(total=9)
        pbar.set_description('read data mesh')
        data_mesh = o3d.io.read_triangle_mesh(args.data)

        vertices = np.asarray(data_mesh.vertices)
        triangles = np.asarray(data_mesh.triangles)
        tri_vert = vertices[triangles]

        pbar.update(1)
        pbar.set_description('sample pcd from mesh')
        v1 = tri_vert[:,1] - tri_vert[:,0]
        v2 = tri_vert[:,2] - tri_vert[:,0]
        l1 = np.linalg.norm(v1, axis=-1, keepdims=True)
        l2 = np.linalg.norm(v2, axis=-1, keepdims=True)
        area2 = np.linalg.norm(np.cross(v1, v2), axis=-1, keepdims=True)
        non_zero_area = (area2 > 0)[:,0]
        l1, l2, area2, v1, v2, tri_vert = [
            arr[non_zero_area] for arr in [l1, l2, area2, v1, v2, tri_vert]
        ]
        thr = thresh * np.sqrt(l1 * l2 / area2)
        n1 = np.floor(l1 / thr)
        n2 = np.floor(l2 / thr)

        with mp.Pool() as mp_pool:
            new_pts = mp_pool.map(sample_single_tri, ((n1[i,0], n2[i,0], v1[i:i+1], v2[i:i+1], tri_vert[i:i+1,0]) for i in range(len(n1))), chunksize=1024)

        new_pts = np.concatenate(new_pts, axis=0)
        data_pcd = np.concatenate([vertices, new_pts], axis=0)
    
    elif args.mode == 'pcd':
        pbar = tqdm(total=8)
        pbar.set_description('read data pcd')
        data_pcd_o3d = o3d.io.read_point_cloud(args.data)
        data_pcd = np.asarray(data_pcd_o3d.points)

    pbar.update(1)
    pbar.set_description('random shuffle pcd index')
    shuffle_rng = np.random.default_rng()
    shuffle_rng.shuffle(data_pcd, axis=0)

    pbar.update(1)
    pbar.set_description('downsample pcd')
    nn_engine = skln.NearestNeighbors(n_neighbors=1, radius=thresh, algorithm='kd_tree', n_jobs=-1)
    nn_engine.fit(data_pcd)
    rnn_idxs = nn_engine.radius_neighbors(data_pcd, radius=thresh, return_distance=False)
    mask = np.ones(data_pcd.shape[0], dtype=np.bool_)
    for curr, idxs in enumerate(rnn_idxs):
        if mask[curr]:
            mask[idxs] = 0
            mask[curr] = 1
    data_down = data_pcd[mask]

    pbar.update(1)
    pbar.set_description('masking data pcd')
    obs_mask_file = loadmat(f'{args.dataset_dir}/ObsMask/ObsMask{args.scan}_10.mat')
    ObsMask, BB, Res = [obs_mask_file[attr] for attr in ['ObsMask', 'BB', 'Res']]
    BB = BB.astype(np.float32)

    patch = args.patch_size
    inbound = ((data_down >= BB[:1]-patch) & (data_down < BB[1:]+patch*2)).sum(axis=-1) ==3
    data_in = data_down[inbound]

    data_grid = np.around((data_in - BB[:1]) / Res).astype(np.int32)
    grid_inbound = ((data_grid >= 0) & (data_grid < np.expand_dims(ObsMask.shape, 0))).sum(axis=-1) ==3
    data_grid_in = data_grid[grid_inbound]
    in_obs = ObsMask[data_grid_in[:,0], data_grid_in[:,1], data_grid_in[:,2]].astype(np.bool_)
    data_in_obs = data_in[grid_inbound][in_obs]

    pbar.update(1)
    pbar.set_description('read STL pcd')
    stl_pcd = o3d.io.read_point_cloud(f'{args.dataset_dir}/Points/stl/stl{args.scan:03}_total.ply')
    stl = np.asarray(stl_pcd.points)

    pbar.update(1)
    pbar.set_description('compute data2stl')
    nn_engine.fit(stl)
    dist_d2s, idx_d2s = nn_engine.kneighbors(data_in_obs, n_neighbors=1, return_distance=True)
    max_dist = args.max_dist
    mean_d2s = dist_d2s[dist_d2s < max_dist].mean()

    pbar.update(1)
    pbar.set_description('compute stl2data')
    ground_plane = loadmat(f'{args.dataset_dir}/ObsMask/Plane{args.scan}.mat')['P']

    stl_hom = np.concatenate([stl, np.ones_like(stl[:,:1])], -1)
    above = (ground_plane.reshape((1,4)) * stl_hom).sum(-1) > 0
    stl_above = stl[above]

    nn_engine.fit(data_in)
    dist_s2d, idx_s2d = nn_engine.kneighbors(stl_above, n_neighbors=1, return_distance=True)
    mean_s2d = dist_s2d[dist_s2d < max_dist].mean()

    pbar.update(1)
    pbar.set_description('visualize error')
    vis_dist = args.visualize_threshold
    R = np.array([[1,0,0]], dtype=np.float64)
    G = np.array([[0,1,0]], dtype=np.float64)
    B = np.array([[0,0,1]], dtype=np.float64)
    W = np.array([[1,1,1]], dtype=np.float64)
    data_color = np.tile(B, (data_down.shape[0], 1))
    data_alpha = dist_d2s.clip(max=vis_dist) / vis_dist
    data_color[ np.where(inbound)[0][grid_inbound][in_obs] ] = R * data_alpha + W * (1-data_alpha)
    data_color[ np.where(inbound)[0][grid_inbound][in_obs][dist_d2s[:,0] >= max_dist] ] = G
    write_vis_pcd(f'{args.vis_out_dir}/vis_{args.scan:03}_d2s.ply', data_down, data_color)
    stl_color = np.tile(B, (stl.shape[0], 1))
    stl_alpha = dist_s2d.clip(max=vis_dist) / vis_dist
    stl_color[ np.where(above)[0] ] = R * stl_alpha + W * (1-stl_alpha)
    stl_color[ np.where(above)[0][dist_s2d[:,0] >= max_dist] ] = G
    write_vis_pcd(f'{args.vis_out_dir}/vis_{args.scan:03}_s2d.ply', stl, stl_color)

    pbar.update(1)
    pbar.set_description('done')
    pbar.close()
    over_all = (mean_d2s + mean_s2d) / 2
    print(mean_d2s, mean_s2d, over_all)
    
    import json
    with open(f'{args.vis_out_dir}/results.json', 'w') as fp:
        json.dump({
            'mean_d2s': mean_d2s,
            'mean_s2d': mean_s2d,
            'overall': over_all,
        }, fp, indent=True)




================================================
FILE: evaluation/eval_dtu/evaluate_single_scene.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
import cv2
import numpy as np
import os
import glob
from skimage.morphology import binary_dilation, disk
import argparse

import trimesh
from pathlib import Path
from tqdm import tqdm

import sys

sys.path.append(os.getcwd())

import evaluation.eval_dtu.render_utils as rend_util


def cull_scan(scan, mesh_path, result_mesh_file, instance_dir):
    
    # load poses
    image_dir = '{0}/images'.format(instance_dir)
    image_paths = sorted(glob.glob(os.path.join(image_dir, "*.png")))
    n_images = len(image_paths)
    cam_file = '{0}/cameras.npz'.format(instance_dir)
    camera_dict = np.load(cam_file)
    scale_mats = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(n_images)]
    world_mats = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(n_images)]

    intrinsics_all = []
    pose_all = []
    for scale_mat, world_mat in zip(scale_mats, world_mats):
        P = world_mat @ scale_mat
        P = P[:3, :4]
        intrinsics, pose = rend_util.load_K_Rt_from_P(None, P)
        intrinsics_all.append(torch.from_numpy(intrinsics).float())
        pose_all.append(torch.from_numpy(pose).float())
    
    # load mask
    mask_dir = '{0}/mask'.format(instance_dir)
    mask_paths = sorted(glob.glob(os.path.join(mask_dir, "*.png")))
    masks = []
    for p in mask_paths:
        mask = cv2.imread(p)
        masks.append(mask)

    # hard-coded image shape
    W, H = 1600, 1200

    # load mesh
    mesh = trimesh.load(mesh_path)
    
    # load transformation matrix

    vertices = mesh.vertices

    # project and filter
    vertices = torch.from_numpy(vertices).cuda()
    vertices = torch.cat((vertices, torch.ones_like(vertices[:, :1])), dim=-1)
    vertices = vertices.permute(1, 0)
    vertices = vertices.float()

    sampled_masks = []
    for i in tqdm(range(n_images),  desc="Culling mesh given masks"):
        pose = pose_all[i]
        w2c = torch.inverse(pose).cuda()
        intrinsic = intrinsics_all[i].cuda()

        with torch.no_grad():
            # transform and project
            cam_points = intrinsic @ w2c @ vertices
            pix_coords = cam_points[:2, :] / (cam_points[2, :].unsqueeze(0) + 1e-6)
            pix_coords = pix_coords.permute(1, 0)
            pix_coords[..., 0] /= W - 1
            pix_coords[..., 1] /= H - 1
            pix_coords = (pix_coords - 0.5) * 2
            valid = ((pix_coords > -1. ) & (pix_coords < 1.)).all(dim=-1).float()
            
            # dialate mask similar to unisurf
            maski = masks[i][:, :, 0].astype(np.float32) / 256.
            maski = torch.from_numpy(binary_dilation(maski, disk(24))).float()[None, None].cuda()

            sampled_mask = F.grid_sample(maski, pix_coords[None, None], mode='nearest', padding_mode='zeros', align_corners=True)[0, -1, 0]

            sampled_mask = sampled_mask + (1. - valid)
            sampled_masks.append(sampled_mask)

    sampled_masks = torch.stack(sampled_masks, -1)
    # filter
    
    mask = (sampled_masks > 0.).all(dim=-1).cpu().numpy()
    face_mask = mask[mesh.faces].all(axis=1)

    mesh.update_vertices(mask)
    mesh.update_faces(face_mask)
    
    # transform vertices to world 
    scale_mat = scale_mats[0]
    mesh.vertices = mesh.vertices * scale_mat[0, 0] + scale_mat[:3, 3][None]
    
    # Taking the biggest connected component
    print("Taking the biggest connected component")
    components = mesh.split(only_watertight=False)
    areas = np.array([c.area for c in components], dtype=np.float32)
    mesh = components[areas.argmax()]
    
    mesh.export(result_mesh_file)
    del mesh
    

if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description='Arguments to evaluate the mesh.'
    )

    parser.add_argument('--input_mesh', type=str,  help='path to the mesh to be evaluated')
    parser.add_argument('--scan_id', type=str,  help='scan id of the input mesh')
    parser.add_argument('--output_dir', type=str, default='evaluation_results_single', help='path to the output folder')
    parser.add_argument('--mask_dir', type=str,  default='mask', help='path to uncropped mask')
    parser.add_argument('--DTU', type=str,  default='Offical_DTU_Dataset', help='path to the GT DTU point clouds')
    args = parser.parse_args()

    Offical_DTU_Dataset = args.DTU
    out_dir = args.output_dir
    Path(out_dir).mkdir(parents=True, exist_ok=True)

    scan = args.scan_id
    ply_file = args.input_mesh
    print("cull mesh ....")
    result_mesh_file = os.path.join(out_dir, "culled_mesh.ply")
    cull_scan(scan, ply_file, result_mesh_file, instance_dir=os.path.join(args.mask_dir, f'scan{args.scan_id}'))

    script_dir = os.path.dirname(os.path.abspath(__file__))
    cmd = f"python {script_dir}/eval.py --data {result_mesh_file} --scan {scan} --mode mesh --dataset_dir {Offical_DTU_Dataset} --vis_out_dir {out_dir}"
    os.system(cmd)

================================================
FILE: evaluation/eval_dtu/render_utils.py
================================================
import numpy as np
import imageio
import skimage
import cv2
import torch
from torch.nn import functional as F


def get_psnr(img1, img2, normalize_rgb=False):
    if normalize_rgb: # [-1,1] --> [0,1]
        img1 = (img1 + 1.) / 2.
        img2 = (img2 + 1. ) / 2.

    mse = torch.mean((img1 - img2) ** 2)
    psnr = -10. * torch.log(mse) / torch.log(torch.Tensor([10.]).cuda())

    return psnr


def load_rgb(path, normalize_rgb = False):
    img = imageio.imread(path)
    img = skimage.img_as_float32(img)

    if normalize_rgb: # [-1,1] --> [0,1]
        img -= 0.5
        img *= 2.
    img = img.transpose(2, 0, 1)
    return img


def load_K_Rt_from_P(filename, P=None):
    if P is None:
        lines = open(filename).read().splitlines()
        if len(lines) == 4:
            lines = lines[1:]
        lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
        P = np.asarray(lines).astype(np.float32).squeeze()

    out = cv2.decomposeProjectionMatrix(P)
    K = out[0]
    R = out[1]
    t = out[2]

    K = K/K[2,2]
    intrinsics = np.eye(4)
    intrinsics[:3, :3] = K

    pose = np.eye(4, dtype=np.float32)
    pose[:3, :3] = R.transpose()
    pose[:3,3] = (t[:3] / t[3])[:,0]

    return intrinsics, pose


def get_camera_params(uv, pose, intrinsics):
    if pose.shape[1] == 7: #In case of quaternion vector representation
        cam_loc = pose[:, 4:]
        R = quat_to_rot(pose[:,:4])
        p = torch.eye(4).repeat(pose.shape[0],1,1).cuda().float()
        p[:, :3, :3] = R
        p[:, :3, 3] = cam_loc
    else: # In case of pose matrix representation
        cam_loc = pose[:, :3, 3]
        p = pose

    batch_size, num_samples, _ = uv.shape

    depth = torch.ones((batch_size, num_samples)).cuda()
    x_cam = uv[:, :, 0].view(batch_size, -1)
    y_cam = uv[:, :, 1].view(batch_size, -1)
    z_cam = depth.view(batch_size, -1)

    pixel_points_cam = lift(x_cam, y_cam, z_cam, intrinsics=intrinsics)

    # permute for batch matrix product
    pixel_points_cam = pixel_points_cam.permute(0, 2, 1)

    world_coords = torch.bmm(p, pixel_points_cam).permute(0, 2, 1)[:, :, :3]
    ray_dirs = world_coords - cam_loc[:, None, :]
    ray_dirs = F.normalize(ray_dirs, dim=2)

    return ray_dirs, cam_loc


def get_camera_for_plot(pose):
    if pose.shape[1] == 7: #In case of quaternion vector representation
        cam_loc = pose[:, 4:].detach()
        R = quat_to_rot(pose[:,:4].detach())
    else: # In case of pose matrix representation
        cam_loc = pose[:, :3, 3]
        R = pose[:, :3, :3]
    cam_dir = R[:, :3, 2]
    return cam_loc, cam_dir


def lift(x, y, z, intrinsics):
    # parse intrinsics
    intrinsics = intrinsics.cuda()
    fx = intrinsics[:, 0, 0]
    fy = intrinsics[:, 1, 1]
    cx = intrinsics[:, 0, 2]
    cy = intrinsics[:, 1, 2]
    sk = intrinsics[:, 0, 1]

    x_lift = (x - cx.unsqueeze(-1) + cy.unsqueeze(-1)*sk.unsqueeze(-1)/fy.unsqueeze(-1) - sk.unsqueeze(-1)*y/fy.unsqueeze(-1)) / fx.unsqueeze(-1) * z
    y_lift = (y - cy.unsqueeze(-1)) / fy.unsqueeze(-1) * z

    # homogeneous
    return torch.stack((x_lift, y_lift, z, torch.ones_like(z).cuda()), dim=-1)


def quat_to_rot(q):
    batch_size, _ = q.shape
    q = F.normalize(q, dim=1)
    R = torch.ones((batch_size, 3,3)).cuda()
    qr=q[:,0]
    qi = q[:, 1]
    qj = q[:, 2]
    qk = q[:, 3]
    R[:, 0, 0]=1-2 * (qj**2 + qk**2)
    R[:, 0, 1] = 2 * (qj *qi -qk*qr)
    R[:, 0, 2] = 2 * (qi * qk + qr * qj)
    R[:, 1, 0] = 2 * (qj * qi + qk * qr)
    R[:, 1, 1] = 1-2 * (qi**2 + qk**2)
    R[:, 1, 2] = 2*(qj*qk - qi*qr)
    R[:, 2, 0] = 2 * (qk * qi-qj * qr)
    R[:, 2, 1] = 2 * (qj*qk + qi*qr)
    R[:, 2, 2] = 1-2 * (qi**2 + qj**2)
    return R


def rot_to_quat(R):
    batch_size, _,_ = R.shape
    q = torch.ones((batch_size, 4)).cuda()

    R00 = R[:, 0,0]
    R01 = R[:, 0, 1]
    R02 = R[:, 0, 2]
    R10 = R[:, 1, 0]
    R11 = R[:, 1, 1]
    R12 = R[:, 1, 2]
    R20 = R[:, 2, 0]
    R21 = R[:, 2, 1]
    R22 = R[:, 2, 2]

    q[:,0]=torch.sqrt(1.0+R00+R11+R22)/2
    q[:, 1]=(R21-R12)/(4*q[:,0])
    q[:, 2] = (R02 - R20) / (4 * q[:, 0])
    q[:, 3] = (R10 - R01) / (4 * q[:, 0])
    return q


def get_sphere_intersections(cam_loc, ray_directions, r = 1.0):
    # Input: n_rays x 3 ; n_rays x 3
    # Output: n_rays x 1, n_rays x 1 (close and far)

    ray_cam_dot = torch.bmm(ray_directions.view(-1, 1, 3),
                            cam_loc.view(-1, 3, 1)).squeeze(-1)
    under_sqrt = ray_cam_dot ** 2 - (cam_loc.norm(2, 1, keepdim=True) ** 2 - r ** 2)

    # sanity check
    if (under_sqrt <= 0).sum() > 0:
        print('BOUNDING SPHERE PROBLEM!')
        exit()

    sphere_intersections = torch.sqrt(under_sqrt) * torch.Tensor([-1, 1]).cuda().float() - ray_cam_dot
    sphere_intersections = sphere_intersections.clamp_min(0.0)

    return sphere_intersections

================================================
FILE: evaluation/eval_tnt.py
================================================
import os
import trimesh
import argparse
import numpy as np
import open3d as o3d
from sklearn.neighbors import KDTree


def nn_correspondance(verts1, verts2):
    indices = []
    distances = []
    if len(verts1) == 0 or len(verts2) == 0:
        return indices, distances

    kdtree = KDTree(verts1)
    distances, indices = kdtree.query(verts2)
    distances = distances.reshape(-1)

    return distances


def evaluate(mesh_pred, mesh_trgt, threshold=.05, down_sample=.02):
    pcd_trgt = o3d.geometry.PointCloud()
    pcd_pred = o3d.geometry.PointCloud()
    
    pcd_trgt.points = o3d.utility.Vector3dVector(mesh_trgt.vertices[:, :3])
    pcd_pred.points = o3d.utility.Vector3dVector(mesh_pred.vertices[:, :3])

    if down_sample:
        pcd_pred = pcd_pred.voxel_down_sample(down_sample)
        pcd_trgt = pcd_trgt.voxel_down_sample(down_sample)

    verts_pred = np.asarray(pcd_pred.points)
    verts_trgt = np.asarray(pcd_trgt.points)

    dist1 = nn_correspondance(verts_pred, verts_trgt)
    dist2 = nn_correspondance(verts_trgt, verts_pred)

    precision = np.mean((dist2 < threshold).astype('float'))
    recal = np.mean((dist1 < threshold).astype('float'))
    fscore = 2 * precision * recal / (precision + recal)
    metrics = {
        'Acc': np.mean(dist2),
        'Comp': np.mean(dist1),
        'Prec': precision,
        'Recal': recal,
        'F-score': fscore,
    }
    return metrics


def main(args):
    assert os.path.exists(args.ply_path), f"PLY file {args.ply_path} does not exist."
    
    mesh_rec = trimesh.load(args.ply_path, process=False)
    mesh_gt = trimesh.load(args.gt_path, process=False)
    
    to_align, _ = trimesh.bounds.oriented_bounds(mesh_gt)
    mesh_gt.vertices = (to_align[:3, :3] @ mesh_gt.vertices.T + to_align[:3, 3:]).T
    mesh_rec.vertices = (to_align[:3, :3] @ mesh_rec.vertices.T + to_align[:3, 3:]).T
    
    min_points = mesh_gt.vertices.min(axis=0)
    max_points = mesh_gt.vertices.max(axis=0)

    mask_min = (mesh_rec.vertices - min_points[None]) > 0
    mask_max = (mesh_rec.vertices - max_points[None]) < 0

    mask = np.concatenate((mask_min, mask_max), axis=1).all(axis=1)
    face_mask = mask[mesh_rec.faces].all(axis=1)

    mesh_rec.update_vertices(mask)
    mesh_rec.update_faces(face_mask)
    
    metrics = evaluate(mesh_rec, mesh_gt)
    
    metrics_path = os.path.join(os.path.dirname(args.ply_path), 'metrics.txt')
    with open(metrics_path, 'w') as f:
        for k, v in metrics.items():
            f.write(f'{k}: {v}\n')
    
    print('Scene: {} F-score: {}'.format(args.scene, metrics['F-score']))
    
    mesh_rec.vertices = (to_align[:3, :3].T @ mesh_rec.vertices.T - to_align[:3, :3].T @ to_align[:3, 3:]).T
    mesh_rec.export(args.ply_path.replace('.ply', '_crop.ply'))
    
    return


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--gt_path",
        type=str,
        default='/your/path//Barn_GT.ply',
        help="path to a dataset/scene directory containing X.json, X.ply, ...",
    )
    parser.add_argument(
        "--ply_path",
        type=str,
        default='/your/path//Barn_lowres.ply',
        help="path to reconstruction ply file",
    )
    parser.add_argument(
        "--scene",
        type=str,
        default='Barn',
        help="path to reconstruction ply file",
    )
    args = parser.parse_args()
    
    main(args)



================================================
FILE: evaluation/full_eval.py
================================================
#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import os
from argparse import ArgumentParser

mipnerf360_outdoor_scenes = ["bicycle", "flowers", "garden", "stump", "treehill"]
mipnerf360_indoor_scenes = ["room", "counter", "kitchen", "bonsai"]
tanks_and_temples_scenes = ["truck", "train"]
deep_blending_scenes = ["drjohnson", "playroom"]

parser = ArgumentParser(description="Full evaluation script parameters")
parser.add_argument("--skip_training", action="store_true")
parser.add_argument("--skip_rendering", action="store_true")
parser.add_argument("--skip_metrics", action="store_true")
parser.add_argument("--output_path", default="./eval")
args, _ = parser.parse_known_args()

all_scenes = []
all_scenes.extend(mipnerf360_outdoor_scenes)
all_scenes.extend(mipnerf360_indoor_scenes)
all_scenes.extend(tanks_and_temples_scenes)
all_scenes.extend(deep_blending_scenes)

if not args.skip_training or not args.skip_rendering:
    parser.add_argument('--mipnerf360', "-m360", required=True, type=str)
    parser.add_argument("--tanksandtemples", "-tat", required=True, type=str)
    parser.add_argument("--deepblending", "-db", required=True, type=str)
    args = parser.parse_args()

if not args.skip_training:
    common_args = " --quiet --eval --test_iterations -1 "
    for scene in mipnerf360_outdoor_scenes:
        source = args.mipnerf360 + "/" + scene
        os.system("python train.py -s " + source + " -i images_4 -m " + args.output_path + "/" + scene + common_args)
    for scene in mipnerf360_indoor_scenes:
        source = args.mipnerf360 + "/" + scene
        os.system("python train.py -s " + source + " -i images_2 -m " + args.output_path + "/" + scene + common_args)
    for scene in tanks_and_temples_scenes:
        source = args.tanksandtemples + "/" + scene
        os.system("python train.py -s " + source + " -m " + args.output_path + "/" + scene + common_args)
    for scene in deep_blending_scenes:
        source = args.deepblending + "/" + scene
        os.system("python train.py -s " + source + " -m " + args.output_path + "/" + scene + common_args)

if not args.skip_rendering:
    all_sources = []
    for scene in mipnerf360_outdoor_scenes:
        all_sources.append(args.mipnerf360 + "/" + scene)
    for scene in mipnerf360_indoor_scenes:
        all_sources.append(args.mipnerf360 + "/" + scene)
    for scene in tanks_and_temples_scenes:
        all_sources.append(args.tanksandtemples + "/" + scene)
    for scene in deep_blending_scenes:
        all_sources.append(args.deepblending + "/" + scene)

    common_args = " --quiet --eval --skip_train"
    for scene, source in zip(all_scenes, all_sources):
        os.system("python render.py --iteration 7000 -s " + source + " -m " + args.output_path + "/" + scene + common_args)
        os.system("python render.py --iteration 30000 -s " + source + " -m " + args.output_path + "/" + scene + common_args)

if not args.skip_metrics:
    scenes_string = ""
    for scene in all_scenes:
        scenes_string += "\"" + args.output_path + "/" + scene + "\" "

    os.system("python metrics.py -m " + scenes_string)

================================================
FILE: evaluation/lpipsPyTorch/__init__.py
================================================
import torch

from .modules.lpips import LPIPS


def lpips(x: torch.Tensor,
          y: torch.Tensor,
          net_type: str = 'alex',
          version: str = '0.1'):
    r"""Function that measures
    Learned Perceptual Image Patch Similarity (LPIPS).

    Arguments:
        x, y (torch.Tensor): the input tensors to compare.
        net_type (str): the network type to compare the features: 
                        'alex' | 'squeeze' | 'vgg'. Default: 'alex'.
        version (str): the version of LPIPS. Default: 0.1.
    """
    device = x.device
    criterion = LPIPS(net_type, version).to(device)
    return criterion(x, y)


================================================
FILE: evaluation/lpipsPyTorch/modules/lpips.py
================================================
import torch
import torch.nn as nn

from .networks import get_network, LinLayers
from .utils import get_state_dict


class LPIPS(nn.Module):
    r"""Creates a criterion that measures
    Learned Perceptual Image Patch Similarity (LPIPS).

    Arguments:
        net_type (str): the network type to compare the features: 
                        'alex' | 'squeeze' | 'vgg'. Default: 'alex'.
        version (str): the version of LPIPS. Default: 0.1.
    """
    def __init__(self, net_type: str = 'alex', version: str = '0.1'):

        assert version in ['0.1'], 'v0.1 is only supported now'

        super(LPIPS, self).__init__()

        # pretrained network
        self.net = get_network(net_type)

        # linear layers
        self.lin = LinLayers(self.net.n_channels_list)
        self.lin.load_state_dict(get_state_dict(net_type, version))

    def forward(self, x: torch.Tensor, y: torch.Tensor):
        feat_x, feat_y = self.net(x), self.net(y)

        diff = [(fx - fy) ** 2 for fx, fy in zip(feat_x, feat_y)]
        res = [l(d).mean((2, 3), True) for d, l in zip(diff, self.lin)]

        return torch.sum(torch.cat(res, 0), 0, True)


================================================
FILE: evaluation/lpipsPyTorch/modules/networks.py
================================================
from typing import Sequence

from itertools import chain

import torch
import torch.nn as nn
from torchvision import models

from .utils import normalize_activation


def get_network(net_type: str):
    if net_type == 'alex':
        return AlexNet()
    elif net_type == 'squeeze':
        return SqueezeNet()
    elif net_type == 'vgg':
        return VGG16()
    else:
        raise NotImplementedError('choose net_type from [alex, squeeze, vgg].')


class LinLayers(nn.ModuleList):
    def __init__(self, n_channels_list: Sequence[int]):
        super(LinLayers, self).__init__([
            nn.Sequential(
                nn.Identity(),
                nn.Conv2d(nc, 1, 1, 1, 0, bias=False)
            ) for nc in n_channels_list
        ])

        for param in self.parameters():
            param.requires_grad = False


class BaseNet(nn.Module):
    def __init__(self):
        super(BaseNet, self).__init__()

        # register buffer
        self.register_buffer(
            'mean', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
        self.register_buffer(
            'std', torch.Tensor([.458, .448, .450])[None, :, None, None])

    def set_requires_grad(self, state: bool):
        for param in chain(self.parameters(), self.buffers()):
            param.requires_grad = state

    def z_score(self, x: torch.Tensor):
        return (x - self.mean) / self.std

    def forward(self, x: torch.Tensor):
        x = self.z_score(x)

        output = []
        for i, (_, layer) in enumerate(self.layers._modules.items(), 1):
            x = layer(x)
            if i in self.target_layers:
                output.append(normalize_activation(x))
            if len(output) == len(self.target_layers):
                break
        return output


class SqueezeNet(BaseNet):
    def __init__(self):
        super(SqueezeNet, self).__init__()

        self.layers = models.squeezenet1_1(True).features
        self.target_layers = [2, 5, 8, 10, 11, 12, 13]
        self.n_channels_list = [64, 128, 256, 384, 384, 512, 512]

        self.set_requires_grad(False)


class AlexNet(BaseNet):
    def __init__(self):
        super(AlexNet, self).__init__()

        self.layers = models.alexnet(True).features
        self.target_layers = [2, 5, 8, 10, 12]
        self.n_channels_list = [64, 192, 384, 256, 256]

        self.set_requires_grad(False)


class VGG16(BaseNet):
    def __init__(self):
        super(VGG16, self).__init__()

        self.layers = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).features
        self.target_layers = [4, 9, 16, 23, 30]
        self.n_channels_list = [64, 128, 256, 512, 512]

        self.set_requires_grad(False)


================================================
FILE: evaluation/lpipsPyTorch/modules/utils.py
================================================
from collections import OrderedDict

import torch


def normalize_activation(x, eps=1e-10):
    norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True))
    return x / (norm_factor + eps)


def get_state_dict(net_type: str = 'alex', version: str = '0.1'):
    # build url
    url = 'https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/' \
        + f'master/lpips/weights/v{version}/{net_type}.pth'

    # download
    old_state_dict = torch.hub.load_state_dict_from_url(
        url, progress=True,
        map_location=None if torch.cuda.is_available() else torch.device('cpu')
    )

    # rename keys
    new_state_dict = OrderedDict()
    for key, val in old_state_dict.items():
        new_key = key
        new_key = new_key.replace('lin', '')
        new_key = new_key.replace('model.', '')
        new_state_dict[new_key] = val

    return new_state_dict


================================================
FILE: evaluation/metrics.py
================================================
#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import os
import sys
import json
import torch
from PIL import Image
from tqdm import tqdm
from pathlib import Path
import torchvision.transforms.functional as tf
sys.path.append(os.getcwd())

from tools.loss_utils import ssim
from lpipsPyTorch import lpips
from tools.image_utils import psnr
from argparse import ArgumentParser
from configs.config import Config
from tools.general_utils import set_random_seed


def readImages(renders_dir, gt_dir):
    renders = []
    gts = []
    image_names = []
    for fname in os.listdir(renders_dir):
        render = Image.open(renders_dir / fname)
        gt = Image.open(gt_dir / fname)
        renders.append(tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda())
        gts.append(tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda())
        image_names.append(fname)
    return renders, gts, image_names

def evaluate(model_paths):

    full_dict = {}
    per_view_dict = {}
    full_dict_polytopeonly = {}
    per_view_dict_polytopeonly = {}
    print("")

    for scene_dir in model_paths:
        try:
            print("Scene:", scene_dir)
            full_dict[scene_dir] = {}
            per_view_dict[scene_dir] = {}
            full_dict_polytopeonly[scene_dir] = {}
            per_view_dict_polytopeonly[scene_dir] = {}

            test_dir = Path(scene_dir) / "test"

            for method in os.listdir(test_dir):
                print("Method:", method)

                full_dict[scene_dir][method] = {}
                per_view_dict[scene_dir][method] = {}
                full_dict_polytopeonly[scene_dir][method] = {}
                per_view_dict_polytopeonly[scene_dir][method] = {}

                method_dir = test_dir / method
                gt_dir = method_dir/ "gt"
                renders_dir = method_dir / "renders"
                renders, gts, image_names = readImages(renders_dir, gt_dir)

                ssims = []
                psnrs = []
                lpipss = []

                for idx in tqdm(range(len(renders)), desc="Metric evaluation progress"):
                    ssims.append(ssim(renders[idx], gts[idx]))
                    psnrs.append(psnr(renders[idx], gts[idx]))
                    lpipss.append(lpips(renders[idx], gts[idx], net_type='vgg'))


                full_dict[scene_dir][method].update({"SSIM": torch.tensor(ssims).mean().item(),
                                                        "PSNR": torch.tensor(psnrs).mean().item(),
                                                        "LPIPS": torch.tensor(lpipss).mean().item()})
                per_view_dict[scene_dir][method].update({"SSIM": {name: ssim for ssim, name in zip(torch.tensor(ssims).tolist(), image_names)},
                                                            "PSNR": {name: psnr for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names)},
                                                            "LPIPS": {name: lp for lp, name in zip(torch.tensor(lpipss).tolist(), image_names)}})

            with open(scene_dir + "/results.json", 'w') as fp:
                json.dump(full_dict[scene_dir], fp, indent=True)
            with open(scene_dir + "/per_view.json", 'w') as fp:
                json.dump(per_view_dict[scene_dir], fp, indent=True)
        except:
            print("Unable to compute metrics for model", scene_dir)

if __name__ == "__main__":
    device = torch.device("cuda:0")
    torch.cuda.set_device(device)

    # Set up command line argument parser
    parser = ArgumentParser(description="Training script parameters")
    parser.add_argument('--cfg_path', type=str, default='configs/config_base.yaml')
    args = parser.parse_args()
    
    cfg = Config(args.cfg_path)
    cfg.model.data_device = 'cpu'
    cfg.model.load_normal = False
    
    set_random_seed(cfg.seed)
    
    evaluate([cfg.model.model_path])


================================================
FILE: evaluation/render.py
================================================
#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import os
import sys
import torch
import torchvision
from tqdm import tqdm
from argparse import ArgumentParser
sys.path.append(os.getcwd())

from scene import Scene
from gaussian_renderer import render, render_fast
from gaussian_renderer import GaussianModel
from configs.config import Config
from tools.general_utils import set_random_seed
from tools.loss_utils import cos_weight


def render_set(model_path, name, iteration, views, gaussians, cfg, background):
    render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")
    gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt")

    os.makedirs(render_path, exist_ok=True)
    os.makedirs(gts_path, exist_ok=True)
    alphas = []

    for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
        outs = render(view, gaussians, cfg, background)
        # outs = render_fast(view, gaussians, cfg, background)
        
        rendering = outs["render"]
        gt = view.original_image[0:3, :, :]
        torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
        torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ".png"))
        alphas.append(outs["alpha"].detach().clone().view(-1).cpu())
        
        if False:
            normal_map = outs["normal"].detach().clone()
            normal_gt = view.normal.cuda()
            cos = cos_weight(normal_gt, normal_map, cfg.optim.exp_t, cfg.optim.cos_thr)
            torchvision.utils.save_image(cos, os.path.join(render_path, '{0:05d}_cosine'.format(idx) + ".png"))
    
    # alphas = torch.cat(alphas, dim=0)
    # print("Alpha min: {}, max: {}".format(alphas.min(), alphas.max()))
    # print("Alpha mean: {}, std: {}".format(alphas.mean(), alphas.std()))
    # print("Alpha median: {}".format(alphas.median()))


def render_sets(cfg, iteration : int, skip_train : bool, skip_test : bool):
    with torch.no_grad():
        gaussians = GaussianModel(cfg.model)
        scene = Scene(cfg.model, gaussians, load_iteration=iteration, shuffle=False)
        # gaussians.extent = scene.cameras_extent

        bg_color = [1,1,1] if cfg.model.white_background else [0, 0, 0]
        background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")

        if not skip_train:
            render_set(cfg.model.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, cfg, background)

        if not skip_test:
            render_set(cfg.model.model_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, cfg, background)
            

if __name__ == "__main__":
    # Set up command line argument parser
    parser = ArgumentParser()
    parser.add_argument('--cfg_path', type=str, default='configs/config_base.yaml')
    parser.add_argument("--iteration", default=-1, type=int)
    parser.add_argument("--skip_train", action="store_true")
    parser.add_argument("--skip_test", action="store_true")
    args = parser.parse_args()
    
    cfg = Config(args.cfg_path)
    cfg.model.data_device = 'cuda'
    cfg.model.load_normal = False
    cfg.model.load_mask = False
    
    set_random_seed(cfg.seed)

    # Initialize system state (RNG)
    # safe_state(args.quiet)

    render_sets(cfg, args.iteration, args.skip_train, args.skip_test)

================================================
FILE: evaluation/tnt_eval/README.md
================================================
# Python Toolbox for Evaluation

This Python script evaluates **training** dataset of TanksAndTemples benchmark.
The script requires ``Open3D`` and a few Python packages such as ``matplotlib``, ``json``, and ``numpy``.

## How to use:
**Step 0**. Reconstruct 3D models and recover camera poses from the training dataset.
The raw videos of the training dataset can be found from:
https://tanksandtemples.org/download/

**Step 1**. Download evaluation data (ground truth geometry + reference reconstruction) using
[this link](https://drive.google.com/open?id=1UoKPiUUsKa0AVHFOrnMRhc5hFngjkE-t). In this example, we regard ``TanksAndTemples/evaluation/data/`` as a dataset folder.

**Step 2**. Install Open3D. Follow instructions in http://open3d.org/docs/getting_started.html

**Step 3**. Run the evaluation script and grab some coffee.
```
python run.py --dataset-dir path/to/TanksAndTemples/evaluation/data/Ignatius --traj-path path/to/TanksAndTemples/evaluation/data/Ignatius/Ignatius_COLMAP_SfM.log --ply-path path/to/TanksAndTemples/evaluation/data/Ignatius/Ignatius_COLMAP.ply
```
Output (evaluation of Ignatius):
```
===========================
Evaluating Ignatius
===========================
path/to/TanksAndTemples/evaluation/data/Ignatius/Ignatius_COLMAP.ply
Reading PLY: [========================================] 100%
Read PointCloud: 6929586 vertices.
path/to/TanksAndTemples/evaluation/data/Ignatius/Ignatius.ply
Reading PLY: [========================================] 100%
:
ICP Iteration #0: Fitness 0.9980, RMSE 0.0044
ICP Iteration #1: Fitness 0.9980, RMSE 0.0043
ICP Iteration #2: Fitness 0.9980, RMSE 0.0043
ICP Iteration #3: Fitness 0.9980, RMSE 0.0043
ICP Iteration #4: Fitness 0.9980, RMSE 0.0042
ICP Iteration #5: Fitness 0.9980, RMSE 0.0042
ICP Iteration #6: Fitness 0.9979, RMSE 0.0042
ICP Iteration #7: Fitness 0.9979, RMSE 0.0042
ICP Iteration #8: Fitness 0.9979, RMSE 0.0042
ICP Iteration #9: Fitness 0.9979, RMSE 0.0042
ICP Iteration #10: Fitness 0.9979, RMSE 0.0042
[EvaluateHisto]
Cropping geometry: [========================================] 100%
Pointcloud down sampled from 6929586 points to 1449840 points.
Pointcloud down sampled from 1449840 points to 1365628 points.
path/to/TanksAndTemples/evaluation/data/Ignatius/evaluation//Ignatius.precision.ply
Cropping geometry: [========================================] 100%
Pointcloud down sampled from 5016769 points to 4957123 points.
Pointcloud down sampled from 4957123 points to 4181506 points.
[compute_point_cloud_to_point_cloud_distance]
[compute_point_cloud_to_point_cloud_distance]
:
[ViewDistances] Add color coding to visualize error
[ViewDistances] Add color coding to visualize error
:
[get_f1_score_histo2]
==============================
evaluation result : Ignatius
==============================
distance tau : 0.003
precision : 0.7679
recall : 0.7937
f-score : 0.7806
==============================
```

**Step 5**. Go to the evaluation folder. ``TanksAndTemples/evaluation/data/Ignatius/evaluation/`` will have the following outputs.

<img src="images/f-score.jpg" width="400">

``PR_Ignatius_@d_th_0_0030.pdf`` (Precision and recall curves with a F-score)

| <img src="images/precision.jpg" width="200"> | <img src="images/recall.jpg" width="200"> |
|--|--|
| ``Ignatius.precision.ply``  | ``Ignatius.recall.ply`` |

(3D visualization of precision and recall. Each mesh is color coded using hot colormap)

# Requirements

- Python 3
- open3d v0.9.0
- matplotlib


================================================
FILE: evaluation/tnt_eval/config.py
================================================
# ----------------------------------------------------------------------------
# -                   TanksAndTemples Website Toolbox                        -
# -                    http://www.tanksandtemples.org                        -
# ----------------------------------------------------------------------------
# The MIT License (MIT)
#
# Copyright (c) 2017
# Arno Knapitsch <arno.knapitsch@gmail.com >
# Jaesik Park <syncle@gmail.com>
# Qian-Yi Zhou <Qianyi.Zhou@gmail.com>
# Vladlen Koltun <vkoltun@gmail.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# ----------------------------------------------------------------------------

# some global parameters - do not modify
scenes_tau_dict = {
    "Barn": 0.01,
    "Caterpillar": 0.005,
    "Church": 0.025,
    "Courthouse": 0.025,
    "Ignatius": 0.003,
    "Meetingroom": 0.01,
    "Truck": 0.005,
}


================================================
FILE: evaluation/tnt_eval/evaluation.py
================================================
# ----------------------------------------------------------------------------
# -                   TanksAndTemples Website Toolbox                        -
# -                    http://www.tanksandtemples.org                        -
# ----------------------------------------------------------------------------
# The MIT License (MIT)
#
# Copyright (c) 2017
# Arno Knapitsch <arno.knapitsch@gmail.com >
# Jaesik Park <syncle@gmail.com>
# Qian-Yi Zhou <Qianyi.Zhou@gmail.com>
# Vladlen Koltun <vkoltun@gmail.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# ----------------------------------------------------------------------------
#
# This python script is for downloading dataset from www.tanksandtemples.org
# The dataset has a different license, please refer to
# https://tanksandtemples.org/license/

import json
import copy
import os
import numpy as np
import open3d as o3d
import matplotlib.pyplot as plt


def read_alignment_transformation(filename):
    with open(filename) as data_file:
        data = json.load(data_file)
    return np.asarray(data["transformation"]).reshape((4, 4)).transpose()


def write_color_distances(path, pcd, distances, max_distance):
    o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Debug)
    # cmap = plt.get_cmap("afmhot")
    cmap = plt.get_cmap("hot_r")
    distances = np.array(distances)
    colors = cmap(np.minimum(distances, max_distance) / max_distance)[:, :3]
    pcd.colors = o3d.utility.Vector3dVector(colors)
    o3d.io.write_point_cloud(path, pcd)


def EvaluateHisto(
    source,
    target,
    trans,
    crop_volume,
    voxel_size,
    threshold,
    filename_mvs,
    plot_stretch,
    scene_name,
    verbose=True,
):
    print("[EvaluateHisto]")
    o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Debug)
    s = copy.deepcopy(source)
    s.transform(trans)
    s = crop_volume.crop_point_cloud(s)
    s = s.voxel_down_sample(voxel_size)
    s.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamKNN(knn=20))
    print(filename_mvs + "/" + scene_name + ".precision.ply")

    t = copy.deepcopy(target)
    t = crop_volume.crop_point_cloud(t)
    t = t.voxel_down_sample(voxel_size)
    t.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamKNN(knn=20))
    print("[compute_point_cloud_to_point_cloud_distance]")
    distance1 = s.compute_point_cloud_distance(t)
    print("[compute_point_cloud_to_point_cloud_distance]")
    distance2 = t.compute_point_cloud_distance(s)

    # write the distances to bin files
    # np.array(distance1).astype("float64").tofile(
    #     filename_mvs + "/" + scene_name + ".precision.bin"
    # )
    # np.array(distance2).astype("float64").tofile(
    #     filename_mvs + "/" + scene_name + ".recall.bin"
    # )

    # Colorize the poincloud files prith the precision and recall values
    # o3d.io.write_point_cloud(
    #     filename_mvs + "/" + scene_name + ".precision.ply", s
    # )
    # o3d.io.write_point_cloud(
    #     filename_mvs + "/" + scene_name + ".precision.ncb.ply", s
    # )
    # o3d.io.write_point_cloud(filename_mvs + "/" + scene_name + ".recall.ply", t)

    source_n_fn = filename_mvs + "/" + scene_name + ".precision.ply"
    target_n_fn = filename_mvs + "/" + scene_name + ".recall.ply"

    print("[ViewDistances] Add color coding to visualize error")
    # eval_str_viewDT = (
    #     OPEN3D_EXPERIMENTAL_BIN_PATH
    #     + "ViewDistances "
    #     + source_n_fn
    #     + " --max_distance "
    #     + str(threshold * 3)
    #     + " --write_color_back --without_gui"
    # )
    # os.system(eval_str_viewDT)
    write_color_distances(source_n_fn, s, distance1, 3 * threshold)

    print("[ViewDistances] Add color coding to visualize error")
    # eval_str_viewDT = (
    #     OPEN3D_EXPERIMENTAL_BIN_PATH
    #     + "ViewDistances "
    #     + target_n_fn
    #     + " --max_distance "
    #     + str(threshold * 3)
    #     + " --write_color_back --without_gui"
    # )
    # os.system(eval_str_viewDT)
    write_color_distances(target_n_fn, t, distance2, 3 * threshold)

    # get histogram and f-score
    [
        precision,
        recall,
        fscore,
        edges_source,
        cum_source,
        edges_target,
        cum_target,
    ] = get_f1_score_histo2(threshold, filename_mvs, plot_stretch, distance1,
                            distance2)
    np.savetxt(filename_mvs + "/" + scene_name + ".recall.txt", cum_target)
    np.savetxt(filename_mvs + "/" + scene_name + ".precision.txt", cum_source)
    np.savetxt(
        filename_mvs + "/" + scene_name + ".prf_tau_plotstr.txt",
        np.array([precision, recall, fscore, threshold, plot_stretch]),
    )

    return [
        precision,
        recall,
        fscore,
        edges_source,
        cum_source,
        edges_target,
        cum_target,
    ]


def get_f1_score_histo2(threshold,
                        filename_mvs,
                        plot_stretch,
                        distance1,
                        distance2,
                        verbose=True):
    print("[get_f1_score_histo2]")
    dist_threshold = threshold
    if len(distance1) and len(distance2):

        recall = float(sum(d < threshold for d in distance2)) / float(
            len(distance2))
        precision = float(sum(d < threshold for d in distance1)) / float(
            len(distance1))
        fscore = 2 * recall * precision / (recall + precision)
        num = len(distance1)
        bins = np.arange(0, dist_threshold * plot_stretch, dist_threshold / 100)
        hist, edges_source = np.histogram(distance1, bins)
        cum_source = np.cumsum(hist).astype(float) / num

        num = len(distance2)
        bins = np.arange(0, dist_threshold * plot_stretch, dist_threshold / 100)
        hist, edges_target = np.histogram(distance2, bins)
        cum_target = np.cumsum(hist).astype(float) / num

    else:
        precision = 0
        recall = 0
        fscore = 0
        edges_source = np.array([0])
        cum_source = np.array([0])
        edges_target = np.array([0])
        cum_target = np.array([0])

    return [
        precision,
        recall,
        fscore,
        edges_source,
        cum_source,
        edges_target,
        cum_target,
    ]


================================================
FILE: evaluation/tnt_eval/plot.py
================================================
# ----------------------------------------------------------------------------
# -                   TanksAndTemples Website Toolbox                        -
# -                    http://www.tanksandtemples.org                        -
# ----------------------------------------------------------------------------
# The MIT License (MIT)
#
# Copyright (c) 2017
# Arno Knapitsch <arno.knapitsch@gmail.com >
# Jaesik Park <syncle@gmail.com>
# Qian-Yi Zhou <Qianyi.Zhou@gmail.com>
# Vladlen Koltun <vkoltun@gmail.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# ----------------------------------------------------------------------------
#
# This python script is for downloading dataset from www.tanksandtemples.org
# The dataset has a different license, please refer to
# https://tanksandtemples.org/license/

import matplotlib.pyplot as plt
from cycler import cycler


def plot_graph(
    scene,
    fscore,
    dist_threshold,
    edges_source,
    cum_source,
    edges_target,
    cum_target,
    plot_stretch,
    mvs_outpath,
    show_figure=False,
):
    f = plt.figure()
    plt_size = [14, 7]
    pfontsize = "medium"

    ax = plt.subplot(111)
    label_str = "precision"
    ax.plot(
        edges_source[1::],
        cum_source * 100,
        c="red",
        label=label_str,
        linewidth=2.0,
    )

    label_str = "recall"
    ax.plot(
        edges_target[1::],
        cum_target * 100,
        c="blue",
        label=label_str,
        linewidth=2.0,
    )

    ax.grid(True)
    plt.rcParams["figure.figsize"] = plt_size
    plt.rc("axes", prop_cycle=cycler("color", ["r", "g", "b", "y"]))
    plt.title("Precision and Recall: " + scene + ", " + "%02.2f f-score" %
              (fscore * 100))
    plt.axvline(x=dist_threshold, c="black", ls="dashed", linewidth=2.0)

    plt.ylabel("# of points (%)", fontsize=15)
    plt.xlabel("Meters", fontsize=15)
    plt.axis([0, dist_threshold * plot_stretch, 0, 100])
    ax.legend(shadow=True, fancybox=True, fontsize=pfontsize)
    # plt.axis([0, dist_threshold*plot_stretch, 0, 100])

    plt.setp(ax.get_legend().get_texts(), fontsize=pfontsize)

    plt.legend(loc=2, borderaxespad=0.0, fontsize=pfontsize)
    plt.legend(loc=4)
    leg = plt.legend(loc="lower right")

    box = ax.get_position()
    ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])

    # Put a legend to the right of the current axis
    ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
    plt.setp(ax.get_legend().get_texts(), fontsize=pfontsize)
    png_name = mvs_outpath + "/PR_{0}_@d_th_0_{1}.png".format(
        scene, "%04d" % (dist_threshold * 10000))
    pdf_name = mvs_outpath + "/PR_{0}_@d_th_0_{1}.pdf".format(
        scene, "%04d" % (dist_threshold * 10000))

    # save figure and display
    f.savefig(png_name, format="png", bbox_inches="tight")
    f.savefig(pdf_name, format="pdf", bbox_inches="tight")
    if show_figure:
        plt.show()


================================================
FILE: evaluation/tnt_eval/registration.py
================================================
# ----------------------------------------------------------------------------
# -                   TanksAndTemples Website Toolbox                        -
# -                    http://www.tanksandtemples.org                        -
# ----------------------------------------------------------------------------
# The MIT License (MIT)
#
# Copyright (c) 2017
# Arno Knapitsch <arno.knapitsch@gmail.com >
# Jaesik Park <syncle@gmail.com>
# Qian-Yi Zhou <Qianyi.Zhou@gmail.com>
# Vladlen Koltun <vkoltun@gmail.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# ----------------------------------------------------------------------------
#
# This python script is for downloading dataset from www.tanksandtemples.org
# The dataset has a different license, please refer to
# https://tanksandtemples.org/license/

from trajectory_io import read_trajectory, convert_trajectory_to_pointcloud
import copy
import numpy as np
import open3d as o3d

MAX_POINT_NUMBER = 4e6


def read_mapping(filename):
    mapping = []
    with open(filename, "r") as f:
        n_sampled_frames = int(f.readline())
        n_total_frames = int(f.readline())
        mapping = np.zeros(shape=(n_sampled_frames, 2))
        metastr = f.readline()
        for iter in range(n_sampled_frames):
            metadata = list(map(int, metastr.split()))
            mapping[iter, :] = metadata
            metastr = f.readline()
    return [n_sampled_frames, n_total_frames, mapping]


def gen_sparse_trajectory(mapping, f_trajectory):
    sparse_traj = []
    for m in mapping:
        sparse_traj.append(f_trajectory[int(m[1] - 1)])
    return sparse_traj


def trajectory_alignment(map_file, traj_to_register, gt_traj_col, gt_trans,
                         scene):
    traj_pcd_col = convert_trajectory_to_pointcloud(gt_traj_col)
    traj_pcd_col.transform(gt_trans)
    corres = o3d.utility.Vector2iVector(
        np.asarray(list(map(lambda x: [x, x], range(len(gt_traj_col))))))
    rr = o3d.registration.RANSACConvergenceCriteria()
    rr.max_iteration = 100000
    rr.max_validation = 100000

    # in this case a log file was used which contains
    # every movie frame (see tutorial for details)
    if len(traj_to_register) > 1600:
        n_sampled_frames, n_total_frames, mapping = read_mapping(map_file)
        traj_col2 = gen_sparse_trajectory(mapping, traj_to_register)
        traj_to_register_pcd = convert_trajectory_to_pointcloud(traj_col2)
    else:
        traj_to_register_pcd = convert_trajectory_to_pointcloud(
            traj_to_register)
    randomvar = 0.0
    nr_of_cam_pos = len(traj_to_register_pcd.points)
    rand_number_added = np.asanyarray(traj_to_register_pcd.points) * (
        np.random.rand(nr_of_cam_pos, 3) * randomvar - randomvar / 2.0 + 1)
    list_rand = list(rand_number_added)
    traj_to_register_pcd_rand = o3d.geometry.PointCloud()
    for elem in list_rand:
        traj_to_register_pcd_rand.points.append(elem)

    # Rough registration based on aligned colmap SfM data
    reg = o3d.registration.registration_ransac_based_on_correspondence(
        traj_to_register_pcd_rand,
        traj_pcd_col,
        corres,
        0.2,
        o3d.registration.TransformationEstimationPointToPoint(True),
        6,
        rr,
    )
    return reg.transformation


def crop_and_downsample(
        pcd,
        crop_volume,
        down_sample_method="voxel",
        voxel_size=0.01,
        trans=np.identity(4),
):
    pcd_copy = copy.deepcopy(pcd)
    pcd_copy.transform(trans)
    pcd_crop = crop_volume.crop_point_cloud(pcd_copy)
    if down_sample_method == "voxel":
        # return voxel_down_sample(pcd_crop, voxel_size)
        return pcd_crop.voxel_down_sample(voxel_size)
    elif down_sample_method == "uniform":
        n_points = len(pcd_crop.points)
        if n_points > MAX_POINT_NUMBER:
            ds_rate = int(round(n_points / float(MAX_POINT_NUMBER)))
            return pcd_crop.uniform_down_sample(ds_rate)
    return pcd_crop


def registration_unif(
    source,
    gt_target,
    init_trans,
    crop_volume,
    threshold,
    max_itr,
    max_size=4 * MAX_POINT_NUMBER,
    verbose=True,
):
    if verbose:
        print("[Registration] threshold: %f" % threshold)
        o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Debug)
    s = crop_and_downsample(source,
                            crop_volume,
                            down_sample_method="uniform",
                            trans=init_trans)
    t = crop_and_downsample(gt_target,
                            crop_volume,
                            down_sample_method="uniform")
    reg = o3d.registration.registration_icp(
        s,
        t,
        threshold,
        np.identity(4),
        o3d.registration.TransformationEstimationPointToPoint(True),
        o3d.registration.ICPConvergenceCriteria(1e-6, max_itr),
    )
    reg.transformation = np.matmul(reg.transformation, init_trans)
    return reg


def registration_vol_ds(
    source,
    gt_target,
    init_trans,
    crop_volume,
    voxel_size,
    threshold,
    max_itr,
    verbose=True,
):
    if verbose:
        print("[Registration] voxel_size: %f, threshold: %f" %
              (voxel_size, threshold))
        o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Debug)
    s = crop_and_downsample(
        source,
        crop_volume,
        down_sample_method="voxel",
        voxel_size=voxel_size,
        trans=init_trans,
    )
    t = crop_and_downsample(
        gt_target,
        crop_volume,
        down_sample_method="voxel",
        voxel_size=voxel_size,
    )
    
    s = crop_based_target(s, t)
    
    reg = o3d.registration.registration_icp(
        s,
        t,
        threshold,
        np.identity(4),
        o3d.registration.TransformationEstimationPointToPoint(True),
        o3d.registration.ICPConvergenceCriteria(1e-6, max_itr),
    )
    reg.transformation = np.matmul(reg.transformation, init_trans)
    return reg


def crop_based_target(s, t):
    bbox_t = t.get_axis_aligned_bounding_box()

    min_bound = bbox_t.get_min_bound()
    max_bound = bbox_t.get_max_bound()

    s_filtered = o3d.geometry.PointCloud()
    
    valid = np.logical_and(np.all(s.points >= min_bound, axis=1), np.all(s.points <= max_bound, axis=1))
    s_filtered.points = o3d.utility.Vector3dVector(np.asarray(s.points)[valid])
    
    return s_filtered

================================================
FILE: evaluation/tnt_eval/requirements.txt
================================================
matplotlib>=1.3
open3d==0.9


================================================
FILE: evaluation/tnt_eval/run.py
================================================
# ----------------------------------------------------------------------------
# -                   TanksAndTemples Website Toolbox                        -
# -                    http://www.tanksandtemples.org                        -
# ----------------------------------------------------------------------------
# The MIT License (MIT)
#
# Copyright (c) 2017
# Arno Knapitsch <arno.knapitsch@gmail.com >
# Jaesik Park <syncle@gmail.com>
# Qian-Yi Zhou <Qianyi.Zhou@gmail.com>
# Vladlen Koltun <vkoltun@gmail.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# ----------------------------------------------------------------------------
#
# This python script is for downloading dataset from www.tanksandtemples.org
# The dataset has a different license, please refer to
# https://tanksandtemples.org/license/

# this script requires Open3D python binding
# please follow the intructions in setup.py before running this script.
import numpy as np
import open3d as o3d
import os
import argparse
import sys
sys.path.append(os.getcwd())

from config import scenes_tau_dict
from registration import (
    trajectory_alignment,
    registration_vol_ds,
    registration_unif,
    read_trajectory,
)
from evaluation import EvaluateHisto
from util import make_dir
from plot import plot_graph



def run_evaluation(dataset_dir, traj_path, ply_path, out_dir):
    scene = os.path.basename(os.path.normpath(dataset_dir))

    if scene not in scenes_tau_dict:
        print(dataset_dir, scene)
        raise Exception("invalid dataset-dir, not in scenes_tau_dict")

    print("")
    print("===========================")
    print("Evaluating %s" % scene)
    print("===========================")

    dTau = scenes_tau_dict[scene]
    # put the crop-file, the GT file, the COLMAP SfM log file and
    # the alignment of the according scene in a folder of
    # the same scene name in the dataset_dir
    colmap_ref_logfile = os.path.join(dataset_dir, scene + "_COLMAP_SfM.log")
    alignment = os.path.join(dataset_dir, scene + "_trans.txt")
    gt_filen = os.path.join(dataset_dir, scene + ".ply")
    # gt_filen = os.path.join(dataset_dir, scene + "_GT.ply")
    cropfile = os.path.join(dataset_dir, scene + ".json")
    map_file = os.path.join(dataset_dir, scene + "_mapping_reference.txt")

    make_dir(out_dir)
    
    assert os.path.exists(ply_path), f"ply_path {ply_path} does not exist"

    # Load reconstruction and according GT
    print(gt_filen)
    gt_pcd = o3d.io.read_point_cloud(gt_filen)
    print(ply_path)
    # pcd = o3d.io.read_point_cloud(ply_path)
    mesh = o3d.io.read_triangle_mesh(ply_path)
    pcd = mesh.sample_points_uniformly(len(gt_pcd.points))

    gt_trans = np.loadtxt(alignment)
    traj_to_register = read_trajectory(traj_path)
    gt_traj_col = read_trajectory(colmap_ref_logfile)

    trajectory_transform = trajectory_alignment(map_file, traj_to_register,
                                                gt_traj_col, gt_trans, scene)

    # Refine alignment by using the actual GT and MVS pointclouds
    vol = o3d.visualization.read_selection_polygon_volume(cropfile)
    # big pointclouds will be downlsampled to this number to speed up alignment
    dist_threshold = dTau

    # Registration refinment in 3 iterations
    r2 = registration_vol_ds(pcd, gt_pcd, trajectory_transform, vol, dTau,
                             dTau * 80, 20)
    r3 = registration_vol_ds(pcd, gt_pcd, r2.transformation, vol, dTau / 2.0,
                             dTau * 20, 20)
    r = registration_unif(pcd, gt_pcd, r3.transformation, vol, 2 * dTau, 20)

    # Histogramms and P/R/F1
    plot_stretch = 5
    [
        precision,
        recall,
        fscore,
        edges_source,
        cum_source,
        edges_target,
        cum_target,
    ] = EvaluateHisto(
        pcd,
        gt_pcd,
        r.transformation,
        vol,
        dTau / 2.0,
        dTau,
        out_dir,
        plot_stretch,
        scene,
    )
    eva = [precision, recall, fscore]
    # eva = [i*100 for i in eva]
    print("==============================")
    print("evaluation result : %s" % scene)
    print("==============================")
    print("distance tau : %.3f" % dTau)
    print("precision : %.4f" % eva[0])
    print("recall : %.4f" % eva[1])
    print("f-score : %.4f" % eva[2])
    print("==============================")
    
    with open(os.path.join(out_dir, "evaluation.txt"), "w") as f:
        f.write("evaluation result : %s\n" % scene)
        f.write("distance tau : %.3f\n" % dTau)
        f.write("precision : %.4f\n" % eva[0])
        f.write("recall : %.4f\n" % eva[1])
        f.write("f-score : %.4f\n" % eva[2])

    # Plotting
    plot_graph(
        scene,
        fscore,
        dist_threshold,
        edges_source,
        cum_source,
        edges_target,
        cum_target,
        plot_stretch,
        out_dir,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dataset-dir",
        type=str,
        required=True,
        help="path to a dataset/scene directory containing X.json, X.ply, ...",
    )
    parser.add_argument(
        "--traj-path",
        type=str,
        required=True,
        help=
        "path to trajectory file. See `convert_to_logfile.py` to create this file.",
    )
    parser.add_argument(
        "--ply-path",
        type=str,
        required=True,
        help="path to reconstruction ply file",
    )
    parser.add_argument(
        "--out-dir",
        type=str,
        default="",
        help=
        "output directory, default: an evaluation directory is created in the directory of the ply file",
    )
    args = parser.parse_args()

    if args.out_dir.strip() == "":
        args.out_dir = os.path.join(os.path.dirname(args.ply_path),
                                    "evaluation")

    run_evaluation(
        dataset_dir=args.dataset_dir,
        traj_path=args.traj_path,
        ply_path=args.ply_path,
        out_dir=args.out_dir,
    )


================================================
FILE: evaluation/tnt_eval/trajectory_io.py
================================================
import numpy as np
import open3d as o3d


class CameraPose:

    def __init__(self, meta, mat):
        self.metadata = meta
        self.pose = mat

    def __str__(self):
        return ("Metadata : " + " ".join(map(str, self.metadata)) + "\n" +
                "Pose : " + "\n" + np.array_str(self.pose))


def convert_trajectory_to_pointcloud(traj):
    pcd = o3d.geometry.PointCloud()
    for t in traj:
        pcd.points.append(t.pose[:3, 3])
    return pcd


def read_trajectory(filename):
    traj = []
    with open(filename, "r") as f:
        metastr = f.readline()
        while metastr:
            metadata = map(int, metastr.split())
            mat = np.zeros(shape=(4, 4))
            for i in range(4):
                matstr = f.readline()
                mat[i, :] = np.fromstring(matstr, dtype=float, sep=" \t")
            traj.append(CameraPose(metadata, mat))
            metastr = f.readline()
    return traj


def write_trajectory(traj, filename):
    with open(filename, "w") as f:
        for x in traj:
            p = x.pose.tolist()
            f.write(" ".join(map(str, x.metadata)) + "\n")
            f.write("\n".join(
                " ".join(map("{0:.12f}".format, p[i])) for i in range(4)))
            f.write("\n")


================================================
FILE: evaluation/tnt_eval/util.py
================================================
import os


def make_dir(path):
    if not os.path.exists(path):
        os.makedirs(path)


================================================
FILE: gaussian_renderer/__init__.py
================================================
#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import math
import torch
import torch.nn.functional as F

from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
from scene.gaussian_model import GaussianModel
from tools.sh_utils import eval_sh
from tools.normal_utils import compute_normals


def render(viewpoint_camera, pc : GaussianModel, cfg, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None, 
           return_normal = True, is_all = True, dirs=None, mask_depth_thr=0.8):
    """
    Render the scene. 
    
    Background tensor (bg_color) must be on GPU!
    """

    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    screenspace_points_densify = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    try:
        screenspace_points.retain_grad()
        screenspace_points_densify.retain_grad()
    except:
        pass

    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=cfg.pipline.debug,
        f_count=0,
    )

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)

    means3D = pc.get_xyz
    means2D = screenspace_points
    means2D_densify = screenspace_points_densify
    opacity = pc.get_opacity

    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
    # scaling / rotation by the rasterizer.
    scales = None
    rotations = None
    cov3D_precomp = None
    if cfg.pipline.compute_cov3D_python:
        cov3D_precomp = pc.get_covariance(scaling_modifier)
    else:
        scales = pc.get_scaling
        rotations = pc.get_rotation

    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
    shs = None
    colors_precomp = None
    if override_color is None:
        if cfg.pipline.convert_SHs_python:
            shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2)
            dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1))
            dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True)
            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
        else:
            shs = pc.get_features
    else:
        colors_precomp = override_color

    normals_precomp = None
    # inside, _ = pc.get_inside_gaus_normalized()
    if return_normal:
        normal = pc.get_normal(is_all=is_all)
        # convert normal direction to the camera; calculate the normal in the camera coordinate
        view_dir = means3D - viewpoint_camera.camera_center
        normal   = normal * ((((view_dir * normal).sum(dim=-1) > 0) * 1 - 0.5) * 2)[..., None]
        R_w2c = torch.tensor(viewpoint_camera.R.T).cuda().to(torch.float32)
        normals_precomp = normal @ R_w2c.transpose(0, 1)        # camera coordinate
    
    sem_feats = pc.get_objects.squeeze(1) if cfg.optim.loss_weight.semantic > 0 else None
    inside = None

    # Rasterize visible Gaussians to image, obtain their radii (on screen). 
    rendered_out, radii = rasterizer(
        means3D = means3D,
        means2D = means2D,
        means2D_densify = means2D_densify,
        shs = shs,
        colors_precomp = colors_precomp,
        normals_precomp = normals_precomp,
        semantics_precomp = sem_feats,
        opacities = opacity,
        scales = scales,
        rotations = rotations,
        cov3D_precomp = cov3D_precomp,
        dirs = dirs,
        inside = inside)
    
    chs = [3, 1, 3, 1]
    rendered_image, rendered_depth, rendered_normal, rendered_alpha = rendered_out[:sum(chs)].split(chs, dim=0)
    
    with torch.no_grad():
        mask = viewpoint_camera.mask.bool() if hasattr(viewpoint_camera, 'mask') else \
                torch.ones_like(rendered_depth, dtype=torch.bool).squeeze(0)
        if cfg.optim.mask_depth_thr > 0:
            mask1 = rendered_depth < (pc.extent * cfg.optim.mask_depth_thr)
            mask1 = mask1.squeeze(0)
            mask = mask & mask1
    
    rendered_normal = rendered_normal.permute(1, 2, 0)
    rendered_normal = F.normalize(rendered_normal, dim = -1)
    
    est_normal = compute_normals(rendered_depth, viewpoint_camera.intr)

    out = {"render": rendered_image,
            "depth": rendered_depth,
            "normal": rendered_normal,
            "est_normal": est_normal,
            "alpha": rendered_alpha,
            "viewspace_points": screenspace_points,
            "viewspace_points_densify": screenspace_points_densify,
            "visibility_filter" : radii > 0,
            "mask": mask,
            "radii": radii,}

    if cfg.optim.loss_weight.semantic > 0:
        rendered_sem = rendered_out[sum(chs):sum(chs)+cfg.model.ch_sem_feat]
        rendered_sem = pc.classifier(rendered_sem[None])[0].permute(1, 2, 0)    # [H, W, cls]
        out.update({"render_sem": rendered_sem})
    
    if hasattr(cfg.optim.loss_weight, 'depth_var') and cfg.optim.loss_weight.depth_var > 0:
        d1 = rendered_out[-2:-1]
        d2 = rendered_out[-1:]
        depth_var = d2 / rendered_alpha - (d1 / rendered_alpha) ** 2
        out.update({"depth_var": depth_var})
    
    if hasattr(cfg.optim.loss_weight, 'distortion') and cfg.optim.loss_weight.distortion > 0:
        rendered_dist = rendered_out[-1:]
        out.update({"distortion": rendered_dist})

    return out


def render_fast(viewpoint_camera, pc : GaussianModel, cfg, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None):
    """
    use the original Gaussian Splatting cuda code!!!!
    """
 
    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    try:
        screenspace_points.retain_grad()
    except:
        pass

    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=cfg.pipline.debug
    )

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)

    means3D = pc.get_xyz
    means2D = screenspace_points
    opacity = pc.get_opacity

    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
    # scaling / rotation by the rasterizer.
    scales = None
    rotations = None
    cov3D_precomp = None
    if cfg.pipline.compute_cov3D_python:
        cov3D_precomp = pc.get_covariance(scaling_modifier)
    else:
        scales = pc.get_scaling
        rotations = pc.get_rotation

    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
    shs = None
    colors_precomp = None
    if override_color is None:
        if cfg.pipline.convert_SHs_python:
            shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2)
            dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1))
            dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True)
            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
        else:
            shs = pc.get_features
    else:
        colors_precomp = override_color

    # Rasterize visible Gaussians to image, obtain their radii (on screen). 
    rendered_image, radii = rasterizer(
        means3D = means3D,
        means2D = means2D,
        shs = shs,
        colors_precomp = colors_precomp,
        opacities = opacity,
        scales = scales,
        rotations = rotations,
        cov3D_precomp = cov3D_precomp)

    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
    # They will be excluded from value updates used in the splitting criteria.
    return {"render": rendered_image,
            "viewspace_points": screenspace_points,
            "visibility_filter" : radii > 0,
            "radii": radii}


def count_render(
    viewpoint_camera,
    pc: GaussianModel,
    pipe,
    bg_color: torch.Tensor,
    scaling_modifier=1.0,
    override_color=None,
):
    """
    Render the scene.

    Background tensor (bg_color) must be on GPU!
    """
    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = (
        torch.zeros_like(
            pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda"
        )
        + 0
    )
    try:
        screenspace_points.retain_grad()
    except:
        pass

    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug,
        f_count=1,
    )

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)
    means3D = pc.get_xyz
    means2D = screenspace_points
    opacity = pc.get_opacity

    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
    # scaling / rotation by the rasterizer.
    scales = None
    rotations = None
    cov3D_precomp = None
    if pipe.compute_cov3D_python:
        cov3D_precomp = pc.get_covariance(scaling_modifier)
    else:
        scales = pc.get_scaling
        rotations = pc.get_rotation

    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
    shs = None
    colors_precomp = None
    if override_color is None:
        if pipe.convert_SHs_python:
            shs_view = pc.get_features.transpose(1, 2).view(
                -1, 3, (pc.max_sh_degree + 1) ** 2
            )
            dir_pp = pc.get_xyz - viewpoint_camera.camera_center.repeat(
                pc.get_features.shape[0], 1
            )
            dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True)
            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
        else:
            shs = pc.get_features
    else:
        colors_precomp = override_color

    # Rasterize visible Gaussians to image, obtain their radii (on screen).
    gaussians_count, important_score, rendered_image, radii = rasterizer(
        means3D=means3D,
        means2D=means2D,
        means2D_densify=None,
        shs=shs,
        colors_precomp=colors_precomp,
        normals_precomp = None,
        semantics_precomp = None,
        opacities=opacity,
        scales=scales,
        rotations=rotations,
        cov3D_precomp=cov3D_precomp,
    )

    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
    # They will be excluded from value updates used in the splitting criteria.
    return {
        "render": rendered_image,
        "viewspace_points": screenspace_points,
        "visibility_filter": radii > 0,
        "radii": radii,
        "gaussians_count": gaussians_count,
        "important_score": important_score,
    }


def visi_render(
    viewpoint_camera,
    pc: GaussianModel,
    pipe,
    bg_color: torch.Tensor,
    scaling_modifier=1.0,
    override_color=None,
):
    """
    Render the scene.

    Background tensor (bg_color) must be on GPU!
    """
    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = (
        torch.zeros_like(
            pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda"
        )
        + 0
    )
    try:
        screenspace_points.retain_grad()
    except:
        pass

    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug,
        f_count=2,
    )

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)
    means3D = pc.get_xyz
    means2D = screenspace_points
    opacity = pc.get_opacity

    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
    # scaling / rotation by the rasterizer.
    scales = None
    rotations = None
    cov3D_precomp = None
    if pipe.compute_cov3D_python:
        cov3D_precomp = pc.get_covariance(scaling_modifier)
    else:
        scales = pc.get_scaling
        rotations = pc.get_rotation

    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
    shs = None
    colors_precomp = None
    if override_color is None:
        if pipe.convert_SHs_python:
            shs_view = pc.get_features.transpose(1, 2).view(
                -1, 3, (pc.max_sh_degree + 1) ** 2
            )
            dir_pp = pc.get_xyz - viewpoint_camera.camera_center.repeat(
                pc.get_features.shape[0], 1
            )
            dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True)
            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
        else:
            shs = pc.get_features
    else:
        colors_precomp = override_color

    # Rasterize visible Gaussians to image, obtain their radii (on screen).
    
    countlist, important_score, rendered_image, radii = rasterizer(
            means3D=means3D,
            means2D=means2D,
            means2D_densify=None,
            shs=shs,
            colors_precomp=colors_precomp,
            normals_precomp = None,
            semantics_precomp = None,
            opacities=opacity,
            scales=scales,
            rotations=rotations,
            cov3D_precomp=cov3D_precomp,
        )

    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
    # They will be excluded from value updates used in the splitting criteria.
    return {
        "render": rendered_image,
        "viewspace_points": screenspace_points,
        "visibility_filter": radii > 0,
        "radii": radii,
        "countlist": countlist,
        "important_score": important_score,
    }


def visi_acc_render(
    viewpoint_camera,
    pc: GaussianModel,
    pipe,
    bg_color: torch.Tensor,
    scaling_modifier=1.0,
    override_color=None,
):
    """
    Render the scene.

    Background tensor (bg_color) must be on GPU!
    """
    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = (
        torch.zeros_like(
            pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda"
        )
        + 0
    )
    try:
        screenspace_points.retain_grad()
    except:
        pass

    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug,
        f_count=3,
    )

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)
    means3D = pc.get_xyz
    means2D = screenspace_points
    opacity = pc.get_opacity

    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
    # scaling / rotation by the rasterizer.
    scales = None
    rotations = None
    cov3D_precomp = None
    if pipe.compute_cov3D_python:
        cov3D_precomp = pc.get_covariance(scaling_modifier)
    else:
        scales = pc.get_scaling
        rotations = pc.get_rotation

    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
    shs = None
    colors_precomp = None
    if override_color is None:
        if pipe.convert_SHs_python:
            shs_view = pc.get_features.transpose(1, 2).view(
                -1, 3, (pc.max_sh_degree + 1) ** 2
            )
            dir_pp = pc.get_xyz - viewpoint_camera.camera_center.repeat(
                pc.get_features.shape[0], 1
            )
            dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True)
            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
        else:
            shs = pc.get_features
    else:
        colors_precomp = override_color

    # Rasterize visible Gaussians to image, obtain their radii (on screen).
    
    countlist, radii = rasterizer(
            means3D=means3D,
            means2D=means2D,
            means2D_densify=None,
            shs=shs,
            colors_precomp=colors_precomp,
            normals_precomp = None,
            semantics_precomp = None,
            opacities=opacity,
            scales=scales,
            rotations=rotations,
            cov3D_precomp=cov3D_precomp,
        )

    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
    # They will be excluded from value updates used in the splitting criteria.
    return {
        "viewspace_points": screenspace_points,
        "visibility_filter": radii > 0,
        "radii": radii,
        "countlist": countlist,
    }



================================================
FILE: gaussian_renderer/network_gui.py
================================================
#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import torch
import traceback
import socket
import json
from scene.cameras import MiniCam

host = "127.0.0.1"
port = 6009

conn = None
addr = None

listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

def init(wish_host, wish_port):
    global host, port, listener
    host = wish_host
    port = wish_port
    listener.bind((host, port))
    listener.listen()
    listener.settimeout(0)

def try_connect():
    global conn, addr, listener
    try:
        conn, addr = listener.accept()
        print(f"\nConnected by {addr}")
        conn.settimeout(None)
    except Exception as inst:
        pass
            
def read():
    global conn
    messageLength = conn.recv(4)
    messageLength = int.from_bytes(messageLength, 'little')
    message = conn.recv(messageLength)
    return json.loads(message.decode("utf-8"))

def send(message_bytes, verify):
    global conn
    if message_bytes != None:
        conn.sendall(message_bytes)
    conn.sendall(len(verify).to_bytes(4, 'little'))
    conn.sendall(bytes(verify, 'ascii'))

def receive():
    message = read()

    width = message["resolution_x"]
    height = message["resolution_y"]

    if width != 0 and height != 0:
        try:
            do_training = bool(message["train"])
            fovy = message["fov_y"]
            fovx = message["fov_x"]
            znear = message["z_near"]
            zfar = message["z_far"]
            do_shs_python = bool(message["shs_python"])
            do_rot_scale_python = bool(message["rot_scale_python"])
            keep_alive = bool(message["keep_alive"])
            scaling_modifier = message["scaling_modifier"]
            world_view_transform = torch.reshape(torch.tensor(message["view_matrix"]), (4, 4)).cuda()
            world_view_transform[:,1] = -world_view_transform[:,1]
            world_view_transform[:,2] = -world_view_transform[:,2]
            full_proj_transform = torch.reshape(torch.tensor(message["view_projection_matrix"]), (4, 4)).cuda()
            full_proj_transform[:,1] = -full_proj_transform[:,1]
            custom_cam = MiniCam(width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform)
        except Exception as e:
            print("")
            traceback.print_exc()
            raise e
        return custom_cam, do_training, do_shs_python, do_rot_scale_python, keep_alive, scaling_modifier
    else:
        return None, None, None, None, None, None

================================================
FILE: process_data/convert.py
================================================
#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import os
import json
import logging
from argparse import ArgumentParser
import shutil
import sys
import importlib

sys.path.append(os.getcwd())


def create_init_files(pinhole_dict_file, db_file, out_dir):
    # Partially adapted from https://github.com/Kai-46/nerfplusplus/blob/master/colmap_runner/run_colmap_posed.py
    # COLMAPDatabase = getattr(importlib.import_module(f'{args.colmap_path}.scripts.python.database'), 'COLMAPDatabase')
    from submodules.colmap.scripts.python.database import COLMAPDatabase  # NOQA

    if not os.path.exists(out_dir):
        os.mkdir(out_dir)

    # create template
    with open(pinhole_dict_file) as fp:
        pinhole_dict = json.load(fp)

    template = {}
    cameras_line_template = '{camera_id} RADIAL {width} {height} {f} {cx} {cy} {k1} {k2}\n'
    images_line_template = '{image_id} {qw} {qx} {qy} {qz} {tx} {ty} {tz} {camera_id} {image_name}\n\n'

    for img_name in pinhole_dict:
        # w, h, fx, fy, cx, cy, qvec, t
        params = pinhole_dict[img_name]
        w = params[0]
        h = params[1]
        fx = params[2]
        # fy = params[3]
        cx = params[4]
        cy = params[5]
        qvec = params[6:10]
        tvec = params[10:13]

        cam_line = cameras_line_template.format(
            camera_id="{camera_id}", width=w, height=h, f=fx, cx=cx, cy=cy, k1=0, k2=0)
        img_line = images_line_template.format(image_id="{image_id}", qw=qvec[0], qx=qvec[1], qy=qvec[2], qz=qvec[3],
                                               tx=tvec[0], ty=tvec[1], tz=tvec[2], camera_id="{camera_id}",
                                               image_name=img_name)
        template[img_name] = (cam_line, img_line)

    # read database
    db = COLMAPDatabase.connect(db_file)
    table_images = db.execute("SELECT * FROM images")
    img_name2id_dict = {}
    for row in table_images:
        img_name2id_dict[row[1]] = row[0]

    cameras_txt_lines = [template[img_name][0].format(camera_id=1)]
    images_txt_lines = []
    for img_name, img_id in img_name2id_dict.items():
        image_line = template[img_name][1].format(image_id=img_id, camera_id=1)
        images_txt_lines.append(image_line)

    with open(os.path.join(out_dir, 'cameras.txt'), 'w') as fp:
        fp.writelines(cameras_txt_lines)

    with open(os.path.join(out_dir, 'images.txt'), 'w') as fp:
        fp.writelines(images_txt_lines)
        fp.write('\n')

    # create an empty points3D.txt
    fp = open(os.path.join(out_dir, 'points3D.txt'), 'w')
    fp.close()


def main(args):
    colmap_command = '"{}"'.format(args.colmap_executable) if len(args.colmap_executable) > 0 else "colmap"
    magick_command = '"{}"'.format(args.magick_executable) if len(args.magick_executable) > 0 else "magick"
    use_gpu = 1 if not args.no_gpu else 0

    if not args.skip_matching:
        os.makedirs(args.source_path + "/distorted/sparse", exist_ok=True)

        ## Feature extraction
        feat_extracton_cmd = colmap_command + " feature_extractor "\
            "--database_path " + args.source_path + "/distorted/database.db \
            --image_path " + args.source_path + "/input \
            --ImageReader.single_camera 1 \
            --ImageReader.camera_model " + args.camera + " \
            --SiftExtraction.use_gpu " + str(use_gpu)
        exit_code = os.system(feat_extracton_cmd)
        if exit_code != 0:
            logging.error(f"Feature extraction failed with code {exit_code}. Exiting.")
            exit(exit_code)

        ## Feature matching
        feat_matching_cmd = colmap_command + " exhaustive_matcher \
            --database_path " + args.source_path + "/distorted/database.db \
            --SiftMatching.use_gpu " + str(use_gpu)
        exit_code = os.system(feat_matching_cmd)
        if exit_code != 0:
            logging.error(f"Feature matching failed with code {exit_code}. Exiting.")
            exit(exit_code)

        if args.existing_pose:
            db_file = os.path.join(args.source_path, 'distorted/database.db')
            sfm_dir = os.path.join(args.source_path, 'distorted/sparse/0')
            pinhole_dict_file = os.path.join(args.source_path, 'pinhole_dict.json')
            create_init_files(pinhole_dict_file, db_file, sfm_dir)
            
        ### Bundle adjustment
        # The default Mapper tolerance is unnecessarily large,
        # decreasing it speeds up bundle adjustment steps.
        mapper_cmd = (colmap_command + " mapper \
            --database_path " + args.source_path + "/distorted/database.db \
            --image_path "  + args.source_path + "/input \
            --output_path "  + args.source_path + "/distorted/sparse \
            --Mapper.ba_global_function_tolerance=0.000001")
        exit_code = os.system(mapper_cmd)
        if exit_code != 0:
            logging.error(f"Mapper failed with code {exit_code}. Exiting.")
            exit(exit_code)

    if not args.skip_distorting:
        ### Image undistortion
        ## We need to undistort our images into ideal pinhole intrinsics.
        img_undist_cmd = (colmap_command + " image_undistorter \
            --image_path " + args.source_path + "/input \
            --input_path " + args.source_path + "/distorted/sparse/0 \
            --output_path " + args.source_path + "\
            --output_type COLMAP")
        exit_code = os.system(img_undist_cmd)
        if exit_code != 0:
            logging.error(f"Mapper failed with code {exit_code}. Exiting.")
            exit(exit_code)

    files = os.listdir(args.source_path + "/distorted/sparse/0")
    os.makedirs(args.source_path + "/sparse/0", exist_ok=True)
    # Copy each file from the source directory to the destination directory
    for file in files:
        source_file = os.path.join(args.source_path, "distorted/sparse/0", file)
        destination_file = os.path.join(args.source_path, "sparse", "0", file)
        shutil.move(source_file, destination_file)

    if(args.resize):
        print("Copying and resizing...")

        # Resize images.
        os.makedirs(args.source_path + "/images_2", exist_ok=True)
        os.makedirs(args.source_path + "/images_4", exist_ok=True)
        os.makedirs(args.source_path + "/images_8", exist_ok=True)
        # Get the list of files in the source directory
        files = os.listdir(args.source_path + "/images")
        # Copy each file from the source directory to the destination directory
        for file in files:
            source_file = os.path.join(args.source_path, "images", file)

            destination_file = os.path.join(args.source_path, "images_2", file)
            shutil.copy2(source_file, destination_file)
            exit_code = os.system(magick_command + " mogrify -resize 50% " + destination_file)
            if exit_code != 0:
                logging.error(f"50% resize failed with code {exit_code}. Exiting.")
                exit(exit_code)

            destination_file = os.path.join(args.source_path, "images_4", file)
            shutil.copy2(source_file, destination_file)
            exit_code = os.system(magick_command + " mogrify -resize 25% " + destination_file)
            if exit_code != 0:
                logging.error(f"25% resize failed with code {exit_code}. Exiting.")
                exit(exit_code)

            destination_file = os.path.join(args.source_path, "images_8", file)
            shutil.copy2(source_file, destination_file)
            exit_code = os.system(magick_command + " mogrify -resize 12.5% " + destination_file)
            if exit_code != 0:
                logging.error(f"12.5% resize failed with code {exit_code}. Exiting.")
                exit(exit_code)

    print("Done.")


if __name__ == '__main__':
    # This Python script is based on the shell converter script provided in the MipNerF 360 repository.
    parser = ArgumentParser("Colmap converter")
    parser.add_argument("--no_gpu", action='store_true')
    parser.add_argument("--skip_matching", action='store_true')
    parser.add_argument("--skip_distorting", action='store_true')
    parser.add_argument("--source_path", "-s", required=True, type=str)
    parser.add_argument("--camera", default="OPENCV", type=str)
    parser.add_argument("--colmap_executable", default="", type=str)
    parser.add_argument("--resize", action="store_true")
    parser.add_argument("--magick_executable", default="", type=str)
    parser.add_argument("--existing_pose", action='store_true')
    parser.add_argument("--colmap_path", default="submodules.colmap", type=str)
    args = parser.parse_args()

    main(args)

================================================
FILE: process_data/convert_360_to_json.py
================================================
import os
import numpy as np
import json
import sys
from pathlib import Path
from argparse import ArgumentParser
import trimesh

dir_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0]
sys.path.append(dir_path.__str__())

from process_data.convert_data_to_json import export_to_json, get_split_dict, bound_by_pose  # NOQA

from submodules.colmap.scripts.python.database import COLMAPDatabase  # NOQA
from submodules.colmap.scripts.python.read_write_model import read_model, rotmat2qvec  # NOQA


def create_init_files(pinhole_dict_file, db_file, out_dir):
    # Partially adapted from https://github.com/Kai-46/nerfplusplus/blob/master/colmap_runner/run_colmap_posed.py

    if not os.path.exists(out_dir):
        os.mkdir(out_dir)

    # create template
    with open(pinhole_dict_file) as fp:
        pinhole_dict = json.load(fp)

    template = {}
    cameras_line_template = '{camera_id} RADIAL {width} {height} {f} {cx} {cy} {k1} {k2}\n'
    images_line_template = '{image_id} {qw} {qx} {qy} {qz} {tx} {ty} {tz} {camera_id} {image_name}\n\n'

    for img_name in pinhole_dict:
        # w, h, fx, fy, cx, cy, qvec, t
        params = pinhole_dict[img_name]
        w = params[0]
        h = params[1]
        fx = params[2]
        # fy = params[3]
        cx = params[4]
        cy = params[5]
        qvec = params[6:10]
        tvec = params[10:13]

        cam_line = cameras_line_template.format(
            camera_id="{camera_id}", width=w, height=h, f=fx, cx=cx, cy=cy, k1=0, k2=0)
        img_line = images_line_template.format(image_id="{image_id}", qw=qvec[0], qx=qvec[1], qy=qvec[2], qz=qvec[3],
                                               tx=tvec[0], ty=tvec[1], tz=tvec[2], camera_id="{camera_id}",
                                               image_name=img_name)
        template[img_name] = (cam_line, img_line)

    # read database
    db = COLMAPDatabase.connect(db_file)
    table_images = db.execute("SELECT * FROM images")
    img_name2id_dict = {}
    for row in table_images:
        img_name2id_dict[row[1]] = row[0]

    cameras_txt_lines = [template[img_name][0].format(camera_id=1)]
    images_txt_lines = []
    for img_name, img_id in img_name2id_dict.items():
        image_line = template[img_name][1].format(image_id=img_id, camera_id=1)
        images_txt_lines.append(image_line)

    with open(os.path.join(out_dir, 'cameras.txt'), 'w') as fp:
        fp.writelines(cameras_txt_lines)

    with open(os.path.join(out_dir, 'images.txt'), 'w') as fp:
        fp.writelines(images_txt_lines)
        fp.write('\n')

    # create an empty points3D.txt
    fp = open(os.path.join(out_dir, 'points3D.txt'), 'w')
    fp.close()


def convert_cam_dict_to_pinhole_dict(cam_dict, pinhole_dict_file):
    # Partially adapted from https://github.com/Kai-46/nerfplusplus/blob/master/colmap_runner/run_colmap_posed.py

    print('Writing pinhole_dict to: ', pinhole_dict_file)
    h = 1080
    w = 1920

    pinhole_dict = {}
    for img_name in cam_dict:
        W2C = cam_dict[img_name]

        # params
        fx = 0.6 * w
        fy = 0.6 * w
        cx = w / 2.0
        cy = h / 2.0

        qvec = rotmat2qvec(W2C[:3, :3])
        tvec = W2C[:3, 3]

        params = [w, h, fx, fy, cx, cy,
                  qvec[0], qvec[1], qvec[2], qvec[3],
                  tvec[0], tvec[1], tvec[2]]
        pinhole_dict[img_name] = params

    with open(pinhole_dict_file, 'w') as fp:
        json.dump(pinhole_dict, fp, indent=2, sort_keys=True)


def load_COLMAP_poses(cam_file, img_dir, tf='w2c'):
    # load img_dir namges
    names = sorted(os.listdir(img_dir))

    with open(cam_file) as f:
        lines = f.readlines()

    # C2W
    poses = {}
    for idx, line in enumerate(lines):
        if idx % 5 == 0:  # header
            img_idx, valid, _ = line.split(' ')
            if valid != '-1':
                poses[int(img_idx)] = np.eye(4)
                poses[int(img_idx)]
        else:
            if int(img_idx) in poses:
                num = np.array([float(n) for n in line.split(' ')])
                poses[int(img_idx)][idx % 5-1, :] = num

    if tf == 'c2w':
        return poses
    else:
        # convert to W2C (follow nerf convention)
        poses_w2c = {}
        for k, v in poses.items():
            poses_w2c[names[k]] = np.linalg.inv(v)
        return poses_w2c


def load_transformation(trans_file):
    with open(trans_file) as f:
        lines = f.readlines()

    trans = np.eye(4)
    for idx, line in enumerate(lines):
        num = np.array([float(n) for n in line.split(' ')])
        trans[idx, :] = num

    return trans


def align_gt_with_cam(pts, trans):
    trans_inv = np.linalg.inv(trans)
    pts_aligned = pts @ trans_inv[:3, :3].transpose(-1, -2) + trans_inv[:3, -1]
    return pts_aligned


def main(args):
    assert args.data_path, "Provide path to 360 dataset"
    scene_list = os.listdir(args.data_path)
    scene_list = sorted(scene_list)

    for scene in scene_list:
        scene_path = os.path.join(args.data_path, scene)
        if not os.path.isdir(scene_path): continue
        
        cameras, images, points3D = read_model(os.path.join(scene_path, "sparse/0"), ext=".bin")

        trans, scale, bounding_box = bound_by_pose(images)
        trans = trans.tolist()
        
        export_to_json(trans, scale, scene_path, 'meta.json')
        print('Writing data to json file: ', os.path.join(scene_path, 'meta.json'))


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--data_path', type=str, default=None, help='Path to tanks and temples dataset')
    parser.add_argument('--run_colmap', action='store_true', help='Run colmap')
    parser.add_argument('--export_json', action='store_true', help='export json')

    args = parser.parse_args()

    main(args)


================================================
FILE: process_data/convert_data_to_json.py
================================================
'''
-----------------------------------------------------------------------------
Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.

NVIDIA CORPORATION and its licensors retain all intellectual property
and proprietary rights in and to this software, related documentation
and any modifications thereto. Any use, reproduction, disclosure or
distribution of this software and related documentation without an express
license agreement from NVIDIA CORPORATION is strictly prohibited.
-----------------------------------------------------------------------------
'''

import numpy as np
from argparse import ArgumentParser
import os
import sys
from pathlib import Path
import json
import trimesh

dir_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0]
sys.path.append(dir_path.__str__())

from submodules.colmap.scripts.python.read_write_model import read_model, qvec2rotmat  # NOQA


def find_closest_point(p1, d1, p2, d2):
    # Calculate the direction vectors of the lines
    d1_norm = d1 / np.linalg.norm(d1)
    d2_norm = d2 / np.linalg.norm(d2)

    # Create the coefficient matrix A and the constant vector b
    A = np.vstack((d1_norm, -d2_norm)).T
    b = p2 - p1

    # Solve the linear system to find the parameters t1 and t2
    t1, t2 = np.linalg.lstsq(A, b, rcond=None)[0]

    # Calculate the closest point on each line
    closest_point1 = p1 + d1_norm * t1
    closest_point2 = p2 + d2_norm * t2

    # Calculate the average of the two closest points
    closest_point = 0.5 * (closest_point1 + closest_point2)

    return closest_point


def bound_by_pose(images):
    poses = []
    for img in images.values():
        rotation = qvec2rotmat(img.qvec)
        translation = img.tvec.reshape(3, 1)
        w2c = np.concatenate([rotation, translation], 1)
        w2c = np.concatenate([w2c, np.array([0, 0, 0, 1])[None]], 0)
        c2w = np.linalg.inv(w2c)
        poses.append(c2w)

    center = np.array([0.0, 0.0, 0.0])
    for f in poses:
        src_frame = f[0:3, :]
        for g in poses:
            tgt_frame = g[0:3, :]
            p = find_closest_point(src_frame[:, 3], src_frame[:, 2], tgt_frame[:, 3], tgt_frame[:, 2])
            center += p
    center /= len(poses) ** 2

    radius = 0.0
    for f in poses:
        radius += np.linalg.norm(f[0:3, 3])
    radius /= len(poses)
    bounding_box = [
        [center[0] - radius, center[0] + radius],
        [center[1] - radius, center[1] + radius],
        [center[2] - radius, center[2] + radius],
    ]
    return center, radius, bounding_box


def bound_by_points(points3D):
    if not isinstance(points3D, np.ndarray):
        xyzs = np.stack([point.xyz for point in points3D.values()])
    else:
        xyzs = points3D
    center = xyzs.mean(axis=0)
    std = xyzs.std(axis=0)
    # radius = float(std.max() * 2)  # use 2*std to define the region, equivalent to 95% percentile
    radius = np.abs(xyzs).max(0) * 1.1
    bounding_box = [
        [center[0] - std[0] * 3, center[0] + std[0] * 3],
        [center[1] - std[1] * 3, center[1] + std[1] * 3],
        [center[2] - std[2] * 3, center[2] + std[2] * 3],
    ]
    return center, radius, bounding_box


def compute_oriented_bound(pts):
    to_align, _ = trimesh.bounds.oriented_bounds(pts)
    
    scale = (np.abs((to_align[:3, :3] @ pts.vertices.T + to_align[:3, 3:]).T).max(0) * 1.2).tolist()
    
    return to_align.tolist(), scale


def split_data(names, split=10):
    split_dict = {'train': [], 'test': []}
    names = sorted(names)
    
    for i, name in enumerate(names):
        if i % split == 0:
            split_dict['test'].append(name)
        else:
            split_dict['train'].append(name)
    
    split_dict['train'] = sorted(split_dict['train'])
    split_dict['test'] = sorted(split_dict['test'])
    return split_dict


def get_split_dict(scene_path):
    split_dict = None
    
    if os.path.exists(os.path.join(scene_path, 'train_test_lists.json')):
        image_names = os.listdir(os.path.join(scene_path, "images"))
        image_names = sorted(['{:06}'.format(int(i.split(".")[0])) for i in image_names])
        
        with open(os.path.join(scene_path, 'train_test_lists.json'), 'r') as fp:
            split_dict = json.load(fp)
            
        test_split = sorted([i.split(".")[0] for i in split_dict['test']])
        train_split = [i for i in image_names if i not in test_split]
        
        assert len(train_split) + len(test_split) == len(image_names), "train and test split do not cover all images"
        
        split_dict = {
            'train': train_split,
            'test': test_split,
        }
    
    return split_dict


def check_concentric(images, ang_tol=np.pi / 6.0, radii_tol=0.5, pose_tol=0.5):
    look_at = []
    cam_loc = []
    for img in images.values():
        rotation = qvec2rotmat(img.qvec)
        translation = img.tvec.reshape(3, 1)
        w2c = np.concatenate([rotation, translation], 1)
        w2c = np.concatenate([w2c, np.array([0, 0, 0, 1])[None]], 0)
        c2w = np.linalg.inv(w2c)
        cam_loc.append(c2w[:3, -1])
        look_at.append(c2w[:3, 2])
    look_at = np.stack(look_at)
    look_at = look_at / np.linalg.norm(look_at, axis=1, keepdims=True)
    cam_loc = np.stack(cam_loc)
    num_images = cam_loc.shape[0]

    center = cam_loc.mean(axis=0)
    vec = center - cam_loc
    radii = np.linalg.norm(vec, axis=1, keepdims=True)
    vec_unit = vec / radii
    ang = np.arccos((look_at * vec_unit).sum(axis=-1, keepdims=True))
    ang_valid = ang < ang_tol
    print(f"Fraction of images looking at the center: {ang_valid.sum()/num_images:.2f}.")

    radius_mean = radii.mean()
    radii_valid = np.isclose(radius_mean, radii, rtol=radii_tol)
    print(f"Fraction of images positioned around the center: {radii_valid.sum()/num_images:.2f}.")

    valid = ang_valid * radii_valid
    print(f"Valid fraction of concentric images: {valid.sum()/num_images:.2f}.")

    return valid.sum() / num_images > pose_tol


def export_to_json(trans, scale, scene_path, file_name, split_dict=None, do_split=False):
    out = {
        "trans": trans,
        "scale": scale,
    }

    if do_split:
        if split_dict is None:
            image_names = os.listdir(os.path.join(scene_path, "images"))
            image_names = ['{:06}'.format(int(i.split(".")[0])) for i in image_names]
            split_dict = split_data(image_names, split=10)
        
        out.update(split_dict)

    with open(os.path.join(scene_path, file_name), "w") as outputfile:
        json.dump(out, outputfile, indent=4)

    return


def data_to_json(args):
    cameras, images, points3D = read_model(os.path.join(args.data_dir, "sparse"), ext=".bin")

    # define bounding regions based on scene type
    if args.scene_type == "outdoor":
        if check_concentric(images):
            center, scale, bounding_box = bound_by_pose(images)
        else:
            center, scale, bounding_box = bound_by_points(points3D)
    elif args.scene_type == "indoor":
        # use sfm points as a proxy to define bounding regions
        center, scale, bounding_box = bound_by_points(points3D)
    elif args.scene_type == "object":
        # use poses as a proxy to define bounding regions
        center, scale, bounding_box = bound_by_pose(images)
    else:
        raise TypeError("Unknown scene type")

    # export json file
    export_to_json(list(center), scale, args.data_dir, "meta.json")
    print("Writing data to json file: ", os.path.join(args.data_dir, "meta.json"))
    return


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--data_dir", type=str, default=None, help="Path to data")
    parser.add_argument(
        "--scene_type",
        type=str,
        default="outdoor",
        choices=["outdoor", "indoor", "object"],
        help="Select scene type. Outdoor for building-scale reconstruction; "
        "indoor for room-scale reconstruction; object for object-centric scene reconstruction.",
    )
    args = parser.parse_args()
    data_to_json(args)


================================================
FILE: process_data/convert_dtu_to_json.py
================================================
'''
-----------------------------------------------------------------------------
Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.

NVIDIA CORPORATION and its licensors retain all intellectual property
and proprietary rights in and to this software, related documentation
and any modifications thereto. Any use, reproduction, disclosure or
distribution of this software and related documentation without an express
license agreement from NVIDIA CORPORATION is strictly prohibited.
-----------------------------------------------------------------------------
'''

import numpy as np
import json
from argparse import ArgumentParser
import os
import cv2
from PIL import Image, ImageFile
from glob import glob
import math
import sys
from pathlib import Path
from tqdm import tqdm
import trimesh


dir_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0]
sys.path.append(dir_path.__str__())
# from process_data.convert_data_to_json import _cv_to_gl  # noqa: E402
from process_data.convert_data_to_json import export_to_json, compute_oriented_bound  # NOQA
from submodules.colmap.scripts.python.database import COLMAPDatabase  # NOQA
from submodules.colmap.scripts.python.read_write_model import rotmat2qvec  # NOQA

ImageFile.LOAD_TRUNCATED_IMAGES = True


def load_K_Rt_from_P(filename, P=None):
    # This function is borrowed from IDR: https://github.com/lioryariv/idr
    if P is None:
        lines = open(filename).read().splitlines()
        if len(lines) == 4:
            lines = lines[1:]
        lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
        P = np.asarray(lines).astype(np.float32).squeeze()

    out = cv2.decomposeProjectionMatrix(P)
    K = out[0]
    R = out[1]
    t = out[2]

    K = K / K[2, 2]
    intrinsics = np.eye(4)
    intrinsics[:3, :3] = K

    pose = np.eye(4, dtype=np.float32)
    pose[:3, :3] = R.transpose()
    pose[:3, 3] = (t[:3] / t[3])[:, 0]

    return intrinsics, pose


def dtu_to_json(args):
    assert args.dtu_path, "Provide path to DTU dataset"
    scene_list = os.listdir(args.dtu_path)

    test_indexes = [8, 13, 16, 21, 26, 31, 34, 56]
    for scene in tqdm(scene_list):
        scene_path = os.path.join(args.dtu_path, scene)
        if not os.path.isdir(scene_path) or 'scan' not in scene:
            continue

        # trans = [0., 0., 0.]
        # scale = 1.
        id = int(scene[4:])
        pts = trimesh.load(os.path.join(args.dtu_path, f'Points/stl/stl{id:03}_total.ply'))
        trans, scale = compute_oriented_bound(pts)
        
        out = {
            "trans": trans,
            "scale": scale,
        }

        # split_dict = None
        if args.split:
            images_names = os.listdir(os.path.join(scene_path, 'images'))
            images_names = sorted([i for i in images_names if 'png' in i])
            
            train_images = [i.split('.')[0] for i in images_names if int(i.split('.')[0]) not in test_indexes]
            test_images = [i.split('.')[0] for i in images_names if int(i.split('.')[0]) in test_indexes]
            
            train_images = sorted(train_images)
            test_images = sorted(test_images)
            
            out.update({
                    'train': train_images,
                    'test': test_images,
                    })
        
            assert len(train_images) + len(test_images) == len(images_names)
        
        file_path = os.path.join(scene_path, 'meta.json')
        with open(file_path, "w") as outputfile:
            json.dump(out, outputfile, indent=4)
        # print('Writing data to json file: ', file_path)


def load_poses(scene_path):
    camera_param = dict(np.load(os.path.join(scene_path, 'cameras_sphere.npz')))
    images_lis = sorted(glob(os.path.join(scene_path, 'image/*.png')))
    c2ws = {}
    for idx, image in enumerate(images_lis):
        image = os.path.basename(image)

        world_mat = camera_param['world_mat_%d' % idx]
        scale_mat = camera_param['scale_mat_%d' % idx]

        # scale and decompose
        P = world_mat @ scale_mat
        P = P[:3, :4]
        intrinsic_param, c2w = load_K_Rt_from_P(None, P)
        c2ws[image] = c2w
    
    w, h = Image.open(os.path.join(scene_path, 'image', image)).size
    
    return c2ws, intrinsic_param, w, h
        

def convert_cam_dict_to_pinhole_dict(scene_path, pinhole_dict_file):
    # Partially adapted from https://github.com/Kai-46/nerfplusplus/blob/master/colmap_runner/run_colmap_posed.py
    
    c2ws, intrinsic_param, w, h = load_poses(scene_path)
    
    fx = intrinsic_param[0][0]
    fy = intrinsic_param[1][1]
    cx = intrinsic_param[0][2]
    cy = intrinsic_param[1][2]
    sk_x = intrinsic_param[0][1]
    sk_y = intrinsic_param[1][0]

    print('Writing pinhole_dict to: ', pinhole_dict_file)

    pinhole_dict = {}
    for img_name in c2ws:
        c2w = c2ws[img_name]
        W2C = np.linalg.inv(c2w)

        # params
        qvec = rotmat2qvec(W2C[:3, :3])
        tvec = W2C[:3, 3]

        params = [w, h, fx, fy, cx, cy, sk_x, sk_y,
                  qvec[0], qvec[1], qvec[2], qvec[3],
                  tvec[0], tvec[1], tvec[2]]
        pinhole_dict[img_name] = params
    
    with open(pinhole_dict_file, 'w') as fp:
        pinhole_dict = {k: [float(x) for x in v] for k, v in pinhole_dict.items()}
        json.dump(pinhole_dict, fp, indent=2, sort_keys=True)


def create_init_files(pinhole_dict_file, db_file, out_dir):
    # Partially adapted from https://github.com/Kai-46/nerfplusplus/blob/master/colmap_runner/run_colmap_posed.py

    if not os.path.exists(out_dir):
        os.mkdir(out_dir)

    # create template
    with open(pinhole_dict_file) as fp:
        pinhole_dict = json.load(fp)

    template = {}
    cameras_line_template = '{camera_id} RADIAL {width} {height} {fx} {fy} {cx} {cy} {k1} {k2}\n'
    images_line_template = '{image_id} {qw} {qx} {qy} {qz} {tx} {ty} {tz} {camera_id} {image_name}\n\n'

    for img_name in pinhole_dict:
        # w, h, fx, fy, cx, cy, qvec, t
        params = pinhole_dict[img_name]
        w = params[0]
        h = params[1]
        fx = params[2]
        fy = params[3]
        cx = params[4]
        cy = params[5]
        sk_x = params[6]
        sk_y = params[7]
        qvec = params[8:12]
        tvec = params[12:15]

        cam_line = cameras_line_template.format(
            camera_id="{camera_id}", width=w, height=h, fx=fx, fy=fy, cx=cx, cy=cy, k1=sk_x, k2=sk_y)
        img_line = images_line_template.format(image_id="{image_id}", qw=qvec[0], qx=qvec[1], qy=qvec[2], qz=qvec[3],
                                               tx=tvec[0], ty=tvec[1], tz=tvec[2], camera_id="{camera_id}",
                                               image_name=img_name)
        template[img_name] = (cam_line, img_line)

    # read database
    db = COLMAPDatabase.connect(db_file)
    table_images = db.execute("SELECT * FROM images")
    img_name2id_dict = {}
    for row in table_images:
        img_name2id_dict[row[1]] = row[0]

    cameras_txt_lines = [template[img_name][0].format(camera_id=1)]
    images_txt_lines = []
    for img_name, img_id in img_name2id_dict.items():
        image_line = template[img_name][1].format(image_id=img_id, camera_id=1)
        images_txt_lines.append(image_line)

    with open(os.path.join(out_dir, 'cameras.txt'), 'w') as fp:
        fp.writelines(cameras_txt_lines)

    with open(os.path.join(out_dir, 'images.txt'), 'w') as fp:
        fp.writelines(images_txt_lines)
        fp.write('\n')

    # create an empty points3D.txt
    fp = open(os.path.join(out_dir, 'points3D.txt'), 'w')
    fp.close()


def init_colmap(args):
    assert args.dtu_path, "Provide path to DTU dataset"
    scene_list = os.listdir(args.dtu_path)
    scene_list = sorted([i for i in scene_list if 'scan' in i])

    pbar = tqdm(total=len(scene_list))
    for scene in scene_list:
        pbar.set_description(desc=f'Scene: {scene}')
        pbar.update(1)
        scene_path = os.path.join(args.dtu_path, scene)

        if not os.path.exists(f"{scene_path}/image"):
            raise Exception(f"'image` folder cannot be found in {scene_path}."
                            "Please check the expected folder structure in DATA_PREPROCESSING.md")

        # extract features
        os.system(f"colmap feature_extractor --database_path {scene_path}/database.db \
                --image_path {scene_path}/image \
                --ImageReader.camera_model=RADIAL \
                --SiftExtraction.use_gpu=true \
                --SiftExtraction.num_threads=32 \
                --ImageReader.single_camera=true"
                  )
                # --ImageReader.camera_model=RADIAL \

        # match features
        os.system(f"colmap sequential_matcher \
                --database_path {scene_path}/database.db \
                --SiftMatching.use_gpu=true"
                  )

        pinhole_dict_file = os.path.join(scene_path, 'pinhole_dict.json')
        convert_cam_dict_to_pinhole_dict(scene_path, pinhole_dict_file)

        db_file = os.path.join(scene_path, 'database.db')
        sfm_dir = os.path.join(scene_path, 'sparse')
        # sfm_dir = os.path.join(scene_path, 'colmap')
        create_init_files(pinhole_dict_file, db_file, sfm_dir)

        # bundle adjustment
        os.system(f"colmap point_triangulator \
                --database_path {scene_path}/database.db \
                --image_path {scene_path}/image \
                --input_path {scene_path}/sparse \
                --output_path {scene_path}/sparse \
                --clear_points 1 \
                --Mapper.tri_ignore_two_view_tracks=true"
                  )
        os.system(f"colmap bundle_adjuster \
                --input_path {scene_path}/sparse \
                --output_path {scene_path}/sparse \
                --BundleAdjustment.refine_extrinsics=false"
                  )
        
        # undistortion
        os.system(f"colmap image_undistorter \
            --image_path {scene_path}/image \
            --input_path {scene_path}/sparse \
            --output_path {scene_path} \
            --output_type COLMAP \
            --max_image_size 1600"
                )

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--dtu_path', type=str, default=None)
    parser.add_argument('--export_json', action='store_true', help='export json')
    parser.add_argument('--run_colmap', action='store_true', help='export json')
    parser.add_argument('--split', action='store_true', help='export json')

    args = parser.parse_args()

    if args.run_colmap:
        init_colmap(args)

    if args.export_json:
        dtu_to_json(args)

================================================
FILE: process_data/convert_tnt_to_json.py
================================================
import os
import numpy as np
import json
import sys
from pathlib import Path
from argparse import ArgumentParser
import trimesh

dir_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0]
sys.path.append(dir_path.__str__())

from process_data.convert_data_to_json import export_to_json, get_split_dict, compute_oriented_bound  # NOQA

from submodules.colmap.scripts.python.database import COLMAPDatabase  # NOQA
from submodules.colmap.scripts.python.read_write_model import rotmat2qvec  # NOQA


def create_init_files(pinhole_dict_file, db_file, out_dir):
    # Partially adapted from https://github.com/Kai-46/nerfplusplus/blob/master/colmap_runner/run_colmap_posed.py

    if not os.path.exists(out_dir):
        os.mkdir(out_dir)

    # create template
    with open(pinhole_dict_file) as fp:
        pinhole_dict = json.load(fp)

    template = {}
    cameras_line_template = '{camera_id} RADIAL {width} {height} {f} {cx} {cy} {k1} {k2}\n'
    images_line_template = '{image_id} {qw} {qx} {qy} {qz} {tx} {ty} {tz} {camera_id} {image_name}\n\n'

    for img_name in pinhole_dict:
        # w, h, fx, fy, cx, cy, qvec, t
        params = pinhole_dict[img_name]
        w = params[0]
        h = params[1]
        fx = params[2]
        # fy = params[3]
        cx = params[4]
        cy = params[5]
        qvec = params[6:10]
        tvec = params[10:13]

        cam_line = cameras_line_template.format(
            camera_id="{camera_id}", width=w, height=h, f=fx, cx=cx, cy=cy, k1=0, k2=0)
        img_line = images_line_template.format(image_id="{image_id}", qw=qvec[0], qx=qvec[1], qy=qvec[2], qz=qvec[3],
                                               tx=tvec[0], ty=tvec[1], tz=tvec[2], camera_id="{camera_id}",
                                               image_name=img_name)
        template[img_name] = (cam_line, img_line)

    # read database
    db = COLMAPDatabase.connect(db_file)
    table_images = db.execute("SELECT * FROM images")
    img_name2id_dict = {}
    for row in table_images:
        img_name2id_dict[row[1]] = row[0]

    cameras_txt_lines = [template[img_name][0].format(camera_id=1)]
    images_txt_lines = []
    for img_name, img_id in img_name2id_dict.items():
        image_line = template[img_name][1].format(image_id=img_id, camera_id=1)
        images_txt_lines.append(image_line)

    with open(os.path.join(out_dir, 'cameras.txt'), 'w') as fp:
        fp.writelines(cameras_txt_lines)

    with open(os.path.join(out_dir, 'images.txt'), 'w') as fp:
        fp.writelines(images_txt_lines)
        fp.write('\n')

    # create an empty points3D.txt
    fp = open(os.path.join(out_dir, 'points3D.txt'), 'w')
    fp.close()


def convert_cam_dict_to_pinhole_dict(cam_dict, pinhole_dict_file):
    # Partially adapted from https://github.com/Kai-46/nerfplusplus/blob/master/colmap_runner/run_colmap_posed.py

    print('Writing pinhole_dict to: ', pinhole_dict_file)
    h = 1080
    w = 1920

    pinhole_dict = {}
    for img_name in cam_dict:
        W2C = cam_dict[img_name]

        # params
        fx = 0.6 * w
        fy = 0.6 * w
        cx = w / 2.0
        cy = h / 2.0

        qvec = rotmat2qvec(W2C[:3, :3])
        tvec = W2C[:3, 3]

        params = [w, h, fx, fy, cx, cy,
                  qvec[0], qvec[1], qvec[2], qvec[3],
                  tvec[0], tvec[1], tvec[2]]
        pinhole_dict[img_name] = params

    with open(pinhole_dict_file, 'w') as fp:
        json.dump(pinhole_dict, fp, indent=2, sort_keys=True)


def load_COLMAP_poses(cam_file, img_dir, tf='w2c'):
    # load img_dir namges
    names = sorted(os.listdir(img_dir))

    with open(cam_file) as f:
        lines = f.readlines()

    # C2W
    poses = {}
    for idx, line in enumerate(lines):
        if idx % 5 == 0:  # header
            img_idx, valid, _ = line.split(' ')
            if valid != '-1':
                poses[int(img_idx)] = np.eye(4)
                poses[int(img_idx)]
        else:
            if int(img_idx) in poses:
                num = np.array([float(n) for n in line.split(' ')])
                poses[int(img_idx)][idx % 5-1, :] = num

    if tf == 'c2w':
        return poses
    else:
        # convert to W2C (follow nerf convention)
        poses_w2c = {}
        for k, v in poses.items():
            poses_w2c[names[k]] = np.linalg.inv(v)
        return poses_w2c


def load_transformation(trans_file):
    with open(trans_file) as f:
        lines = f.readlines()

    trans = np.eye(4)
    for idx, line in enumerate(lines):
        num = np.array([float(n) for n in line.split(' ')])
        trans[idx, :] = num

    return trans


def align_gt_with_cam(pts, trans):
    trans_inv = np.linalg.inv(trans)
    pts_aligned = pts @ trans_inv[:3, :3].transpose(-1, -2) + trans_inv[:3, -1]
    return pts_aligned


def compute_bound(pts):
    bounding_box = np.array([pts.min(axis=0), pts.max(axis=0)])
    center = bounding_box.mean(axis=0)
    # sphere radius
    # scale = np.max(np.linalg.norm(pts - center, axis=-1)) * 1.01
    # cube 
    # scale = (np.abs(pts - center).max(0) * 1.2).tolist() # cuboid for street
    scale = (np.abs(pts - center).max(0) * 1.).tolist() # cuboid for street
    return center, scale, bounding_box.T.tolist()


def init_colmap(args):
    assert args.tnt_path, "Provide path to Tanks and Temples dataset"
    scene_list = os.listdir(args.tnt_path)
    if 'Church' in scene_list: scene_list.remove('Church')
    scene_list = sorted(scene_list)

    for scene in scene_list:
        scene_path = os.path.join(args.tnt_path, scene)

        if args.run_colmap:
            if not os.path.exists(f"{scene_path}/images_raw"):
                raise Exception(f"'images_raw` folder cannot be found in {scene_path}."
                                "Please check the expected folder structure in DATA_PREPROCESSING.md")


            # extract features
            os.system(f"colmap feature_extractor --database_path {scene_path}/database.db \
                    --image_path {scene_path}/images_raw \
                    --ImageReader.camera_model=RADIAL \
                    --SiftExtraction.use_gpu=true \
                    --SiftExtraction.num_threads=32 \
                    --ImageReader.single_camera=true"
                    )

            # match features
            os.system(f"colmap sequential_matcher \
                    --database_path {scene_path}/database.db \
                    --SiftMatching.use_gpu=true"
                  )

            # read poses
            poses = load_COLMAP_poses(os.path.join(scene_path, f'{scene}_COLMAP_SfM.log'),
                                    os.path.join(scene_path, 'images_raw'))

            # convert to colmap files
            pinhole_dict_file = os.path.join(scene_path, 'pinhole_dict.json')
            convert_cam_dict_to_pinhole_dict(poses, pinhole_dict_file)

            db_file = os.path.join(scene_path, 'database.db')
            sfm_dir = os.path.join(scene_path, 'sparse')
            create_init_files(pinhole_dict_file, db_file, sfm_dir)

            # bundle adjustment
            os.system(f"colmap point_triangulator \
                    --database_path {scene_path}/database.db \
                    --image_path {scene_path}/images_raw \
                    --input_path {scene_path}/sparse \
                    --output_path {scene_path}/sparse \
                    --Mapper.tri_ignore_two_view_tracks=true"
                    )
            os.system(f"colmap bundle_adjuster \
                    --input_path {scene_path}/sparse \
                    --output_path {scene_path}/sparse \
                    --BundleAdjustment.refine_extrinsics=false"
                    )

            # undistortion
            os.system(f"colmap image_undistorter \
                --image_path {scene_path}/images_raw \
                --input_path {scene_path}/sparse \
                --output_path {scene_path} \
                --output_type COLMAP \
                --max_image_size 1500"
                    )

        if args.export_json:
            # read for bounding information
            trans = load_transformation(os.path.join(scene_path, f'{scene}_trans.txt'))
            pts = trimesh.load(os.path.join(scene_path, f'{scene}.ply'))
            # pts = pts.vertices
            # pts_aligned = align_gt_with_cam(pts, trans)
            # center, scale, bounding_box = compute_bound(pts_aligned[::100])
            pts.vertices = align_gt_with_cam(pts.vertices, trans)
            # pts = pts.sample(20000)
            pts.vertices = pts.vertices[::100]
            trans, scale = compute_oriented_bound(pts)
            
            split_dict = get_split_dict(scene_path)
            
            export_to_json(trans, scale, scene_path, 'meta.json', split_dict=split_dict)
            print('Writing data to json file: ', os.path.join(scene_path, 'meta.json'))


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--tnt_path', type=str, default=None, help='Path to tanks and temples dataset')
    parser.add_argument('--run_colmap', action='store_true', help='Run colmap')
    parser.add_argument('--export_json', action='store_true', help='export json')

    args = parser.parse_args()

    init_colmap(args)


================================================
FILE: process_data/extract_mask.py
================================================
import argparse
import os
import gc
import sys

import numpy as np
import json
import torch
from PIL import Image
from tqdm import tqdm
import torch.nn.functional as F

# segment anything
from segment_anything import (
    sam_model_registry,
    sam_hq_model_registry,
    SamPredictor
)
import cv2
import numpy as np
import matplotlib.pyplot as plt

sys.path.append(os.getcwd())
from tools.semantic_id import text_label_dict


text_prompt_dict = {
    'indoor': 'window.floor.',
    'outdoor': 'sky.',
}


def load_image(image_path):
    # load image
    image_pil = Image.open(image_path).convert("RGB")  # load image

    transform = T.Compose(
        [
            T.RandomResize([800], max_size=1333),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )
    image, _ = transform(image_pil, None)  # 3, h, w
    return image_pil, image


def print_(a):
    pass


def load_model(model_config_path, model_checkpoint_path, device):
    args = SLConfig.fromfile(model_config_path)
    args.device = device
    model = build_model(args)
    checkpoint = torch.load(model_checkpoint_path, map_location="cpu")
    load_res = model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False)
    print(load_res)
    _ = model.eval()
    return model


def get_grounding_output(model, image, caption, box_threshold, text_threshold, with_logits=True, device="cpu"):
    caption = caption.lower()
    caption = caption.strip()
    if not caption.endswith("."):
        caption = caption + "."
    model = model.to(device)
    image = image.to(device)
    with torch.no_grad():
        outputs = model(image[None], captions=[caption])
    logits = outputs["pred_logits"].cpu().sigmoid()[0]  # (nq, 256)
    boxes = outputs["pred_boxes"].cpu()[0]  # (nq, 4)
    logits.shape[0]

    # filter output
    logits_filt = logits.clone()
    boxes_filt = boxes.clone()
    filt_mask = logits_filt.max(dim=1)[0] > box_threshold
    logits_filt = logits_filt[filt_mask]  # num_filt, 256
    boxes_filt = boxes_filt[filt_mask]  # num_filt, 4
    logits_filt.shape[0]

    # get phrase
    tokenlizer = model.tokenizer
    tokenized = tokenlizer(caption)
    # build pred
    pred_phrases = []
    for logit, box in zip(logits_filt, boxes_filt):
        pred_phrase = get_phrases_from_posmap(logit > text_threshold, tokenized, tokenlizer)
        if with_logits:
            pred_phrases.append(pred_phrase + f"({str(logit.max().item())[:4]})")
        else:
            pred_phrases.append(pred_phrase)

    return boxes_filt, pred_phrases


def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)


def show_box(box, ax, label):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))
    ax.text(x0, y0, label)


def save_mask_data(output_dir, mask_list, box_list, label_list, name):
    value = 1

    mask_img = torch.ones(mask_list.shape[-2:]) * value
    for idx, mask in enumerate(mask_list):
        if len(label_list) == 0: break
        sem = label_list[idx].split('(')[0]
        try:
            mask_img[mask.cpu().numpy()[0] == True] = text_label_dict.get(sem, value)
        except KeyError:
            import pdb; pdb.set_trace()
    
    mask_img = mask_img.numpy().astype(np.uint8)
    cv2.imwrite(os.path.join(output_dir, f'{name}.png'), mask_img)


def morphology_open(x, k1=21, k2=21):
    out = x.float()[None]
    p1 = (k1 - 1) // 2
    out = -F.max_pool2d(-out, kernel_size=k1, stride=1, padding=p1)
    out = F.max_pool2d(out, kernel_size=k1, stride=1, padding=p1)
    return out


def process_image(image_name):
    name = image_name.split('.')[0]
    image_path = os.path.join(image_dir, image_name)
    # load image
    image_pil, image = load_image(image_path)
    # visualize raw image
    # image_pil.save(os.path.join(output_dir, "raw_image.jpg"))

    # run grounding dino model
    boxes_filt, pred_phrases = get_grounding_output(
        model, image, text_prompt, box_threshold, text_threshold, device=device
    )

    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    predictor.set_image(image)

    size = image_pil.size
    H, W = size[1], size[0]
    for i in range(boxes_filt.size(0)):
        boxes_filt[i] = boxes_filt[i] * torch.Tensor([W, H, W, H])
        boxes_filt[i][:2] -= boxes_filt[i][2:] / 2
        boxes_filt[i][2:] += boxes_filt[i][:2]

    boxes_filt = boxes_filt.cpu()
    transformed_boxes = predictor.transform.apply_boxes_torch(boxes_filt, image.shape[:2]).to(device)

    with torch.no_grad():
        try:
            masks, _, _ = predictor.predict_torch(
                point_coords = None,
                point_labels = None,
                boxes = transformed_boxes.to(device),
                multimask_output = False,
            )
        except RuntimeError:
            print(f"Error in {name}")
            masks = torch.zeros([1, 1, H, W]).to(device).bool()

    masks = masks.cpu()

    if args.vis:
        # draw output image
        plt.figure(figsize=(10, 10))
        plt.imshow(image)
        for mask in masks:
            show_mask(mask.cpu().numpy(), plt.gca(), random_color=True)
        for box, label in zip(boxes_filt, pred_phrases):
            show_box(box.numpy(), plt.gca(), label)

        plt.axis('off')
        plt.savefig(
            os.path.join(output_dir, f"{name}_output.png"),
            bbox_inches="tight", dpi=100, pad_inches=0.0
        )
        plt.close()             # important!!! close the plot to release memory

    save_mask_data(output_dir, masks, boxes_filt, pred_phrases, name)


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Grounded-Segment-Anything Demo", add_help=True)
    parser.add_argument("--config", type=str, required=True, help="path to config file")
    parser.add_argument(
        "--grounded_checkpoint", type=str, required=True, help="path to checkpoint file"
    )
    parser.add_argument(
        "--sam_version", type=str, default="vit_h", required=False, help="SAM ViT version: vit_b / vit_l / vit_h"
    )
    parser.add_argument(
        "--sam_checkpoint", type=str, required=False, help="path to sam checkpoint file"
    )
    parser.add_argument(
        "--sam_hq_checkpoint", type=str, default=None, help="path to sam-hq checkpoint file"
    )
    parser.add_argument(
        "--use_sam_hq", action="store_true", help="using sam-hq for prediction"
    )
    parser.add_argument("--input_image", type=str, required=True, help="path to image file")
    parser.add_argument("--text_prompt", type=str, default=None, help="text prompt")
    parser.add_argument("--scene_type", type=str, choices=['indoor', 'outdoor'], help="text prompt")
    parser.add_argument("--scene", type=str, default=None, help="text prompt")
    parser.add_argument(
        "--output_dir", "-o", type=str, default="outputs", required=True, help="output directory"
    )

    parser.add_argument("--box_threshold", type=float, default=0.3, help="box threshold")
    parser.add_argument("--text_threshold", type=float, default=0.25, help="text threshold")
    parser.add_argument("--gsam_path", dest="gsam_path", help="path to gsam")
    parser.add_argument('--vis', action='store_true', help='visualize the output')

    parser.add_argument("--device", type=str, default="cpu", help="running on cpu only!, default=False")
    args = parser.parse_args()

    gsam_path = args.gsam_path

    sys.path.append(args.gsam_path)
    sys.path.append(os.path.join(gsam_path, "GroundingDINO"))
    sys.path.append(os.path.join(gsam_path, "segment_anything"))

    # Grounding DINO
    import GroundingDINO.groundingdino.datasets.transforms as T
    from GroundingDINO.groundingdino.models import build_model
    from GroundingDINO.groundingdino.util.slconfig import SLConfig
    from GroundingDINO.groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap

    # print = print_
    seed = 0
    np.random.seed(seed)
    torch.manual_seed(seed)         # sets seed on the current CPU & all GPUs
    # cfg
    config_file = args.config  # change the path of the model config file
    grounded_checkpoint = args.grounded_checkpoint  # change the path of the model
    sam_version = args.sam_version
    sam_checkpoint = args.sam_checkpoint
    sam_hq_checkpoint = args.sam_hq_checkpoint
    use_sam_hq = args.use_sam_hq
    image_dir = args.input_image
    if args.text_prompt is not None:
        text_prompt = args.text_prompt
    else:
        text_prompt = text_prompt_dict[args.scene_type]
        if args.scene is not None:
            text_prompt = text_prompt_dict.get(args.scene, text_prompt_dict[args.scene_type])
            
    output_dir = args.output_dir
    box_threshold = args.box_threshold
    text_threshold = args.text_threshold
    device = args.device

    # make dir
    os.makedirs(output_dir, exist_ok=True)
    # load model
    model = load_model(config_file, grounded_checkpoint, device=device)

    image_names = os.listdir(image_dir)
    image_names = sorted([i for i in image_names if i.endswith(".jpg") or i.endswith(".png")])
    # initialize SAM
    if use_sam_hq:
        predictor = SamPredictor(sam_hq_model_registry[sam_version](checkpoint=sam_hq_checkpoint).to(device))
    else:
        predictor = SamPredictor(sam_model_registry[sam_version](checkpoint=sam_checkpoint).to(device))

    for image_name in tqdm(image_names):
        process_image(image_name)

================================================
FILE: process_data/extract_normal.py
================================================
import os
import sys
import glob
import math
import struct
import argparse
import numpy as np
import collections

import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image, ImageFile
from tqdm import tqdm
ImageFile.LOAD_TRUNCATED_IMAGES = True

sys.path.append(os.getcwd())
from tools.general_utils import set_random_seed


Camera = collections.namedtuple(
    "Camera", ["id", "model", "width", "height", "params"])
CameraModel = collections.namedtuple(
    "CameraModel", ["model_id", "model_name", "num_params"])
CAMERA_MODELS = {
    CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3),
    CameraModel(model_id=1, model_name="PINHOLE", num_params=4),
    CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4),
    CameraModel(model_id=3, model_name="RADIAL", num_params=5),
    CameraModel(model_id=4, model_name="OPENCV", num_params=8),
    CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8),
    CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12),
    CameraModel(model_id=7, model_name="FOV", num_params=5),
    CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4),
    CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),
    CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12)
}
CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model)
                         for camera_model in CAMERA_MODELS])
CAMERA_MODEL_NAMES = dict([(camera_model.model_name, camera_model)
                           for camera_model in CAMERA_MODELS])



def get_args(test=False):
    parser = get_default_parser()

    #↓↓↓↓
    #NOTE: project-specific args
    parser.add_argument('--NNET_architecture', type=str, default='v02')
    parser.add_argument('--NNET_output_dim', type=int, default=3, help='{3, 4}')
    parser.add_argument('--NNET_output_type', type=str, default='R', help='{R, G}')
    parser.add_argument('--NNET_feature_dim', type=int, default=64)
    parser.add_argument('--NNET_hidden_dim', type=int, default=64)

    parser.add_argument('--NNET_encoder_B', type=int, default=5)

    parser.add_argument('--NNET_decoder_NF', type=int, default=2048)
    parser.add_argument('--NNET_decoder_BN', default=False, action="store_true")
    parser.add_argument('--NNET_decoder_down', type=int, default=8)
    parser.add_argument('--NNET_learned_upsampling', default=False, action="store_true")

    parser.add_argument('--NRN_prop_ps', type=int, default=5)
    parser.add_argument('--NRN_num_iter_train', type=int, default=5)
    parser.add_argument('--NRN_num_iter_test', type=int, default=5)
    parser.add_argument('--NRN_ray_relu', default=False, action="store_true")

    parser.add_argument('--loss_fn', type=str, default='AL')
    parser.add_argument('--loss_gamma', type=float, default=0.8)
    parser.add_argument('--outdir', type=str, default='/your/log/path/')
    #↑↑↑↑

    # read arguments from txt file
    assert '.txt' in sys.argv[1]
    arg_filename_with_prefix = '@' + sys.argv[1]
    args = parser.parse_args([arg_filename_with_prefix] + sys.argv[2:])

    #↓↓↓↓
    #NOTE: update args
    args.exp_root = os.path.join(args.outdir, 'dsine')
    args.load_normal = True
    args.load_intrins = True
    #↑↑↑↑

    # set working dir
    exp_dir = os.path.join(args.exp_root, args.exp_name)

    args.output_dir = os.path.join(exp_dir, args.exp_id)
    return args


def focal2fov(focal, pixels):
    return 2*math.atan(pixels/(2*focal))


def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
    """Read and unpack the next bytes from a binary file.
    :param fid:
    :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.
    :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
    :param endian_character: Any of {@, =, <, >, !}
    :return: Tuple of read and unpacked values.
    """
    data = fid.read(num_bytes)
    return struct.unpack(endian_character + format_char_sequence, data)


def read_intrinsics_binary(path_to_model_file):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::WriteCamerasBinary(const std::string& path)
        void Reconstruction::ReadCamerasBinary(const std::string& path)
    """
    cameras = {}
    with open(path_to_model_file, "rb") as fid:
        num_cameras = read_next_bytes(fid, 8, "Q")[0]
        for _ in range(num_cameras):
            camera_properties = read_next_bytes(
                fid, num_bytes=24, format_char_sequence="iiQQ")
            camera_id = camera_properties[0]
            model_id = camera_properties[1]
            model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name
            width = camera_properties[2]
            height = camera_properties[3]
            num_params = CAMERA_MODEL_IDS[model_id].num_params
            params = read_next_bytes(fid, num_bytes=8*num_params,
                                     format_char_sequence="d"*num_params)
            cameras[camera_id] = Camera(id=camera_id,
                                        model=model_name,
                                        width=width,
                                        height=height,
                                        params=np.array(params))
        assert len(cameras) == num_cameras
    return cameras


def read_intrinsics_text(path):
    """
    Taken from https://github.com/colmap/colmap/blob/dev/scripts/python/read_write_model.py
    """
    cameras = {}
    with open(path, "r") as fid:
        while True:
            line = fid.readline()
            if not line:
                break
            line = line.strip()
            if len(line) > 0 and line[0] != "#":
                elems = line.split()
                camera_id = int(elems[0])
                model = elems[1]
                assert model == "PINHOLE", "While the loader support other types, the rest of the code assumes PINHOLE"
                width = int(elems[2])
                height = int(elems[3])
                params = np.array(tuple(map(float, elems[4:])))
                cameras[camera_id] = Camera(id=camera_id, model=model,
                                            width=width, height=height,
                                            params=params)
    return cameras


def load_intrinsic_colmap(path):
    intr_dir = os.path.join(path, "sparse", "0")
    if not os.path.exists(intr_dir):
        intr_dir = os.path.join(path, "sparse")
    # support only one camera for now
    try:
        cameras_intrinsic_file = os.path.join(intr_dir, "cameras.bin")
        cam_intrinsics = read_intrinsics_binary(cameras_intrinsic_file)
    except:
        cameras_intrinsic_file = os.path.join(intr_dir, "cameras.txt")
        cam_intrinsics = read_intrinsics_text(cameras_intrinsic_file)

    intrinsics = []
    for idx, key in enumerate(cam_intrinsics):
        intrinsic = np.eye(3)
        intrinsic = torch.eye(3, dtype=torch.float32)

        intr = cam_intrinsics[key]
        height = intr.height
        width = intr.width

        if intr.model=="SIMPLE_PINHOLE":
            focal_length_x = intr.params[0]
            FovY = focal2fov(focal_length_x, height)
            FovX = focal2fov(focal_length_x, width)
        elif intr.model=="PINHOLE":
            focal_length_x = intr.params[0]
            focal_length_y = intr.params[1]
            FovY = focal2fov(focal_length_y, height)
            FovX = focal2fov(focal_length_x, width)
        else:
            assert False, "Colmap camera model not handled: only undistorted datasets (PINHOLE or SIMPLE_PINHOLE cameras) supported!"

    
        intrinsic[0, 0] = focal_length_x # FovX
        intrinsic[1, 1] = focal_length_y # FovY
        intrinsic[0, 2] = width / 2
        intrinsic[1, 2] = height / 2
        
        intrinsics.append(intrinsic)
    
    intrinsics = torch.stack(intrinsics, axis=0)

    return intrinsics


def test_samples(args, model, intrins=None, device='cpu'):
    img_paths = glob.glob(f'{args.img_path}/*.png') + glob.glob(f'{args.img_path}/*.jpg') + glob.glob(f'{args.img_path}/*.JPG')
    img_paths.sort()

    # normalize
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

    intrin = load_intrinsic_colmap(args.intrins_path).to(device)
    os.makedirs(args.output_path, exist_ok=True)
    
    with torch.no_grad():
        for img_path in tqdm(img_paths):
            ext = os.path.splitext(img_path)[1]
            img = Image.open(img_
Download .txt
gitextract_oa0m8dnw/

├── .gitmodules
├── LICENSE.md
├── README.md
├── arguments/
│   └── __init__.py
├── bash_scripts/
│   ├── 0_train.sh
│   ├── 1_preprocess_tnt.sh
│   ├── 2_extract_normal_dsine.sh
│   ├── 3_extract_mask.sh
│   ├── 4_extract_normal_geow.sh
│   ├── convert.sh
│   └── install.sh
├── configs/
│   ├── 360_v2/
│   │   └── base.yaml
│   ├── config.py
│   ├── config_base.yaml
│   ├── dtu/
│   │   ├── base.yaml
│   │   └── dtu_scan24.yaml
│   ├── reconstruct.yaml
│   ├── scannetpp/
│   │   └── base.yaml
│   └── tnt/
│       ├── Barn.yaml
│       ├── Caterpillar.yaml
│       ├── Courthouse.yaml
│       ├── Ignatius.yaml
│       ├── Meetingroom.yaml
│       ├── Truck.yaml
│       └── base.yaml
├── environment.yml
├── evaluation/
│   ├── crop_mesh.py
│   ├── eval_dtu/
│   │   ├── eval.py
│   │   ├── evaluate_single_scene.py
│   │   └── render_utils.py
│   ├── eval_tnt.py
│   ├── full_eval.py
│   ├── lpipsPyTorch/
│   │   ├── __init__.py
│   │   └── modules/
│   │       ├── lpips.py
│   │       ├── networks.py
│   │       └── utils.py
│   ├── metrics.py
│   ├── render.py
│   └── tnt_eval/
│       ├── README.md
│       ├── config.py
│       ├── evaluation.py
│       ├── plot.py
│       ├── registration.py
│       ├── requirements.txt
│       ├── run.py
│       ├── trajectory_io.py
│       └── util.py
├── gaussian_renderer/
│   ├── __init__.py
│   └── network_gui.py
├── process_data/
│   ├── convert.py
│   ├── convert_360_to_json.py
│   ├── convert_data_to_json.py
│   ├── convert_dtu_to_json.py
│   ├── convert_tnt_to_json.py
│   ├── extract_mask.py
│   ├── extract_normal.py
│   ├── extract_normal_geo.py
│   ├── visualize_colmap.ipynb
│   └── visualize_transforms.ipynb
├── pyproject.toml
├── python_scripts/
│   ├── run_base.py
│   ├── run_dtu.py
│   ├── run_mipnerf360.py
│   ├── run_tnt.py
│   ├── show_360.py
│   ├── show_dtu.py
│   └── show_tnt.py
├── requirements.txt
├── scene/
│   ├── __init__.py
│   ├── appearance_network.py
│   ├── cameras.py
│   ├── colmap_loader.py
│   ├── dataset_readers.py
│   └── gaussian_model.py
├── tools/
│   ├── __init__.py
│   ├── camera.py
│   ├── camera_utils.py
│   ├── crop_mesh.py
│   ├── denoise_pcd.py
│   ├── depth2mesh.py
│   ├── distributed.py
│   ├── general_utils.py
│   ├── graphics_utils.py
│   ├── image_utils.py
│   ├── loss_utils.py
│   ├── math_utils.py
│   ├── mcube_utils.py
│   ├── mesh_utils.py
│   ├── normal_utils.py
│   ├── prune.py
│   ├── render_utils.py
│   ├── semantic_id.py
│   ├── sh_utils.py
│   ├── system_utils.py
│   ├── termcolor.py
│   ├── visualization.py
│   └── visualize.py
├── train.py
└── trainer.py
Download .txt
SYMBOL INDEX (492 symbols across 64 files)

FILE: arguments/__init__.py
  class GroupParams (line 16) | class GroupParams:
  class ParamGroup (line 19) | class ParamGroup:
    method __init__ (line 20) | def __init__(self, parser: ArgumentParser, name : str, fill_none = Fal...
    method extract (line 40) | def extract(self, args):
  class ModelParams (line 47) | class ModelParams(ParamGroup):
    method __init__ (line 48) | def __init__(self, parser, sentinel=False):
    method extract (line 59) | def extract(self, args):
  class PipelineParams (line 64) | class PipelineParams(ParamGroup):
    method __init__ (line 65) | def __init__(self, parser):
  class OptimizationParams (line 71) | class OptimizationParams(ParamGroup):
    method __init__ (line 72) | def __init__(self, parser):
  function get_combined_args (line 92) | def get_combined_args(parser : ArgumentParser):

FILE: configs/config.py
  class AttrDict (line 26) | class AttrDict(dict):
    method __init__ (line 29) | def __init__(self, *args, **kwargs):
    method yaml (line 41) | def yaml(self):
    method __repr__ (line 59) | def __repr__(self):
  class Config (line 83) | class Config(AttrDict):
    method __init__ (line 87) | def __init__(self, filename=None, verbose=False):
    method load_config (line 107) | def load_config(self, filename):
    method print_config (line 136) | def print_config(self, level=0):
    method save_config (line 145) | def save_config(self, logdir):
  function rsetattr (line 151) | def rsetattr(obj, attr, val):
  function rgetattr (line 157) | def rgetattr(obj, attr, *args):
  function recursive_update (line 167) | def recursive_update(d, u):
  function recursive_update_strict (line 182) | def recursive_update_strict(d, u, stack=[]):
  function parse_cmdline_arguments (line 200) | def parse_cmdline_arguments(args):

FILE: evaluation/crop_mesh.py
  function align_gt_with_cam (line 12) | def align_gt_with_cam(pts, trans):
  function main (line 18) | def main(args):

FILE: evaluation/eval_dtu/eval.py
  function sample_single_tri (line 10) | def sample_single_tri(input_):
  function write_vis_pcd (line 21) | def write_vis_pcd(file, points, colors):

FILE: evaluation/eval_dtu/evaluate_single_scene.py
  function cull_scan (line 22) | def cull_scan(scan, mesh_path, result_mesh_file, instance_dir):

FILE: evaluation/eval_dtu/render_utils.py
  function get_psnr (line 9) | def get_psnr(img1, img2, normalize_rgb=False):
  function load_rgb (line 20) | def load_rgb(path, normalize_rgb = False):
  function load_K_Rt_from_P (line 31) | def load_K_Rt_from_P(filename, P=None):
  function get_camera_params (line 55) | def get_camera_params(uv, pose, intrinsics):
  function get_camera_for_plot (line 85) | def get_camera_for_plot(pose):
  function lift (line 96) | def lift(x, y, z, intrinsics):
  function quat_to_rot (line 112) | def quat_to_rot(q):
  function rot_to_quat (line 132) | def rot_to_quat(R):
  function get_sphere_intersections (line 153) | def get_sphere_intersections(cam_loc, ray_directions, r = 1.0):

FILE: evaluation/eval_tnt.py
  function nn_correspondance (line 9) | def nn_correspondance(verts1, verts2):
  function evaluate (line 22) | def evaluate(mesh_pred, mesh_trgt, threshold=.05, down_sample=.02):
  function main (line 52) | def main(args):

FILE: evaluation/lpipsPyTorch/__init__.py
  function lpips (line 6) | def lpips(x: torch.Tensor,

FILE: evaluation/lpipsPyTorch/modules/lpips.py
  class LPIPS (line 8) | class LPIPS(nn.Module):
    method __init__ (line 17) | def __init__(self, net_type: str = 'alex', version: str = '0.1'):
    method forward (line 30) | def forward(self, x: torch.Tensor, y: torch.Tensor):

FILE: evaluation/lpipsPyTorch/modules/networks.py
  function get_network (line 12) | def get_network(net_type: str):
  class LinLayers (line 23) | class LinLayers(nn.ModuleList):
    method __init__ (line 24) | def __init__(self, n_channels_list: Sequence[int]):
  class BaseNet (line 36) | class BaseNet(nn.Module):
    method __init__ (line 37) | def __init__(self):
    method set_requires_grad (line 46) | def set_requires_grad(self, state: bool):
    method z_score (line 50) | def z_score(self, x: torch.Tensor):
    method forward (line 53) | def forward(self, x: torch.Tensor):
  class SqueezeNet (line 66) | class SqueezeNet(BaseNet):
    method __init__ (line 67) | def __init__(self):
  class AlexNet (line 77) | class AlexNet(BaseNet):
    method __init__ (line 78) | def __init__(self):
  class VGG16 (line 88) | class VGG16(BaseNet):
    method __init__ (line 89) | def __init__(self):

FILE: evaluation/lpipsPyTorch/modules/utils.py
  function normalize_activation (line 6) | def normalize_activation(x, eps=1e-10):
  function get_state_dict (line 11) | def get_state_dict(net_type: str = 'alex', version: str = '0.1'):

FILE: evaluation/metrics.py
  function readImages (line 30) | def readImages(renders_dir, gt_dir):
  function evaluate (line 42) | def evaluate(model_paths):

FILE: evaluation/render.py
  function render_set (line 28) | def render_set(model_path, name, iteration, views, gaussians, cfg, backg...
  function render_sets (line 58) | def render_sets(cfg, iteration : int, skip_train : bool, skip_test : bool):

FILE: evaluation/tnt_eval/evaluation.py
  function read_alignment_transformation (line 44) | def read_alignment_transformation(filename):
  function write_color_distances (line 50) | def write_color_distances(path, pcd, distances, max_distance):
  function EvaluateHisto (line 60) | def EvaluateHisto(
  function get_f1_score_histo2 (line 163) | def get_f1_score_histo2(threshold,

FILE: evaluation/tnt_eval/plot.py
  function plot_graph (line 40) | def plot_graph(

FILE: evaluation/tnt_eval/registration.py
  function read_mapping (line 44) | def read_mapping(filename):
  function gen_sparse_trajectory (line 58) | def gen_sparse_trajectory(mapping, f_trajectory):
  function trajectory_alignment (line 65) | def trajectory_alignment(map_file, traj_to_register, gt_traj_col, gt_trans,
  function crop_and_downsample (line 106) | def crop_and_downsample(
  function registration_unif (line 127) | def registration_unif(
  function registration_vol_ds (line 159) | def registration_vol_ds(
  function crop_based_target (line 201) | def crop_based_target(s, t):

FILE: evaluation/tnt_eval/run.py
  function run_evaluation (line 58) | def run_evaluation(dataset_dir, traj_path, ply_path, out_dir):

FILE: evaluation/tnt_eval/trajectory_io.py
  class CameraPose (line 5) | class CameraPose:
    method __init__ (line 7) | def __init__(self, meta, mat):
    method __str__ (line 11) | def __str__(self):
  function convert_trajectory_to_pointcloud (line 16) | def convert_trajectory_to_pointcloud(traj):
  function read_trajectory (line 23) | def read_trajectory(filename):
  function write_trajectory (line 38) | def write_trajectory(traj, filename):

FILE: evaluation/tnt_eval/util.py
  function make_dir (line 4) | def make_dir(path):

FILE: gaussian_renderer/__init__.py
  function render (line 22) | def render(viewpoint_camera, pc : GaussianModel, cfg, bg_color : torch.T...
  function render_fast (line 167) | def render_fast(viewpoint_camera, pc : GaussianModel, cfg, bg_color : to...
  function count_render (line 250) | def count_render(
  function visi_render (line 358) | def visi_render(
  function visi_acc_render (line 467) | def visi_acc_render(

FILE: gaussian_renderer/network_gui.py
  function init (line 26) | def init(wish_host, wish_port):
  function try_connect (line 34) | def try_connect():
  function read (line 43) | def read():
  function send (line 50) | def send(message_bytes, verify):
  function receive (line 57) | def receive():

FILE: process_data/convert.py
  function create_init_files (line 23) | def create_init_files(pinhole_dict_file, db_file, out_dir):
  function main (line 83) | def main(args):

FILE: process_data/convert_360_to_json.py
  function create_init_files (line 18) | def create_init_files(pinhole_dict_file, db_file, out_dir):
  function convert_cam_dict_to_pinhole_dict (line 76) | def convert_cam_dict_to_pinhole_dict(cam_dict, pinhole_dict_file):
  function load_COLMAP_poses (line 105) | def load_COLMAP_poses(cam_file, img_dir, tf='w2c'):
  function load_transformation (line 135) | def load_transformation(trans_file):
  function align_gt_with_cam (line 147) | def align_gt_with_cam(pts, trans):
  function main (line 153) | def main(args):

FILE: process_data/convert_data_to_json.py
  function find_closest_point (line 27) | def find_closest_point(p1, d1, p2, d2):
  function bound_by_pose (line 49) | def bound_by_pose(images):
  function bound_by_points (line 80) | def bound_by_points(points3D):
  function compute_oriented_bound (line 97) | def compute_oriented_bound(pts):
  function split_data (line 105) | def split_data(names, split=10):
  function get_split_dict (line 120) | def get_split_dict(scene_path):
  function check_concentric (line 143) | def check_concentric(images, ang_tol=np.pi / 6.0, radii_tol=0.5, pose_to...
  function export_to_json (line 177) | def export_to_json(trans, scale, scene_path, file_name, split_dict=None,...
  function data_to_json (line 197) | def data_to_json(args):

FILE: process_data/convert_dtu_to_json.py
  function load_K_Rt_from_P (line 37) | def load_K_Rt_from_P(filename, P=None):
  function dtu_to_json (line 62) | def dtu_to_json(args):
  function load_poses (line 107) | def load_poses(scene_path):
  function convert_cam_dict_to_pinhole_dict (line 128) | def convert_cam_dict_to_pinhole_dict(scene_path, pinhole_dict_file):
  function create_init_files (line 161) | def create_init_files(pinhole_dict_file, db_file, out_dir):
  function init_colmap (line 221) | def init_colmap(args):

FILE: process_data/convert_tnt_to_json.py
  function create_init_files (line 18) | def create_init_files(pinhole_dict_file, db_file, out_dir):
  function convert_cam_dict_to_pinhole_dict (line 76) | def convert_cam_dict_to_pinhole_dict(cam_dict, pinhole_dict_file):
  function load_COLMAP_poses (line 105) | def load_COLMAP_poses(cam_file, img_dir, tf='w2c'):
  function load_transformation (line 135) | def load_transformation(trans_file):
  function align_gt_with_cam (line 147) | def align_gt_with_cam(pts, trans):
  function compute_bound (line 153) | def compute_bound(pts):
  function init_colmap (line 164) | def init_colmap(args):

FILE: process_data/extract_mask.py
  function load_image (line 33) | def load_image(image_path):
  function print_ (line 48) | def print_(a):
  function load_model (line 52) | def load_model(model_config_path, model_checkpoint_path, device):
  function get_grounding_output (line 63) | def get_grounding_output(model, image, caption, box_threshold, text_thre...
  function show_mask (line 99) | def show_mask(mask, ax, random_color=False):
  function show_box (line 109) | def show_box(box, ax, label):
  function save_mask_data (line 116) | def save_mask_data(output_dir, mask_list, box_list, label_list, name):
  function morphology_open (line 132) | def morphology_open(x, k1=21, k2=21):
  function process_image (line 140) | def process_image(image_name):

FILE: process_data/extract_normal.py
  function get_args (line 45) | def get_args(test=False):
  function focal2fov (line 92) | def focal2fov(focal, pixels):
  function read_next_bytes (line 96) | def read_next_bytes(fid, num_bytes, format_char_sequence, endian_charact...
  function read_intrinsics_binary (line 108) | def read_intrinsics_binary(path_to_model_file):
  function read_intrinsics_text (line 137) | def read_intrinsics_text(path):
  function load_intrinsic_colmap (line 162) | def load_intrinsic_colmap(path):
  function test_samples (line 208) | def test_samples(args, model, intrins=None, device='cpu'):

FILE: python_scripts/run_base.py
  function worker (line 6) | def worker(gpu, scene, factor, fn):
  function dispatch_jobs (line 13) | def dispatch_jobs(jobs, executor, excluded_gpus, fn):
  function check_finish (line 46) | def check_finish(scene, path, type='mesh'):

FILE: python_scripts/run_dtu.py
  function train_scene (line 58) | def train_scene(gpu, scene, factor):

FILE: python_scripts/run_mipnerf360.py
  function train_scene (line 61) | def train_scene(gpu, scene, factor):

FILE: python_scripts/run_tnt.py
  function train_scene (line 61) | def train_scene(gpu, scene, factor):

FILE: python_scripts/show_360.py
  function show_matrix (line 16) | def show_matrix(scenes, output_dirs, TRIAL_NAME):

FILE: python_scripts/show_dtu.py
  function show_matrix_old (line 10) | def show_matrix_old(scenes, output_dirs, TRIAL_NAME):
  function show_matrix (line 38) | def show_matrix(scenes, output_dirs, TRIAL_NAME):

FILE: python_scripts/show_tnt.py
  function show_matrix (line 17) | def show_matrix(scenes, output_dirs, TRIAL_NAME):

FILE: scene/__init__.py
  class Scene (line 24) | class Scene:
    method __init__ (line 28) | def __init__(self, args : ModelParams, gaussians : GaussianModel, load...
    method save (line 105) | def save(self, iteration, visi=None, surf=None, save_splat=False):
    method getTrainCameras (line 119) | def getTrainCameras(self, scale=1.0):
    method getTestCameras (line 122) | def getTestCameras(self, scale=1.0):
    method getFullCameras (line 125) | def getFullCameras(self, scale=1.0):
    method getUpCameras (line 131) | def getUpCameras(self):
    method getAroundCameras (line 134) | def getAroundCameras(self):
    method getRandCameras (line 137) | def getRandCameras(self, n, up=False, around=True, sample_mode='unifor...

FILE: scene/appearance_network.py
  class UpsampleBlock (line 5) | class UpsampleBlock(nn.Module):
    method __init__ (line 6) | def __init__(self, num_input_channels, num_output_channels):
    method forward (line 12) | def forward(self, x):
  class AppearanceNetwork (line 19) | class AppearanceNetwork(nn.Module):
    method __init__ (line 20) | def __init__(self, num_input_channels, num_output_channels):
    method forward (line 34) | def forward(self, x):

FILE: scene/cameras.py
  class Camera (line 19) | class Camera(nn.Module):
    method __init__ (line 20) | def __init__(self, colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask,
  class MiniCam (line 76) | class MiniCam:
    method __init__ (line 77) | def __init__(self, width, height, fovy, fovx, znear, zfar, world_view_...
  class SampleCam (line 90) | class SampleCam(nn.Module):
    method __init__ (line 91) | def __init__(self, w2c, width, height, FoVx, FoVy, device='cuda'):
  class MiniCam2 (line 115) | class MiniCam2:
    method __init__ (line 116) | def __init__(self, c2w, width, height, fovy, fovx, znear, zfar):

FILE: scene/colmap_loader.py
  function qvec2rotmat (line 43) | def qvec2rotmat(qvec):
  function rotmat2qvec (line 55) | def rotmat2qvec(R):
  class Image (line 68) | class Image(BaseImage):
    method qvec2rotmat (line 69) | def qvec2rotmat(self):
  function read_next_bytes (line 72) | def read_next_bytes(fid, num_bytes, format_char_sequence, endian_charact...
  function read_points3D_text (line 83) | def read_points3D_text(path):
  function read_points3D_binary (line 125) | def read_points3D_binary(path_to_model_file):
  function read_intrinsics_text (line 156) | def read_intrinsics_text(path):
  function read_extrinsics_binary (line 180) | def read_extrinsics_binary(path_to_model_file):
  function read_intrinsics_binary (line 215) | def read_intrinsics_binary(path_to_model_file):
  function read_extrinsics_text (line 244) | def read_extrinsics_text(path):
  function read_colmap_bin_array (line 273) | def read_colmap_bin_array(path):

FILE: scene/dataset_readers.py
  class CameraInfo (line 32) | class CameraInfo(NamedTuple):
  class SceneInfo (line 47) | class SceneInfo(NamedTuple):
  function getNerfppNorm (line 57) | def getNerfppNorm(cam_info):
  function readColmapCameras (line 80) | def readColmapCameras(cam_extrinsics, cam_intrinsics, images_folder, loa...
  function fetchPly (line 149) | def fetchPly(path):
  function storePly (line 157) | def storePly(path, xyz, rgb, normals=None):
  function get_inside_mask (line 174) | def get_inside_mask(pts, trans, scale):
  function filter_point_cloud (line 180) | def filter_point_cloud(trans, scale, xyz, rgb, nb_points=5, radius=0.1):
  function readColmapSceneInfo (line 201) | def readColmapSceneInfo(path, images, eval, llffhold=8, ratio=0, split=F...
  function readCamerasFromTransforms (line 300) | def readCamerasFromTransforms(path, transformsfile, white_background, ex...
  function readNerfSyntheticInfo (line 342) | def readNerfSyntheticInfo(path, white_background, eval, extension=".png"):

FILE: scene/gaussian_model.py
  class GaussianModel (line 36) | class GaussianModel:
    method setup_functions (line 37) | def setup_functions(self):
    method __init__ (line 54) | def __init__(self, cfg):
    method capture (line 88) | def capture(self):
    method restore (line 105) | def restore(self, model_args, training_args):
    method get_scaling (line 126) | def get_scaling(self):
    method get_rotation (line 131) | def get_rotation(self):
    method get_xyz (line 135) | def get_xyz(self):
    method get_features (line 139) | def get_features(self):
    method get_objects (line 145) | def get_objects(self):
    method get_cls (line 148) | def get_cls(self, idx=None):
    method logits_2_label (line 154) | def logits_2_label(self, logits):
    method logits2prob (line 157) | def logits2prob(self, logits):
    method get_opacity (line 161) | def get_opacity(self):
    method get_apperance_embedding (line 164) | def get_apperance_embedding(self, idx):
    method get_normal (line 168) | def get_normal(self, valid=None, idx=None, refine_sign=True, is_all=Fa...
    method get_covariance (line 194) | def get_covariance(self, scaling_modifier = 1):
    method oneupSHdegree (line 197) | def oneupSHdegree(self):
    method create_from_pcd (line 201) | def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : fl...
    method training_setup (line 232) | def training_setup(self, training_args, neural_sdf_params=None):
    method update_learning_rate (line 264) | def update_learning_rate(self, iteration):
    method construct_list_of_attributes (line 272) | def construct_list_of_attributes(self):
    method save_ply (line 289) | def save_ply(self, path):
    method save_inside_ply (line 323) | def save_inside_ply(self, path, inside=None):
    method save_visi_ply (line 355) | def save_visi_ply(self, path, visi):
    method reset_opacity (line 361) | def reset_opacity(self):
    method load_ply (line 366) | def load_ply(self, path):
    method replace_tensor_to_optimizer (line 425) | def replace_tensor_to_optimizer(self, tensor, name):
    method _prune_optimizer (line 442) | def _prune_optimizer(self, mask):
    method prune_points (line 462) | def prune_points(self, mask):
    method cat_tensors_to_optimizer (line 480) | def cat_tensors_to_optimizer(self, tensors_dict):
    method densification_postfix (line 504) | def densification_postfix(self, new_xyz, new_features_dc, new_features...
    method densify_and_split (line 533) | def densify_and_split(self, grads, grad_threshold, scene_extent, visi=...
    method get_dir_max_scaling (line 569) | def get_dir_max_scaling(self, scaling, rots):
    method densify_and_split_along_maxscaling (line 579) | def densify_and_split_along_maxscaling(self, grads, grad_threshold, sc...
    method densify_and_clone (line 629) | def densify_and_clone(self, grads, grad_threshold, scene_extent):
    method densify_and_prune (line 645) | def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_...
    method prune_gaussians (line 661) | def prune_gaussians(self, percent, import_score: list):
    method add_densification_stats (line 669) | def add_densification_stats(self, viewspace_point_tensor, update_filter):
    method get_inside_gaus_normalized (line 673) | def get_inside_gaus_normalized(self):
    method normalize_pts (line 677) | def normalize_pts(self, pts):
    method filter_points (line 681) | def filter_points(self, nb_points=5, radius=0.01, std_ratio=0.01):
    method prune_outside (line 689) | def prune_outside(self):
    method prune_outliers (line 693) | def prune_outliers(self):
    method prune_semantics (line 699) | def prune_semantics(self, cls=BACKGROUND):

FILE: tools/camera.py
  class Pose (line 17) | class Pose():
    method __call__ (line 23) | def __call__(self, R=None, t=None):
    method invert (line 46) | def invert(self, pose, use_inverse=False):
    method compose (line 54) | def compose(self, pose_list):
    method compose_pair (line 62) | def compose_pair(self, pose_a, pose_b):
    method scale_center (line 70) | def scale_center(self, pose, scale):
    method interpolate (line 79) | def interpolate(self, pose_a, pose_b, alpha):
  class Lie (line 99) | class Lie():
    method so3_to_SO3 (line 104) | def so3_to_SO3(self, w):  # [..., 3]
    method SO3_to_so3 (line 113) | def SO3_to_so3(self, R, eps=1e-7):  # [..., 3, 3]
    method se3_to_SE3 (line 122) | def se3_to_SE3(self, wu):  # [...,3]
    method SE3_to_se3 (line 135) | def SE3_to_se3(self, Rt, eps=1e-8):  # [...,3,4]
    method skew_symmetric (line 148) | def skew_symmetric(self, w):
    method taylor_A (line 156) | def taylor_A(self, x, nth=10):
    method taylor_B (line 166) | def taylor_B(self, x, nth=10):
    method taylor_C (line 175) | def taylor_C(self, x, nth=10):
  class Quaternion (line 185) | class Quaternion():
    method q_to_R (line 187) | def q_to_R(self, q):  # [...,4]
    method R_to_q (line 197) | def R_to_q(self, R, eps=1e-6):  # [...,3,3]
    method invert (line 212) | def invert(self, q):  # [...,4]
    method product (line 218) | def product(self, q1, q2):  # [...,4]
    method interpolate (line 227) | def interpolate(self, q1, q2, alpha):  # [...,4],[...,4],[...,1]
  function to_hom (line 242) | def to_hom(X):
  function world2cam (line 249) | def world2cam(X, pose):  # [B,N,3]
  function cam2img (line 254) | def cam2img(X, cam_intr):
  function img2cam (line 258) | def img2cam(X, cam_intr):
  function cam2world (line 262) | def cam2world(X, pose):
  function angle_to_rotation_matrix (line 268) | def angle_to_rotation_matrix(a, axis):
  function get_center_and_ray (line 282) | def get_center_and_ray(pose, intr, image_size):
  function get_3D_points_from_dist (line 313) | def get_3D_points_from_dist(center, ray_unit, dist, multi=True):
  function convert_NDC (line 322) | def convert_NDC(center, ray, intr, near=1):
  function convert_NDC2 (line 342) | def convert_NDC2(center, ray, intr):
  function rotation_distance (line 363) | def rotation_distance(R1, R2, eps=1e-7):
  function get_oscil_novel_view_poses (line 371) | def get_oscil_novel_view_poses(N=60, angle=0.05, dist=5):
  function cross_product_matrix (line 382) | def cross_product_matrix(x):
  function essential_matrix (line 395) | def essential_matrix(poses):
  function fundamental_matrix (line 407) | def fundamental_matrix(poses, intr1, intr2):
  function get_ray_depth_plane_intersection (line 418) | def get_ray_depth_plane_intersection(center, ray, depths):
  function unit_view_vector_to_rotation_matrix (line 437) | def unit_view_vector_to_rotation_matrix(v, axes="ZYZ"):
  function sample_on_spherical_cap (line 456) | def sample_on_spherical_cap(anchor, N, max_angle):
  function sample_on_spherical_cap_northern (line 486) | def sample_on_spherical_cap_northern(anchor, N, max_angle, away_from=Non...

FILE: tools/camera_utils.py
  function loadCam (line 30) | def loadCam(args, id, cam_info, resolution_scale):
  function cameraList_from_camInfos (line 79) | def cameraList_from_camInfos(cam_infos, resolution_scale, args):
  function camera_to_JSON (line 88) | def camera_to_JSON(id, camera):
  function find_up_axis (line 111) | def find_up_axis(R):
  function find_axis (line 124) | def find_axis(R, axis_name='up'):
  function dot (line 144) | def dot(x, y):
  function length (line 151) | def length(x, eps=1e-20):
  function safe_normalize (line 158) | def safe_normalize(x, eps=1e-20):
  function look_at_np (line 162) | def look_at_np(campos, target, opengl=True):
  function look_at (line 182) | def look_at(campos, target, opengl=True):
  function orbit_camera (line 203) | def orbit_camera(elevation, azimuth, radius=1, is_degree=True, target=No...
  function cubic_camera (line 223) | def cubic_camera(n, trans, scale, target=None, opengl=False):
  function check_tensor (line 243) | def check_tensor(x):
  function up_camera (line 249) | def up_camera(n, trans, scale, target=None, opengl=False): # colmap
  function around_camera (line 278) | def around_camera(n, trans, scale, height=None, target=None, opengl=False):
  function bb_camera (line 315) | def bb_camera(n, trans, scale, height=None, target=None, opengl=False, u...
  function around_grid_posi (line 404) | def around_grid_posi(num_points, scale, right_axis, up_axis, front_axis,...
  function up_grid_posi (line 463) | def up_grid_posi(num_points, scale, right_axis, up_axis, front_axis):
  function grid_camera (line 484) | def grid_camera(trans, scale, opengl=False):
  function sample_cameras (line 519) | def sample_cameras(model, n, up=False, around=True, look_mode='target', ...
  class OrbitCamera (line 540) | class OrbitCamera:
    method __init__ (line 541) | def __init__(self, W, H, r=2, fovy=60, near=0.01, far=100):
    method fovx (line 553) | def fovx(self):
    method campos (line 557) | def campos(self):
    method pose (line 562) | def pose(self):
    method view (line 576) | def view(self):
    method perspective (line 581) | def perspective(self):
    method intrinsics (line 601) | def intrinsics(self):
    method mvp (line 606) | def mvp(self):
    method orbit (line 609) | def orbit(self, dx, dy):
    method scale (line 616) | def scale(self, delta):
    method pan (line 619) | def pan(self, dx, dy, dz=0):

FILE: tools/crop_mesh.py
  function align_gt_with_cam (line 7) | def align_gt_with_cam(pts, trans):
  function filter_largest_cc (line 13) | def filter_largest_cc(mesh):
  function main (line 23) | def main(args):

FILE: tools/denoise_pcd.py
  function remove_radius_outlier (line 4) | def remove_radius_outlier(xyz, nb_points=5, radius=0.1):
  function remove_statistical_outlier (line 12) | def remove_statistical_outlier(xyz, nb_points=20, std_ratio=20.):

FILE: tools/depth2mesh.py
  function tsdf_fusion (line 22) | def tsdf_fusion(args, cfg, model, cameras, dirs, bg, outdir, mesh_name='...
  function tsdf_cpu (line 103) | def tsdf_cpu(args, cfg, model, cameras, dirs, bg, outdir, mesh_name='fus...
  function main (line 123) | def main(args):

FILE: tools/distributed.py
  function init_dist (line 21) | def init_dist(local_rank, backend='nccl', **kwargs):
  function get_rank (line 38) | def get_rank():
  function get_world_size (line 47) | def get_world_size():
  function broadcast_object_list (line 56) | def broadcast_object_list(message, src=0):
  function master_only (line 65) | def master_only(func):
  function is_master (line 77) | def is_master():
  function is_dist (line 82) | def is_dist():
  function barrier (line 86) | def barrier():
  function master_first (line 92) | def master_first():
  function is_local_master (line 100) | def is_local_master():
  function master_only_print (line 105) | def master_only_print(*args):
  function dist_reduce_tensor (line 110) | def dist_reduce_tensor(tensor, rank=0, reduce='mean'):
  function dist_all_reduce_tensor (line 127) | def dist_all_reduce_tensor(tensor, reduce='mean'):
  function dist_all_gather_tensor (line 143) | def dist_all_gather_tensor(tensor):

FILE: tools/general_utils.py
  function inverse_sigmoid (line 22) | def inverse_sigmoid(x):
  function PILtoTorch (line 26) | def PILtoTorch(pil_image, resolution):
  function NumpytoTorch (line 35) | def NumpytoTorch(image, resolution):
  function get_expon_lr_func (line 49) | def get_expon_lr_func(
  function strip_lowerdiag (line 84) | def strip_lowerdiag(L):
  function strip_symmetric (line 95) | def strip_symmetric(sym):
  function build_rotation (line 98) | def build_rotation(r):
  function build_scaling_rotation (line 121) | def build_scaling_rotation(s, r):
  function safe_state (line 132) | def safe_state(silent):
  function set_random_seed (line 151) | def set_random_seed(seed):

FILE: tools/graphics_utils.py
  class BasicPointCloud (line 17) | class BasicPointCloud(NamedTuple):
  function geom_transform_points (line 22) | def geom_transform_points(points, transf_matrix):
  function getWorld2View (line 31) | def getWorld2View(R, t):
  function getWorld2View2 (line 38) | def getWorld2View2(R, t, translate=np.array([.0, .0, .0]), scale=1.0):
  function getView2World (line 51) | def getView2World(R, t):
  function getProjectionMatrix (line 63) | def getProjectionMatrix(znear, zfar, fovX, fovY):
  function getIntrinsic (line 89) | def getIntrinsic(fovX, fovY, h, w):
  function fov2focal (line 104) | def fov2focal(fov, pixels):
  function focal2fov (line 107) | def focal2fov(focal, pixels):
  function ndc_2_cam (line 111) | def ndc_2_cam(ndc_xyz, intrinsic, W, H):
  function depth2point_cam (line 120) | def depth2point_cam(sampled_depth, ref_intrinsic):
  function depth2point (line 134) | def depth2point(depth_image, intrinsic_matrix, extrinsic_matrix):
  function get_all_px_dir (line 144) | def get_all_px_dir(intrinsics, height, width):

FILE: tools/image_utils.py
  function mse (line 14) | def mse(img1, img2):
  function psnr (line 17) | def psnr(img1, img2):

FILE: tools/loss_utils.py
  function entropy_loss (line 30) | def entropy_loss(opacity):
  function l1_loss (line 36) | def l1_loss(network_output, gt):
  function log_l1_loss (line 40) | def log_l1_loss(network_output, gt):
  function l2_loss (line 45) | def l2_loss(network_output, gt):
  function gaussian (line 49) | def gaussian(window_size, sigma):
  function create_window (line 54) | def create_window(window_size, channel):
  function ssim (line 61) | def ssim(img1, img2, window_size=11, size_average=True):
  function _ssim (line 72) | def _ssim(img1, img2, window, window_size, channel, size_average=True):
  function eikonal_loss (line 95) | def eikonal_loss(gradients):
  function curvature_loss (line 102) | def curvature_loss(hessian):
  function compute_normal_loss (line 109) | def compute_normal_loss(normal_pred, normal_gt, weight=None):
  function monosdf_normal_loss (line 122) | def monosdf_normal_loss(normal_pred: torch.Tensor, normal_gt: torch.Tens...
  function cos_weight (line 135) | def cos_weight(render_normal, gt_normal, exp_t=1.0):
  function compute_scale_and_shift (line 146) | def compute_scale_and_shift(prediction, target, mask):
  function reduction_batch_based (line 169) | def reduction_batch_based(image_loss, M):
  function reduction_image_based (line 181) | def reduction_image_based(image_loss, M):
  function mse_loss (line 192) | def mse_loss(prediction, target, mask, reduction=reduction_batch_based):
  function gradient_loss (line 201) | def gradient_loss(prediction, target, mask, reduction=reduction_batch_ba...
  class MSELoss (line 221) | class MSELoss(nn.Module):
    method __init__ (line 222) | def __init__(self, reduction='batch-based'):
    method forward (line 230) | def forward(self, prediction, target, mask):
  class GradientLoss (line 234) | class GradientLoss(nn.Module):
    method __init__ (line 235) | def __init__(self, scales=4, reduction='batch-based'):
    method forward (line 245) | def forward(self, prediction, target, mask):
  class ScaleAndShiftInvariantLoss (line 257) | class ScaleAndShiftInvariantLoss(nn.Module):
    method __init__ (line 258) | def __init__(self, alpha=0.5, scales=1, reduction='batch-based'):
    method forward (line 267) | def forward(self, prediction, target, mask=None):
    method __get_prediction_ssi (line 280) | def __get_prediction_ssi(self):
  function normal2curv (line 287) | def normal2curv(normal, mask = None):
  function L1_loss_appearance (line 303) | def L1_loss_appearance(image, gt_image, gaussians, view_idx, return_tran...

FILE: tools/math_utils.py
  function eps_sqrt (line 4) | def eps_sqrt(squared, eps=1e-17):
  function ndc_to_pix (line 12) | def ndc_to_pix(p, resolution):
  function decompose_to_R_and_t (line 25) | def decompose_to_R_and_t(transform_mat, row_major=True):
  function to_homogen (line 39) | def to_homogen(x, dim=-1):
  function normalize_pts (line 50) | def normalize_pts(pts, trans, scale):
  function inv_normalize_pts (line 61) | def inv_normalize_pts(pts, trans, scale):
  function get_inside_normalized (line 70) | def get_inside_normalized(xyz, trans, scale):

FILE: tools/mcube_utils.py
  function marching_cubes_with_contraction (line 17) | def marching_cubes_with_contraction(

FILE: tools/mesh_utils.py
  function post_process_mesh (line 15) | def post_process_mesh(mesh, cluster_to_keep=1000):
  function to_cam_open3d (line 38) | def to_cam_open3d(viewpoint_stack):
  class GaussianExtractor (line 57) | class GaussianExtractor(object):
    method __init__ (line 58) | def __init__(self, gaussians, render, cfg, bg_color=None, dirs=None, p...
    method clean (line 78) | def clean(self):
    method reconstruction (line 87) | def reconstruction(self, viewpoint_stack):
    method extract_mesh_bounded (line 125) | def extract_mesh_bounded(self, voxel_size=0.004, sdf_trunc=0.02, depth...
    method extract_mesh_unbounded (line 169) | def extract_mesh_unbounded(self, resolution=1024):
    method export_image (line 279) | def export_image(self, path):

FILE: tools/normal_utils.py
  function get_normal_sign (line 7) | def get_normal_sign(normals, begin=None, end=None, trans=None, mode='ori...
  function compute_gradient (line 24) | def compute_gradient(img):
  function compute_normals (line 30) | def compute_normals(depth_map, K):
  function compute_edge (line 44) | def compute_edge(image, k=11, thr=0.01):
  function get_edge_aware_distortion_map (line 57) | def get_edge_aware_distortion_map(gt_image, distortion_map):

FILE: tools/prune.py
  function calculate_v_imp_score (line 6) | def calculate_v_imp_score(gaussians, imp_list, v_pow):
  function prune_list (line 25) | def prune_list(gaussians, viewpoint_stack, pipe, background):
  function get_visi_list (line 51) | def get_visi_list(gaussians, viewpoint_stack, pipe, background):

FILE: tools/render_utils.py
  function normalize (line 26) | def normalize(x: np.ndarray) -> np.ndarray:
  function pad_poses (line 30) | def pad_poses(p: np.ndarray) -> np.ndarray:
  function unpad_poses (line 36) | def unpad_poses(p: np.ndarray) -> np.ndarray:
  function recenter_poses (line 41) | def recenter_poses(poses: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
  function average_pose (line 49) | def average_pose(poses: np.ndarray) -> np.ndarray:
  function viewmatrix (line 57) | def viewmatrix(lookdir: np.ndarray, up: np.ndarray,
  function focus_point_fn (line 66) | def focus_point_fn(poses: np.ndarray) -> np.ndarray:
  function transform_poses_pca (line 74) | def transform_poses_pca(poses: np.ndarray) -> Tuple[np.ndarray, np.ndarr...
  function generate_ellipse_path (line 108) | def generate_ellipse_path(poses: np.ndarray,
  function generate_path (line 153) | def generate_path(viewpoint_cameras, n_frames=480):
  function load_img (line 176) | def load_img(pth: str) -> np.ndarray:
  function create_videos (line 183) | def create_videos(base_dir, input_dir, out_name, num_frames=480):
  function save_img_u8 (line 249) | def save_img_u8(img, pth):
  function save_img_f32 (line 257) | def save_img_f32(depthmap, pth):

FILE: tools/sh_utils.py
  function eval_sh (line 57) | def eval_sh(deg, sh, dirs):
  function RGB2SH (line 114) | def RGB2SH(rgb):
  function SH2RGB (line 117) | def SH2RGB(sh):

FILE: tools/system_utils.py
  function mkdir_p (line 16) | def mkdir_p(folder_path):
  function searchForMaxIteration (line 26) | def searchForMaxIteration(folder):

FILE: tools/termcolor.py
  function red (line 18) | def red(x): return termcolor.colored(str(x), color="red")
  function green (line 19) | def green(x): return termcolor.colored(str(x), color="green")
  function blue (line 20) | def blue(x): return termcolor.colored(str(x), color="blue")
  function cyan (line 21) | def cyan(x): return termcolor.colored(str(x), color="cyan")
  function yellow (line 22) | def yellow(x): return termcolor.colored(str(x), color="yellow")
  function magenta (line 23) | def magenta(x): return termcolor.colored(str(x), color="magenta")
  function grey (line 24) | def grey(x): return termcolor.colored(str(x), color="grey")
  function PP (line 32) | def PP(x):
  function alert (line 39) | def alert(x, color='red'):

FILE: tools/visualization.py
  function wandb_image (line 20) | def wandb_image(images, from_range=(0, 1)):
  function preprocess_image (line 26) | def preprocess_image(images, from_range=(0, 1), cmap="viridis"):
  function wandb_sem (line 36) | def wandb_sem(image, palette=PALETTE):
  function tensor2pil (line 44) | def tensor2pil(images):
  function get_heatmap (line 50) | def get_heatmap(gray, cmap):  # [N,H,W]
  function save_render (line 56) | def save_render(render, path):

FILE: tools/visualize.py
  function get_camera_mesh (line 22) | def get_camera_mesh(pose, depth=1):
  function merge_meshes (line 39) | def merge_meshes(vertices, faces):
  function merge_wireframes_k3d (line 46) | def merge_wireframes_k3d(wireframe):
  function merge_wireframes_plotly (line 52) | def merge_wireframes_plotly(wireframe):
  function get_xyz_indicators (line 58) | def get_xyz_indicators(pose, length=0.1):
  function merge_xyz_indicators_k3d (line 64) | def merge_xyz_indicators_k3d(xyz):  # [N,4,3]
  function merge_xyz_indicators_plotly (line 72) | def merge_xyz_indicators_plotly(xyz):  # [N,4,3]
  function k3d_visualize_pose (line 81) | def k3d_visualize_pose(poses, vis_depth=0.5, xyz_length=0.1, center_size...
  function plotly_visualize_pose (line 120) | def plotly_visualize_pose(poses, vis_depth=0.5, xyz_length=0.5, center_s...

FILE: train.py
  function parse_args (line 10) | def parse_args():
  function main (line 19) | def main():

FILE: trainer.py
  class Trainer (line 42) | class Trainer(object):
    method __init__ (line 43) | def __init__(self, cfg):
    method setup_model (line 68) | def setup_model(self, cfg):
    method setup_dataset (line 71) | def setup_dataset(self, cfg):
    method init_writer (line 78) | def init_writer(self, cfg):
    method init_wandb (line 98) | def init_wandb(self, cfg, wandb_id=None, project="", run_name=None, mo...
    method init_losses (line 138) | def init_losses(self):
    method setup_optimizer (line 147) | def setup_optimizer(self, cfg):
    method init_attributes (line 150) | def init_attributes(self):
    method train (line 199) | def train(self):
    method get_center_scale (line 212) | def get_center_scale(self):
    method model_forward (line 225) | def model_forward(self, data, mode):
    method _compute_loss (line 233) | def _compute_loss(self, data, mode=None):
    method _get_total_loss (line 310) | def _get_total_loss(self):
    method train_step (line 323) | def train_step(self, mode='train'):
    method start_of_iteration (line 394) | def start_of_iteration(self):
    method end_of_iteration (line 406) | def end_of_iteration(self, output, render, progress_bar):
    method log_wandb_scalars (line 438) | def log_wandb_scalars(self, output, mode=None):
    method log_wandb_images (line 452) | def log_wandb_images(self, data, mode=None):
    method log_hist (line 496) | def log_hist(self, tensor, name, num_bin=10):
    method test (line 505) | def test(self, renderFunc):
    method finalize (line 589) | def finalize(self):
    method log_writer (line 593) | def log_writer(self, mode=None):
    method save_vis (line 598) | def save_vis(self, data, name, mode='train'):
    method sample_cameras (line 621) | def sample_cameras(self, n, up=False, around=True, look_mode='target',...
    method get_visi_mask (line 637) | def get_visi_mask(self, n=500, up=False, around=True, denoise_after=Fa...
    method get_visi_mask_acc (line 688) | def get_visi_mask_acc(self, n=500, up=False, around=True, sample_mode=...
    method save_gaussians (line 705) | def save_gaussians(self):
Condensed preview — 99 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (509K chars).
[
  {
    "path": ".gitmodules",
    "chars": 482,
    "preview": "[submodule \"submodules/simple-knn\"]\n\tpath = submodules/simple-knn\n\turl = https://gitlab.inria.fr/bkerbl/simple-knn.git\n["
  },
  {
    "path": "LICENSE.md",
    "chars": 4275,
    "preview": "Gaussian-Splatting License  \n===========================  \n\n**Inria** and **the Max Planck Institut for Informatik (MPII"
  },
  {
    "path": "README.md",
    "chars": 10830,
    "preview": "<p align=\"center\">\n\n  <h1 align=\"center\">VCR-GauS: View Consistent Depth-Normal Regularizer for Gaussian Surface Reconst"
  },
  {
    "path": "arguments/__init__.py",
    "chars": 3821,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "bash_scripts/0_train.sh",
    "chars": 465,
    "preview": "GPU=0\nexport CUDA_VISIBLE_DEVICES=${GPU}\nls\n\n\nDATASET=tnt\nSCENE=Barn\nNAME=${SCENE}\n\nPROJECT=vcr_gaus\n\nTRIAL_NAME=vcr_gau"
  },
  {
    "path": "bash_scripts/1_preprocess_tnt.sh",
    "chars": 218,
    "preview": "echo \"Compute intrinsics, undistort images and generate json files. This may take a while\"\npython process_data/convert_t"
  },
  {
    "path": "bash_scripts/2_extract_normal_dsine.sh",
    "chars": 540,
    "preview": "export CUDA_VISIBLE_DEVICES=0\n\nDOMAIN_TYPE=indoor\nDATADIR=/your/data/path\n\nCODE_PATH=/your/dsine/code/path\nCKPT=/your/ds"
  },
  {
    "path": "bash_scripts/3_extract_mask.sh",
    "chars": 1026,
    "preview": "export CUDA_VISIBLE_DEVICES=0\n\nDATADIR=/your/data/path\nGSAM_PATH=~/code/gsam\nCKPT_PATH=${GSAM_PATH}\n\nfor SCENE in Barn C"
  },
  {
    "path": "bash_scripts/4_extract_normal_geow.sh",
    "chars": 616,
    "preview": "export CUDA_VISIBLE_DEVICES=0\n\n# DOMAIN_TYPE=outdoor\n# DOMAIN_TYPE=indoor\nDOMAIN_TYPE=object\nDATADIR=/your/data/path/DTU"
  },
  {
    "path": "bash_scripts/convert.sh",
    "chars": 139,
    "preview": "SCENE=Truck\nDATA_ROOT=/your/data/path/${SCENE}\n\npython convert.py -s $DATA_ROOT # [--resize] #If not resizing, ImageMagi"
  },
  {
    "path": "bash_scripts/install.sh",
    "chars": 159,
    "preview": "env=vcr\nconda create -n $env -y python=3.10\nconda activate $env\npip install -e \".[train]\"\nexport CUDA_HOME=/usr/local/cu"
  },
  {
    "path": "configs/360_v2/base.yaml",
    "chars": 292,
    "preview": "_parent_: configs/reconstruct.yaml\n\nmodel:\n    eval: True\n    llffhold: 8\n    split: False\n\noptim:\n    mask_depth_thr: 1"
  },
  {
    "path": "configs/config.py",
    "chars": 8333,
    "preview": "'''\n-----------------------------------------------------------------------------\nCopyright (c) 2023, NVIDIA CORPORATION"
  },
  {
    "path": "configs/config_base.yaml",
    "chars": 1984,
    "preview": "logdir: \"/your/log/path/debug/\"\nip: 127.0.0.1\nport: -1\ndetect_anomaly: False\nsilent: 0\nseed: 0\n\nmodel:\n    sh_degree: 3\n"
  },
  {
    "path": "configs/dtu/base.yaml",
    "chars": 640,
    "preview": "_parent_: configs/reconstruct.yaml\n\nmodel:\n    use_decoupled_appearance: False\n    use_decoupled_dnormal: False\n    norm"
  },
  {
    "path": "configs/dtu/dtu_scan24.yaml",
    "chars": 32,
    "preview": "_parent_: configs/dtu/base.yaml\n"
  },
  {
    "path": "configs/reconstruct.yaml",
    "chars": 1062,
    "preview": "_parent_: configs/config_base.yaml\n\n\nmodel:\n    load_mask: False\n    use_decoupled_appearance: False\n    use_decoupled_d"
  },
  {
    "path": "configs/scannetpp/base.yaml",
    "chars": 385,
    "preview": "_parent_: configs/reconstruct.yaml\n\nmodel:\n    split: True\n    eval: True\n    use_decoupled_appearance: False\n    use_de"
  },
  {
    "path": "configs/tnt/Barn.yaml",
    "chars": 32,
    "preview": "_parent_: configs/tnt/base.yaml\n"
  },
  {
    "path": "configs/tnt/Caterpillar.yaml",
    "chars": 32,
    "preview": "_parent_: configs/tnt/base.yaml\n"
  },
  {
    "path": "configs/tnt/Courthouse.yaml",
    "chars": 32,
    "preview": "_parent_: configs/tnt/base.yaml\n"
  },
  {
    "path": "configs/tnt/Ignatius.yaml",
    "chars": 32,
    "preview": "_parent_: configs/tnt/base.yaml\n"
  },
  {
    "path": "configs/tnt/Meetingroom.yaml",
    "chars": 266,
    "preview": "_parent_: configs/tnt/base.yaml\n\noptim:\n    exp_t: 1e-3\n    mask_depth_thr: 0\n    densify_large:\n        percent_dense: "
  },
  {
    "path": "configs/tnt/Truck.yaml",
    "chars": 32,
    "preview": "_parent_: configs/tnt/base.yaml\n"
  },
  {
    "path": "configs/tnt/base.yaml",
    "chars": 255,
    "preview": "_parent_: configs/reconstruct.yaml\n\nmodel:\n    use_decoupled_appearance: True\n    use_decoupled_dnormal: False\n    eval:"
  },
  {
    "path": "environment.yml",
    "chars": 548,
    "preview": "name: fast_render\nchannels:\n  - pytorch\n  - nvidia\n  - conda-forge\n  - defaults\ndependencies:\n  - python=3.10\n  - pytorc"
  },
  {
    "path": "evaluation/crop_mesh.py",
    "chars": 2975,
    "preview": "import os\nimport json\nimport plyfile\nimport argparse\n# import open3d as o3d\nimport numpy as np\n# from tqdm import tqdm\ni"
  },
  {
    "path": "evaluation/eval_dtu/eval.py",
    "chars": 6719,
    "preview": "# adapted from https://github.com/jzhangbs/DTUeval-python\nimport numpy as np\nimport open3d as o3d\nimport sklearn.neighbo"
  },
  {
    "path": "evaluation/eval_dtu/evaluate_single_scene.py",
    "chars": 4964,
    "preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport cv2\nimport numpy as np\nimport os\nimport glob\nf"
  },
  {
    "path": "evaluation/eval_dtu/render_utils.py",
    "chars": 4835,
    "preview": "import numpy as np\nimport imageio\nimport skimage\nimport cv2\nimport torch\nfrom torch.nn import functional as F\n\n\ndef get_"
  },
  {
    "path": "evaluation/eval_tnt.py",
    "chars": 3410,
    "preview": "import os\nimport trimesh\nimport argparse\nimport numpy as np\nimport open3d as o3d\nfrom sklearn.neighbors import KDTree\n\n\n"
  },
  {
    "path": "evaluation/full_eval.py",
    "chars": 3340,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "evaluation/lpipsPyTorch/__init__.py",
    "chars": 635,
    "preview": "import torch\n\nfrom .modules.lpips import LPIPS\n\n\ndef lpips(x: torch.Tensor,\n          y: torch.Tensor,\n          net_typ"
  },
  {
    "path": "evaluation/lpipsPyTorch/modules/lpips.py",
    "chars": 1151,
    "preview": "import torch\nimport torch.nn as nn\n\nfrom .networks import get_network, LinLayers\nfrom .utils import get_state_dict\n\n\ncla"
  },
  {
    "path": "evaluation/lpipsPyTorch/modules/networks.py",
    "chars": 2692,
    "preview": "from typing import Sequence\n\nfrom itertools import chain\n\nimport torch\nimport torch.nn as nn\nfrom torchvision import mod"
  },
  {
    "path": "evaluation/lpipsPyTorch/modules/utils.py",
    "chars": 885,
    "preview": "from collections import OrderedDict\n\nimport torch\n\n\ndef normalize_activation(x, eps=1e-10):\n    norm_factor = torch.sqrt"
  },
  {
    "path": "evaluation/metrics.py",
    "chars": 4125,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "evaluation/render.py",
    "chars": 3611,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "evaluation/tnt_eval/README.md",
    "chars": 3463,
    "preview": "# Python Toolbox for Evaluation\n\nThis Python script evaluates **training** dataset of TanksAndTemples benchmark.\nThe scr"
  },
  {
    "path": "evaluation/tnt_eval/config.py",
    "chars": 1875,
    "preview": "# ----------------------------------------------------------------------------\n# -                   TanksAndTemples Web"
  },
  {
    "path": "evaluation/tnt_eval/evaluation.py",
    "chars": 7264,
    "preview": "# ----------------------------------------------------------------------------\n# -                   TanksAndTemples Web"
  },
  {
    "path": "evaluation/tnt_eval/plot.py",
    "chars": 3944,
    "preview": "# ----------------------------------------------------------------------------\n# -                   TanksAndTemples Web"
  },
  {
    "path": "evaluation/tnt_eval/registration.py",
    "chars": 7392,
    "preview": "# ----------------------------------------------------------------------------\n# -                   TanksAndTemples Web"
  },
  {
    "path": "evaluation/tnt_eval/requirements.txt",
    "chars": 28,
    "preview": "matplotlib>=1.3\nopen3d==0.9\n"
  },
  {
    "path": "evaluation/tnt_eval/run.py",
    "chars": 7024,
    "preview": "# ----------------------------------------------------------------------------\n# -                   TanksAndTemples Web"
  },
  {
    "path": "evaluation/tnt_eval/trajectory_io.py",
    "chars": 1257,
    "preview": "import numpy as np\nimport open3d as o3d\n\n\nclass CameraPose:\n\n    def __init__(self, meta, mat):\n        self.metadata = "
  },
  {
    "path": "evaluation/tnt_eval/util.py",
    "chars": 91,
    "preview": "import os\n\n\ndef make_dir(path):\n    if not os.path.exists(path):\n        os.makedirs(path)\n"
  },
  {
    "path": "gaussian_renderer/__init__.py",
    "chars": 20508,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "gaussian_renderer/network_gui.py",
    "chars": 2716,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "process_data/convert.py",
    "chars": 8854,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "process_data/convert_360_to_json.py",
    "chars": 5815,
    "preview": "import os\nimport numpy as np\nimport json\nimport sys\nfrom pathlib import Path\nfrom argparse import ArgumentParser\nimport "
  },
  {
    "path": "process_data/convert_data_to_json.py",
    "chars": 8059,
    "preview": "'''\n-----------------------------------------------------------------------------\nCopyright (c) 2023, NVIDIA CORPORATION"
  },
  {
    "path": "process_data/convert_dtu_to_json.py",
    "chars": 10702,
    "preview": "'''\n-----------------------------------------------------------------------------\nCopyright (c) 2023, NVIDIA CORPORATION"
  },
  {
    "path": "process_data/convert_tnt_to_json.py",
    "chars": 9255,
    "preview": "import os\nimport numpy as np\nimport json\nimport sys\nfrom pathlib import Path\nfrom argparse import ArgumentParser\nimport "
  },
  {
    "path": "process_data/extract_mask.py",
    "chars": 9820,
    "preview": "import argparse\nimport os\nimport gc\nimport sys\n\nimport numpy as np\nimport json\nimport torch\nfrom PIL import Image\nfrom t"
  },
  {
    "path": "process_data/extract_normal.py",
    "chars": 11757,
    "preview": "import os\nimport sys\nimport glob\nimport math\nimport struct\nimport argparse\nimport numpy as np\nimport collections\n\nimport"
  },
  {
    "path": "process_data/extract_normal_geo.py",
    "chars": 8715,
    "preview": "# A reimplemented version in public environments by Xiao Fu and Mu Hu\n\nimport os\nimport sys\nimport logging\nimport argpar"
  },
  {
    "path": "process_data/visualize_colmap.ipynb",
    "chars": 10346,
    "preview": "{\n \"cells\": [\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"id\": \"8b8d7b17-af50-42cd-b531-ef61c49c9e61\",\n "
  },
  {
    "path": "process_data/visualize_transforms.ipynb",
    "chars": 5481,
    "preview": "{\n \"cells\": [\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"id\": \"8b8d7b17-af50-42cd-b531-ef61c49c9e61\",\n "
  },
  {
    "path": "pyproject.toml",
    "chars": 1556,
    "preview": "[tool.black]\nline-length = 240\n\n[build-system]\nrequires = [\"setuptools>=61.0\"]\nbuild-backend = \"setuptools.build_meta\"\n\n"
  },
  {
    "path": "python_scripts/run_base.py",
    "chars": 4876,
    "preview": "import os\nimport time\nimport GPUtil\n\n\ndef worker(gpu, scene, factor, fn):\n    print(f\"Starting job on GPU {gpu} with sce"
  },
  {
    "path": "python_scripts/run_dtu.py",
    "chars": 3411,
    "preview": "# training scripts for the TNT datasets\nimport os\nimport sys\nimport time\nfrom concurrent.futures import ThreadPoolExecut"
  },
  {
    "path": "python_scripts/run_mipnerf360.py",
    "chars": 3509,
    "preview": "# Training script for the Mip-NeRF 360 dataset\nimport os\nimport sys\nimport time\nfrom concurrent.futures import ThreadPoo"
  },
  {
    "path": "python_scripts/run_tnt.py",
    "chars": 3449,
    "preview": "# training scripts for the TNT datasets\nimport os\nimport sys\nimport time\nfrom concurrent.futures import ThreadPoolExecut"
  },
  {
    "path": "python_scripts/show_360.py",
    "chars": 3099,
    "preview": "import json\nimport numpy as np\n\nscenes = ['bicycle', 'flowers', 'garden', 'stump', 'treehill', 'room', 'counter', 'kitch"
  },
  {
    "path": "python_scripts/show_dtu.py",
    "chars": 2334,
    "preview": "import os\nimport json\nimport numpy as np\n\nscenes = [24, 37, 40, 55, 63, 65, 69, 83, 97, 105, 106, 110, 114, 118, 122]\nou"
  },
  {
    "path": "python_scripts/show_tnt.py",
    "chars": 1875,
    "preview": "import os\nimport numpy as np\n\ntraining_list = [\n    'Barn', 'Caterpillar', 'Courthouse', 'Ignatius', 'Meetingroom', 'Tru"
  },
  {
    "path": "requirements.txt",
    "chars": 122,
    "preview": "submodules/diff-gaussian-rasterization\nsubmodules/simple-knn/\ngit+https://github.com/facebookresearch/pytorch3d.git@stab"
  },
  {
    "path": "scene/__init__.py",
    "chars": 6943,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "scene/appearance_network.py",
    "chars": 1868,
    "preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass UpsampleBlock(nn.Module):\n    def __init__(sel"
  },
  {
    "path": "scene/cameras.py",
    "chars": 5198,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "scene/colmap_loader.py",
    "chars": 11859,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "scene/dataset_readers.py",
    "chars": 15203,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "scene/gaussian_model.py",
    "chars": 34578,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "tools/__init__.py",
    "chars": 0,
    "preview": ""
  },
  {
    "path": "tools/camera.py",
    "chars": 20621,
    "preview": "'''\n-----------------------------------------------------------------------------\nCopyright (c) 2023, NVIDIA CORPORATION"
  },
  {
    "path": "tools/camera_utils.py",
    "chars": 21549,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "tools/crop_mesh.py",
    "chars": 3184,
    "preview": "import os\nimport argparse\nimport numpy as np\nimport trimesh\n\n\ndef align_gt_with_cam(pts, trans):\n    trans_inv = np.lina"
  },
  {
    "path": "tools/denoise_pcd.py",
    "chars": 1221,
    "preview": "from pytorch3d.ops import ball_query, knn_points\n\n\ndef remove_radius_outlier(xyz, nb_points=5, radius=0.1):\n    if xyz.d"
  },
  {
    "path": "tools/depth2mesh.py",
    "chars": 8565,
    "preview": "import os\nimport sys\nimport math\nimport torch\nimport argparse\nimport numpy as np\nimport open3d as o3d\nimport open3d.core"
  },
  {
    "path": "tools/distributed.py",
    "chars": 4204,
    "preview": "'''\n-----------------------------------------------------------------------------\nCopyright (c) 2023, NVIDIA CORPORATION"
  },
  {
    "path": "tools/general_utils.py",
    "chars": 5032,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "tools/graphics_utils.py",
    "chars": 4794,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "tools/image_utils.py",
    "chars": 554,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "tools/loss_utils.py",
    "chars": 10702,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "tools/math_utils.py",
    "chars": 1879,
    "preview": "import torch\n\n\ndef eps_sqrt(squared, eps=1e-17):\n    \"\"\"\n    Prepare for the input for sqrt, make sure the input positiv"
  },
  {
    "path": "tools/mcube_utils.py",
    "chars": 3461,
    "preview": "#\n# Copyright (C) 2024, ShanghaiTech\n# SVIP research group, https://github.com/svip-lab\n# All rights reserved.\n#\n# This "
  },
  {
    "path": "tools/mesh_utils.py",
    "chars": 13706,
    "preview": "import torch\nimport numpy as np\nimport os\nimport math\nfrom tqdm import tqdm\nfrom functools import partial\nimport open3d "
  },
  {
    "path": "tools/normal_utils.py",
    "chars": 2307,
    "preview": "import torch\nimport torch.nn.functional as F\n\nfrom tools.graphics_utils import depth2point_cam\n\n\ndef get_normal_sign(nor"
  },
  {
    "path": "tools/prune.py",
    "chars": 2418,
    "preview": "import torch\n\nfrom gaussian_renderer import count_render, visi_acc_render\n\n\ndef calculate_v_imp_score(gaussians, imp_lis"
  },
  {
    "path": "tools/render_utils.py",
    "chars": 9650,
    "preview": "# Copyright 2022 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this f"
  },
  {
    "path": "tools/semantic_id.py",
    "chars": 163,
    "preview": "\nBACKGROUND = 0\ntext_label_dict = {\n    'window': BACKGROUND,\n    'sky': BACKGROUND,\n    'sky window': BACKGROUND,\n    '"
  },
  {
    "path": "tools/sh_utils.py",
    "chars": 4371,
    "preview": "#  Copyright 2021 The PlenOctree Authors.\n#  Redistribution and use in source and binary forms, with or without\n#  modif"
  },
  {
    "path": "tools/system_utils.py",
    "chars": 785,
    "preview": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# Thi"
  },
  {
    "path": "tools/termcolor.py",
    "chars": 1449,
    "preview": "'''\n-----------------------------------------------------------------------------\nCopyright (c) 2023, NVIDIA CORPORATION"
  },
  {
    "path": "tools/visualization.py",
    "chars": 1931,
    "preview": "import wandb\nimport imageio\nimport torch\nimport torchvision\n\nfrom matplotlib import pyplot as plt\nfrom torchvision.trans"
  },
  {
    "path": "tools/visualize.py",
    "chars": 7048,
    "preview": "'''\n-----------------------------------------------------------------------------\nCopyright (c) 2023, NVIDIA CORPORATION"
  },
  {
    "path": "train.py",
    "chars": 1157,
    "preview": "import os\nimport sys\nimport argparse\nsys.path.append(os.getcwd())\n\nfrom configs.config import Config, recursive_update_s"
  },
  {
    "path": "trainer.py",
    "chars": 35316,
    "preview": "import os\nimport json\nimport uuid\nimport math\nimport wandb\nimport imageio\nimport numpy as np\nfrom torch import nn\nfrom t"
  }
]

About this extraction

This page contains the full source code of the HLinChen/VCR-GauS GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 99 files (474.8 KB), approximately 130.4k tokens, and a symbol index with 492 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.

Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.

Copied to clipboard!