Upload 133 files

One-2-3-45-master 2/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

One-2-3-45-master 2/.gitignore

@@ -0,0 +1,11 @@
+__pycache__/
+exp/
+src/
+*.DS_Store
+*.ipynb
+*.egg-info/
+*.ckpt
+*.pth
+
+!example.ipynb
+!reconstruction/exp

One-2-3-45-master 2/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

One-2-3-45-master 2/README.md

@@ -0,0 +1,221 @@
+<p align="center" width="100%">
+<img src="https://github.com/Dustinpro/Dustinpro/assets/23076389/0fbdb69a-0fb4-4b42-b9da-e0b28532bdfd"  width="80%" height="80%">
+</p>
+
+
+<p align="center">
+  [<a href="https://arxiv.org/pdf/2306.16928.pdf"><strong>Paper</strong></a>]
+  [<a href="http://one-2-3-45.com"><strong>Project</strong></a>]
+  [<a href="https://huggingface.co/spaces/One-2-3-45/One-2-3-45"><strong>Demo</strong></a>]
+  [<a href="#citation"><strong>BibTeX</strong></a>]
+</p>
+
+<p align="center">
+  <a href="https://huggingface.co/spaces/One-2-3-45/One-2-3-45">
+    <img alt="Hugging Face Spaces" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Space_of_the_Week_%F0%9F%94%A5-blue">
+  </a>
+</p>
+
+One-2-3-45 rethinks how to leverage 2D diffusion models for 3D AIGC and introduces a novel forward-only paradigm that avoids the time-consuming optimization.
+
+https://github.com/One-2-3-45/One-2-3-45/assets/16759292/a81d6e32-8d29-43a5-b044-b5112b9f9664
+
+
+
+https://github.com/One-2-3-45/One-2-3-45/assets/16759292/5ecd45ef-8fd3-4643-af4c-fac3050a0428
+
+
+## News
+**[09/21/2023]**
+One-2-3-45 is accepted by NeurIPS 2023. See you in New Orleans!
+
+**[09/11/2023]**
+Training code released.
+
+**[08/18/2023]**
+Inference code released.
+
+**[07/24/2023]**
+Our demo reached the HuggingFace top 4 trending and was featured in 🤗 Spaces of the Week 🔥! Special thanks to HuggingFace 🤗 for sponsoring this demo!!
+
+**[07/11/2023]**
+[Online interactive demo](https://huggingface.co/spaces/One-2-3-45/One-2-3-45) released! Explore it and create your own 3D models in just 45 seconds! 
+
+**[06/29/2023]**
+Check out our [paper](https://arxiv.org/pdf/2306.16928.pdf). [[X](https://twitter.com/_akhaliq/status/1674617785119305728)]
+
+## Installation
+Hardware requirement: an NVIDIA GPU with memory >=18GB (_e.g._, RTX 3090 or A10). Tested on Ubuntu.
+
+We offer two ways to setup the environment:
+
+### Traditional Installation
+<details>
+<summary>Step 1: Install Debian packages. </summary> 
+
+```bash
+sudo apt update && sudo apt install git-lfs libsparsehash-dev build-essential
+```
+</details>
+
+<details>
+<summary>Step 2: Create and activate a conda environment. </summary>
+
+```bash
+conda create -n One2345 python=3.10
+conda activate One2345
+```
+</details>
+
+<details>
+<summary>Step 3: Clone the repository to the local machine. </summary>
+
+```bash
+# Make sure you have git-lfs installed.
+git lfs install
+git clone https://github.com/One-2-3-45/One-2-3-45
+cd One-2-3-45
+```
+</details>
+
+<details>
+<summary>Step 4: Install project dependencies using pip. </summary>
+
+```bash
+# Ensure that the installed CUDA version matches the torch's cuda version.
+# Example: CUDA 11.8 installation
+wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run
+sudo sh cuda_11.8.0_520.61.05_linux.run
+export PATH="/usr/local/cuda-11.8/bin:$PATH"
+export LD_LIBRARY_PATH="/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH"
+# Install PyTorch 2.0
+pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
+# Install dependencies
+pip install -r requirements.txt
+# Install inplace_abn and torchsparse
+export TORCH_CUDA_ARCH_LIST="7.0;7.2;8.0;8.6+PTX" # CUDA architectures. Modify according to your hardware.
+export IABN_FORCE_CUDA=1
+pip install inplace_abn
+FORCE_CUDA=1 pip install --no-cache-dir git+https://github.com/mit-han-lab/torchsparse.git@v1.4.0
+```
+</details>
+
+<details>
+<summary>Step 5: Download model checkpoints. </summary>
+
+```bash
+python download_ckpt.py
+```
+</details>
+
+
+### Installation by Docker Images
+<details>
+<summary>Option 1: Pull and Play (environment and checkpoints). (~22.3G)</summary> 
+
+```bash
+# Pull the Docker image that contains the full repository.
+docker pull chaoxu98/one2345:demo_1.0
+# An interactive demo will be launched automatically upon running the container.
+# This will provide a public URL like XXXXXXX.gradio.live
+docker run --name One-2-3-45_demo --gpus all -it chaoxu98/one2345:demo_1.0
+```
+</details>
+
+<details>
+<summary>Option 2: Environment Only. (~7.3G)</summary> 
+
+```bash
+# Pull the Docker image that installed all project dependencies.
+docker pull chaoxu98/one2345:1.0
+# Start a Docker container named One2345.
+docker run --name One-2-3-45 --gpus all -it chaoxu98/one2345:1.0
+# Get a bash shell in the container.
+docker exec -it One-2-3-45 /bin/bash
+# Clone the repository to the local machine.
+git clone https://github.com/One-2-3-45/One-2-3-45
+cd One-2-3-45
+# Download model checkpoints. 
+python download_ckpt.py
+# Refer to getting started for inference.
+```
+</details>
+
+## Getting Started (Inference)
+
+First-time running will take longer time to compile the models.
+
+Expected time cost per image: 40s on an NVIDIA A6000.
+```bash
+# 1. Script
+python run.py --img_path PATH_TO_INPUT_IMG --half_precision
+
+# 2. Interactive demo (Gradio) with a friendly web interface
+#    An URL will be provided in the output 
+#    (Local: 127.0.0.1:7860; Public: XXXXXXX.gradio.live)
+cd demo/
+python app.py
+
+# 3. Jupyter Notebook
+example.ipynb
+```
+
+## Training Your Own Model
+
+### Data Preparation
+We use Objaverse-LVIS dataset for training and render the selected shapes (with CC-BY license) into 2D images with Blender. 
+#### Download the training images.
+Download all One2345.zip.part-* files (5 files in total) from <a href="https://huggingface.co/datasets/One-2-3-45/training_data/tree/main">here</a> and then cat them into a single .zip file using the following command:
+```bash
+cat One2345.zip.part-* > One2345.zip
+```
+
+#### Unzip the training images zip file.
+Unzip the zip file into a folder specified by yourself (`YOUR_BASE_FOLDER`) with the following command:
+
+```bash
+unzip One2345.zip -d YOUR_BASE_FOLDER
+```
+
+#### Download meta files.
+
+Download `One2345_training_pose.json` and `lvis_split_cc_by.json` from <a href="https://huggingface.co/datasets/One-2-3-45/training_data/tree/main">here</a> and put them into the same folder as the training images (`YOUR_BASE_FOLDER`).
+
+Your file structure should look like this:
+```
+# One2345 is your base folder used in the previous steps
+
+One2345
+├── One2345_training_pose.json
+├── lvis_split_cc_by.json
+└── zero12345_narrow
+    ├── 000-000
+    ├── 000-001
+    ├── 000-002
+    ...
+    └── 000-159
+    
+```
+
+### Training
+Specify the `trainpath`, `valpath`, and `testpath` in the config file `./reconstruction/confs/one2345_lod_train.conf` to be `YOUR_BASE_FOLDER` used in data preparation steps and run the following command:
+```bash
+cd reconstruction
+python exp_runner_generic_blender_train.py --mode train --conf confs/one2345_lod_train.conf
+```
+Experiment logs and checkpoints will be saved in `./reconstruction/exp/`.
+
+## Citation
+
+If you find our code helpful, please cite our paper:
+
+```
+@misc{liu2023one2345,
+      title={One-2-3-45: Any Single Image to 3D Mesh in 45 Seconds without Per-Shape Optimization}, 
+      author={Minghua Liu and Chao Xu and Haian Jin and Linghao Chen and Mukund Varma T and Zexiang Xu and Hao Su},
+      year={2023},
+      eprint={2306.16928},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

One-2-3-45-master 2/configs/sd-objaverse-finetune-c_concat-256.yaml

@@ -0,0 +1,117 @@
+model:
+  base_learning_rate: 1.0e-04
+  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "image_target"
+    cond_stage_key: "image_cond"
+    image_size: 32
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: hybrid
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+
+    scheduler_config: # 10000 warmup steps
+      target: ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 100 ]
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 8
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenCLIPImageEmbedder
+
+
+data:
+  target: ldm.data.simple.ObjaverseDataModuleFromConfig
+  params:
+    root_dir: 'views_whole_sphere'
+    batch_size: 192
+    num_workers: 16
+    total_view: 4
+    train:
+      validation: False
+      image_transforms:
+        size: 256
+
+    validation:
+      validation: True
+      image_transforms:
+        size: 256
+
+
+lightning:
+  find_unused_parameters: false
+  metrics_over_trainsteps_checkpoint: True
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 5000
+  callbacks:
+    image_logger:
+      target: main.ImageLogger
+      params:
+        batch_frequency: 500
+        max_images: 32
+        increase_log_steps: False
+        log_first_step: True
+        log_images_kwargs:
+          use_ema_scope: False
+          inpaint: False
+          plot_progressive_rows: False
+          plot_diffusion_rows: False
+          N: 32
+          unconditional_guidance_scale: 3.0
+          unconditional_guidance_label: [""]
+
+  trainer:
+    benchmark: True
+    val_check_interval: 5000000 # really sorry
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1

One-2-3-45-master 2/demo/.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

One-2-3-45-master 2/demo/.gitignore

@@ -0,0 +1,4 @@
+weights/
+data/
+*.ipynb
+demo_examples_*

One-2-3-45-master 2/demo/app.py

@@ -0,0 +1,639 @@
+import os
+import sys
+import shutil
+import torch
+import fire
+import gradio as gr
+import numpy as np
+import cv2
+from PIL import Image
+import plotly.graph_objects as go
+from functools import partial
+import trimesh
+import tempfile
+from rembg import remove
+
+code_dir = "../"
+sys.path.append(code_dir)
+from utils.zero123_utils import init_model, predict_stage1_gradio, zero123_infer
+from utils.sam_utils import sam_init, sam_out_nosave
+from utils.utils import image_preprocess_nosave, gen_poses
+from elevation_estimate.estimate_wild_imgs import estimate_elev
+
+_GPU_INDEX = 0
+_HALF_PRECISION = True
+_MESH_RESOLUTION = 256
+
+_TITLE = '''One-2-3-45: Any Single Image to 3D Mesh in 45 Seconds without Per-Shape Optimization'''
+_DESCRIPTION = '''
+<div>
+<a style="display:inline-block" href="http://one-2-3-45.com"><img src="https://img.shields.io/badge/Project_Homepage-f9f7f7?logo=data:image/webp;base64,UklGRmIRAABXRUJQVlA4IFYRAABQPwCdASrIAMgAPm00lEekpiolqDvpMIANiWJu3pE7maI+vTDkhN5f7PfmGT7nS6p8nKBr0I+YBzr/ML+2/rG/8j1Sf3/1AP6v/sOsW/bn2GP2W9Zv/zeyf/cf+Z+63tReoB/6uCx2p5E/iUkPwG1FO9t/XgHPL7RH7TzI8EvEA4Mr1T2CP0J6xH+d5J/rz2F/LG9h37l+x9+4xQ3m86D2Te/zeVV/tWyTw7s85XZ0ABD4N2CpzWHt8feKiWkqdTkRjojREWrbUDAKXlYsV7EGU9rWR2gCxVXnstqpNVhwra603swvYRlMyRzKc5nJHEEeLuV8EDee/MpPVIq2DNUcXpCZXvFLHgXBWvZWzSZCFo4iub8df+Yu9q7rw5qemOe2Nt1IIoyjBmXdjCunMherehhPjQIQGiI6PDcriy/zhhwHE7O+0gmpUsYmcDR+ixOfLPY0yjnQosZkIoK1pfttDGtirMbSMqndDVi73JMxcSlNb0MNFtgdAAXNk5Z77wgsPz9RRj9oWO/KRpXn5ra4gUt+mMgSCvFG86zgghSehTRD54z10sNxqnG3/rpKDifOvT4EQU1uA9ZckUZcUt5L5C0+dOdj1I56uLJEAsn0432gHD5wRG7dgSfYusXvhGl2uMaczlXSJ0JfX+Z0e9q7sHywvyEkWC+iJREwvtWi1K+NQAD+/JSRhGP+QTeW9xU73vXKZO+JaR/TAb6vV9dNzIjket6jYZdxK0qCcaf95QeouegLeSQL/9WeH5l2/DE2AKdYjhEYzzIefp7c6cTfM3D3q3kSFxAF/xP/f/3YUFjjOzfzl5xrD3XaWz0TAehn6+ze5pANq6t5GDX8ZOfpIBGUplJj6UZXd76ropLkDdM+d/F2Megl53hry7QvtcUGNlKgjLd7/txvzvkYIPre5sKVvAJzj9DEml706Piekk2NTtBnCMQtQAPO7/Soo3p3QbqLnMIY2PKCq3jFUkeMDAB6uvaHy7e8G/yi+LlFCfYgju+h+ha+jj6NYh6xUx/9TpQoQ1VFrpEw7pCAaQ2NbzVcj/EfBLQUWQBwliZd6FG70L3ATK7AS/cu+Pm/ASndDhIDTx08uveDvY2kW7Mqproq8D4ImWzJ7ZwM8JfrvyN9/wh0Iu00O3UbTDU58dYfWzxI1gDb2Yt6+AyvgjRY/WUM8aikx5MTFi6ZEWeffMc8ruwWeKmfwJtpDxNYhJgSN5gZoOS+XedZmwoYfiuaf9hhPdDtJCM429liA9mZQ2GNfMOPtcLJV/4xTUuWJx4/d43remtcIdsy1GlD79SNNSlfWuSCF8LIopGEcQwZnBVOBmJ7O2bQvBHNAQ6dlz+Pc8zL7MgsN7uff5PDGyXHqV4lCh2Q/xbeZaYVv1agSO4QU5FvEX/6AQpNxQqgtvrXM+XsXTWYJpiO7+ucPvPudldDswT7/ITpp7AdSJ9OjPYr3cMRBVy5sXXkyY8SVv0z//QqGJbxMA3IV81dfN5dUNNlMvl+EBv6Qrnq42ZAEXMEbW/zcmuvIO+539I0BKM+COuGTuEmhulQBlMmdlBNVaII5lFuENjHGpIPNULKARA/PhM9jOvRG2xs8SFCjLZ9ZNLyWznJxq3sZWeGUNXcyJPKDr3WAWCmP8tNyjKQk6aKOD1s/+2MCNQ9b4Zb2XJJLW4kBOS6P10n42Scz8D1K6GTeouELbGApoVNEYjw2jdbNqsNLZiJi6XfSs7Kz5ozupOLJsYOzfLQTjDz7BZOzEYFoB+acySl5Qs3mO84Mw6pIb9MqHhzFqtmrsD4V6uwssP9SUvJTMA4orRWZAMdCA9lMHQi6XBPR9CBidQdudyeYtUj5gWMyzdcSkG8l/2zbzpz8THi23y/8+Ijg5naj6GtYnpJna+6QpI6dlY+E2KF7bOK2ctKwBjClUvAjHXa1162i6DsToHLQE4stmDJdvI1POb9Hj0Mq+azo1wrOfqVFcAS5XNc37IJeYBs/cQYZ08mg2vXWWJYVWz648jTHABHf+LiHsy4WRaVo4oOOSyeampoUYSM9WUJ3iOlTMis5U2DCrGoAiATOAyyuwMcYgTni5FGSpdE5BnoS6ORUiYapPetM/XmsvikTkKNn4z4jhiLFFcU+bH1pZ2DseVK9vCgY5s9ZDjNb9Ky+8fwn9dJtsZ6M7opvXhqde9Ljos6KWQ/8hj3pswa2lLZ7WRc9xaxTjq1sytCxfOd+J+VrsXedNuuYDMwumYIzF1Xsbz1VCURDw6C1twAPizF49s4JfToWsGhgG4wtBE5NAU4KvnGleFGzv54AwBR9qqPVD9bcN7ZmZhphTcAGnR2oOcvT98FmknnVXqJYdHSeP9nWG6A8YUUgDmy7rYLtbxgpid5ysrxvLeyMOeTaQXuNZXn5kQeqDGELMfQ5U2PAk+/DhnbWTsirG5NAP0tgGbSNGM5cv9+trgSk6sXdw1lZOLrfqLGZ8Dt19DWcxmjsbDU30CoSc1alYxX5G+uIHy72tQxjzsot1O4iZeNO34PItzBoeg0Fq+YQZGsHdknwJkAbffRl96xFsDxM6l4g22OaMEHxLMC9uFFE8Ee/xf+krkjv7YCfJrCM3Nw6xfyrhtxN3x2GxSg4YTu2dtjb3zVI/6jYNwgGiaDWoh5I29uQ8ZvES8Ros5jgxDzeKB1tJ3HtDM9SFGNJfQiLiSyYZQLBjCcGbi3+vlythB3k6af+P5rDqah2oPFl29Ngnw/tmpkmRIvri5i55FPeY9J4nXfvWYHTHdoB0oVA2NEk2nropP+T7GXhAxA2NgyGtzHaVU2yxiSju87w8MLIo1eac26wOnbEo/oD6Zcb8vyu0x7ug9iERQ5FlppDnIktT6QC6Kk3qBxovLzOPdEvYQoytaN256n2dmkxAaq78klv6PnU7HiH3e/I9RC27VOP0j6JDW19KvC9/uN9tfOi6WMr0IGKpTsZAUZXm+Ukyk/Rpu9ZPIH5/3CL+yfj3ROts+BWIZNj8lpFHfmYhmN/J0+/lDIGmbRVMbvmif9tqr53fqb8EkFpdMHnK8jc0oIYu2Io5SWOzHc7GMdwt5RB8LR5jUjV6Xv+rR7N4IcTQRphe7WarFsxHmihpNr8sLroqsVxBH+6dOjC5DPhRV6aJB9ZB0NjpLtbjRsEKe1By1huo8rJa+DS73fTUfxWPaJjQsasBOcc6xwuob3OBjTFjUOxfiBbhMDNUFcamlMphrkbmTe2smHz0hrScXZjoHxphV537e8PNenBpI//N58bUOcmV4Lo1H1BLLjNTw1gK+rKFgaU/WOZQ0DZ1kRRqCa86XYnFposmkLgDNooS/yeW/RGfvopRDH40d2TeW8t1+2fDHQcwocSXolq+dxC6JMGsu2rCrhdjzhqd1KPMp5EVGQuCyLc8LfjUhQ8fSs63P9aVDYZBDhO8oWSI2Lbk7cRpKJ38ww9dD0b0OjvucHkJl1zIwyQFqKKEfIN7RPvV8Q1Xxot6Y5f8/UqOCOVZRt+IM1JFcJ4AstPMOXs9hAyZzEs1EY9lv3976/18LNNvL8K7RPNH1uz3qwAajMXLaOTEK7IzCjex3YZQ0LCICPzWVKMNbkSFpmy5ow1A54fK4F45T0apL1FE8dc/Jy6ERymiJ8ZvT+BJHUtbS5oB72w8NeIb0zTuqTzYwMQiKeCI+DlJTd6R3dgbvDETb7XtLT1L5quVxBiyJLxgARoeWU1DY3eWTFJkicFp/UqIFCYgLUhQgGm/1gAxylWf4wZmbQy6RGlY3/pfn5qxqFq8Xmza7Unght3AckydGZ6u6yWcooxZwILsHaklA/Bu2HRlCLzLer57IQWfvHUjJ8pqEoZ/TE0WqZc4SF6CBVC4KGEIqyPnH/+chaIQRfGuKg0rKAAc5tB+7vGl4ck72A+dA9iW0UUwXqD6Y333q9MEdov8hbXuiRkRMv1CEm0h1N8yhxOEe1SLWxlCmvUHcVvhojM6S4ODYr2rxHxOqx63MVVCk6PpQAB2gn4D/9+QHVBBqAxLV8Pggh6aRlEmPuHNEc+b/1Zqh4lvwxUgyMFngOgTAhqZAZqBpRRD41KfU7wEbzruYQhOIxxPMbGIe93LRFgYZLz21sLmS/02JhQ1eY6cSu/iumOWleWujzjCcB9qxDpOBlsugpiveHwunOO9Lu66uGNeKuw6Aqo+zTcdEX+BOlserNtSeyYmhQrwLA+mnEqaAtKv5eAyTC03krdlSEI+++xVMU+kqsGF+6H9yNBQj5aZxmOzd7BejBdBBInEjlj868zR80jlgVKb+yQ7XkdiFIvQl/XvaFpPGqYb4UR70U0jNe/I9UuFggu8M90wyOi7Ihm3t9FZTBPv4zmee4ue5pVKpdsOOZLwBSODTdpUJb6ctU3I9n0KAUBM4+RkNoQzCyb+iXoXl22CL2eQWlOlBi8IG84Y2bMIiLnPs5qeUth9zlniN14oQNTtVJibuIgkylT50ExHyuqz1ra2+wW3QDltErT6yyrKnL8rmkPesI3aPAL880z4U6TWXqcU6hkryL8W5gdI94KYuDTBEim0GM6IAAKf8JZNX3sM/OIB9h3XbFUuNXRocJY9iqQAGdinm3YPLbRBxP5S5EWwlTdIVK5yjUpV+tCN0HXOVf7xj9pnyIMPDz/Znf4zufz+0zonywFQLgAdKiVwC5a6EDC2rmxYC4L82QIO17SKc8NCAJZPTWwwrPGgb0nhQdi3g32QzHUAqE2qhq2jyM7WINI34P28PN5IE50uRx/XFn2a8h2Qgla55PnsIT7KbDBo0Nd4XUkCWINxReQK0/NZEDZrUuZghyZYnnIuIi0pTpecJWliTLvfxyiRkIsb9t2mT6VzM8H2HN8nq0rF7BC27r0JoLl/5YgZQZmw763cQ625wkmPOX0vr1M35fZYv06zKm1ux/L+W6O3ju3VdFudKgEgRIeT+bIOQKoKaT+knRugmPDGt1JAt6bKTT2bvIYnf5OvZs9id5x+qy5UeotL3uxYiBj7SyGxTCHdovbak2BG5hGmuVWxRojEJS9IEqUKwy133zg24keiFy0bXsG125D1XgQ/uI3IM8dijJ4N6jHObWneJl3zHvKb+cX97XFAv5VV5ySEfm0Iglkir9QaTOXP9SuTeCao7Q3fULO6Jcp+sOU6x+jCjlsmHiF7oDCAb/RYITB9oIcAGGQJkT6ccTNJPlyPiNf4/nefe5V4j5RWLTDm7Wb/kt426CIGzE2ekBjrvExlch914MGYjMJcBllUj/LTfXTYLSqPwPzU/xBVSUOR9o7wFnFBZTaI0ApgY11rRsoEgTu9yRgBxw0h71O/RjpN5Ku/U/er87C9/jHzucfXpRDcP1JxOOxoJziSE01YxnPjmyDigmCcur63bY/xXdZeNQNprWvE3mAIP14fFkdJ4+0vwkAP+BXokPPQBkZuEWFAUEz1H/YQf4Q9bCQZXl/WSpUpG/TjBo8EpZLTJ2Jwa1G3H2hVIUlifUnV/SvKDYbpUvl6mKuwdgglJxkJOXjtf84FjvjeHUOzf8ZhLw3PH53rUrDz0INySaGJ/n4a/iuvLMaL146Ii98kND4sM0nElTIxnJe+LJF/8aimynAAiTshwnKc7MHqCtuDaUFfEQCGw0tmys5yKZM5zzawr6LW7MdQs9XUDiyTrX+YcI0uPZZ43oMnO737u5Tmc/sAeKCNGIWt8kw87EMQ+BP5NMrf8b8wDvSD2XVEZu7xqwNCizeSYQJGVJVSAdJ27XwXrFfHtdHHrlojW+3BFzE5rOzDsUsA00zYHxt+e9zo9Yn0sImcxGhbDFBGD892Rgz9G+eor3huRF8h4p1qYpjTe/ykVkhWyvHRjNNevOV7Gk1jhgiwOajzrXwNsIJNUvAQQB017GRYgey7MAEeBoAx5RuxYU+oMH6DNk5eYcrJDxo48XGbO4QhCMRgAA"></a>
+<a style="display:inline-block; margin-left: .5em" href="https://arxiv.org/abs/2306.16928"><img src="https://img.shields.io/badge/2306.16928-f9f7f7?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAADcAAABMCAYAAADJPi9EAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAAuIwAALiMBeKU/dgAAABl0RVh0U29mdHdhcmUAd3d3Lmlua3NjYXBlLm9yZ5vuPBoAAAa2SURBVHja3Zt7bBRFGMAXUCDGF4rY7m7bAwuhlggKStFgLBgFEkCIIRJEEoOBYHwRFYKilUgEReVNJEGCJJpehHI3M9vZvd3bUP1DjNhEIRQQsQgSHiJgQZ5dv7krWEvvdmZ7d7vHJN+ft/f99pv5XvOtJMFCqvoCUpTdIEeRLC+L9Ox5i3Q9LACaCeK0kXoSChVcD3C/tQPHpAEsquQ73IkUcEz2kcLCknyGW5MGjkljRFVL8xJOKyi4CwCOuQAeAkfTP1+tNxLkogvgEbDgffkJqKqvuMA5ifOpqg/5qWecRstNg7xoUTI1Fovdxg8oy2s5AP8CGeYHmGngeZaOL4I4LXLcpHg4149/GDz4xqgsb+UAbMKKUpkrqHA43MUyyJpWUK0EHeG2YKRXr7tB+QMcgGewLD+ebTDbtrtbBt7UPlhS4rV4IvcDI7J8P1OeA/AcAI7LHljN7aB8XTowJmZt9EFRD/o0SDMH4HlwMhMyDWZZSAHFf3YDs3RS49WDLuaAY3IJq+qzmQKLxXAZKN7oDoYbdV3v5elPqiSpMyiOuAEVZVqHXb1OhloUH+MA+ztO0cAO/RkrfyBE7OAEbAZvO8vzVtTRWFD6DAfY5biBM3PWiaL0a4lvXICwnV8WjmE6ntYmhqX2jjp5LbMZjCw/wbYeN6CizOa2GMVzQOlmHjB4Ceuyk6LJ8huccEmR5Xddg7OOV/NAtchW+E3XbOag60QA4Qwuarca0bRuEJyr+cFQwzcY98huxhAKdQelt4kAQpj4qJ3gvFXAYn+aJumXk1yPlpQUgtIHhbYoFMUstNRRWgjnpl4A7IKlayNymqFHFaWCpV9CFry3LGxR1CgA5kB5M8OX2goApwpaz6mdOMGxtAgXWJySxb4WuQD4qTDgU+N5AAnzpr7ChSWpCyisiQJqY0Y7FtmSKpbV23b45kC0KHBxcQ9QeI8w4KgnHRPVtIU7rOtbioLVg5Hl/qDwSVFAMqLSMSObroCdZYlzIJtMRFVHCaRo/wFWPgaAXzdbBpkc2A4aKzCNd97+URQuESYGDDhIVfWOQIKZJu4D2+oXlgDTV1865gUQZDts756BArMNMoR1oa46BYqbyPixZz1ZUFV3sgwoGBajuBKATl3btIn8QYYMuezRgrsiRUWyr2BxA40EkPMpA/Hm6gbUu7fjEXA3azP6AsbKD9bxdUuhjM9W7fII52BF+daRpE4+WA3P501+jbfmHvQKyFqMuXf7Ot4mkN2fr50y+bRH61X7AXdUpHSxaPQ4GVbR5AGw3g+434XgQGKfr72I+vQRhfsu92dOx7WicInzt3CBg1RVpMm0NveWo2SqFzgmdNZMbriILD+S+zoueWf2vSdAipzacWN5nMl6XxNlUHa/J8DoJodUDE0HR8Ll5V0lPxcrLEHZPV4AzS83OLis7FowVa3RSku7BSNxJqQAlN3hBTC2apmDSkpaw22wJemGQFUG7J4MlP3JC6A+f96V7vRyX9It3nzT/GrjIU8edM7rMSnIi10f476lzbE1K7yEiEuWro0OJBguLCwDuFOJc1Na6sRWL/cCeMIwUN9ggSVbe3v/5/EgzTKWLvEAiBrYRUkgwNI2ZaFQNT75UDxEUEx97zYnzpmiLEmbaYCbNxYtFAb0/Z4AztgUrhyxuNgxPnhfHFDHz/vTgFWUQZxTRkkJhQ6YNdVUEPAfO6ZV5BRss6LcCVb7VaAma9giy0XJZBt9IQh42NY0NSdgbLIPlLUF6rEdrdt0CUCK1wsCbkcI3ZSLc7ZSwGLbmJXbPsNxnE5xilYKAobZ77LpGZ8TAIun+/iCKQoF71IxQDI3K2CCd+ARNvXg9sykBcnHAoCZG4u66hlDoQLe6QV4CRtFSxZQ+D0BwNO2jgdkzoGoah1nj3FVlSR19taTSYxI8QLut23U8dsgzqHulJNCQpcqBnpTALCuQ6NSYLHpmR5i42gZzuIdcrMMvMJbQlxe3jXxyZnLACl7ARm/FjPIDOY8ODtpM71sxwfcZpvBeUzKWmfNINM5AS+wO0Khh7dMqKccu4+qatarZjYAwDlgetzStHtEt+XedsBOQtU9XMrRgjg4KTnc5nr+dmqadit/4C4uLm8DuA9koJTj1TL7fI5nDL+qqoo/FLGAzL7dYT17PzvAcQONYSUQRxW/QMrHZVIyik0ZuQA2mzp+Ji8BW4YM3Mbzm9inaHkJCGfrUZZjujiYailfFwA8DHIy3acwUj4v9vUVa+SmgNsl5fuyDTKovW9/IAmfLV0Pi2UncA515kjYdrwC9i9rpuHiq3JwtAAAAABJRU5ErkJggg=="></a>
+<a style="display:inline-block; margin-left: .5em" href='https://github.com/One-2-3-45/One-2-3-45'><img src='https://img.shields.io/github/stars/One-2-3-45/One-2-3-45?style=social' /></a>
+</div>
+We reconstruct a 3D textured mesh from a single image by initially predicting multi-view images and then lifting them to 3D. 
+'''
+_USER_GUIDE = "Please upload an image in the block above (or choose an example above) and click **Run Generation**." 
+_BBOX_1 = "Predicting bounding box for the input image..."
+_BBOX_2 = "Bounding box adjusted. Continue adjusting or **Run Generation**."
+_BBOX_3 = "Bounding box predicted. Adjust it using sliders or **Run Generation**."
+_SAM = "Preprocessing the input image... (safety check, SAM segmentation, *etc*.)"
+_GEN_1 = "Predicting multi-view images... (may take \~13 seconds) <br> Images will be shown in the bottom right blocks."
+_GEN_2 = "Predicting nearby views and generating mesh... (may take \~33 seconds) <br> Mesh will be shown on the right."
+_DONE = "Done! Mesh is shown on the right. <br> If it is not satisfactory, please select **Retry view** checkboxes for inaccurate views and click **Regenerate selected view(s)** at the bottom."
+_REGEN_1 = "Selected view(s) are regenerated. You can click **Regenerate nearby views and mesh**. <br> Alternatively, if the regenerated view(s) are still not satisfactory, you can repeat the previous step (select the view and regenerate)."
+_REGEN_2 = "Regeneration done. Mesh is shown on the right."
+
+
+def calc_cam_cone_pts_3d(polar_deg, azimuth_deg, radius_m, fov_deg):
+    '''
+    :param polar_deg (float).
+    :param azimuth_deg (float).
+    :param radius_m (float).
+    :param fov_deg (float).
+    :return (5, 3) array of float with (x, y, z).
+    '''
+    polar_rad = np.deg2rad(polar_deg)
+    azimuth_rad = np.deg2rad(azimuth_deg)
+    fov_rad = np.deg2rad(fov_deg)
+    polar_rad = -polar_rad  # NOTE: Inverse of how used_x relates to x.
+
+    # Camera pose center:
+    cam_x = radius_m * np.cos(azimuth_rad) * np.cos(polar_rad)
+    cam_y = radius_m * np.sin(azimuth_rad) * np.cos(polar_rad)
+    cam_z = radius_m * np.sin(polar_rad)
+
+    # Obtain four corners of camera frustum, assuming it is looking at origin.
+    # First, obtain camera extrinsics (rotation matrix only):
+    camera_R = np.array([[np.cos(azimuth_rad) * np.cos(polar_rad),
+                          -np.sin(azimuth_rad),
+                          -np.cos(azimuth_rad) * np.sin(polar_rad)],
+                         [np.sin(azimuth_rad) * np.cos(polar_rad),
+                          np.cos(azimuth_rad),
+                          -np.sin(azimuth_rad) * np.sin(polar_rad)],
+                         [np.sin(polar_rad),
+                          0.0,
+                          np.cos(polar_rad)]])
+
+    # Multiply by corners in camera space to obtain go to space:
+    corn1 = [-1.0, np.tan(fov_rad / 2.0), np.tan(fov_rad / 2.0)]
+    corn2 = [-1.0, -np.tan(fov_rad / 2.0), np.tan(fov_rad / 2.0)]
+    corn3 = [-1.0, -np.tan(fov_rad / 2.0), -np.tan(fov_rad / 2.0)]
+    corn4 = [-1.0, np.tan(fov_rad / 2.0), -np.tan(fov_rad / 2.0)]
+    corn1 = np.dot(camera_R, corn1)
+    corn2 = np.dot(camera_R, corn2)
+    corn3 = np.dot(camera_R, corn3)
+    corn4 = np.dot(camera_R, corn4)
+
+    # Now attach as offset to actual 3D camera position:
+    corn1 = np.array(corn1) / np.linalg.norm(corn1, ord=2)
+    corn_x1 = cam_x + corn1[0]
+    corn_y1 = cam_y + corn1[1]
+    corn_z1 = cam_z + corn1[2]
+    corn2 = np.array(corn2) / np.linalg.norm(corn2, ord=2)
+    corn_x2 = cam_x + corn2[0]
+    corn_y2 = cam_y + corn2[1]
+    corn_z2 = cam_z + corn2[2]
+    corn3 = np.array(corn3) / np.linalg.norm(corn3, ord=2)
+    corn_x3 = cam_x + corn3[0]
+    corn_y3 = cam_y + corn3[1]
+    corn_z3 = cam_z + corn3[2]
+    corn4 = np.array(corn4) / np.linalg.norm(corn4, ord=2)
+    corn_x4 = cam_x + corn4[0]
+    corn_y4 = cam_y + corn4[1]
+    corn_z4 = cam_z + corn4[2]
+
+    xs = [cam_x, corn_x1, corn_x2, corn_x3, corn_x4]
+    ys = [cam_y, corn_y1, corn_y2, corn_y3, corn_y4]
+    zs = [cam_z, corn_z1, corn_z2, corn_z3, corn_z4]
+
+    return np.array([xs, ys, zs]).T
+
+class CameraVisualizer:
+    def __init__(self, gradio_plot):
+        self._gradio_plot = gradio_plot
+        self._fig = None
+        self._polar = 0.0
+        self._azimuth = 0.0
+        self._radius = 0.0
+        self._raw_image = None
+        self._8bit_image = None
+        self._image_colorscale = None
+
+    def encode_image(self, raw_image, elev=90):
+        '''
+        :param raw_image (H, W, 3) array of uint8 in [0, 255].
+        '''
+        # https://stackoverflow.com/questions/60685749/python-plotly-how-to-add-an-image-to-a-3d-scatter-plot
+
+        dum_img = Image.fromarray(np.ones((3, 3, 3), dtype='uint8')).convert('P', palette='WEB')
+        idx_to_color = np.array(dum_img.getpalette()).reshape((-1, 3))
+
+        self._raw_image = raw_image
+        self._8bit_image = Image.fromarray(raw_image).convert('P', palette='WEB', dither=None)
+        # self._8bit_image = Image.fromarray(raw_image.clip(0, 254)).convert(
+        #     'P', palette='WEB', dither=None)
+        self._image_colorscale = [
+            [i / 255.0, 'rgb({}, {}, {})'.format(*rgb)] for i, rgb in enumerate(idx_to_color)]
+        self._elev = elev
+        # return self.update_figure()
+
+    def update_figure(self):
+        fig = go.Figure()
+
+        if self._raw_image is not None:
+            (H, W, C) = self._raw_image.shape
+
+            x = np.zeros((H, W))
+            (y, z) = np.meshgrid(np.linspace(-1.0, 1.0, W), np.linspace(1.0, -1.0, H) * H / W)
+            
+            angle_deg = self._elev-90
+            angle = np.radians(90-self._elev)
+            rotation_matrix = np.array([
+                [np.cos(angle), 0, np.sin(angle)],
+                [0, 1, 0],
+                [-np.sin(angle), 0, np.cos(angle)]
+            ])
+            # Assuming x, y, z are the original 3D coordinates of the image
+            coordinates = np.stack((x, y, z), axis=-1)  # Combine x, y, z into a single array
+            # Apply the rotation matrix
+            rotated_coordinates = np.matmul(coordinates, rotation_matrix)
+            # Extract the new x, y, z coordinates from the rotated coordinates
+            x, y, z = rotated_coordinates[..., 0], rotated_coordinates[..., 1], rotated_coordinates[..., 2]
+
+            fig.add_trace(go.Surface(
+                x=x, y=y, z=z,
+                surfacecolor=self._8bit_image,
+                cmin=0,
+                cmax=255,
+                colorscale=self._image_colorscale,
+                showscale=False,
+                lighting_diffuse=1.0,
+                lighting_ambient=1.0,
+                lighting_fresnel=1.0,
+                lighting_roughness=1.0,
+                lighting_specular=0.3))
+
+            scene_bounds = 3.5
+            base_radius = 2.5
+            zoom_scale = 1.5  # Note that input radius offset is in [-0.5, 0.5].
+            fov_deg = 50.0
+            edges = [(0, 1), (0, 2), (0, 3), (0, 4), (1, 2), (2, 3), (3, 4), (4, 1)]
+
+            input_cone = calc_cam_cone_pts_3d(
+                angle_deg, 0.0, base_radius, fov_deg)  # (5, 3).
+            output_cone = calc_cam_cone_pts_3d(
+                self._polar, self._azimuth, base_radius + self._radius * zoom_scale, fov_deg)  # (5, 3).
+            output_cones = []
+            for i in range(1,4):
+                output_cones.append(calc_cam_cone_pts_3d(
+                    angle_deg, i*90, base_radius + self._radius * zoom_scale, fov_deg))
+            delta_deg = 30 if angle_deg <= -15 else -30
+            for i in range(4):
+                output_cones.append(calc_cam_cone_pts_3d(
+                    angle_deg+delta_deg, 30+i*90, base_radius + self._radius * zoom_scale, fov_deg))
+
+            cones = [(input_cone, 'rgb(174, 54, 75)', 'Input view (Predicted view 1)')]
+            for i in range(len(output_cones)):
+                cones.append((output_cones[i], 'rgb(32, 77, 125)', f'Predicted view {i+2}'))
+
+            for idx, (cone, clr, legend) in enumerate(cones):
+
+                for (i, edge) in enumerate(edges):
+                    (x1, x2) = (cone[edge[0], 0], cone[edge[1], 0])
+                    (y1, y2) = (cone[edge[0], 1], cone[edge[1], 1])
+                    (z1, z2) = (cone[edge[0], 2], cone[edge[1], 2])
+                    fig.add_trace(go.Scatter3d(
+                        x=[x1, x2], y=[y1, y2], z=[z1, z2], mode='lines',
+                        line=dict(color=clr, width=3),
+                        name=legend, showlegend=(i == 1) and (idx <= 1)))
+
+                # Add label.
+                if cone[0, 2] <= base_radius / 2.0:
+                    fig.add_trace(go.Scatter3d(
+                        x=[cone[0, 0]], y=[cone[0, 1]], z=[cone[0, 2] - 0.05], showlegend=False,
+                        mode='text', text=legend, textposition='bottom center'))
+                else:
+                    fig.add_trace(go.Scatter3d(
+                        x=[cone[0, 0]], y=[cone[0, 1]], z=[cone[0, 2] + 0.05], showlegend=False,
+                        mode='text', text=legend, textposition='top center'))
+
+            # look at center of scene
+            fig.update_layout(
+                # width=640,
+                # height=480,
+                # height=400,
+                height=450,
+                autosize=True,
+                hovermode=False,
+                margin=go.layout.Margin(l=0, r=0, b=0, t=0),
+                showlegend=False,
+                legend=dict(
+                    yanchor='bottom',
+                    y=0.01,
+                    xanchor='right',
+                    x=0.99,
+                ),
+                scene=dict(
+                    aspectmode='manual',
+                    aspectratio=dict(x=1, y=1, z=1.0),
+                    camera=dict(
+                        eye=dict(x=base_radius - 1.6, y=0.0, z=0.6),
+                        center=dict(x=0.0, y=0.0, z=0.0),
+                        up=dict(x=0.0, y=0.0, z=1.0)),
+                    xaxis_title='',
+                    yaxis_title='',
+                    zaxis_title='',
+                    xaxis=dict(
+                        range=[-scene_bounds, scene_bounds],
+                        showticklabels=False,
+                        showgrid=True,
+                        zeroline=False,
+                        showbackground=True,
+                        showspikes=False,
+                        showline=False,
+                        ticks=''),
+                    yaxis=dict(
+                        range=[-scene_bounds, scene_bounds],
+                        showticklabels=False,
+                        showgrid=True,
+                        zeroline=False,
+                        showbackground=True,
+                        showspikes=False,
+                        showline=False,
+                        ticks=''),
+                    zaxis=dict(
+                        range=[-scene_bounds, scene_bounds],
+                        showticklabels=False,
+                        showgrid=True,
+                        zeroline=False,
+                        showbackground=True,
+                        showspikes=False,
+                        showline=False,
+                        ticks='')))
+
+        self._fig = fig
+        return fig
+    
+
+def stage1_run(models, device, cam_vis, tmp_dir,
+               input_im, scale, ddim_steps, elev=None, rerun_all=[],
+               *btn_retrys):
+    is_rerun = True if cam_vis is None else False
+    model = models['turncam']
+
+    stage1_dir = os.path.join(tmp_dir, "stage1_8")
+    if not is_rerun:
+        os.makedirs(stage1_dir, exist_ok=True)
+        output_ims = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=list(range(4)), device=device, ddim_steps=ddim_steps, scale=scale)
+        stage2_steps = 50 # ddim_steps
+        zero123_infer(model, tmp_dir, indices=[0], device=device, ddim_steps=stage2_steps, scale=scale)
+        try:
+            elev_output = estimate_elev(tmp_dir)
+        except:
+            print("Failed to estimate polar angle")
+            elev_output = 90
+        print("Estimated polar angle:", elev_output)
+        gen_poses(tmp_dir, elev_output)
+        show_in_im1 = np.asarray(input_im, dtype=np.uint8)
+        cam_vis.encode_image(show_in_im1, elev=elev_output)
+        new_fig = cam_vis.update_figure()
+
+        flag_lower_cam = elev_output <= 75
+        if flag_lower_cam:
+            output_ims_2 = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=list(range(4,8)), device=device, ddim_steps=ddim_steps, scale=scale)
+        else:
+            output_ims_2 = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=list(range(8,12)), device=device, ddim_steps=ddim_steps, scale=scale)
+        torch.cuda.empty_cache()
+        return (90-elev_output, new_fig, *output_ims, *output_ims_2)
+    else:
+        rerun_idx = [i for i in range(len(btn_retrys)) if btn_retrys[i]]
+        if 90-int(elev["label"]) > 75:
+            rerun_idx_in = [i if i < 4 else i+4 for i in rerun_idx]
+        else:
+            rerun_idx_in = rerun_idx
+        for idx in rerun_idx_in:
+            if idx not in rerun_all:
+                rerun_all.append(idx)
+        print("rerun_idx", rerun_all)
+        output_ims = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=rerun_idx_in, device=device, ddim_steps=ddim_steps, scale=scale)
+        outputs = [gr.update(visible=True)] * 8
+        for idx, view_idx in enumerate(rerun_idx):
+            outputs[view_idx] = output_ims[idx]
+        reset = [gr.update(value=False)] * 8
+        torch.cuda.empty_cache()
+        return (rerun_all, *reset, *outputs)
+    
+def stage2_run(models, device, tmp_dir,
+               elev, scale, is_glb=False, rerun_all=[], stage2_steps=50):
+    flag_lower_cam = 90-int(elev["label"]) <= 75
+    is_rerun = True if rerun_all else False
+    model = models['turncam']
+    if not is_rerun:
+        if flag_lower_cam:
+            zero123_infer(model, tmp_dir, indices=list(range(1,8)), device=device, ddim_steps=stage2_steps, scale=scale)
+        else:
+            zero123_infer(model, tmp_dir, indices=list(range(1,4))+list(range(8,12)), device=device, ddim_steps=stage2_steps, scale=scale)
+    else:
+        print("rerun_idx", rerun_all)
+        zero123_infer(model, tmp_dir, indices=rerun_all, device=device, ddim_steps=stage2_steps, scale=scale)
+    
+    dataset = tmp_dir
+    main_dir_path = os.path.dirname(__file__)
+    torch.cuda.empty_cache()
+    os.chdir(os.path.join(code_dir, 'reconstruction/'))
+
+    bash_script = f'CUDA_VISIBLE_DEVICES={_GPU_INDEX} python exp_runner_generic_blender_val.py \
+                    --specific_dataset_name {dataset} \
+                    --mode export_mesh \
+                    --conf confs/one2345_lod0_val_demo.conf \
+                    --resolution {_MESH_RESOLUTION}'
+    print(bash_script)
+    os.system(bash_script)
+    os.chdir(main_dir_path)
+
+    ply_path = os.path.join(tmp_dir, f"mesh.ply")
+    mesh_ext = ".glb" if is_glb else ".obj"
+    mesh_path = os.path.join(tmp_dir, f"mesh{mesh_ext}")
+    # Read the textured mesh from .ply file
+    mesh = trimesh.load_mesh(ply_path)
+    rotation_matrix = trimesh.transformations.rotation_matrix(np.pi/2, [1, 0, 0])
+    mesh.apply_transform(rotation_matrix)
+    rotation_matrix = trimesh.transformations.rotation_matrix(np.pi, [0, 0, 1])
+    mesh.apply_transform(rotation_matrix)
+    # flip x
+    mesh.vertices[:, 0] = -mesh.vertices[:, 0]
+    mesh.faces = np.fliplr(mesh.faces)
+    # Export the mesh as .obj file with colors
+    if not is_glb:
+        mesh.export(mesh_path, file_type='obj', include_color=True)
+    else:
+        mesh.export(mesh_path, file_type='glb')
+    torch.cuda.empty_cache()
+
+    if not is_rerun:
+        return (mesh_path)
+    else:
+        return (mesh_path, gr.update(value=[]), gr.update(visible=False), gr.update(visible=False))
+
+def nsfw_check(models, raw_im, device='cuda'):
+    safety_checker_input = models['clip_fe'](raw_im, return_tensors='pt').to(device)
+    (_, has_nsfw_concept) = models['nsfw'](
+        images=np.ones((1, 3)), clip_input=safety_checker_input.pixel_values)
+    del safety_checker_input
+    if np.any(has_nsfw_concept):
+        print('NSFW content detected.')
+        return Image.open("unsafe.png")
+    else:
+        print('Safety check passed.')
+        return False
+
+def preprocess_run(predictor, models, raw_im, lower_contrast, *bbox_sliders):
+    raw_im.thumbnail([512, 512], Image.Resampling.LANCZOS)
+    check_results = nsfw_check(models, raw_im, device=predictor.device)
+    if check_results:
+        return check_results
+    image_sam = sam_out_nosave(predictor, raw_im.convert("RGB"), *bbox_sliders)
+    input_256 = image_preprocess_nosave(image_sam, lower_contrast=lower_contrast, rescale=True)
+    torch.cuda.empty_cache()
+    return input_256
+
+def on_coords_slider(image, x_min, y_min, x_max, y_max, color=(88, 191, 131, 255)):
+    """Draw a bounding box annotation for an image."""
+    print("Slider adjusted, drawing bbox...")
+    image.thumbnail([512, 512], Image.Resampling.LANCZOS)
+    image_size = image.size
+    if max(image_size) > 224:
+        image.thumbnail([224, 224], Image.Resampling.LANCZOS)
+        shrink_ratio = max(image.size) / max(image_size)
+        x_min = int(x_min * shrink_ratio)
+        y_min = int(y_min * shrink_ratio)
+        x_max = int(x_max * shrink_ratio)
+        y_max = int(y_max * shrink_ratio)
+    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGBA2BGRA)
+    image = cv2.rectangle(image, (x_min, y_min), (x_max, y_max), color, int(max(max(image.shape) / 400*2, 2)))
+    return cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA) # image[:, :, ::-1]
+
+def init_bbox(image):
+    image.thumbnail([512, 512], Image.Resampling.LANCZOS)
+    width, height = image.size
+    image_rem = image.convert('RGBA')
+    image_nobg = remove(image_rem, alpha_matting=True)
+    arr = np.asarray(image_nobg)[:,:,-1]
+    x_nonzero = np.nonzero(arr.sum(axis=0))
+    y_nonzero = np.nonzero(arr.sum(axis=1))
+    x_min = int(x_nonzero[0].min())
+    y_min = int(y_nonzero[0].min())
+    x_max = int(x_nonzero[0].max())
+    y_max = int(y_nonzero[0].max())
+    image_mini = image.copy()
+    image_mini.thumbnail([224, 224], Image.Resampling.LANCZOS)
+    shrink_ratio = max(image_mini.size) / max(width, height)
+    x_min_shrink = int(x_min * shrink_ratio)
+    y_min_shrink = int(y_min * shrink_ratio)
+    x_max_shrink = int(x_max * shrink_ratio)
+    y_max_shrink = int(y_max * shrink_ratio)
+
+    return [on_coords_slider(image_mini, x_min_shrink, y_min_shrink, x_max_shrink, y_max_shrink),
+            gr.update(value=x_min, maximum=width),
+            gr.update(value=y_min, maximum=height),
+            gr.update(value=x_max, maximum=width),
+            gr.update(value=y_max, maximum=height)]
+
+
+def run_demo(
+        device_idx=_GPU_INDEX,
+        ckpt='zero123-xl.ckpt'):
+
+    device = f"cuda:{device_idx}" if torch.cuda.is_available() else "cpu"
+    models = init_model(device, os.path.join(code_dir, 'zero123-xl.ckpt'), half_precision=_HALF_PRECISION)
+
+    # init sam model
+    predictor = sam_init(device_idx)
+
+    with open('instructions_12345.md', 'r') as f:
+        article = f.read()
+
+    # NOTE: Examples must match inputs
+    example_folder = os.path.join(os.path.dirname(__file__), 'demo_examples')
+    example_fns = os.listdir(example_folder)
+    example_fns.sort()
+    examples_full = [os.path.join(example_folder, x) for x in example_fns if x.endswith('.png')]
+
+    # Compose demo layout & data flow.
+    with gr.Blocks(title=_TITLE, css="style.css") as demo:
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown('# ' + _TITLE)
+            with gr.Column(scale=0):
+                gr.DuplicateButton(value='Duplicate Space for private use',
+                            elem_id='duplicate-button')
+        gr.Markdown(_DESCRIPTION)
+
+        with gr.Row(variant='panel'):
+            with gr.Column(scale=1.2):
+                image_block = gr.Image(type='pil', image_mode='RGBA', height=290, label='Input image', tool=None)
+
+                gr.Examples(
+                    examples=examples_full,  # NOTE: elements must match inputs list!
+                    inputs=[image_block],
+                    outputs=[image_block],
+                    cache_examples=False,
+                    label='Examples (click one of the images below to start)',
+                    examples_per_page=40
+                )
+                preprocess_chk = gr.Checkbox(
+                        False, label='Reduce image contrast (mitigate shadows on the backside)')
+                with gr.Accordion('Advanced options', open=False):
+                    scale_slider = gr.Slider(0, 30, value=3, step=1,
+                                             label='Diffusion guidance scale')
+                    steps_slider = gr.Slider(5, 200, value=75, step=5,
+                                             label='Number of diffusion inference steps')
+                    glb_chk = gr.Checkbox(
+                        False, label='Export the mesh in .glb format')
+
+                run_btn = gr.Button('Run Generation', variant='primary', interactive=False)
+                guide_text = gr.Markdown(_USER_GUIDE, visible=True)
+                
+            with gr.Column(scale=.8):
+                with gr.Row():
+                    bbox_block = gr.Image(type='pil', label="Bounding box", height=290, interactive=False)
+                    sam_block = gr.Image(type='pil', label="SAM output", interactive=False)
+                max_width = max_height = 256
+                with gr.Row():
+                    x_min_slider = gr.Slider(label="X min", interactive=True, value=0, minimum=0, maximum=max_width, step=1)
+                    y_min_slider = gr.Slider(label="Y min", interactive=True, value=0, minimum=0, maximum=max_height, step=1)
+                with gr.Row():
+                    x_max_slider = gr.Slider(label="X max", interactive=True, value=max_width, minimum=0, maximum=max_width, step=1)
+                    y_max_slider = gr.Slider(label="Y max", interactive=True, value=max_height, minimum=0, maximum=max_height, step=1)
+                bbox_sliders = [x_min_slider, y_min_slider, x_max_slider, y_max_slider]
+
+                mesh_output = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="One-2-3-45's Textured Mesh", elem_id="model-3d-out")
+        
+        with gr.Row(variant='panel'):
+            with gr.Column(scale=0.85):
+                elev_output = gr.Label(label='Estimated elevation (degree, w.r.t. the horizontal plane)')
+                vis_output = gr.Plot(label='Camera poses of the input view (red) and predicted views (blue)', elem_id="plot-out")
+                
+            with gr.Column(scale=1.15):
+                gr.Markdown('Predicted multi-view images')
+                with gr.Row():
+                    view_1 = gr.Image(interactive=False, height=200, show_label=False)
+                    view_2 = gr.Image(interactive=False, height=200, show_label=False)
+                    view_3 = gr.Image(interactive=False, height=200, show_label=False)
+                    view_4 = gr.Image(interactive=False, height=200, show_label=False)
+                with gr.Row():
+                    btn_retry_1 = gr.Checkbox(label='Retry view 1')
+                    btn_retry_2 = gr.Checkbox(label='Retry view 2')
+                    btn_retry_3 = gr.Checkbox(label='Retry view 3')
+                    btn_retry_4 = gr.Checkbox(label='Retry view 4')
+                with gr.Row():
+                    view_5 = gr.Image(interactive=False, height=200, show_label=False)
+                    view_6 = gr.Image(interactive=False, height=200, show_label=False)
+                    view_7 = gr.Image(interactive=False, height=200, show_label=False)
+                    view_8 = gr.Image(interactive=False, height=200, show_label=False)
+                with gr.Row():
+                    btn_retry_5 = gr.Checkbox(label='Retry view 5')
+                    btn_retry_6 = gr.Checkbox(label='Retry view 6')
+                    btn_retry_7 = gr.Checkbox(label='Retry view 7')
+                    btn_retry_8 = gr.Checkbox(label='Retry view 8')
+                with gr.Row():
+                    regen_view_btn = gr.Button('1. Regenerate selected view(s)', variant='secondary', visible=False)
+                    regen_mesh_btn = gr.Button('2. Regenerate nearby views and mesh', variant='secondary', visible=False)
+        
+        gr.Markdown(article)
+        gr.HTML("""
+            <div class="footer">
+                <p> 
+                One-2-3-45 Demo by <a style="text-decoration:none" href="https://chaoxu.xyz" target="_blank">Chao Xu</a>
+                </p>
+            </div>
+        """)
+
+        update_guide = lambda GUIDE_TEXT: gr.update(value=GUIDE_TEXT)
+
+        views = [view_1, view_2, view_3, view_4, view_5, view_6, view_7, view_8]
+        btn_retrys = [btn_retry_1, btn_retry_2, btn_retry_3, btn_retry_4, btn_retry_5, btn_retry_6, btn_retry_7, btn_retry_8]
+        
+        rerun_idx = gr.State([])
+        tmp_dir = gr.State('./demo_tmp/tmp_dir')
+
+        def refresh(tmp_dir):
+            if os.path.exists(tmp_dir):
+                shutil.rmtree(tmp_dir)
+            tmp_dir = tempfile.TemporaryDirectory(dir=os.path.join(os.path.dirname(__file__), 'demo_tmp'))
+            print("create tmp_dir", tmp_dir.name)
+            clear = [gr.update(value=[])] + [None] * 5 + [gr.update(visible=False)] * 2 + [None] * 8 + [gr.update(value=False)] * 8
+            return (tmp_dir.name, *clear)
+        
+        placeholder = gr.Image(visible=False)
+        tmp_func = lambda x: False if not x else gr.update(visible=False)
+        disable_func = lambda x: gr.update(interactive=False)
+        enable_func = lambda x: gr.update(interactive=True)
+        image_block.change(disable_func, inputs=run_btn, outputs=run_btn, queue=False
+                           ).success(fn=refresh,
+                           inputs=[tmp_dir],
+                           outputs=[tmp_dir, rerun_idx, bbox_block, sam_block, elev_output, vis_output, mesh_output, regen_view_btn, regen_mesh_btn, *views, *btn_retrys],
+                           queue=False
+                           ).success(fn=tmp_func, inputs=[image_block], outputs=[placeholder], queue=False
+                           ).success(fn=partial(update_guide, _BBOX_1), outputs=[guide_text], queue=False
+                           ).success(fn=init_bbox,
+                                     inputs=[image_block],
+                                     outputs=[bbox_block, *bbox_sliders], queue=False
+                           ).success(fn=partial(update_guide, _BBOX_3), outputs=[guide_text], queue=False
+                           ).success(enable_func, inputs=run_btn, outputs=run_btn, queue=False)
+
+
+        for bbox_slider in bbox_sliders:
+            bbox_slider.release(fn=on_coords_slider,
+                               inputs=[image_block, *bbox_sliders],
+                               outputs=[bbox_block], 
+                               queue=False
+                               ).success(fn=partial(update_guide, _BBOX_2), outputs=[guide_text], queue=False)
+
+        cam_vis = CameraVisualizer(vis_output)
+
+        # Define the function to be called when any of the btn_retry buttons are clicked
+        def on_retry_button_click(*btn_retrys):
+            any_checked = any([btn_retry for btn_retry in btn_retrys])
+            print('any_checked:', any_checked, [btn_retry for btn_retry in btn_retrys])
+            if any_checked:
+                return (gr.update(visible=True), gr.update(visible=True))
+            else:
+                return (gr.update(), gr.update())
+        # make regen_btn visible when any of the btn_retry is checked
+        for btn_retry in btn_retrys:
+            # Add the event handlers to the btn_retry buttons
+            btn_retry.change(fn=on_retry_button_click, inputs=[*btn_retrys], outputs=[regen_view_btn, regen_mesh_btn], queue=False)
+
+
+        run_btn.click(fn=partial(update_guide, _SAM), outputs=[guide_text], queue=False
+                      ).success(fn=partial(preprocess_run, predictor, models), 
+                                inputs=[image_block, preprocess_chk, *bbox_sliders], 
+                                outputs=[sam_block]
+                      ).success(fn=partial(update_guide, _GEN_1), outputs=[guide_text], queue=False
+                      ).success(fn=partial(stage1_run, models, device, cam_vis),
+                                inputs=[tmp_dir, sam_block, scale_slider, steps_slider],
+                                outputs=[elev_output, vis_output, *views]
+                      ).success(fn=partial(update_guide, _GEN_2), outputs=[guide_text], queue=False
+                      ).success(fn=partial(stage2_run, models, device),
+                                inputs=[tmp_dir, elev_output, scale_slider, glb_chk],
+                                outputs=[mesh_output]
+                      ).success(fn=partial(update_guide, _DONE), outputs=[guide_text], queue=False)
+    
+
+        regen_view_btn.click(fn=partial(stage1_run, models, device, None),
+                             inputs=[tmp_dir, sam_block, scale_slider, steps_slider, elev_output, rerun_idx, *btn_retrys],
+                             outputs=[rerun_idx, *btn_retrys, *views]
+                            ).success(fn=partial(update_guide, _REGEN_1), outputs=[guide_text], queue=False)
+        regen_mesh_btn.click(fn=partial(stage2_run, models, device),
+                             inputs=[tmp_dir, elev_output, scale_slider, glb_chk, rerun_idx],
+                             outputs=[mesh_output, rerun_idx, regen_view_btn, regen_mesh_btn]
+                            ).success(fn=partial(update_guide, _REGEN_2), outputs=[guide_text], queue=False)
+
+
+    demo.queue().launch(share=True, max_threads=80) # auth=("admin", os.environ['PASSWD'])
+
+
+if __name__ == '__main__':
+    fire.Fire(run_demo)

One-2-3-45-master 2/demo/demo_tmp/.gitignore

@@ -0,0 +1 @@
+tmp*

One-2-3-45-master 2/demo/instructions_12345.md

@@ -0,0 +1,10 @@
+## Tuning Tips:
+
+1. The multi-view prediction module (Zero123) operates probabilistically. If some of the predicted views are not satisfactory, you may select and regenerate them.
+
+2. In “advanced options”, you can tune two parameters as in other common diffusion models:
+  - Diffusion Guidance Scale determines how much you want the model to respect the input information (input image + viewpoints). Increasing the scale typically results in better adherence, less diversity, and also higher image distortion.
+  
+  - Number of diffusion inference steps controls the number of diffusion steps applied to generate each image. Generally, a higher value yields better results but with diminishing returns.
+
+Enjoy creating your 3D asset!

One-2-3-45-master 2/demo/memora/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

One-2-3-45-master 2/demo/memora/README.md

@@ -0,0 +1,12 @@
+---
+title: Memora
+emoji: 🐨
+colorFrom: purple
+colorTo: green
+sdk: gradio
+sdk_version: 3.47.1
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

One-2-3-45-master 2/demo/style.css

@@ -0,0 +1,33 @@
+#model-3d-out {
+    height: 400px;
+} 
+
+#plot-out {
+    height: 450px;
+}
+
+#duplicate-button {
+    margin-left: auto;
+    color: #fff;
+    background: #1565c0;
+  }
+
+.footer {
+    margin-bottom: 45px;
+    margin-top: 10px;
+    text-align: center;
+    border-bottom: 1px solid #e5e5e5;
+}
+.footer>p {
+    font-size: .8rem;
+    display: inline-block;
+    padding: 0 10px;
+    transform: translateY(10px);
+    background: white;
+}
+.dark .footer {
+    border-color: #303030;
+}
+.dark .footer>p {
+    background: #0b0f19;
+}

One-2-3-45-master 2/download_ckpt.py

@@ -0,0 +1,30 @@
+import urllib.request
+from tqdm import tqdm
+
+def download_checkpoint(url, save_path):
+    try:
+        with urllib.request.urlopen(url) as response, open(save_path, 'wb') as file:
+            file_size = int(response.info().get('Content-Length', -1))
+            chunk_size = 8192
+            num_chunks = file_size // chunk_size if file_size > chunk_size else 1
+
+            with tqdm(total=file_size, unit='B', unit_scale=True, desc='Downloading', ncols=100) as pbar:
+                for chunk in iter(lambda: response.read(chunk_size), b''):
+                    file.write(chunk)
+                    pbar.update(len(chunk))
+        
+        print(f"Checkpoint downloaded and saved to: {save_path}")
+    except Exception as e:
+        print(f"Error downloading checkpoint: {e}")
+
+if __name__ == "__main__":
+    ckpts = {
+        "sam_vit_h_4b8939.pth": "https://huggingface.co/One-2-3-45/code/resolve/main/sam_vit_h_4b8939.pth",
+        "zero123-xl.ckpt": "https://huggingface.co/One-2-3-45/code/resolve/main/zero123-xl.ckpt",
+        "elevation_estimate/utils/weights/indoor_ds_new.ckpt" : "https://huggingface.co/One-2-3-45/code/resolve/main/one2345_elev_est/tools/weights/indoor_ds_new.ckpt",
+        "reconstruction/exp/lod0/checkpoints/ckpt_215000.pth": "https://huggingface.co/One-2-3-45/code/resolve/main/SparseNeuS_demo_v1/exp/lod0/checkpoints/ckpt_215000.pth"
+    }
+    for ckpt_name, ckpt_url in ckpts.items():
+        print(f"Downloading checkpoint: {ckpt_name}")
+        download_checkpoint(ckpt_url, ckpt_name)
+

One-2-3-45-master 2/elevation_estimate/.gitignore

@@ -0,0 +1,3 @@
+build/
+.idea/
+*.egg-info/

One-2-3-45-master 2/elevation_estimate/estimate_wild_imgs.py

@@ -0,0 +1,10 @@
+import os.path as osp
+from .utils.elev_est_api import elev_est_api
+
+def estimate_elev(root_dir):
+    img_dir = osp.join(root_dir, "stage2_8")
+    img_paths = []
+    for i in range(4):
+        img_paths.append(f"{img_dir}/0_{i}.png")
+    elev = elev_est_api(img_paths)
+    return elev

One-2-3-45-master 2/elevation_estimate/loftr/__init__.py

@@ -0,0 +1,2 @@
+from .loftr import LoFTR
+from .utils.cvpr_ds_config import default_cfg

One-2-3-45-master 2/elevation_estimate/loftr/backbone/__init__.py

@@ -0,0 +1,11 @@
+from .resnet_fpn import ResNetFPN_8_2, ResNetFPN_16_4
+
+
+def build_backbone(config):
+    if config['backbone_type'] == 'ResNetFPN':
+        if config['resolution'] == (8, 2):
+            return ResNetFPN_8_2(config['resnetfpn'])
+        elif config['resolution'] == (16, 4):
+            return ResNetFPN_16_4(config['resnetfpn'])
+    else:
+        raise ValueError(f"LOFTR.BACKBONE_TYPE {config['backbone_type']} not supported.")

One-2-3-45-master 2/elevation_estimate/loftr/backbone/resnet_fpn.py

@@ -0,0 +1,199 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution without padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False)
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_planes, planes, stride=1):
+        super().__init__()
+        self.conv1 = conv3x3(in_planes, planes, stride)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+
+        if stride == 1:
+            self.downsample = None
+        else:
+            self.downsample = nn.Sequential(
+                conv1x1(in_planes, planes, stride=stride),
+                nn.BatchNorm2d(planes)
+            )
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.bn1(self.conv1(y)))
+        y = self.bn2(self.conv2(y))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x+y)
+
+
+class ResNetFPN_8_2(nn.Module):
+    """
+    ResNet+FPN, output resolution are 1/8 and 1/2.
+    Each block has 2 layers.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        # Config
+        block = BasicBlock
+        initial_dim = config['initial_dim']
+        block_dims = config['block_dims']
+
+        # Class Variable
+        self.in_planes = initial_dim
+
+        # Networks
+        self.conv1 = nn.Conv2d(1, initial_dim, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(initial_dim)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.layer1 = self._make_layer(block, block_dims[0], stride=1)  # 1/2
+        self.layer2 = self._make_layer(block, block_dims[1], stride=2)  # 1/4
+        self.layer3 = self._make_layer(block, block_dims[2], stride=2)  # 1/8
+
+        # 3. FPN upsample
+        self.layer3_outconv = conv1x1(block_dims[2], block_dims[2])
+        self.layer2_outconv = conv1x1(block_dims[1], block_dims[2])
+        self.layer2_outconv2 = nn.Sequential(
+            conv3x3(block_dims[2], block_dims[2]),
+            nn.BatchNorm2d(block_dims[2]),
+            nn.LeakyReLU(),
+            conv3x3(block_dims[2], block_dims[1]),
+        )
+        self.layer1_outconv = conv1x1(block_dims[0], block_dims[1])
+        self.layer1_outconv2 = nn.Sequential(
+            conv3x3(block_dims[1], block_dims[1]),
+            nn.BatchNorm2d(block_dims[1]),
+            nn.LeakyReLU(),
+            conv3x3(block_dims[1], block_dims[0]),
+        )
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, block, dim, stride=1):
+        layer1 = block(self.in_planes, dim, stride=stride)
+        layer2 = block(dim, dim, stride=1)
+        layers = (layer1, layer2)
+
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        # ResNet Backbone
+        x0 = self.relu(self.bn1(self.conv1(x)))
+        x1 = self.layer1(x0)  # 1/2
+        x2 = self.layer2(x1)  # 1/4
+        x3 = self.layer3(x2)  # 1/8
+
+        # FPN
+        x3_out = self.layer3_outconv(x3)
+
+        x3_out_2x = F.interpolate(x3_out, scale_factor=2., mode='bilinear', align_corners=True)
+        x2_out = self.layer2_outconv(x2)
+        x2_out = self.layer2_outconv2(x2_out+x3_out_2x)
+
+        x2_out_2x = F.interpolate(x2_out, scale_factor=2., mode='bilinear', align_corners=True)
+        x1_out = self.layer1_outconv(x1)
+        x1_out = self.layer1_outconv2(x1_out+x2_out_2x)
+
+        return [x3_out, x1_out]
+
+
+class ResNetFPN_16_4(nn.Module):
+    """
+    ResNet+FPN, output resolution are 1/16 and 1/4.
+    Each block has 2 layers.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        # Config
+        block = BasicBlock
+        initial_dim = config['initial_dim']
+        block_dims = config['block_dims']
+
+        # Class Variable
+        self.in_planes = initial_dim
+
+        # Networks
+        self.conv1 = nn.Conv2d(1, initial_dim, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(initial_dim)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.layer1 = self._make_layer(block, block_dims[0], stride=1)  # 1/2
+        self.layer2 = self._make_layer(block, block_dims[1], stride=2)  # 1/4
+        self.layer3 = self._make_layer(block, block_dims[2], stride=2)  # 1/8
+        self.layer4 = self._make_layer(block, block_dims[3], stride=2)  # 1/16
+
+        # 3. FPN upsample
+        self.layer4_outconv = conv1x1(block_dims[3], block_dims[3])
+        self.layer3_outconv = conv1x1(block_dims[2], block_dims[3])
+        self.layer3_outconv2 = nn.Sequential(
+            conv3x3(block_dims[3], block_dims[3]),
+            nn.BatchNorm2d(block_dims[3]),
+            nn.LeakyReLU(),
+            conv3x3(block_dims[3], block_dims[2]),
+        )
+
+        self.layer2_outconv = conv1x1(block_dims[1], block_dims[2])
+        self.layer2_outconv2 = nn.Sequential(
+            conv3x3(block_dims[2], block_dims[2]),
+            nn.BatchNorm2d(block_dims[2]),
+            nn.LeakyReLU(),
+            conv3x3(block_dims[2], block_dims[1]),
+        )
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, block, dim, stride=1):
+        layer1 = block(self.in_planes, dim, stride=stride)
+        layer2 = block(dim, dim, stride=1)
+        layers = (layer1, layer2)
+
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        # ResNet Backbone
+        x0 = self.relu(self.bn1(self.conv1(x)))
+        x1 = self.layer1(x0)  # 1/2
+        x2 = self.layer2(x1)  # 1/4
+        x3 = self.layer3(x2)  # 1/8
+        x4 = self.layer4(x3)  # 1/16
+
+        # FPN
+        x4_out = self.layer4_outconv(x4)
+
+        x4_out_2x = F.interpolate(x4_out, scale_factor=2., mode='bilinear', align_corners=True)
+        x3_out = self.layer3_outconv(x3)
+        x3_out = self.layer3_outconv2(x3_out+x4_out_2x)
+
+        x3_out_2x = F.interpolate(x3_out, scale_factor=2., mode='bilinear', align_corners=True)
+        x2_out = self.layer2_outconv(x2)
+        x2_out = self.layer2_outconv2(x2_out+x3_out_2x)
+
+        return [x4_out, x2_out]

One-2-3-45-master 2/elevation_estimate/loftr/loftr.py

@@ -0,0 +1,81 @@
+import torch
+import torch.nn as nn
+from einops.einops import rearrange
+
+from .backbone import build_backbone
+from .utils.position_encoding import PositionEncodingSine
+from .loftr_module import LocalFeatureTransformer, FinePreprocess
+from .utils.coarse_matching import CoarseMatching
+from .utils.fine_matching import FineMatching
+
+
+class LoFTR(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        # Misc
+        self.config = config
+
+        # Modules
+        self.backbone = build_backbone(config)
+        self.pos_encoding = PositionEncodingSine(
+            config['coarse']['d_model'],
+            temp_bug_fix=config['coarse']['temp_bug_fix'])
+        self.loftr_coarse = LocalFeatureTransformer(config['coarse'])
+        self.coarse_matching = CoarseMatching(config['match_coarse'])
+        self.fine_preprocess = FinePreprocess(config)
+        self.loftr_fine = LocalFeatureTransformer(config["fine"])
+        self.fine_matching = FineMatching()
+
+    def forward(self, data):
+        """ 
+        Update:
+            data (dict): {
+                'image0': (torch.Tensor): (N, 1, H, W)
+                'image1': (torch.Tensor): (N, 1, H, W)
+                'mask0'(optional) : (torch.Tensor): (N, H, W) '0' indicates a padded position
+                'mask1'(optional) : (torch.Tensor): (N, H, W)
+            }
+        """
+        # 1. Local Feature CNN
+        data.update({
+            'bs': data['image0'].size(0),
+            'hw0_i': data['image0'].shape[2:], 'hw1_i': data['image1'].shape[2:]
+        })
+
+        if data['hw0_i'] == data['hw1_i']:  # faster & better BN convergence
+            feats_c, feats_f = self.backbone(torch.cat([data['image0'], data['image1']], dim=0))
+            (feat_c0, feat_c1), (feat_f0, feat_f1) = feats_c.split(data['bs']), feats_f.split(data['bs'])
+        else:  # handle different input shapes
+            (feat_c0, feat_f0), (feat_c1, feat_f1) = self.backbone(data['image0']), self.backbone(data['image1'])
+
+        data.update({
+            'hw0_c': feat_c0.shape[2:], 'hw1_c': feat_c1.shape[2:],
+            'hw0_f': feat_f0.shape[2:], 'hw1_f': feat_f1.shape[2:]
+        })
+
+        # 2. coarse-level loftr module
+        # add featmap with positional encoding, then flatten it to sequence [N, HW, C]
+        feat_c0 = rearrange(self.pos_encoding(feat_c0), 'n c h w -> n (h w) c')
+        feat_c1 = rearrange(self.pos_encoding(feat_c1), 'n c h w -> n (h w) c')
+
+        mask_c0 = mask_c1 = None  # mask is useful in training
+        if 'mask0' in data:
+            mask_c0, mask_c1 = data['mask0'].flatten(-2), data['mask1'].flatten(-2)
+        feat_c0, feat_c1 = self.loftr_coarse(feat_c0, feat_c1, mask_c0, mask_c1)
+
+        # 3. match coarse-level
+        self.coarse_matching(feat_c0, feat_c1, data, mask_c0=mask_c0, mask_c1=mask_c1)
+
+        # 4. fine-level refinement
+        feat_f0_unfold, feat_f1_unfold = self.fine_preprocess(feat_f0, feat_f1, feat_c0, feat_c1, data)
+        if feat_f0_unfold.size(0) != 0:  # at least one coarse level predicted
+            feat_f0_unfold, feat_f1_unfold = self.loftr_fine(feat_f0_unfold, feat_f1_unfold)
+
+        # 5. match fine-level
+        self.fine_matching(feat_f0_unfold, feat_f1_unfold, data)
+
+    def load_state_dict(self, state_dict, *args, **kwargs):
+        for k in list(state_dict.keys()):
+            if k.startswith('matcher.'):
+                state_dict[k.replace('matcher.', '', 1)] = state_dict.pop(k)
+        return super().load_state_dict(state_dict, *args, **kwargs)

One-2-3-45-master 2/elevation_estimate/loftr/loftr_module/__init__.py

@@ -0,0 +1,2 @@
+from .transformer import LocalFeatureTransformer
+from .fine_preprocess import FinePreprocess

One-2-3-45-master 2/elevation_estimate/loftr/loftr_module/fine_preprocess.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops.einops import rearrange, repeat
+
+
+class FinePreprocess(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.config = config
+        self.cat_c_feat = config['fine_concat_coarse_feat']
+        self.W = self.config['fine_window_size']
+
+        d_model_c = self.config['coarse']['d_model']
+        d_model_f = self.config['fine']['d_model']
+        self.d_model_f = d_model_f
+        if self.cat_c_feat:
+            self.down_proj = nn.Linear(d_model_c, d_model_f, bias=True)
+            self.merge_feat = nn.Linear(2*d_model_f, d_model_f, bias=True)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.kaiming_normal_(p, mode="fan_out", nonlinearity="relu")
+
+    def forward(self, feat_f0, feat_f1, feat_c0, feat_c1, data):
+        W = self.W
+        stride = data['hw0_f'][0] // data['hw0_c'][0]
+
+        data.update({'W': W})
+        if data['b_ids'].shape[0] == 0:
+            feat0 = torch.empty(0, self.W**2, self.d_model_f, device=feat_f0.device)
+            feat1 = torch.empty(0, self.W**2, self.d_model_f, device=feat_f0.device)
+            return feat0, feat1
+
+        # 1. unfold(crop) all local windows
+        feat_f0_unfold = F.unfold(feat_f0, kernel_size=(W, W), stride=stride, padding=W//2)
+        feat_f0_unfold = rearrange(feat_f0_unfold, 'n (c ww) l -> n l ww c', ww=W**2)
+        feat_f1_unfold = F.unfold(feat_f1, kernel_size=(W, W), stride=stride, padding=W//2)
+        feat_f1_unfold = rearrange(feat_f1_unfold, 'n (c ww) l -> n l ww c', ww=W**2)
+
+        # 2. select only the predicted matches
+        feat_f0_unfold = feat_f0_unfold[data['b_ids'], data['i_ids']]  # [n, ww, cf]
+        feat_f1_unfold = feat_f1_unfold[data['b_ids'], data['j_ids']]
+
+        # option: use coarse-level loftr feature as context: concat and linear
+        if self.cat_c_feat:
+            feat_c_win = self.down_proj(torch.cat([feat_c0[data['b_ids'], data['i_ids']],
+                                                   feat_c1[data['b_ids'], data['j_ids']]], 0))  # [2n, c]
+            feat_cf_win = self.merge_feat(torch.cat([
+                torch.cat([feat_f0_unfold, feat_f1_unfold], 0),  # [2n, ww, cf]
+                repeat(feat_c_win, 'n c -> n ww c', ww=W**2),  # [2n, ww, cf]
+            ], -1))
+            feat_f0_unfold, feat_f1_unfold = torch.chunk(feat_cf_win, 2, dim=0)
+
+        return feat_f0_unfold, feat_f1_unfold

One-2-3-45-master 2/elevation_estimate/loftr/loftr_module/linear_attention.py

@@ -0,0 +1,81 @@
+"""
+Linear Transformer proposed in "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention"
+Modified from: https://github.com/idiap/fast-transformers/blob/master/fast_transformers/attention/linear_attention.py
+"""
+
+import torch
+from torch.nn import Module, Dropout
+
+
+def elu_feature_map(x):
+    return torch.nn.functional.elu(x) + 1
+
+
+class LinearAttention(Module):
+    def __init__(self, eps=1e-6):
+        super().__init__()
+        self.feature_map = elu_feature_map
+        self.eps = eps
+
+    def forward(self, queries, keys, values, q_mask=None, kv_mask=None):
+        """ Multi-Head linear attention proposed in "Transformers are RNNs"
+        Args:
+            queries: [N, L, H, D]
+            keys: [N, S, H, D]
+            values: [N, S, H, D]
+            q_mask: [N, L]
+            kv_mask: [N, S]
+        Returns:
+            queried_values: (N, L, H, D)
+        """
+        Q = self.feature_map(queries)
+        K = self.feature_map(keys)
+
+        # set padded position to zero
+        if q_mask is not None:
+            Q = Q * q_mask[:, :, None, None]
+        if kv_mask is not None:
+            K = K * kv_mask[:, :, None, None]
+            values = values * kv_mask[:, :, None, None]
+
+        v_length = values.size(1)
+        values = values / v_length  # prevent fp16 overflow
+        KV = torch.einsum("nshd,nshv->nhdv", K, values)  # (S,D)' @ S,V
+        Z = 1 / (torch.einsum("nlhd,nhd->nlh", Q, K.sum(dim=1)) + self.eps)
+        queried_values = torch.einsum("nlhd,nhdv,nlh->nlhv", Q, KV, Z) * v_length
+
+        return queried_values.contiguous()
+
+
+class FullAttention(Module):
+    def __init__(self, use_dropout=False, attention_dropout=0.1):
+        super().__init__()
+        self.use_dropout = use_dropout
+        self.dropout = Dropout(attention_dropout)
+
+    def forward(self, queries, keys, values, q_mask=None, kv_mask=None):
+        """ Multi-head scaled dot-product attention, a.k.a full attention.
+        Args:
+            queries: [N, L, H, D]
+            keys: [N, S, H, D]
+            values: [N, S, H, D]
+            q_mask: [N, L]
+            kv_mask: [N, S]
+        Returns:
+            queried_values: (N, L, H, D)
+        """
+
+        # Compute the unnormalized attention and apply the masks
+        QK = torch.einsum("nlhd,nshd->nlsh", queries, keys)
+        if kv_mask is not None:
+            QK.masked_fill_(~(q_mask[:, :, None, None] * kv_mask[:, None, :, None]), float('-inf'))
+
+        # Compute the attention and the weighted average
+        softmax_temp = 1. / queries.size(3)**.5  # sqrt(D)
+        A = torch.softmax(softmax_temp * QK, dim=2)
+        if self.use_dropout:
+            A = self.dropout(A)
+
+        queried_values = torch.einsum("nlsh,nshd->nlhd", A, values)
+
+        return queried_values.contiguous()

One-2-3-45-master 2/elevation_estimate/loftr/loftr_module/transformer.py

@@ -0,0 +1,101 @@
+import copy
+import torch
+import torch.nn as nn
+from .linear_attention import LinearAttention, FullAttention
+
+
+class LoFTREncoderLayer(nn.Module):
+    def __init__(self,
+                 d_model,
+                 nhead,
+                 attention='linear'):
+        super(LoFTREncoderLayer, self).__init__()
+
+        self.dim = d_model // nhead
+        self.nhead = nhead
+
+        # multi-head attention
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, d_model, bias=False)
+        self.v_proj = nn.Linear(d_model, d_model, bias=False)
+        self.attention = LinearAttention() if attention == 'linear' else FullAttention()
+        self.merge = nn.Linear(d_model, d_model, bias=False)
+
+        # feed-forward network
+        self.mlp = nn.Sequential(
+            nn.Linear(d_model*2, d_model*2, bias=False),
+            nn.ReLU(True),
+            nn.Linear(d_model*2, d_model, bias=False),
+        )
+
+        # norm and dropout
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+
+    def forward(self, x, source, x_mask=None, source_mask=None):
+        """
+        Args:
+            x (torch.Tensor): [N, L, C]
+            source (torch.Tensor): [N, S, C]
+            x_mask (torch.Tensor): [N, L] (optional)
+            source_mask (torch.Tensor): [N, S] (optional)
+        """
+        bs = x.size(0)
+        query, key, value = x, source, source
+
+        # multi-head attention
+        query = self.q_proj(query).view(bs, -1, self.nhead, self.dim)  # [N, L, (H, D)]
+        key = self.k_proj(key).view(bs, -1, self.nhead, self.dim)  # [N, S, (H, D)]
+        value = self.v_proj(value).view(bs, -1, self.nhead, self.dim)
+        message = self.attention(query, key, value, q_mask=x_mask, kv_mask=source_mask)  # [N, L, (H, D)]
+        message = self.merge(message.view(bs, -1, self.nhead*self.dim))  # [N, L, C]
+        message = self.norm1(message)
+
+        # feed-forward network
+        message = self.mlp(torch.cat([x, message], dim=2))
+        message = self.norm2(message)
+
+        return x + message
+
+
+class LocalFeatureTransformer(nn.Module):
+    """A Local Feature Transformer (LoFTR) module."""
+
+    def __init__(self, config):
+        super(LocalFeatureTransformer, self).__init__()
+
+        self.config = config
+        self.d_model = config['d_model']
+        self.nhead = config['nhead']
+        self.layer_names = config['layer_names']
+        encoder_layer = LoFTREncoderLayer(config['d_model'], config['nhead'], config['attention'])
+        self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(len(self.layer_names))])
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, feat0, feat1, mask0=None, mask1=None):
+        """
+        Args:
+            feat0 (torch.Tensor): [N, L, C]
+            feat1 (torch.Tensor): [N, S, C]
+            mask0 (torch.Tensor): [N, L] (optional)
+            mask1 (torch.Tensor): [N, S] (optional)
+        """
+
+        assert self.d_model == feat0.size(2), "the feature number of src and transformer must be equal"
+
+        for layer, name in zip(self.layers, self.layer_names):
+            if name == 'self':
+                feat0 = layer(feat0, feat0, mask0, mask0)
+                feat1 = layer(feat1, feat1, mask1, mask1)
+            elif name == 'cross':
+                feat0 = layer(feat0, feat1, mask0, mask1)
+                feat1 = layer(feat1, feat0, mask1, mask0)
+            else:
+                raise KeyError
+
+        return feat0, feat1

One-2-3-45-master 2/elevation_estimate/loftr/utils/coarse_matching.py

@@ -0,0 +1,261 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops.einops import rearrange
+
+INF = 1e9
+
+def mask_border(m, b: int, v):
+    """ Mask borders with value
+    Args:
+        m (torch.Tensor): [N, H0, W0, H1, W1]
+        b (int)
+        v (m.dtype)
+    """
+    if b <= 0:
+        return
+
+    m[:, :b] = v
+    m[:, :, :b] = v
+    m[:, :, :, :b] = v
+    m[:, :, :, :, :b] = v
+    m[:, -b:] = v
+    m[:, :, -b:] = v
+    m[:, :, :, -b:] = v
+    m[:, :, :, :, -b:] = v
+
+
+def mask_border_with_padding(m, bd, v, p_m0, p_m1):
+    if bd <= 0:
+        return
+
+    m[:, :bd] = v
+    m[:, :, :bd] = v
+    m[:, :, :, :bd] = v
+    m[:, :, :, :, :bd] = v
+
+    h0s, w0s = p_m0.sum(1).max(-1)[0].int(), p_m0.sum(-1).max(-1)[0].int()
+    h1s, w1s = p_m1.sum(1).max(-1)[0].int(), p_m1.sum(-1).max(-1)[0].int()
+    for b_idx, (h0, w0, h1, w1) in enumerate(zip(h0s, w0s, h1s, w1s)):
+        m[b_idx, h0 - bd:] = v
+        m[b_idx, :, w0 - bd:] = v
+        m[b_idx, :, :, h1 - bd:] = v
+        m[b_idx, :, :, :, w1 - bd:] = v
+
+
+def compute_max_candidates(p_m0, p_m1):
+    """Compute the max candidates of all pairs within a batch
+    
+    Args:
+        p_m0, p_m1 (torch.Tensor): padded masks
+    """
+    h0s, w0s = p_m0.sum(1).max(-1)[0], p_m0.sum(-1).max(-1)[0]
+    h1s, w1s = p_m1.sum(1).max(-1)[0], p_m1.sum(-1).max(-1)[0]
+    max_cand = torch.sum(
+        torch.min(torch.stack([h0s * w0s, h1s * w1s], -1), -1)[0])
+    return max_cand
+
+
+class CoarseMatching(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        # general config
+        self.thr = config['thr']
+        self.border_rm = config['border_rm']
+        # -- # for trainig fine-level LoFTR
+        self.train_coarse_percent = config['train_coarse_percent']
+        self.train_pad_num_gt_min = config['train_pad_num_gt_min']
+
+        # we provide 2 options for differentiable matching
+        self.match_type = config['match_type']
+        if self.match_type == 'dual_softmax':
+            self.temperature = config['dsmax_temperature']
+        elif self.match_type == 'sinkhorn':
+            try:
+                from .superglue import log_optimal_transport
+            except ImportError:
+                raise ImportError("download superglue.py first!")
+            self.log_optimal_transport = log_optimal_transport
+            self.bin_score = nn.Parameter(
+                torch.tensor(config['skh_init_bin_score'], requires_grad=True))
+            self.skh_iters = config['skh_iters']
+            self.skh_prefilter = config['skh_prefilter']
+        else:
+            raise NotImplementedError()
+
+    def forward(self, feat_c0, feat_c1, data, mask_c0=None, mask_c1=None):
+        """
+        Args:
+            feat0 (torch.Tensor): [N, L, C]
+            feat1 (torch.Tensor): [N, S, C]
+            data (dict)
+            mask_c0 (torch.Tensor): [N, L] (optional)
+            mask_c1 (torch.Tensor): [N, S] (optional)
+        Update:
+            data (dict): {
+                'b_ids' (torch.Tensor): [M'],
+                'i_ids' (torch.Tensor): [M'],
+                'j_ids' (torch.Tensor): [M'],
+                'gt_mask' (torch.Tensor): [M'],
+                'mkpts0_c' (torch.Tensor): [M, 2],
+                'mkpts1_c' (torch.Tensor): [M, 2],
+                'mconf' (torch.Tensor): [M]}
+            NOTE: M' != M during training.
+        """
+        N, L, S, C = feat_c0.size(0), feat_c0.size(1), feat_c1.size(1), feat_c0.size(2)
+
+        # normalize
+        feat_c0, feat_c1 = map(lambda feat: feat / feat.shape[-1]**.5,
+                               [feat_c0, feat_c1])
+
+        if self.match_type == 'dual_softmax':
+            sim_matrix = torch.einsum("nlc,nsc->nls", feat_c0,
+                                      feat_c1) / self.temperature
+            if mask_c0 is not None:
+                sim_matrix.masked_fill_(
+                    ~(mask_c0[..., None] * mask_c1[:, None]).bool(),
+                    -INF)
+            conf_matrix = F.softmax(sim_matrix, 1) * F.softmax(sim_matrix, 2)
+
+        elif self.match_type == 'sinkhorn':
+            # sinkhorn, dustbin included
+            sim_matrix = torch.einsum("nlc,nsc->nls", feat_c0, feat_c1)
+            if mask_c0 is not None:
+                sim_matrix[:, :L, :S].masked_fill_(
+                    ~(mask_c0[..., None] * mask_c1[:, None]).bool(),
+                    -INF)
+
+            # build uniform prior & use sinkhorn
+            log_assign_matrix = self.log_optimal_transport(
+                sim_matrix, self.bin_score, self.skh_iters)
+            assign_matrix = log_assign_matrix.exp()
+            conf_matrix = assign_matrix[:, :-1, :-1]
+
+            # filter prediction with dustbin score (only in evaluation mode)
+            if not self.training and self.skh_prefilter:
+                filter0 = (assign_matrix.max(dim=2)[1] == S)[:, :-1]  # [N, L]
+                filter1 = (assign_matrix.max(dim=1)[1] == L)[:, :-1]  # [N, S]
+                conf_matrix[filter0[..., None].repeat(1, 1, S)] = 0
+                conf_matrix[filter1[:, None].repeat(1, L, 1)] = 0
+
+            if self.config['sparse_spvs']:
+                data.update({'conf_matrix_with_bin': assign_matrix.clone()})
+
+        data.update({'conf_matrix': conf_matrix})
+
+        # predict coarse matches from conf_matrix
+        data.update(**self.get_coarse_match(conf_matrix, data))
+
+    @torch.no_grad()
+    def get_coarse_match(self, conf_matrix, data):
+        """
+        Args:
+            conf_matrix (torch.Tensor): [N, L, S]
+            data (dict): with keys ['hw0_i', 'hw1_i', 'hw0_c', 'hw1_c']
+        Returns:
+            coarse_matches (dict): {
+                'b_ids' (torch.Tensor): [M'],
+                'i_ids' (torch.Tensor): [M'],
+                'j_ids' (torch.Tensor): [M'],
+                'gt_mask' (torch.Tensor): [M'],
+                'm_bids' (torch.Tensor): [M],
+                'mkpts0_c' (torch.Tensor): [M, 2],
+                'mkpts1_c' (torch.Tensor): [M, 2],
+                'mconf' (torch.Tensor): [M]}
+        """
+        axes_lengths = {
+            'h0c': data['hw0_c'][0],
+            'w0c': data['hw0_c'][1],
+            'h1c': data['hw1_c'][0],
+            'w1c': data['hw1_c'][1]
+        }
+        _device = conf_matrix.device
+        # 1. confidence thresholding
+        mask = conf_matrix > self.thr
+        mask = rearrange(mask, 'b (h0c w0c) (h1c w1c) -> b h0c w0c h1c w1c',
+                         **axes_lengths)
+        if 'mask0' not in data:
+            mask_border(mask, self.border_rm, False)
+        else:
+            mask_border_with_padding(mask, self.border_rm, False,
+                                     data['mask0'], data['mask1'])
+        mask = rearrange(mask, 'b h0c w0c h1c w1c -> b (h0c w0c) (h1c w1c)',
+                         **axes_lengths)
+
+        # 2. mutual nearest
+        mask = mask \
+            * (conf_matrix == conf_matrix.max(dim=2, keepdim=True)[0]) \
+            * (conf_matrix == conf_matrix.max(dim=1, keepdim=True)[0])
+
+        # 3. find all valid coarse matches
+        # this only works when at most one `True` in each row
+        mask_v, all_j_ids = mask.max(dim=2)
+        b_ids, i_ids = torch.where(mask_v)
+        j_ids = all_j_ids[b_ids, i_ids]
+        mconf = conf_matrix[b_ids, i_ids, j_ids]
+
+        # 4. Random sampling of training samples for fine-level LoFTR
+        # (optional) pad samples with gt coarse-level matches
+        if self.training:
+            # NOTE:
+            # The sampling is performed across all pairs in a batch without manually balancing
+            # #samples for fine-level increases w.r.t. batch_size
+            if 'mask0' not in data:
+                num_candidates_max = mask.size(0) * max(
+                    mask.size(1), mask.size(2))
+            else:
+                num_candidates_max = compute_max_candidates(
+                    data['mask0'], data['mask1'])
+            num_matches_train = int(num_candidates_max *
+                                    self.train_coarse_percent)
+            num_matches_pred = len(b_ids)
+            assert self.train_pad_num_gt_min < num_matches_train, "min-num-gt-pad should be less than num-train-matches"
+
+            # pred_indices is to select from prediction
+            if num_matches_pred <= num_matches_train - self.train_pad_num_gt_min:
+                pred_indices = torch.arange(num_matches_pred, device=_device)
+            else:
+                pred_indices = torch.randint(
+                    num_matches_pred,
+                    (num_matches_train - self.train_pad_num_gt_min, ),
+                    device=_device)
+
+            # gt_pad_indices is to select from gt padding. e.g. max(3787-4800, 200)
+            gt_pad_indices = torch.randint(
+                    len(data['spv_b_ids']),
+                    (max(num_matches_train - num_matches_pred,
+                        self.train_pad_num_gt_min), ),
+                    device=_device)
+            mconf_gt = torch.zeros(len(data['spv_b_ids']), device=_device)  # set conf of gt paddings to all zero
+
+            b_ids, i_ids, j_ids, mconf = map(
+                lambda x, y: torch.cat([x[pred_indices], y[gt_pad_indices]],
+                                       dim=0),
+                *zip([b_ids, data['spv_b_ids']], [i_ids, data['spv_i_ids']],
+                     [j_ids, data['spv_j_ids']], [mconf, mconf_gt]))
+
+        # These matches select patches that feed into fine-level network
+        coarse_matches = {'b_ids': b_ids, 'i_ids': i_ids, 'j_ids': j_ids}
+
+        # 4. Update with matches in original image resolution
+        scale = data['hw0_i'][0] / data['hw0_c'][0]
+        scale0 = scale * data['scale0'][b_ids] if 'scale0' in data else scale
+        scale1 = scale * data['scale1'][b_ids] if 'scale1' in data else scale
+        mkpts0_c = torch.stack(
+            [i_ids % data['hw0_c'][1], i_ids // data['hw0_c'][1]],
+            dim=1) * scale0
+        mkpts1_c = torch.stack(
+            [j_ids % data['hw1_c'][1], j_ids // data['hw1_c'][1]],
+            dim=1) * scale1
+
+        # These matches is the current prediction (for visualization)
+        coarse_matches.update({
+            'gt_mask': mconf == 0,
+            'm_bids': b_ids[mconf != 0],  # mconf == 0 => gt matches
+            'mkpts0_c': mkpts0_c[mconf != 0],
+            'mkpts1_c': mkpts1_c[mconf != 0],
+            'mconf': mconf[mconf != 0]
+        })
+
+        return coarse_matches

One-2-3-45-master 2/elevation_estimate/loftr/utils/cvpr_ds_config.py

@@ -0,0 +1,50 @@
+from yacs.config import CfgNode as CN
+
+
+def lower_config(yacs_cfg):
+    if not isinstance(yacs_cfg, CN):
+        return yacs_cfg
+    return {k.lower(): lower_config(v) for k, v in yacs_cfg.items()}
+
+
+_CN = CN()
+_CN.BACKBONE_TYPE = 'ResNetFPN'
+_CN.RESOLUTION = (8, 2)  # options: [(8, 2), (16, 4)]
+_CN.FINE_WINDOW_SIZE = 5  # window_size in fine_level, must be odd
+_CN.FINE_CONCAT_COARSE_FEAT = True
+
+# 1. LoFTR-backbone (local feature CNN) config
+_CN.RESNETFPN = CN()
+_CN.RESNETFPN.INITIAL_DIM = 128
+_CN.RESNETFPN.BLOCK_DIMS = [128, 196, 256]  # s1, s2, s3
+
+# 2. LoFTR-coarse module config
+_CN.COARSE = CN()
+_CN.COARSE.D_MODEL = 256
+_CN.COARSE.D_FFN = 256
+_CN.COARSE.NHEAD = 8
+_CN.COARSE.LAYER_NAMES = ['self', 'cross'] * 4
+_CN.COARSE.ATTENTION = 'linear'  # options: ['linear', 'full']
+_CN.COARSE.TEMP_BUG_FIX = False
+
+# 3. Coarse-Matching config
+_CN.MATCH_COARSE = CN()
+_CN.MATCH_COARSE.THR = 0.2
+_CN.MATCH_COARSE.BORDER_RM = 2
+_CN.MATCH_COARSE.MATCH_TYPE = 'dual_softmax'  # options: ['dual_softmax, 'sinkhorn']
+_CN.MATCH_COARSE.DSMAX_TEMPERATURE = 0.1
+_CN.MATCH_COARSE.SKH_ITERS = 3
+_CN.MATCH_COARSE.SKH_INIT_BIN_SCORE = 1.0
+_CN.MATCH_COARSE.SKH_PREFILTER = True
+_CN.MATCH_COARSE.TRAIN_COARSE_PERCENT = 0.4  # training tricks: save GPU memory
+_CN.MATCH_COARSE.TRAIN_PAD_NUM_GT_MIN = 200  # training tricks: avoid DDP deadlock
+
+# 4. LoFTR-fine module config
+_CN.FINE = CN()
+_CN.FINE.D_MODEL = 128
+_CN.FINE.D_FFN = 128
+_CN.FINE.NHEAD = 8
+_CN.FINE.LAYER_NAMES = ['self', 'cross'] * 1
+_CN.FINE.ATTENTION = 'linear'
+
+default_cfg = lower_config(_CN)

One-2-3-45-master 2/elevation_estimate/loftr/utils/fine_matching.py

@@ -0,0 +1,74 @@
+import math
+import torch
+import torch.nn as nn
+
+from kornia.geometry.subpix import dsnt
+from kornia.utils.grid import create_meshgrid
+
+
+class FineMatching(nn.Module):
+    """FineMatching with s2d paradigm"""
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, feat_f0, feat_f1, data):
+        """
+        Args:
+            feat0 (torch.Tensor): [M, WW, C]
+            feat1 (torch.Tensor): [M, WW, C]
+            data (dict)
+        Update:
+            data (dict):{
+                'expec_f' (torch.Tensor): [M, 3],
+                'mkpts0_f' (torch.Tensor): [M, 2],
+                'mkpts1_f' (torch.Tensor): [M, 2]}
+        """
+        M, WW, C = feat_f0.shape
+        W = int(math.sqrt(WW))
+        scale = data['hw0_i'][0] / data['hw0_f'][0]
+        self.M, self.W, self.WW, self.C, self.scale = M, W, WW, C, scale
+
+        # corner case: if no coarse matches found
+        if M == 0:
+            assert self.training == False, "M is always >0, when training, see coarse_matching.py"
+            # logger.warning('No matches found in coarse-level.')
+            data.update({
+                'expec_f': torch.empty(0, 3, device=feat_f0.device),
+                'mkpts0_f': data['mkpts0_c'],
+                'mkpts1_f': data['mkpts1_c'],
+            })
+            return
+
+        feat_f0_picked = feat_f0_picked = feat_f0[:, WW//2, :]
+        sim_matrix = torch.einsum('mc,mrc->mr', feat_f0_picked, feat_f1)
+        softmax_temp = 1. / C**.5
+        heatmap = torch.softmax(softmax_temp * sim_matrix, dim=1).view(-1, W, W)
+
+        # compute coordinates from heatmap
+        coords_normalized = dsnt.spatial_expectation2d(heatmap[None], True)[0]  # [M, 2]
+        grid_normalized = create_meshgrid(W, W, True, heatmap.device).reshape(1, -1, 2)  # [1, WW, 2]
+
+        # compute std over <x, y>
+        var = torch.sum(grid_normalized**2 * heatmap.view(-1, WW, 1), dim=1) - coords_normalized**2  # [M, 2]
+        std = torch.sum(torch.sqrt(torch.clamp(var, min=1e-10)), -1)  # [M]  clamp needed for numerical stability
+        
+        # for fine-level supervision
+        data.update({'expec_f': torch.cat([coords_normalized, std.unsqueeze(1)], -1)})
+
+        # compute absolute kpt coords
+        self.get_fine_match(coords_normalized, data)
+
+    @torch.no_grad()
+    def get_fine_match(self, coords_normed, data):
+        W, WW, C, scale = self.W, self.WW, self.C, self.scale
+
+        # mkpts0_f and mkpts1_f
+        mkpts0_f = data['mkpts0_c']
+        scale1 = scale * data['scale1'][data['b_ids']] if 'scale0' in data else scale
+        mkpts1_f = data['mkpts1_c'] + (coords_normed * (W // 2) * scale1)[:len(data['mconf'])]
+
+        data.update({
+            "mkpts0_f": mkpts0_f,
+            "mkpts1_f": mkpts1_f
+        })

One-2-3-45-master 2/elevation_estimate/loftr/utils/geometry.py

@@ -0,0 +1,54 @@
+import torch
+
+
+@torch.no_grad()
+def warp_kpts(kpts0, depth0, depth1, T_0to1, K0, K1):
+    """ Warp kpts0 from I0 to I1 with depth, K and Rt
+    Also check covisibility and depth consistency.
+    Depth is consistent if relative error < 0.2 (hard-coded).
+    
+    Args:
+        kpts0 (torch.Tensor): [N, L, 2] - <x, y>,
+        depth0 (torch.Tensor): [N, H, W],
+        depth1 (torch.Tensor): [N, H, W],
+        T_0to1 (torch.Tensor): [N, 3, 4],
+        K0 (torch.Tensor): [N, 3, 3],
+        K1 (torch.Tensor): [N, 3, 3],
+    Returns:
+        calculable_mask (torch.Tensor): [N, L]
+        warped_keypoints0 (torch.Tensor): [N, L, 2] <x0_hat, y1_hat>
+    """
+    kpts0_long = kpts0.round().long()
+
+    # Sample depth, get calculable_mask on depth != 0
+    kpts0_depth = torch.stack(
+        [depth0[i, kpts0_long[i, :, 1], kpts0_long[i, :, 0]] for i in range(kpts0.shape[0])], dim=0
+    )  # (N, L)
+    nonzero_mask = kpts0_depth != 0
+
+    # Unproject
+    kpts0_h = torch.cat([kpts0, torch.ones_like(kpts0[:, :, [0]])], dim=-1) * kpts0_depth[..., None]  # (N, L, 3)
+    kpts0_cam = K0.inverse() @ kpts0_h.transpose(2, 1)  # (N, 3, L)
+
+    # Rigid Transform
+    w_kpts0_cam = T_0to1[:, :3, :3] @ kpts0_cam + T_0to1[:, :3, [3]]    # (N, 3, L)
+    w_kpts0_depth_computed = w_kpts0_cam[:, 2, :]
+
+    # Project
+    w_kpts0_h = (K1 @ w_kpts0_cam).transpose(2, 1)  # (N, L, 3)
+    w_kpts0 = w_kpts0_h[:, :, :2] / (w_kpts0_h[:, :, [2]] + 1e-4)  # (N, L, 2), +1e-4 to avoid zero depth
+
+    # Covisible Check
+    h, w = depth1.shape[1:3]
+    covisible_mask = (w_kpts0[:, :, 0] > 0) * (w_kpts0[:, :, 0] < w-1) * \
+        (w_kpts0[:, :, 1] > 0) * (w_kpts0[:, :, 1] < h-1)
+    w_kpts0_long = w_kpts0.long()
+    w_kpts0_long[~covisible_mask, :] = 0
+
+    w_kpts0_depth = torch.stack(
+        [depth1[i, w_kpts0_long[i, :, 1], w_kpts0_long[i, :, 0]] for i in range(w_kpts0_long.shape[0])], dim=0
+    )  # (N, L)
+    consistent_mask = ((w_kpts0_depth - w_kpts0_depth_computed) / w_kpts0_depth).abs() < 0.2
+    valid_mask = nonzero_mask * covisible_mask * consistent_mask
+
+    return valid_mask, w_kpts0

One-2-3-45-master 2/elevation_estimate/loftr/utils/position_encoding.py

@@ -0,0 +1,42 @@
+import math
+import torch
+from torch import nn
+
+
+class PositionEncodingSine(nn.Module):
+    """
+    This is a sinusoidal position encoding that generalized to 2-dimensional images
+    """
+
+    def __init__(self, d_model, max_shape=(256, 256), temp_bug_fix=True):
+        """
+        Args:
+            max_shape (tuple): for 1/8 featmap, the max length of 256 corresponds to 2048 pixels
+            temp_bug_fix (bool): As noted in this [issue](https://github.com/zju3dv/LoFTR/issues/41),
+                the original implementation of LoFTR includes a bug in the pos-enc impl, which has little impact
+                on the final performance. For now, we keep both impls for backward compatability.
+                We will remove the buggy impl after re-training all variants of our released models.
+        """
+        super().__init__()
+
+        pe = torch.zeros((d_model, *max_shape))
+        y_position = torch.ones(max_shape).cumsum(0).float().unsqueeze(0)
+        x_position = torch.ones(max_shape).cumsum(1).float().unsqueeze(0)
+        if temp_bug_fix:
+            div_term = torch.exp(torch.arange(0, d_model//2, 2).float() * (-math.log(10000.0) / (d_model//2)))
+        else:  # a buggy implementation (for backward compatability only)
+            div_term = torch.exp(torch.arange(0, d_model//2, 2).float() * (-math.log(10000.0) / d_model//2))
+        div_term = div_term[:, None, None]  # [C//4, 1, 1]
+        pe[0::4, :, :] = torch.sin(x_position * div_term)
+        pe[1::4, :, :] = torch.cos(x_position * div_term)
+        pe[2::4, :, :] = torch.sin(y_position * div_term)
+        pe[3::4, :, :] = torch.cos(y_position * div_term)
+
+        self.register_buffer('pe', pe.unsqueeze(0), persistent=False)  # [1, C, H, W]
+
+    def forward(self, x):
+        """
+        Args:
+            x: [N, C, H, W]
+        """
+        return x + self.pe[:, :, :x.size(2), :x.size(3)]

One-2-3-45-master 2/elevation_estimate/loftr/utils/supervision.py

@@ -0,0 +1,151 @@
+from math import log
+from loguru import logger
+
+import torch
+from einops import repeat
+from kornia.utils import create_meshgrid
+
+from .geometry import warp_kpts
+
+##############  ↓  Coarse-Level supervision  ↓  ##############
+
+
+@torch.no_grad()
+def mask_pts_at_padded_regions(grid_pt, mask):
+    """For megadepth dataset, zero-padding exists in images"""
+    mask = repeat(mask, 'n h w -> n (h w) c', c=2)
+    grid_pt[~mask.bool()] = 0
+    return grid_pt
+
+
+@torch.no_grad()
+def spvs_coarse(data, config):
+    """
+    Update:
+        data (dict): {
+            "conf_matrix_gt": [N, hw0, hw1],
+            'spv_b_ids': [M]
+            'spv_i_ids': [M]
+            'spv_j_ids': [M]
+            'spv_w_pt0_i': [N, hw0, 2], in original image resolution
+            'spv_pt1_i': [N, hw1, 2], in original image resolution
+        }
+        
+    NOTE:
+        - for scannet dataset, there're 3 kinds of resolution {i, c, f}
+        - for megadepth dataset, there're 4 kinds of resolution {i, i_resize, c, f}
+    """
+    # 1. misc
+    device = data['image0'].device
+    N, _, H0, W0 = data['image0'].shape
+    _, _, H1, W1 = data['image1'].shape
+    scale = config['LOFTR']['RESOLUTION'][0]
+    scale0 = scale * data['scale0'][:, None] if 'scale0' in data else scale
+    scale1 = scale * data['scale1'][:, None] if 'scale0' in data else scale
+    h0, w0, h1, w1 = map(lambda x: x // scale, [H0, W0, H1, W1])
+
+    # 2. warp grids
+    # create kpts in meshgrid and resize them to image resolution
+    grid_pt0_c = create_meshgrid(h0, w0, False, device).reshape(1, h0*w0, 2).repeat(N, 1, 1)    # [N, hw, 2]
+    grid_pt0_i = scale0 * grid_pt0_c
+    grid_pt1_c = create_meshgrid(h1, w1, False, device).reshape(1, h1*w1, 2).repeat(N, 1, 1)
+    grid_pt1_i = scale1 * grid_pt1_c
+
+    # mask padded region to (0, 0), so no need to manually mask conf_matrix_gt
+    if 'mask0' in data:
+        grid_pt0_i = mask_pts_at_padded_regions(grid_pt0_i, data['mask0'])
+        grid_pt1_i = mask_pts_at_padded_regions(grid_pt1_i, data['mask1'])
+
+    # warp kpts bi-directionally and resize them to coarse-level resolution
+    # (no depth consistency check, since it leads to worse results experimentally)
+    # (unhandled edge case: points with 0-depth will be warped to the left-up corner)
+    _, w_pt0_i = warp_kpts(grid_pt0_i, data['depth0'], data['depth1'], data['T_0to1'], data['K0'], data['K1'])
+    _, w_pt1_i = warp_kpts(grid_pt1_i, data['depth1'], data['depth0'], data['T_1to0'], data['K1'], data['K0'])
+    w_pt0_c = w_pt0_i / scale1
+    w_pt1_c = w_pt1_i / scale0
+
+    # 3. check if mutual nearest neighbor
+    w_pt0_c_round = w_pt0_c[:, :, :].round().long()
+    nearest_index1 = w_pt0_c_round[..., 0] + w_pt0_c_round[..., 1] * w1
+    w_pt1_c_round = w_pt1_c[:, :, :].round().long()
+    nearest_index0 = w_pt1_c_round[..., 0] + w_pt1_c_round[..., 1] * w0
+
+    # corner case: out of boundary
+    def out_bound_mask(pt, w, h):
+        return (pt[..., 0] < 0) + (pt[..., 0] >= w) + (pt[..., 1] < 0) + (pt[..., 1] >= h)
+    nearest_index1[out_bound_mask(w_pt0_c_round, w1, h1)] = 0
+    nearest_index0[out_bound_mask(w_pt1_c_round, w0, h0)] = 0
+
+    loop_back = torch.stack([nearest_index0[_b][_i] for _b, _i in enumerate(nearest_index1)], dim=0)
+    correct_0to1 = loop_back == torch.arange(h0*w0, device=device)[None].repeat(N, 1)
+    correct_0to1[:, 0] = False  # ignore the top-left corner
+
+    # 4. construct a gt conf_matrix
+    conf_matrix_gt = torch.zeros(N, h0*w0, h1*w1, device=device)
+    b_ids, i_ids = torch.where(correct_0to1 != 0)
+    j_ids = nearest_index1[b_ids, i_ids]
+
+    conf_matrix_gt[b_ids, i_ids, j_ids] = 1
+    data.update({'conf_matrix_gt': conf_matrix_gt})
+
+    # 5. save coarse matches(gt) for training fine level
+    if len(b_ids) == 0:
+        logger.warning(f"No groundtruth coarse match found for: {data['pair_names']}")
+        # this won't affect fine-level loss calculation
+        b_ids = torch.tensor([0], device=device)
+        i_ids = torch.tensor([0], device=device)
+        j_ids = torch.tensor([0], device=device)
+
+    data.update({
+        'spv_b_ids': b_ids,
+        'spv_i_ids': i_ids,
+        'spv_j_ids': j_ids
+    })
+
+    # 6. save intermediate results (for fast fine-level computation)
+    data.update({
+        'spv_w_pt0_i': w_pt0_i,
+        'spv_pt1_i': grid_pt1_i
+    })
+
+
+def compute_supervision_coarse(data, config):
+    assert len(set(data['dataset_name'])) == 1, "Do not support mixed datasets training!"
+    data_source = data['dataset_name'][0]
+    if data_source.lower() in ['scannet', 'megadepth']:
+        spvs_coarse(data, config)
+    else:
+        raise ValueError(f'Unknown data source: {data_source}')
+
+
+##############  ↓  Fine-Level supervision  ↓  ##############
+
+@torch.no_grad()
+def spvs_fine(data, config):
+    """
+    Update:
+        data (dict):{
+            "expec_f_gt": [M, 2]}
+    """
+    # 1. misc
+    # w_pt0_i, pt1_i = data.pop('spv_w_pt0_i'), data.pop('spv_pt1_i')
+    w_pt0_i, pt1_i = data['spv_w_pt0_i'], data['spv_pt1_i']
+    scale = config['LOFTR']['RESOLUTION'][1]
+    radius = config['LOFTR']['FINE_WINDOW_SIZE'] // 2
+
+    # 2. get coarse prediction
+    b_ids, i_ids, j_ids = data['b_ids'], data['i_ids'], data['j_ids']
+
+    # 3. compute gt
+    scale = scale * data['scale1'][b_ids] if 'scale0' in data else scale
+    # `expec_f_gt` might exceed the window, i.e. abs(*) > 1, which would be filtered later
+    expec_f_gt = (w_pt0_i[b_ids, i_ids] - pt1_i[b_ids, j_ids]) / scale / radius  # [M, 2]
+    data.update({"expec_f_gt": expec_f_gt})
+
+
+def compute_supervision_fine(data, config):
+    data_source = data['dataset_name'][0]
+    if data_source.lower() in ['scannet', 'megadepth']:
+        spvs_fine(data, config)
+    else:
+        raise NotImplementedError

One-2-3-45-master 2/elevation_estimate/pyproject.toml

@@ -0,0 +1,7 @@
+[project]
+name = "elevation_estimate"
+version = "0.1"
+
+[tool.setuptools.packages.find]
+exclude = ["configs", "tests"]  # empty by default
+namespaces = false  # true by default

One-2-3-45-master 2/elevation_estimate/utils/elev_est_api.py

@@ -0,0 +1,205 @@
+import os
+import cv2
+import numpy as np
+import os.path as osp
+import imageio
+from copy import deepcopy
+
+import loguru
+import torch
+import matplotlib.cm as cm
+import matplotlib.pyplot as plt
+
+from ..loftr import LoFTR, default_cfg
+from . import plt_utils
+from .plotting import make_matching_figure
+from .utils3d import rect_to_img, canonical_to_camera, calc_pose
+
+
+class ElevEstHelper:
+    _feature_matcher = None
+
+    @classmethod
+    def get_feature_matcher(cls):
+        if cls._feature_matcher is None:
+            loguru.logger.info("Loading feature matcher...")
+            _default_cfg = deepcopy(default_cfg)
+            _default_cfg['coarse']['temp_bug_fix'] = True  # set to False when using the old ckpt
+            matcher = LoFTR(config=_default_cfg)
+            current_dir = os.path.dirname(os.path.abspath(__file__))
+            ckpt_path = os.path.join(current_dir, "weights/indoor_ds_new.ckpt")
+            if not osp.exists(ckpt_path):
+                loguru.logger.info("Downloading feature matcher...")
+                os.makedirs("weights", exist_ok=True)
+                import gdown
+                gdown.cached_download(url="https://drive.google.com/uc?id=19s3QvcCWQ6g-N1PrYlDCg-2mOJZ3kkgS",
+                                      path=ckpt_path)
+            matcher.load_state_dict(torch.load(ckpt_path)['state_dict'])
+            matcher = matcher.eval().cuda()
+            cls._feature_matcher = matcher
+        return cls._feature_matcher
+
+
+def mask_out_bkgd(img_path, dbg=False):
+    img = imageio.imread_v2(img_path)
+    if img.shape[-1] == 4:
+        fg_mask = img[:, :, :3]
+    else:
+        loguru.logger.info("Image has no alpha channel, using thresholding to mask out background")
+        fg_mask = ~(img > 245).all(axis=-1)
+        if dbg:
+            plt.imshow(plt_utils.vis_mask(img, fg_mask.astype(np.uint8), color=[0, 255, 0]))
+            plt.show()
+    return fg_mask
+
+
+def get_feature_matching(img_paths, dbg=False):
+    assert len(img_paths) == 4
+    matcher = ElevEstHelper.get_feature_matcher()
+    feature_matching = {}
+    masks = []
+    for i in range(4):
+        mask = mask_out_bkgd(img_paths[i], dbg=dbg)
+        masks.append(mask)
+    for i in range(0, 4):
+        for j in range(i + 1, 4):
+            img0_pth = img_paths[i]
+            img1_pth = img_paths[j]
+            mask0 = masks[i]
+            mask1 = masks[j]
+            img0_raw = cv2.imread(img0_pth, cv2.IMREAD_GRAYSCALE)
+            img1_raw = cv2.imread(img1_pth, cv2.IMREAD_GRAYSCALE)
+            original_shape = img0_raw.shape
+            img0_raw_resized = cv2.resize(img0_raw, (480, 480))
+            img1_raw_resized = cv2.resize(img1_raw, (480, 480))
+
+            img0 = torch.from_numpy(img0_raw_resized)[None][None].cuda() / 255.
+            img1 = torch.from_numpy(img1_raw_resized)[None][None].cuda() / 255.
+            batch = {'image0': img0, 'image1': img1}
+
+            # Inference with LoFTR and get prediction
+            with torch.no_grad():
+                matcher(batch)
+                mkpts0 = batch['mkpts0_f'].cpu().numpy()
+                mkpts1 = batch['mkpts1_f'].cpu().numpy()
+                mconf = batch['mconf'].cpu().numpy()
+            mkpts0[:, 0] = mkpts0[:, 0] * original_shape[1] / 480
+            mkpts0[:, 1] = mkpts0[:, 1] * original_shape[0] / 480
+            mkpts1[:, 0] = mkpts1[:, 0] * original_shape[1] / 480
+            mkpts1[:, 1] = mkpts1[:, 1] * original_shape[0] / 480
+            keep0 = mask0[mkpts0[:, 1].astype(int), mkpts1[:, 0].astype(int)]
+            keep1 = mask1[mkpts1[:, 1].astype(int), mkpts1[:, 0].astype(int)]
+            keep = np.logical_and(keep0, keep1)
+            mkpts0 = mkpts0[keep]
+            mkpts1 = mkpts1[keep]
+            mconf = mconf[keep]
+            if dbg:
+                # Draw visualization
+                color = cm.jet(mconf)
+                text = [
+                    'LoFTR',
+                    'Matches: {}'.format(len(mkpts0)),
+                ]
+                fig = make_matching_figure(img0_raw, img1_raw, mkpts0, mkpts1, color, text=text)
+                fig.show()
+            feature_matching[f"{i}_{j}"] = np.concatenate([mkpts0, mkpts1, mconf[:, None]], axis=1)
+
+    return feature_matching
+
+
+def gen_pose_hypothesis(center_elevation):
+    elevations = np.radians(
+        [center_elevation, center_elevation - 10, center_elevation + 10, center_elevation, center_elevation])  # 45~120
+    azimuths = np.radians([30, 30, 30, 20, 40])
+    input_poses = calc_pose(elevations, azimuths, len(azimuths))
+    input_poses = input_poses[1:]
+    input_poses[..., 1] *= -1
+    input_poses[..., 2] *= -1
+    return input_poses
+
+
+def ba_error_general(K, matches, poses):
+    projmat0 = K @ poses[0].inverse()[:3, :4]
+    projmat1 = K @ poses[1].inverse()[:3, :4]
+    match_01 = matches[0]
+    pts0 = match_01[:, :2]
+    pts1 = match_01[:, 2:4]
+    Xref = cv2.triangulatePoints(projmat0.cpu().numpy(), projmat1.cpu().numpy(),
+                                 pts0.cpu().numpy().T, pts1.cpu().numpy().T)
+    Xref = Xref[:3] / Xref[3:]
+    Xref = Xref.T
+    Xref = torch.from_numpy(Xref).cuda().float()
+    reproj_error = 0
+    for match, cp in zip(matches[1:], poses[2:]):
+        dist = (torch.norm(match_01[:, :2][:, None, :] - match[:, :2][None, :, :], dim=-1))
+        if dist.numel() > 0:
+            # print("dist.shape", dist.shape)
+            m0to2_index = dist.argmin(1)
+            keep = dist[torch.arange(match_01.shape[0]), m0to2_index] < 1
+            if keep.sum() > 0:
+                xref_in2 = rect_to_img(K, canonical_to_camera(Xref, cp.inverse()))
+                reproj_error2 = torch.norm(match[m0to2_index][keep][:, 2:4] - xref_in2[keep], dim=-1)
+                conf02 = match[m0to2_index][keep][:, -1]
+                reproj_error += (reproj_error2 * conf02).sum() / (conf02.sum())
+
+    return reproj_error
+
+
+def find_optim_elev(elevs, nimgs, matches, K, dbg=False):
+    errs = []
+    for elev in elevs:
+        err = 0
+        cam_poses = gen_pose_hypothesis(elev)
+        for start in range(nimgs - 1):
+            batch_matches, batch_poses = [], []
+            for i in range(start, nimgs + start):
+                ci = i % nimgs
+                batch_poses.append(cam_poses[ci])
+            for j in range(nimgs - 1):
+                key = f"{start}_{(start + j + 1) % nimgs}"
+                match = matches[key]
+                batch_matches.append(match)
+            err += ba_error_general(K, batch_matches, batch_poses)
+        errs.append(err)
+    errs = torch.tensor(errs)
+    if dbg:
+        plt.plot(elevs, errs)
+        plt.show()
+    optim_elev = elevs[torch.argmin(errs)].item()
+    return optim_elev
+
+
+def get_elev_est(feature_matching, min_elev=30, max_elev=150, K=None, dbg=False):
+    flag = True
+    matches = {}
+    for i in range(4):
+        for j in range(i + 1, 4):
+            match_ij = feature_matching[f"{i}_{j}"]
+            if len(match_ij) == 0:
+                flag = False
+            match_ji = np.concatenate([match_ij[:, 2:4], match_ij[:, 0:2], match_ij[:, 4:5]], axis=1)
+            matches[f"{i}_{j}"] = torch.from_numpy(match_ij).float().cuda()
+            matches[f"{j}_{i}"] = torch.from_numpy(match_ji).float().cuda()
+    if not flag:
+        loguru.logger.info("0 matches, could not estimate elevation")
+        return None
+    interval = 10
+    elevs = np.arange(min_elev, max_elev, interval)
+    optim_elev1 = find_optim_elev(elevs, 4, matches, K)
+
+    elevs = np.arange(optim_elev1 - 10, optim_elev1 + 10, 1)
+    optim_elev2 = find_optim_elev(elevs, 4, matches, K)
+
+    return optim_elev2
+
+
+def elev_est_api(img_paths, min_elev=30, max_elev=150, K=None, dbg=False):
+    feature_matching = get_feature_matching(img_paths, dbg=dbg)
+    if K is None:
+        loguru.logger.warning("K is not provided, using default K")
+        K = np.array([[280.0, 0, 128.0],
+                      [0, 280.0, 128.0],
+                      [0, 0, 1]])
+    K = torch.from_numpy(K).cuda().float()
+    elev = get_elev_est(feature_matching, min_elev, max_elev, K, dbg=dbg)
+    return elev

One-2-3-45-master 2/elevation_estimate/utils/plotting.py

@@ -0,0 +1,154 @@
+import bisect
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib
+
+
+def _compute_conf_thresh(data):
+    dataset_name = data['dataset_name'][0].lower()
+    if dataset_name == 'scannet':
+        thr = 5e-4
+    elif dataset_name == 'megadepth':
+        thr = 1e-4
+    else:
+        raise ValueError(f'Unknown dataset: {dataset_name}')
+    return thr
+
+
+# --- VISUALIZATION --- #
+
+def make_matching_figure(
+        img0, img1, mkpts0, mkpts1, color,
+        kpts0=None, kpts1=None, text=[], dpi=75, path=None):
+    # draw image pair
+    assert mkpts0.shape[0] == mkpts1.shape[0], f'mkpts0: {mkpts0.shape[0]} v.s. mkpts1: {mkpts1.shape[0]}'
+    fig, axes = plt.subplots(1, 2, figsize=(10, 6), dpi=dpi)
+    axes[0].imshow(img0, cmap='gray')
+    axes[1].imshow(img1, cmap='gray')
+    for i in range(2):   # clear all frames
+        axes[i].get_yaxis().set_ticks([])
+        axes[i].get_xaxis().set_ticks([])
+        for spine in axes[i].spines.values():
+            spine.set_visible(False)
+    plt.tight_layout(pad=1)
+    
+    if kpts0 is not None:
+        assert kpts1 is not None
+        axes[0].scatter(kpts0[:, 0], kpts0[:, 1], c='w', s=2)
+        axes[1].scatter(kpts1[:, 0], kpts1[:, 1], c='w', s=2)
+
+    # draw matches
+    if mkpts0.shape[0] != 0 and mkpts1.shape[0] != 0:
+        fig.canvas.draw()
+        transFigure = fig.transFigure.inverted()
+        fkpts0 = transFigure.transform(axes[0].transData.transform(mkpts0))
+        fkpts1 = transFigure.transform(axes[1].transData.transform(mkpts1))
+        fig.lines = [matplotlib.lines.Line2D((fkpts0[i, 0], fkpts1[i, 0]),
+                                            (fkpts0[i, 1], fkpts1[i, 1]),
+                                            transform=fig.transFigure, c=color[i], linewidth=1)
+                                        for i in range(len(mkpts0))]
+        
+        axes[0].scatter(mkpts0[:, 0], mkpts0[:, 1], c=color, s=4)
+        axes[1].scatter(mkpts1[:, 0], mkpts1[:, 1], c=color, s=4)
+
+    # put txts
+    txt_color = 'k' if img0[:100, :200].mean() > 200 else 'w'
+    fig.text(
+        0.01, 0.99, '\n'.join(text), transform=fig.axes[0].transAxes,
+        fontsize=15, va='top', ha='left', color=txt_color)
+
+    # save or return figure
+    if path:
+        plt.savefig(str(path), bbox_inches='tight', pad_inches=0)
+        plt.close()
+    else:
+        return fig
+
+
+def _make_evaluation_figure(data, b_id, alpha='dynamic'):
+    b_mask = data['m_bids'] == b_id
+    conf_thr = _compute_conf_thresh(data)
+    
+    img0 = (data['image0'][b_id][0].cpu().numpy() * 255).round().astype(np.int32)
+    img1 = (data['image1'][b_id][0].cpu().numpy() * 255).round().astype(np.int32)
+    kpts0 = data['mkpts0_f'][b_mask].cpu().numpy()
+    kpts1 = data['mkpts1_f'][b_mask].cpu().numpy()
+    
+    # for megadepth, we visualize matches on the resized image
+    if 'scale0' in data:
+        kpts0 = kpts0 / data['scale0'][b_id].cpu().numpy()[[1, 0]]
+        kpts1 = kpts1 / data['scale1'][b_id].cpu().numpy()[[1, 0]]
+
+    epi_errs = data['epi_errs'][b_mask].cpu().numpy()
+    correct_mask = epi_errs < conf_thr
+    precision = np.mean(correct_mask) if len(correct_mask) > 0 else 0
+    n_correct = np.sum(correct_mask)
+    n_gt_matches = int(data['conf_matrix_gt'][b_id].sum().cpu())
+    recall = 0 if n_gt_matches == 0 else n_correct / (n_gt_matches)
+    # recall might be larger than 1, since the calculation of conf_matrix_gt
+    # uses groundtruth depths and camera poses, but epipolar distance is used here.
+
+    # matching info
+    if alpha == 'dynamic':
+        alpha = dynamic_alpha(len(correct_mask))
+    color = error_colormap(epi_errs, conf_thr, alpha=alpha)
+    
+    text = [
+        f'#Matches {len(kpts0)}',
+        f'Precision({conf_thr:.2e}) ({100 * precision:.1f}%): {n_correct}/{len(kpts0)}',
+        f'Recall({conf_thr:.2e}) ({100 * recall:.1f}%): {n_correct}/{n_gt_matches}'
+    ]
+    
+    # make the figure
+    figure = make_matching_figure(img0, img1, kpts0, kpts1,
+                                  color, text=text)
+    return figure
+
+def _make_confidence_figure(data, b_id):
+    # TODO: Implement confidence figure
+    raise NotImplementedError()
+
+
+def make_matching_figures(data, config, mode='evaluation'):
+    """ Make matching figures for a batch.
+    
+    Args:
+        data (Dict): a batch updated by PL_LoFTR.
+        config (Dict): matcher config
+    Returns:
+        figures (Dict[str, List[plt.figure]]
+    """
+    assert mode in ['evaluation', 'confidence']  # 'confidence'
+    figures = {mode: []}
+    for b_id in range(data['image0'].size(0)):
+        if mode == 'evaluation':
+            fig = _make_evaluation_figure(
+                data, b_id,
+                alpha=config.TRAINER.PLOT_MATCHES_ALPHA)
+        elif mode == 'confidence':
+            fig = _make_confidence_figure(data, b_id)
+        else:
+            raise ValueError(f'Unknown plot mode: {mode}')
+    figures[mode].append(fig)
+    return figures
+
+
+def dynamic_alpha(n_matches,
+                  milestones=[0, 300, 1000, 2000],
+                  alphas=[1.0, 0.8, 0.4, 0.2]):
+    if n_matches == 0:
+        return 1.0
+    ranges = list(zip(alphas, alphas[1:] + [None]))
+    loc = bisect.bisect_right(milestones, n_matches) - 1
+    _range = ranges[loc]
+    if _range[1] is None:
+        return _range[0]
+    return _range[1] + (milestones[loc + 1] - n_matches) / (
+        milestones[loc + 1] - milestones[loc]) * (_range[0] - _range[1])
+
+
+def error_colormap(err, thr, alpha=1.0):
+    assert alpha <= 1.0 and alpha > 0, f"Invaid alpha value: {alpha}"
+    x = 1 - np.clip(err / (thr * 2), 0, 1)
+    return np.clip(
+        np.stack([2-x*2, x*2, np.zeros_like(x), np.ones_like(x)*alpha], -1), 0, 1)

One-2-3-45-master 2/elevation_estimate/utils/plt_utils.py

@@ -0,0 +1,318 @@
+import os.path as osp
+import os
+import matplotlib.pyplot as plt
+import torch
+import cv2
+import math
+
+import numpy as np
+import tqdm
+from cv2 import findContours
+from dl_ext.primitive import safe_zip
+from dl_ext.timer import EvalTime
+
+
+def plot_confidence(confidence):
+    n = len(confidence)
+    plt.plot(np.arange(n), confidence)
+    plt.show()
+
+
+def image_grid(
+        images,
+        rows=None,
+        cols=None,
+        fill: bool = True,
+        show_axes: bool = False,
+        rgb=None,
+        show=True,
+        label=None,
+        **kwargs
+):
+    """
+    A util function for plotting a grid of images.
+    Args:
+        images: (N, H, W, 4) array of RGBA images
+        rows: number of rows in the grid
+        cols: number of columns in the grid
+        fill: boolean indicating if the space between images should be filled
+        show_axes: boolean indicating if the axes of the plots should be visible
+        rgb: boolean, If True, only RGB channels are plotted.
+            If False, only the alpha channel is plotted.
+    Returns:
+        None
+    """
+    evaltime = EvalTime(disable=True)
+    evaltime('')
+    if isinstance(images, torch.Tensor):
+        images = images.detach().cpu()
+    if len(images[0].shape) == 2:
+        rgb = False
+    if images[0].shape[-1] == 2:
+        # flow
+        images = [flow_to_image(im) for im in images]
+    if (rows is None) != (cols is None):
+        raise ValueError("Specify either both rows and cols or neither.")
+
+    if rows is None:
+        rows = int(len(images) ** 0.5)
+        cols = math.ceil(len(images) / rows)
+
+    gridspec_kw = {"wspace": 0.0, "hspace": 0.0} if fill else {}
+    if len(images) < 50:
+        figsize = (10, 10)
+    else:
+        figsize = (15, 15)
+    evaltime('0.5')
+    plt.figure(figsize=figsize)
+    # fig, axarr = plt.subplots(rows, cols, gridspec_kw=gridspec_kw, figsize=figsize)
+    if label:
+        # fig.suptitle(label, fontsize=30)
+        plt.suptitle(label, fontsize=30)
+    # bleed = 0
+    # fig.subplots_adjust(left=bleed, bottom=bleed, right=(1 - bleed), top=(1 - bleed))
+    evaltime('subplots')
+
+    # for i, (ax, im) in enumerate(tqdm.tqdm(zip(axarr.ravel(), images), leave=True, total=len(images))):
+    for i in range(len(images)):
+        # evaltime(f'{i} begin')
+        plt.subplot(rows, cols, i + 1)
+        if rgb:
+            # only render RGB channels
+            plt.imshow(images[i][..., :3], **kwargs)
+            # ax.imshow(im[..., :3], **kwargs)
+        else:
+            # only render Alpha channel
+            plt.imshow(images[i], **kwargs)
+            # ax.imshow(im, **kwargs)
+        if not show_axes:
+            plt.axis('off')
+            # ax.set_axis_off()
+        # ax.set_title(f'{i}')
+        plt.title(f'{i}')
+        # evaltime(f'{i} end')
+    evaltime('2')
+    if show:
+        plt.show()
+    # return fig
+
+
+def depth_grid(
+        depths,
+        rows=None,
+        cols=None,
+        fill: bool = True,
+        show_axes: bool = False,
+):
+    """
+    A util function for plotting a grid of images.
+    Args:
+        images: (N, H, W, 4) array of RGBA images
+        rows: number of rows in the grid
+        cols: number of columns in the grid
+        fill: boolean indicating if the space between images should be filled
+        show_axes: boolean indicating if the axes of the plots should be visible
+        rgb: boolean, If True, only RGB channels are plotted.
+            If False, only the alpha channel is plotted.
+    Returns:
+        None
+    """
+    if (rows is None) != (cols is None):
+        raise ValueError("Specify either both rows and cols or neither.")
+
+    if rows is None:
+        rows = len(depths)
+        cols = 1
+
+    gridspec_kw = {"wspace": 0.0, "hspace": 0.0} if fill else {}
+    fig, axarr = plt.subplots(rows, cols, gridspec_kw=gridspec_kw, figsize=(15, 9))
+    bleed = 0
+    fig.subplots_adjust(left=bleed, bottom=bleed, right=(1 - bleed), top=(1 - bleed))
+
+    for ax, im in zip(axarr.ravel(), depths):
+        ax.imshow(im)
+        if not show_axes:
+            ax.set_axis_off()
+    plt.show()
+
+
+def hover_masks_on_imgs(images, masks):
+    masks = np.array(masks)
+    new_imgs = []
+    tids = list(range(1, masks.max() + 1))
+    colors = colormap(rgb=True, lighten=True)
+    for im, mask in tqdm.tqdm(safe_zip(images, masks), total=len(images)):
+        for tid in tids:
+            im = vis_mask(
+                im,
+                (mask == tid).astype(np.uint8),
+                color=colors[tid],
+                alpha=0.5,
+                border_alpha=0.5,
+                border_color=[255, 255, 255],
+                border_thick=3)
+        new_imgs.append(im)
+    return new_imgs
+
+
+def vis_mask(img,
+             mask,
+             color=[255, 255, 255],
+             alpha=0.4,
+             show_border=True,
+             border_alpha=0.5,
+             border_thick=1,
+             border_color=None):
+    """Visualizes a single binary mask."""
+    if isinstance(mask, torch.Tensor):
+        from anypose.utils.pn_utils import to_array
+        mask = to_array(mask > 0).astype(np.uint8)
+    img = img.astype(np.float32)
+    idx = np.nonzero(mask)
+
+    img[idx[0], idx[1], :] *= 1.0 - alpha
+    img[idx[0], idx[1], :] += [alpha * x for x in color]
+
+    if show_border:
+        contours, _ = findContours(
+            mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+        # contours = [c for c in contours if c.shape[0] > 10]
+        if border_color is None:
+            border_color = color
+        if not isinstance(border_color, list):
+            border_color = border_color.tolist()
+        if border_alpha < 1:
+            with_border = img.copy()
+            cv2.drawContours(with_border, contours, -1, border_color,
+                             border_thick, cv2.LINE_AA)
+            img = (1 - border_alpha) * img + border_alpha * with_border
+        else:
+            cv2.drawContours(img, contours, -1, border_color, border_thick,
+                             cv2.LINE_AA)
+
+    return img.astype(np.uint8)
+
+
+def colormap(rgb=False, lighten=True):
+    """Copied from Detectron codebase."""
+    color_list = np.array(
+        [
+            0.000, 0.447, 0.741,
+            0.850, 0.325, 0.098,
+            0.929, 0.694, 0.125,
+            0.494, 0.184, 0.556,
+            0.466, 0.674, 0.188,
+            0.301, 0.745, 0.933,
+            0.635, 0.078, 0.184,
+            0.300, 0.300, 0.300,
+            0.600, 0.600, 0.600,
+            1.000, 0.000, 0.000,
+            1.000, 0.500, 0.000,
+            0.749, 0.749, 0.000,
+            0.000, 1.000, 0.000,
+            0.000, 0.000, 1.000,
+            0.667, 0.000, 1.000,
+            0.333, 0.333, 0.000,
+            0.333, 0.667, 0.000,
+            0.333, 1.000, 0.000,
+            0.667, 0.333, 0.000,
+            0.667, 0.667, 0.000,
+            0.667, 1.000, 0.000,
+            1.000, 0.333, 0.000,
+            1.000, 0.667, 0.000,
+            1.000, 1.000, 0.000,
+            0.000, 0.333, 0.500,
+            0.000, 0.667, 0.500,
+            0.000, 1.000, 0.500,
+            0.333, 0.000, 0.500,
+            0.333, 0.333, 0.500,
+            0.333, 0.667, 0.500,
+            0.333, 1.000, 0.500,
+            0.667, 0.000, 0.500,
+            0.667, 0.333, 0.500,
+            0.667, 0.667, 0.500,
+            0.667, 1.000, 0.500,
+            1.000, 0.000, 0.500,
+            1.000, 0.333, 0.500,
+            1.000, 0.667, 0.500,
+            1.000, 1.000, 0.500,
+            0.000, 0.333, 1.000,
+            0.000, 0.667, 1.000,
+            0.000, 1.000, 1.000,
+            0.333, 0.000, 1.000,
+            0.333, 0.333, 1.000,
+            0.333, 0.667, 1.000,
+            0.333, 1.000, 1.000,
+            0.667, 0.000, 1.000,
+            0.667, 0.333, 1.000,
+            0.667, 0.667, 1.000,
+            0.667, 1.000, 1.000,
+            1.000, 0.000, 1.000,
+            1.000, 0.333, 1.000,
+            1.000, 0.667, 1.000,
+            0.167, 0.000, 0.000,
+            0.333, 0.000, 0.000,
+            0.500, 0.000, 0.000,
+            0.667, 0.000, 0.000,
+            0.833, 0.000, 0.000,
+            1.000, 0.000, 0.000,
+            0.000, 0.167, 0.000,
+            0.000, 0.333, 0.000,
+            0.000, 0.500, 0.000,
+            0.000, 0.667, 0.000,
+            0.000, 0.833, 0.000,
+            0.000, 1.000, 0.000,
+            0.000, 0.000, 0.167,
+            0.000, 0.000, 0.333,
+            0.000, 0.000, 0.500,
+            0.000, 0.000, 0.667,
+            0.000, 0.000, 0.833,
+            0.000, 0.000, 1.000,
+            0.000, 0.000, 0.000,
+            0.143, 0.143, 0.143,
+            0.286, 0.286, 0.286,
+            0.429, 0.429, 0.429,
+            0.571, 0.571, 0.571,
+            0.714, 0.714, 0.714,
+            0.857, 0.857, 0.857,
+            1.000, 1.000, 1.000
+        ]
+    ).astype(np.float32)
+    color_list = color_list.reshape((-1, 3))
+    if not rgb:
+        color_list = color_list[:, ::-1]
+
+    if lighten:
+        # Make all the colors a little lighter / whiter. This is copied
+        # from the detectron visualization code (search for 'w_ratio').
+        w_ratio = 0.4
+        color_list = (color_list * (1 - w_ratio) + w_ratio)
+    return color_list * 255
+
+
+def vis_layer_mask(masks, save_path=None):
+    masks = torch.as_tensor(masks)
+    tids = masks.unique().tolist()
+    tids.remove(0)
+    for tid in tqdm.tqdm(tids):
+        show = save_path is None
+        image_grid(masks == tid, label=f'{tid}', show=show)
+        if save_path:
+            os.makedirs(osp.dirname(save_path), exist_ok=True)
+            plt.savefig(save_path % tid)
+            plt.close('all')
+
+
+def show(x, **kwargs):
+    if isinstance(x, torch.Tensor):
+        x = x.detach().cpu()
+    plt.imshow(x, **kwargs)
+    plt.show()
+
+
+def vis_title(rgb, text, shift_y=30):
+    tmp = rgb.copy()
+    shift_x = rgb.shape[1] // 2
+    cv2.putText(tmp, text,
+                (shift_x, shift_y), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), thickness=2, lineType=cv2.LINE_AA)
+    return tmp

One-2-3-45-master 2/elevation_estimate/utils/utils3d.py

@@ -0,0 +1,62 @@
+import numpy as np
+import torch
+
+
+def cart_to_hom(pts):
+    """
+    :param pts: (N, 3 or 2)
+    :return pts_hom: (N, 4 or 3)
+    """
+    if isinstance(pts, np.ndarray):
+        pts_hom = np.concatenate((pts, np.ones([*pts.shape[:-1], 1], dtype=np.float32)), -1)
+    else:
+        ones = torch.ones([*pts.shape[:-1], 1], dtype=torch.float32, device=pts.device)
+        pts_hom = torch.cat((pts, ones), dim=-1)
+    return pts_hom
+
+
+def hom_to_cart(pts):
+    return pts[..., :-1] / pts[..., -1:]
+
+
+def canonical_to_camera(pts, pose):
+    pts = cart_to_hom(pts)
+    pts = pts @ pose.transpose(-1, -2)
+    pts = hom_to_cart(pts)
+    return pts
+
+
+def rect_to_img(K, pts_rect):
+    from dl_ext.vision_ext.datasets.kitti.structures import Calibration
+    pts_2d_hom = pts_rect @ K.t()
+    pts_img = Calibration.hom_to_cart(pts_2d_hom)
+    return pts_img
+
+
+def calc_pose(phis, thetas, size, radius=1.2):
+    import torch
+    def normalize(vectors):
+        return vectors / (torch.norm(vectors, dim=-1, keepdim=True) + 1e-10)
+
+    device = torch.device('cuda')
+    thetas = torch.FloatTensor(thetas).to(device)
+    phis = torch.FloatTensor(phis).to(device)
+
+    centers = torch.stack([
+        radius * torch.sin(thetas) * torch.sin(phis),
+        -radius * torch.cos(thetas) * torch.sin(phis),
+        radius * torch.cos(phis),
+    ], dim=-1)  # [B, 3]
+
+    # lookat
+    forward_vector = normalize(centers).squeeze(0)
+    up_vector = torch.FloatTensor([0, 0, 1]).to(device).unsqueeze(0).repeat(size, 1)
+    right_vector = normalize(torch.cross(up_vector, forward_vector, dim=-1))
+    if right_vector.pow(2).sum() < 0.01:
+        right_vector = torch.FloatTensor([0, 1, 0]).to(device).unsqueeze(0).repeat(size, 1)
+    up_vector = normalize(torch.cross(forward_vector, right_vector, dim=-1))
+
+    poses = torch.eye(4, dtype=torch.float, device=device).unsqueeze(0).repeat(size, 1, 1)
+    poses[:, :3, :3] = torch.stack((right_vector, up_vector, forward_vector), dim=-1)
+    poses[:, :3, 3] = centers
+    return poses

One-2-3-45-master 2/ldm/data/base.py

@@ -0,0 +1,40 @@
+import os
+import numpy as np
+from abc import abstractmethod
+from torch.utils.data import Dataset, ConcatDataset, ChainDataset, IterableDataset
+
+
+class Txt2ImgIterableBaseDataset(IterableDataset):
+    '''
+    Define an interface to make the IterableDatasets for text2img data chainable
+    '''
+    def __init__(self, num_records=0, valid_ids=None, size=256):
+        super().__init__()
+        self.num_records = num_records
+        self.valid_ids = valid_ids
+        self.sample_ids = valid_ids
+        self.size = size
+
+        print(f'{self.__class__.__name__} dataset contains {self.__len__()} examples.')
+
+    def __len__(self):
+        return self.num_records
+
+    @abstractmethod
+    def __iter__(self):
+        pass
+
+
+class PRNGMixin(object):
+    """
+    Adds a prng property which is a numpy RandomState which gets
+    reinitialized whenever the pid changes to avoid synchronized sampling
+    behavior when used in conjunction with multiprocessing.
+    """
+    @property
+    def prng(self):
+        currentpid = os.getpid()
+        if getattr(self, "_initpid", None) != currentpid:
+            self._initpid = currentpid
+            self._prng = np.random.RandomState()
+        return self._prng

One-2-3-45-master 2/ldm/data/coco.py

@@ -0,0 +1,253 @@
+import os
+import json
+import albumentations
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from torch.utils.data import Dataset
+from abc import abstractmethod
+
+
+class CocoBase(Dataset):
+    """needed for (image, caption, segmentation) pairs"""
+    def __init__(self, size=None, dataroot="", datajson="", onehot_segmentation=False, use_stuffthing=False,
+                 crop_size=None, force_no_crop=False, given_files=None, use_segmentation=True,crop_type=None):
+        self.split = self.get_split()
+        self.size = size
+        if crop_size is None:
+            self.crop_size = size
+        else:
+            self.crop_size = crop_size
+
+        assert crop_type in [None, 'random', 'center']
+        self.crop_type = crop_type
+        self.use_segmenation = use_segmentation
+        self.onehot = onehot_segmentation       # return segmentation as rgb or one hot
+        self.stuffthing = use_stuffthing        # include thing in segmentation
+        if self.onehot and not self.stuffthing:
+            raise NotImplemented("One hot mode is only supported for the "
+                                 "stuffthings version because labels are stored "
+                                 "a bit different.")
+
+        data_json = datajson
+        with open(data_json) as json_file:
+            self.json_data = json.load(json_file)
+            self.img_id_to_captions = dict()
+            self.img_id_to_filepath = dict()
+            self.img_id_to_segmentation_filepath = dict()
+
+        assert data_json.split("/")[-1] in [f"captions_train{self.year()}.json",
+                                            f"captions_val{self.year()}.json"]
+        # TODO currently hardcoded paths, would be better to follow logic in
+        # cocstuff pixelmaps
+        if self.use_segmenation:
+            if self.stuffthing:
+                self.segmentation_prefix = (
+                    f"data/cocostuffthings/val{self.year()}" if
+                    data_json.endswith(f"captions_val{self.year()}.json") else
+                    f"data/cocostuffthings/train{self.year()}")
+            else:
+                self.segmentation_prefix = (
+                    f"data/coco/annotations/stuff_val{self.year()}_pixelmaps" if
+                    data_json.endswith(f"captions_val{self.year()}.json") else
+                    f"data/coco/annotations/stuff_train{self.year()}_pixelmaps")
+
+        imagedirs = self.json_data["images"]
+        self.labels = {"image_ids": list()}
+        for imgdir in tqdm(imagedirs, desc="ImgToPath"):
+            self.img_id_to_filepath[imgdir["id"]] = os.path.join(dataroot, imgdir["file_name"])
+            self.img_id_to_captions[imgdir["id"]] = list()
+            pngfilename = imgdir["file_name"].replace("jpg", "png")
+            if self.use_segmenation:
+                self.img_id_to_segmentation_filepath[imgdir["id"]] = os.path.join(
+                    self.segmentation_prefix, pngfilename)
+            if given_files is not None:
+                if pngfilename in given_files:
+                    self.labels["image_ids"].append(imgdir["id"])
+            else:
+                self.labels["image_ids"].append(imgdir["id"])
+
+        capdirs = self.json_data["annotations"]
+        for capdir in tqdm(capdirs, desc="ImgToCaptions"):
+            # there are in average 5 captions per image
+            #self.img_id_to_captions[capdir["image_id"]].append(np.array([capdir["caption"]]))
+            self.img_id_to_captions[capdir["image_id"]].append(capdir["caption"])
+
+        self.rescaler = albumentations.SmallestMaxSize(max_size=self.size)
+        if self.split=="validation":
+            self.cropper = albumentations.CenterCrop(height=self.crop_size, width=self.crop_size)
+        else:
+            # default option for train is random crop
+            if self.crop_type in [None, 'random']:
+                self.cropper = albumentations.RandomCrop(height=self.crop_size, width=self.crop_size)
+            else:
+                self.cropper = albumentations.CenterCrop(height=self.crop_size, width=self.crop_size)
+        self.preprocessor = albumentations.Compose(
+            [self.rescaler, self.cropper],
+            additional_targets={"segmentation": "image"})
+        if force_no_crop:
+            self.rescaler = albumentations.Resize(height=self.size, width=self.size)
+            self.preprocessor = albumentations.Compose(
+                [self.rescaler],
+                additional_targets={"segmentation": "image"})
+
+    @abstractmethod
+    def year(self):
+        raise NotImplementedError()
+
+    def __len__(self):
+        return len(self.labels["image_ids"])
+
+    def preprocess_image(self, image_path, segmentation_path=None):
+        image = Image.open(image_path)
+        if not image.mode == "RGB":
+            image = image.convert("RGB")
+        image = np.array(image).astype(np.uint8)
+        if segmentation_path:
+            segmentation = Image.open(segmentation_path)
+            if not self.onehot and not segmentation.mode == "RGB":
+                segmentation = segmentation.convert("RGB")
+            segmentation = np.array(segmentation).astype(np.uint8)
+            if self.onehot:
+                assert self.stuffthing
+                # stored in caffe format: unlabeled==255. stuff and thing from
+                # 0-181. to be compatible with the labels in
+                # https://github.com/nightrome/cocostuff/blob/master/labels.txt
+                # we shift stuffthing one to the right and put unlabeled in zero
+                # as long as segmentation is uint8 shifting to right handles the
+                # latter too
+                assert segmentation.dtype == np.uint8
+                segmentation = segmentation + 1
+
+            processed = self.preprocessor(image=image, segmentation=segmentation)
+
+            image, segmentation = processed["image"], processed["segmentation"]
+        else:
+            image = self.preprocessor(image=image,)['image']
+
+        image = (image / 127.5 - 1.0).astype(np.float32)
+        if segmentation_path:
+            if self.onehot:
+                assert segmentation.dtype == np.uint8
+                # make it one hot
+                n_labels = 183
+                flatseg = np.ravel(segmentation)
+                onehot = np.zeros((flatseg.size, n_labels), dtype=np.bool)
+                onehot[np.arange(flatseg.size), flatseg] = True
+                onehot = onehot.reshape(segmentation.shape + (n_labels,)).astype(int)
+                segmentation = onehot
+            else:
+                segmentation = (segmentation / 127.5 - 1.0).astype(np.float32)
+            return image, segmentation
+        else:
+            return image
+
+    def __getitem__(self, i):
+        img_path = self.img_id_to_filepath[self.labels["image_ids"][i]]
+        if self.use_segmenation:
+            seg_path = self.img_id_to_segmentation_filepath[self.labels["image_ids"][i]]
+            image, segmentation = self.preprocess_image(img_path, seg_path)
+        else:
+            image = self.preprocess_image(img_path)
+        captions = self.img_id_to_captions[self.labels["image_ids"][i]]
+        # randomly draw one of all available captions per image
+        caption = captions[np.random.randint(0, len(captions))]
+        example = {"image": image,
+                   #"caption": [str(caption[0])],
+                   "caption": caption,
+                   "img_path": img_path,
+                   "filename_": img_path.split(os.sep)[-1]
+                    }
+        if self.use_segmenation:
+            example.update({"seg_path": seg_path, 'segmentation': segmentation})
+        return example
+
+
+class CocoImagesAndCaptionsTrain2017(CocoBase):
+    """returns a pair of (image, caption)"""
+    def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,):
+        super().__init__(size=size,
+                         dataroot="data/coco/train2017",
+                         datajson="data/coco/annotations/captions_train2017.json",
+                         onehot_segmentation=onehot_segmentation,
+                         use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop)
+
+    def get_split(self):
+        return "train"
+
+    def year(self):
+        return '2017'
+
+
+class CocoImagesAndCaptionsValidation2017(CocoBase):
+    """returns a pair of (image, caption)"""
+    def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,
+                 given_files=None):
+        super().__init__(size=size,
+                         dataroot="data/coco/val2017",
+                         datajson="data/coco/annotations/captions_val2017.json",
+                         onehot_segmentation=onehot_segmentation,
+                         use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop,
+                         given_files=given_files)
+
+    def get_split(self):
+        return "validation"
+
+    def year(self):
+        return '2017'
+
+
+
+class CocoImagesAndCaptionsTrain2014(CocoBase):
+    """returns a pair of (image, caption)"""
+    def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,crop_type='random'):
+        super().__init__(size=size,
+                         dataroot="data/coco/train2014",
+                         datajson="data/coco/annotations2014/annotations/captions_train2014.json",
+                         onehot_segmentation=onehot_segmentation,
+                         use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop,
+                         use_segmentation=False,
+                         crop_type=crop_type)
+
+    def get_split(self):
+        return "train"
+
+    def year(self):
+        return '2014'
+
+class CocoImagesAndCaptionsValidation2014(CocoBase):
+    """returns a pair of (image, caption)"""
+    def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,
+                 given_files=None,crop_type='center',**kwargs):
+        super().__init__(size=size,
+                         dataroot="data/coco/val2014",
+                         datajson="data/coco/annotations2014/annotations/captions_val2014.json",
+                         onehot_segmentation=onehot_segmentation,
+                         use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop,
+                         given_files=given_files,
+                         use_segmentation=False,
+                         crop_type=crop_type)
+
+    def get_split(self):
+        return "validation"
+
+    def year(self):
+        return '2014'
+
+if __name__ == '__main__':
+    with open("data/coco/annotations2014/annotations/captions_val2014.json", "r") as json_file:
+        json_data = json.load(json_file)
+        capdirs = json_data["annotations"]
+        import pudb; pudb.set_trace()
+    #d2 = CocoImagesAndCaptionsTrain2014(size=256)
+    d2 = CocoImagesAndCaptionsValidation2014(size=256)
+    print("constructed dataset.")
+    print(f"length of {d2.__class__.__name__}: {len(d2)}")
+
+    ex2 = d2[0]
+    # ex3 = d3[0]
+    # print(ex1["image"].shape)
+    print(ex2["image"].shape)
+    # print(ex3["image"].shape)
+    # print(ex1["segmentation"].shape)
+    print(ex2["caption"].__class__.__name__)

One-2-3-45-master 2/ldm/data/dummy.py

@@ -0,0 +1,34 @@
+import numpy as np
+import random
+import string
+from torch.utils.data import Dataset, Subset
+
+class DummyData(Dataset):
+    def __init__(self, length, size):
+        self.length = length
+        self.size = size
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, i):
+        x = np.random.randn(*self.size)
+        letters = string.ascii_lowercase
+        y = ''.join(random.choice(string.ascii_lowercase) for i in range(10))
+        return {"jpg": x, "txt": y}
+
+
+class DummyDataWithEmbeddings(Dataset):
+    def __init__(self, length, size, emb_size):
+        self.length = length
+        self.size = size
+        self.emb_size = emb_size
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, i):
+        x = np.random.randn(*self.size)
+        y = np.random.randn(*self.emb_size).astype(np.float32)
+        return {"jpg": x, "txt": y}
+

One-2-3-45-master 2/ldm/data/imagenet.py

@@ -0,0 +1,394 @@
+import os, yaml, pickle, shutil, tarfile, glob
+import cv2
+import albumentations
+import PIL
+import numpy as np
+import torchvision.transforms.functional as TF
+from omegaconf import OmegaConf
+from functools import partial
+from PIL import Image
+from tqdm import tqdm
+from torch.utils.data import Dataset, Subset
+
+import taming.data.utils as tdu
+from taming.data.imagenet import str_to_indices, give_synsets_from_indices, download, retrieve
+from taming.data.imagenet import ImagePaths
+
+from ldm.modules.image_degradation import degradation_fn_bsr, degradation_fn_bsr_light
+
+
+def synset2idx(path_to_yaml="data/index_synset.yaml"):
+    with open(path_to_yaml) as f:
+        di2s = yaml.load(f)
+    return dict((v,k) for k,v in di2s.items())
+
+
+class ImageNetBase(Dataset):
+    def __init__(self, config=None):
+        self.config = config or OmegaConf.create()
+        if not type(self.config)==dict:
+            self.config = OmegaConf.to_container(self.config)
+        self.keep_orig_class_label = self.config.get("keep_orig_class_label", False)
+        self.process_images = True  # if False we skip loading & processing images and self.data contains filepaths
+        self._prepare()
+        self._prepare_synset_to_human()
+        self._prepare_idx_to_synset()
+        self._prepare_human_to_integer_label()
+        self._load()
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, i):
+        return self.data[i]
+
+    def _prepare(self):
+        raise NotImplementedError()
+
+    def _filter_relpaths(self, relpaths):
+        ignore = set([
+            "n06596364_9591.JPEG",
+        ])
+        relpaths = [rpath for rpath in relpaths if not rpath.split("/")[-1] in ignore]
+        if "sub_indices" in self.config:
+            indices = str_to_indices(self.config["sub_indices"])
+            synsets = give_synsets_from_indices(indices, path_to_yaml=self.idx2syn)  # returns a list of strings
+            self.synset2idx = synset2idx(path_to_yaml=self.idx2syn)
+            files = []
+            for rpath in relpaths:
+                syn = rpath.split("/")[0]
+                if syn in synsets:
+                    files.append(rpath)
+            return files
+        else:
+            return relpaths
+
+    def _prepare_synset_to_human(self):
+        SIZE = 2655750
+        URL = "https://heibox.uni-heidelberg.de/f/9f28e956cd304264bb82/?dl=1"
+        self.human_dict = os.path.join(self.root, "synset_human.txt")
+        if (not os.path.exists(self.human_dict) or
+                not os.path.getsize(self.human_dict)==SIZE):
+            download(URL, self.human_dict)
+
+    def _prepare_idx_to_synset(self):
+        URL = "https://heibox.uni-heidelberg.de/f/d835d5b6ceda4d3aa910/?dl=1"
+        self.idx2syn = os.path.join(self.root, "index_synset.yaml")
+        if (not os.path.exists(self.idx2syn)):
+            download(URL, self.idx2syn)
+
+    def _prepare_human_to_integer_label(self):
+        URL = "https://heibox.uni-heidelberg.de/f/2362b797d5be43b883f6/?dl=1"
+        self.human2integer = os.path.join(self.root, "imagenet1000_clsidx_to_labels.txt")
+        if (not os.path.exists(self.human2integer)):
+            download(URL, self.human2integer)
+        with open(self.human2integer, "r") as f:
+            lines = f.read().splitlines()
+            assert len(lines) == 1000
+            self.human2integer_dict = dict()
+            for line in lines:
+                value, key = line.split(":")
+                self.human2integer_dict[key] = int(value)
+
+    def _load(self):
+        with open(self.txt_filelist, "r") as f:
+            self.relpaths = f.read().splitlines()
+            l1 = len(self.relpaths)
+            self.relpaths = self._filter_relpaths(self.relpaths)
+            print("Removed {} files from filelist during filtering.".format(l1 - len(self.relpaths)))
+
+        self.synsets = [p.split("/")[0] for p in self.relpaths]
+        self.abspaths = [os.path.join(self.datadir, p) for p in self.relpaths]
+
+        unique_synsets = np.unique(self.synsets)
+        class_dict = dict((synset, i) for i, synset in enumerate(unique_synsets))
+        if not self.keep_orig_class_label:
+            self.class_labels = [class_dict[s] for s in self.synsets]
+        else:
+            self.class_labels = [self.synset2idx[s] for s in self.synsets]
+
+        with open(self.human_dict, "r") as f:
+            human_dict = f.read().splitlines()
+            human_dict = dict(line.split(maxsplit=1) for line in human_dict)
+
+        self.human_labels = [human_dict[s] for s in self.synsets]
+
+        labels = {
+            "relpath": np.array(self.relpaths),
+            "synsets": np.array(self.synsets),
+            "class_label": np.array(self.class_labels),
+            "human_label": np.array(self.human_labels),
+        }
+
+        if self.process_images:
+            self.size = retrieve(self.config, "size", default=256)
+            self.data = ImagePaths(self.abspaths,
+                                   labels=labels,
+                                   size=self.size,
+                                   random_crop=self.random_crop,
+                                   )
+        else:
+            self.data = self.abspaths
+
+
+class ImageNetTrain(ImageNetBase):
+    NAME = "ILSVRC2012_train"
+    URL = "http://www.image-net.org/challenges/LSVRC/2012/"
+    AT_HASH = "a306397ccf9c2ead27155983c254227c0fd938e2"
+    FILES = [
+        "ILSVRC2012_img_train.tar",
+    ]
+    SIZES = [
+        147897477120,
+    ]
+
+    def __init__(self, process_images=True, data_root=None, **kwargs):
+        self.process_images = process_images
+        self.data_root = data_root
+        super().__init__(**kwargs)
+
+    def _prepare(self):
+        if self.data_root:
+            self.root = os.path.join(self.data_root, self.NAME)
+        else:
+            cachedir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
+            self.root = os.path.join(cachedir, "autoencoders/data", self.NAME)
+
+        self.datadir = os.path.join(self.root, "data")
+        self.txt_filelist = os.path.join(self.root, "filelist.txt")
+        self.expected_length = 1281167
+        self.random_crop = retrieve(self.config, "ImageNetTrain/random_crop",
+                                    default=True)
+        if not tdu.is_prepared(self.root):
+            # prep
+            print("Preparing dataset {} in {}".format(self.NAME, self.root))
+
+            datadir = self.datadir
+            if not os.path.exists(datadir):
+                path = os.path.join(self.root, self.FILES[0])
+                if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
+                    import academictorrents as at
+                    atpath = at.get(self.AT_HASH, datastore=self.root)
+                    assert atpath == path
+
+                print("Extracting {} to {}".format(path, datadir))
+                os.makedirs(datadir, exist_ok=True)
+                with tarfile.open(path, "r:") as tar:
+                    tar.extractall(path=datadir)
+
+                print("Extracting sub-tars.")
+                subpaths = sorted(glob.glob(os.path.join(datadir, "*.tar")))
+                for subpath in tqdm(subpaths):
+                    subdir = subpath[:-len(".tar")]
+                    os.makedirs(subdir, exist_ok=True)
+                    with tarfile.open(subpath, "r:") as tar:
+                        tar.extractall(path=subdir)
+
+            filelist = glob.glob(os.path.join(datadir, "**", "*.JPEG"))
+            filelist = [os.path.relpath(p, start=datadir) for p in filelist]
+            filelist = sorted(filelist)
+            filelist = "\n".join(filelist)+"\n"
+            with open(self.txt_filelist, "w") as f:
+                f.write(filelist)
+
+            tdu.mark_prepared(self.root)
+
+
+class ImageNetValidation(ImageNetBase):
+    NAME = "ILSVRC2012_validation"
+    URL = "http://www.image-net.org/challenges/LSVRC/2012/"
+    AT_HASH = "5d6d0df7ed81efd49ca99ea4737e0ae5e3a5f2e5"
+    VS_URL = "https://heibox.uni-heidelberg.de/f/3e0f6e9c624e45f2bd73/?dl=1"
+    FILES = [
+        "ILSVRC2012_img_val.tar",
+        "validation_synset.txt",
+    ]
+    SIZES = [
+        6744924160,
+        1950000,
+    ]
+
+    def __init__(self, process_images=True, data_root=None, **kwargs):
+        self.data_root = data_root
+        self.process_images = process_images
+        super().__init__(**kwargs)
+
+    def _prepare(self):
+        if self.data_root:
+            self.root = os.path.join(self.data_root, self.NAME)
+        else:
+            cachedir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
+            self.root = os.path.join(cachedir, "autoencoders/data", self.NAME)
+        self.datadir = os.path.join(self.root, "data")
+        self.txt_filelist = os.path.join(self.root, "filelist.txt")
+        self.expected_length = 50000
+        self.random_crop = retrieve(self.config, "ImageNetValidation/random_crop",
+                                    default=False)
+        if not tdu.is_prepared(self.root):
+            # prep
+            print("Preparing dataset {} in {}".format(self.NAME, self.root))
+
+            datadir = self.datadir
+            if not os.path.exists(datadir):
+                path = os.path.join(self.root, self.FILES[0])
+                if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
+                    import academictorrents as at
+                    atpath = at.get(self.AT_HASH, datastore=self.root)
+                    assert atpath == path
+
+                print("Extracting {} to {}".format(path, datadir))
+                os.makedirs(datadir, exist_ok=True)
+                with tarfile.open(path, "r:") as tar:
+                    tar.extractall(path=datadir)
+
+                vspath = os.path.join(self.root, self.FILES[1])
+                if not os.path.exists(vspath) or not os.path.getsize(vspath)==self.SIZES[1]:
+                    download(self.VS_URL, vspath)
+
+                with open(vspath, "r") as f:
+                    synset_dict = f.read().splitlines()
+                    synset_dict = dict(line.split() for line in synset_dict)
+
+                print("Reorganizing into synset folders")
+                synsets = np.unique(list(synset_dict.values()))
+                for s in synsets:
+                    os.makedirs(os.path.join(datadir, s), exist_ok=True)
+                for k, v in synset_dict.items():
+                    src = os.path.join(datadir, k)
+                    dst = os.path.join(datadir, v)
+                    shutil.move(src, dst)
+
+            filelist = glob.glob(os.path.join(datadir, "**", "*.JPEG"))
+            filelist = [os.path.relpath(p, start=datadir) for p in filelist]
+            filelist = sorted(filelist)
+            filelist = "\n".join(filelist)+"\n"
+            with open(self.txt_filelist, "w") as f:
+                f.write(filelist)
+
+            tdu.mark_prepared(self.root)
+
+
+
+class ImageNetSR(Dataset):
+    def __init__(self, size=None,
+                 degradation=None, downscale_f=4, min_crop_f=0.5, max_crop_f=1.,
+                 random_crop=True):
+        """
+        Imagenet Superresolution Dataloader
+        Performs following ops in order:
+        1.  crops a crop of size s from image either as random or center crop
+        2.  resizes crop to size with cv2.area_interpolation
+        3.  degrades resized crop with degradation_fn
+
+        :param size: resizing to size after cropping
+        :param degradation: degradation_fn, e.g. cv_bicubic or bsrgan_light
+        :param downscale_f: Low Resolution Downsample factor
+        :param min_crop_f: determines crop size s,
+          where s = c * min_img_side_len with c sampled from interval (min_crop_f, max_crop_f)
+        :param max_crop_f: ""
+        :param data_root:
+        :param random_crop:
+        """
+        self.base = self.get_base()
+        assert size
+        assert (size / downscale_f).is_integer()
+        self.size = size
+        self.LR_size = int(size / downscale_f)
+        self.min_crop_f = min_crop_f
+        self.max_crop_f = max_crop_f
+        assert(max_crop_f <= 1.)
+        self.center_crop = not random_crop
+
+        self.image_rescaler = albumentations.SmallestMaxSize(max_size=size, interpolation=cv2.INTER_AREA)
+
+        self.pil_interpolation = False # gets reset later if incase interp_op is from pillow
+
+        if degradation == "bsrgan":
+            self.degradation_process = partial(degradation_fn_bsr, sf=downscale_f)
+
+        elif degradation == "bsrgan_light":
+            self.degradation_process = partial(degradation_fn_bsr_light, sf=downscale_f)
+
+        else:
+            interpolation_fn = {
+            "cv_nearest": cv2.INTER_NEAREST,
+            "cv_bilinear": cv2.INTER_LINEAR,
+            "cv_bicubic": cv2.INTER_CUBIC,
+            "cv_area": cv2.INTER_AREA,
+            "cv_lanczos": cv2.INTER_LANCZOS4,
+            "pil_nearest": PIL.Image.NEAREST,
+            "pil_bilinear": PIL.Image.BILINEAR,
+            "pil_bicubic": PIL.Image.BICUBIC,
+            "pil_box": PIL.Image.BOX,
+            "pil_hamming": PIL.Image.HAMMING,
+            "pil_lanczos": PIL.Image.LANCZOS,
+            }[degradation]
+
+            self.pil_interpolation = degradation.startswith("pil_")
+
+            if self.pil_interpolation:
+                self.degradation_process = partial(TF.resize, size=self.LR_size, interpolation=interpolation_fn)
+
+            else:
+                self.degradation_process = albumentations.SmallestMaxSize(max_size=self.LR_size,
+                                                                          interpolation=interpolation_fn)
+
+    def __len__(self):
+        return len(self.base)
+
+    def __getitem__(self, i):
+        example = self.base[i]
+        image = Image.open(example["file_path_"])
+
+        if not image.mode == "RGB":
+            image = image.convert("RGB")
+
+        image = np.array(image).astype(np.uint8)
+
+        min_side_len = min(image.shape[:2])
+        crop_side_len = min_side_len * np.random.uniform(self.min_crop_f, self.max_crop_f, size=None)
+        crop_side_len = int(crop_side_len)
+
+        if self.center_crop:
+            self.cropper = albumentations.CenterCrop(height=crop_side_len, width=crop_side_len)
+
+        else:
+            self.cropper = albumentations.RandomCrop(height=crop_side_len, width=crop_side_len)
+
+        image = self.cropper(image=image)["image"]
+        image = self.image_rescaler(image=image)["image"]
+
+        if self.pil_interpolation:
+            image_pil = PIL.Image.fromarray(image)
+            LR_image = self.degradation_process(image_pil)
+            LR_image = np.array(LR_image).astype(np.uint8)
+
+        else:
+            LR_image = self.degradation_process(image=image)["image"]
+
+        example["image"] = (image/127.5 - 1.0).astype(np.float32)
+        example["LR_image"] = (LR_image/127.5 - 1.0).astype(np.float32)
+        example["caption"] = example["human_label"]  # dummy caption
+        return example
+
+
+class ImageNetSRTrain(ImageNetSR):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def get_base(self):
+        with open("data/imagenet_train_hr_indices.p", "rb") as f:
+            indices = pickle.load(f)
+        dset = ImageNetTrain(process_images=False,)
+        return Subset(dset, indices)
+
+
+class ImageNetSRValidation(ImageNetSR):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def get_base(self):
+        with open("data/imagenet_val_hr_indices.p", "rb") as f:
+            indices = pickle.load(f)
+        dset = ImageNetValidation(process_images=False,)
+        return Subset(dset, indices)

One-2-3-45-master 2/ldm/data/inpainting/synthetic_mask.py

@@ -0,0 +1,166 @@
+from PIL import Image, ImageDraw
+import numpy as np
+
+settings = {
+    "256narrow": {
+        "p_irr": 1,
+        "min_n_irr": 4,
+        "max_n_irr": 50,
+        "max_l_irr": 40,
+        "max_w_irr": 10,
+        "min_n_box": None,
+        "max_n_box": None,
+        "min_s_box": None,
+        "max_s_box": None,
+        "marg": None,
+    },
+    "256train": {
+        "p_irr": 0.5,
+        "min_n_irr": 1,
+        "max_n_irr": 5,
+        "max_l_irr": 200,
+        "max_w_irr": 100,
+        "min_n_box": 1,
+        "max_n_box": 4,
+        "min_s_box": 30,
+        "max_s_box": 150,
+        "marg": 10,
+    },
+    "512train": {    # TODO: experimental
+            "p_irr": 0.5,
+            "min_n_irr": 1,
+            "max_n_irr": 5,
+            "max_l_irr": 450,
+            "max_w_irr": 250,
+            "min_n_box": 1,
+            "max_n_box": 4,
+            "min_s_box": 30,
+            "max_s_box": 300,
+            "marg": 10,
+        },
+    "512train-large": {    # TODO: experimental
+            "p_irr": 0.5,
+            "min_n_irr": 1,
+            "max_n_irr": 5,
+            "max_l_irr": 450,
+            "max_w_irr": 400,
+            "min_n_box": 1,
+            "max_n_box": 4,
+            "min_s_box": 75,
+            "max_s_box": 450,
+            "marg": 10,
+        },
+}
+
+
+def gen_segment_mask(mask, start, end, brush_width):
+    mask = mask > 0
+    mask = (255 * mask).astype(np.uint8)
+    mask = Image.fromarray(mask)
+    draw = ImageDraw.Draw(mask)
+    draw.line([start, end], fill=255, width=brush_width, joint="curve")
+    mask = np.array(mask) / 255
+    return mask
+
+
+def gen_box_mask(mask, masked):
+    x_0, y_0, w, h = masked
+    mask[y_0:y_0 + h, x_0:x_0 + w] = 1
+    return mask
+
+
+def gen_round_mask(mask, masked, radius):
+    x_0, y_0, w, h = masked
+    xy = [(x_0, y_0), (x_0 + w, y_0 + w)]
+
+    mask = mask > 0
+    mask = (255 * mask).astype(np.uint8)
+    mask = Image.fromarray(mask)
+    draw = ImageDraw.Draw(mask)
+    draw.rounded_rectangle(xy, radius=radius, fill=255)
+    mask = np.array(mask) / 255
+    return mask
+
+
+def gen_large_mask(prng, img_h, img_w,
+                   marg, p_irr, min_n_irr, max_n_irr, max_l_irr, max_w_irr,
+                   min_n_box, max_n_box, min_s_box, max_s_box):
+    """
+    img_h: int, an image height
+    img_w: int, an image width
+    marg: int, a margin for a box starting coordinate
+    p_irr: float, 0 <= p_irr <= 1, a probability of a polygonal chain mask
+
+    min_n_irr: int, min number of segments
+    max_n_irr: int, max number of segments
+    max_l_irr: max length of a segment in polygonal chain
+    max_w_irr: max width of a segment in polygonal chain
+
+    min_n_box: int, min bound for the number of box primitives
+    max_n_box: int, max bound for the number of box primitives
+    min_s_box: int, min length of a box side
+    max_s_box: int, max length of a box side
+    """
+
+    mask = np.zeros((img_h, img_w))
+    uniform = prng.randint
+
+    if np.random.uniform(0, 1) < p_irr:  # generate polygonal chain
+        n = uniform(min_n_irr, max_n_irr)  # sample number of segments
+
+        for _ in range(n):
+            y = uniform(0, img_h)  # sample a starting point
+            x = uniform(0, img_w)
+
+            a = uniform(0, 360)  # sample angle
+            l = uniform(10, max_l_irr)  # sample segment length
+            w = uniform(5, max_w_irr)  # sample a segment width
+
+            # draw segment starting from (x,y) to (x_,y_) using brush of width w
+            x_ = x + l * np.sin(a)
+            y_ = y + l * np.cos(a)
+
+            mask = gen_segment_mask(mask, start=(x, y), end=(x_, y_), brush_width=w)
+            x, y = x_, y_
+    else:  # generate Box masks
+        n = uniform(min_n_box, max_n_box)  # sample number of rectangles
+
+        for _ in range(n):
+            h = uniform(min_s_box, max_s_box)  # sample box shape
+            w = uniform(min_s_box, max_s_box)
+
+            x_0 = uniform(marg, img_w - marg - w)  # sample upper-left coordinates of box
+            y_0 = uniform(marg, img_h - marg - h)
+
+            if np.random.uniform(0, 1) < 0.5:
+                mask = gen_box_mask(mask, masked=(x_0, y_0, w, h))
+            else:
+                r = uniform(0, 60)  # sample radius
+                mask = gen_round_mask(mask, masked=(x_0, y_0, w, h), radius=r)
+    return mask
+
+
+make_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["256train"])
+make_narrow_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["256narrow"])
+make_512_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["512train"])
+make_512_lama_mask_large = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["512train-large"])
+
+
+MASK_MODES = {
+    "256train": make_lama_mask,
+    "256narrow": make_narrow_lama_mask,
+    "512train": make_512_lama_mask,
+    "512train-large": make_512_lama_mask_large
+}
+
+if __name__ == "__main__":
+    import sys
+
+    out = sys.argv[1]
+
+    prng = np.random.RandomState(1)
+    kwargs = settings["256train"]
+    mask = gen_large_mask(prng, 256, 256, **kwargs)
+    mask = (255 * mask).astype(np.uint8)
+    mask = Image.fromarray(mask)
+    mask.save(out)

One-2-3-45-master 2/ldm/data/laion.py

@@ -0,0 +1,537 @@
+import webdataset as wds
+import kornia
+from PIL import Image
+import io
+import os
+import torchvision
+from PIL import Image
+import glob
+import random
+import numpy as np
+import pytorch_lightning as pl
+from tqdm import tqdm
+from omegaconf import OmegaConf
+from einops import rearrange
+import torch
+from webdataset.handlers import warn_and_continue
+
+
+from ldm.util import instantiate_from_config
+from ldm.data.inpainting.synthetic_mask import gen_large_mask, MASK_MODES
+from ldm.data.base import PRNGMixin
+
+
+class DataWithWings(torch.utils.data.IterableDataset):
+    def __init__(self, min_size, transform=None, target_transform=None):
+        self.min_size = min_size
+        self.transform = transform if transform is not None else nn.Identity()
+        self.target_transform = target_transform if target_transform is not None else nn.Identity()
+        self.kv = OnDiskKV(file='/home/ubuntu/laion5B-watermark-safety-ordered', key_format='q', value_format='ee')
+        self.kv_aesthetic = OnDiskKV(file='/home/ubuntu/laion5B-aesthetic-tags-kv', key_format='q', value_format='e')
+        self.pwatermark_threshold = 0.8
+        self.punsafe_threshold = 0.5
+        self.aesthetic_threshold = 5.
+        self.total_samples = 0
+        self.samples = 0
+        location = 'pipe:aws s3 cp --quiet s3://s-datasets/laion5b/laion2B-data/{000000..231349}.tar -'
+
+        self.inner_dataset = wds.DataPipeline(
+            wds.ResampledShards(location),
+            wds.tarfile_to_samples(handler=wds.warn_and_continue),
+            wds.shuffle(1000, handler=wds.warn_and_continue),
+            wds.decode('pilrgb', handler=wds.warn_and_continue),
+            wds.map(self._add_tags, handler=wds.ignore_and_continue),
+            wds.select(self._filter_predicate),
+            wds.map_dict(jpg=self.transform, txt=self.target_transform, punsafe=self._punsafe_to_class, handler=wds.warn_and_continue),
+            wds.to_tuple('jpg', 'txt', 'punsafe', handler=wds.warn_and_continue),
+        )
+
+    @staticmethod
+    def _compute_hash(url, text):
+        if url is None:
+            url = ''
+        if text is None:
+            text = ''
+        total = (url + text).encode('utf-8')
+        return mmh3.hash64(total)[0]
+
+    def _add_tags(self, x):
+        hsh = self._compute_hash(x['json']['url'], x['txt'])
+        pwatermark, punsafe = self.kv[hsh]
+        aesthetic = self.kv_aesthetic[hsh][0]
+        return {**x, 'pwatermark': pwatermark, 'punsafe': punsafe, 'aesthetic': aesthetic}
+
+    def _punsafe_to_class(self, punsafe):
+        return torch.tensor(punsafe >= self.punsafe_threshold).long()
+
+    def _filter_predicate(self, x):
+        try:
+            return x['pwatermark'] < self.pwatermark_threshold and x['aesthetic'] >= self.aesthetic_threshold and x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size
+        except:
+            return False
+
+    def __iter__(self):
+        return iter(self.inner_dataset)
+
+
+def dict_collation_fn(samples, combine_tensors=True, combine_scalars=True):
+    """Take a list  of samples (as dictionary) and create a batch, preserving the keys.
+    If `tensors` is True, `ndarray` objects are combined into
+    tensor batches.
+    :param dict samples: list of samples
+    :param bool tensors: whether to turn lists of ndarrays into a single ndarray
+    :returns: single sample consisting of a batch
+    :rtype: dict
+    """
+    keys = set.intersection(*[set(sample.keys()) for sample in samples])
+    batched = {key: [] for key in keys}
+
+    for s in samples:
+        [batched[key].append(s[key]) for key in batched]
+
+    result = {}
+    for key in batched:
+        if isinstance(batched[key][0], (int, float)):
+            if combine_scalars:
+                result[key] = np.array(list(batched[key]))
+        elif isinstance(batched[key][0], torch.Tensor):
+            if combine_tensors:
+                result[key] = torch.stack(list(batched[key]))
+        elif isinstance(batched[key][0], np.ndarray):
+            if combine_tensors:
+                result[key] = np.array(list(batched[key]))
+        else:
+            result[key] = list(batched[key])
+    return result
+
+
+class WebDataModuleFromConfig(pl.LightningDataModule):
+    def __init__(self, tar_base, batch_size, train=None, validation=None,
+                 test=None, num_workers=4, multinode=True, min_size=None,
+                 max_pwatermark=1.0,
+                 **kwargs):
+        super().__init__(self)
+        print(f'Setting tar base to {tar_base}')
+        self.tar_base = tar_base
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.train = train
+        self.validation = validation
+        self.test = test
+        self.multinode = multinode
+        self.min_size = min_size  # filter out very small images
+        self.max_pwatermark = max_pwatermark # filter out watermarked images
+
+    def make_loader(self, dataset_config, train=True):
+        if 'image_transforms' in dataset_config:
+            image_transforms = [instantiate_from_config(tt) for tt in dataset_config.image_transforms]
+        else:
+            image_transforms = []
+
+        image_transforms.extend([torchvision.transforms.ToTensor(),
+                                 torchvision.transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
+        image_transforms = torchvision.transforms.Compose(image_transforms)
+
+        if 'transforms' in dataset_config:
+            transforms_config = OmegaConf.to_container(dataset_config.transforms)
+        else:
+            transforms_config = dict()
+
+        transform_dict = {dkey: load_partial_from_config(transforms_config[dkey])
+                if transforms_config[dkey] != 'identity' else identity
+                for dkey in transforms_config}
+        img_key = dataset_config.get('image_key', 'jpeg')
+        transform_dict.update({img_key: image_transforms})
+
+        if 'postprocess' in dataset_config:
+            postprocess = instantiate_from_config(dataset_config['postprocess'])
+        else:
+            postprocess = None
+
+        shuffle = dataset_config.get('shuffle', 0)
+        shardshuffle = shuffle > 0
+
+        nodesplitter = wds.shardlists.split_by_node if self.multinode else wds.shardlists.single_node_only
+
+        if self.tar_base == "__improvedaesthetic__":
+            print("## Warning, loading the same improved aesthetic dataset "
+                    "for all splits and ignoring shards parameter.")
+            tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{000000..060207}.tar -"
+        else:
+            tars = os.path.join(self.tar_base, dataset_config.shards)
+
+        dset = wds.WebDataset(
+                tars,
+                nodesplitter=nodesplitter,
+                shardshuffle=shardshuffle,
+                handler=wds.warn_and_continue).repeat().shuffle(shuffle)
+        print(f'Loading webdataset with {len(dset.pipeline[0].urls)} shards.')
+
+        dset = (dset
+                .select(self.filter_keys)
+                .decode('pil', handler=wds.warn_and_continue)
+                .select(self.filter_size)
+                .map_dict(**transform_dict, handler=wds.warn_and_continue)
+                )
+        if postprocess is not None:
+            dset = dset.map(postprocess)
+        dset = (dset
+                .batched(self.batch_size, partial=False,
+                    collation_fn=dict_collation_fn)
+                )
+
+        loader = wds.WebLoader(dset, batch_size=None, shuffle=False,
+                               num_workers=self.num_workers)
+
+        return loader
+
+    def filter_size(self, x):
+        try:
+            valid = True
+            if self.min_size is not None and self.min_size > 1:
+                try:
+                    valid = valid and x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size
+                except Exception:
+                    valid = False
+            if self.max_pwatermark is not None and self.max_pwatermark < 1.0:
+                try:
+                    valid = valid and  x['json']['pwatermark'] <= self.max_pwatermark
+                except Exception:
+                    valid = False
+            return valid
+        except Exception:
+            return False
+
+    def filter_keys(self, x):
+        try:
+            return ("jpg" in x) and ("txt" in x)
+        except Exception:
+            return False
+
+    def train_dataloader(self):
+        return self.make_loader(self.train)
+
+    def val_dataloader(self):
+        return self.make_loader(self.validation, train=False)
+
+    def test_dataloader(self):
+        return self.make_loader(self.test, train=False)
+
+
+from ldm.modules.image_degradation import degradation_fn_bsr_light
+import cv2
+
+class AddLR(object):
+    def __init__(self, factor, output_size, initial_size=None, image_key="jpg"):
+        self.factor = factor
+        self.output_size = output_size
+        self.image_key = image_key
+        self.initial_size = initial_size
+
+    def pt2np(self, x):
+        x = ((x+1.0)*127.5).clamp(0, 255).to(dtype=torch.uint8).detach().cpu().numpy()
+        return x
+
+    def np2pt(self, x):
+        x = torch.from_numpy(x)/127.5-1.0
+        return x
+
+    def __call__(self, sample):
+        # sample['jpg'] is tensor hwc in [-1, 1] at this point
+        x = self.pt2np(sample[self.image_key])
+        if self.initial_size is not None:
+            x = cv2.resize(x, (self.initial_size, self.initial_size), interpolation=2)
+        x = degradation_fn_bsr_light(x, sf=self.factor)['image']
+        x = cv2.resize(x, (self.output_size, self.output_size), interpolation=2)
+        x = self.np2pt(x)
+        sample['lr'] = x
+        return sample
+
+class AddBW(object):
+    def __init__(self, image_key="jpg"):
+        self.image_key = image_key
+
+    def pt2np(self, x):
+        x = ((x+1.0)*127.5).clamp(0, 255).to(dtype=torch.uint8).detach().cpu().numpy()
+        return x
+
+    def np2pt(self, x):
+        x = torch.from_numpy(x)/127.5-1.0
+        return x
+
+    def __call__(self, sample):
+        # sample['jpg'] is tensor hwc in [-1, 1] at this point
+        x = sample[self.image_key]
+        w = torch.rand(3, device=x.device)
+        w /= w.sum()
+        out = torch.einsum('hwc,c->hw', x, w)
+
+        # Keep as 3ch so we can pass to encoder, also we might want to add hints
+        sample['lr'] = out.unsqueeze(-1).tile(1,1,3)
+        return sample
+
+class AddMask(PRNGMixin):
+    def __init__(self, mode="512train", p_drop=0.):
+        super().__init__()
+        assert mode in list(MASK_MODES.keys()), f'unknown mask generation mode "{mode}"'
+        self.make_mask = MASK_MODES[mode]
+        self.p_drop = p_drop
+
+    def __call__(self, sample):
+        # sample['jpg'] is tensor hwc in [-1, 1] at this point
+        x = sample['jpg']
+        mask = self.make_mask(self.prng, x.shape[0], x.shape[1])
+        if self.prng.choice(2, p=[1 - self.p_drop, self.p_drop]):
+            mask = np.ones_like(mask)
+        mask[mask < 0.5] = 0
+        mask[mask > 0.5] = 1
+        mask = torch.from_numpy(mask[..., None])
+        sample['mask'] = mask
+        sample['masked_image'] = x * (mask < 0.5)
+        return sample
+
+
+class AddEdge(PRNGMixin):
+    def __init__(self, mode="512train", mask_edges=True):
+        super().__init__()
+        assert mode in list(MASK_MODES.keys()), f'unknown mask generation mode "{mode}"'
+        self.make_mask = MASK_MODES[mode]
+        self.n_down_choices = [0]
+        self.sigma_choices = [1, 2]
+        self.mask_edges = mask_edges
+
+    @torch.no_grad()
+    def __call__(self, sample):
+        # sample['jpg'] is tensor hwc in [-1, 1] at this point
+        x = sample['jpg']
+
+        mask = self.make_mask(self.prng, x.shape[0], x.shape[1])
+        mask[mask < 0.5] = 0
+        mask[mask > 0.5] = 1
+        mask = torch.from_numpy(mask[..., None])
+        sample['mask'] = mask
+
+        n_down_idx = self.prng.choice(len(self.n_down_choices))
+        sigma_idx = self.prng.choice(len(self.sigma_choices))
+
+        n_choices = len(self.n_down_choices)*len(self.sigma_choices)
+        raveled_idx = np.ravel_multi_index((n_down_idx, sigma_idx),
+                                           (len(self.n_down_choices), len(self.sigma_choices)))
+        normalized_idx = raveled_idx/max(1, n_choices-1)
+
+        n_down = self.n_down_choices[n_down_idx]
+        sigma = self.sigma_choices[sigma_idx]
+
+        kernel_size = 4*sigma+1
+        kernel_size = (kernel_size, kernel_size)
+        sigma = (sigma, sigma)
+        canny = kornia.filters.Canny(
+                low_threshold=0.1,
+                high_threshold=0.2,
+                kernel_size=kernel_size,
+                sigma=sigma,
+                hysteresis=True,
+                )
+        y = (x+1.0)/2.0 # in 01
+        y = y.unsqueeze(0).permute(0, 3, 1, 2).contiguous()
+
+        # down
+        for i_down in range(n_down):
+            size = min(y.shape[-2], y.shape[-1])//2
+            y = kornia.geometry.transform.resize(y, size, antialias=True)
+
+        # edge
+        _, y = canny(y)
+
+        if n_down > 0:
+            size = x.shape[0], x.shape[1]
+            y = kornia.geometry.transform.resize(y, size, interpolation="nearest")
+
+        y = y.permute(0, 2, 3, 1)[0].expand(-1, -1, 3).contiguous()
+        y = y*2.0-1.0
+
+        if self.mask_edges:
+            sample['masked_image'] = y * (mask < 0.5)
+        else:
+            sample['masked_image'] = y
+            sample['mask'] = torch.zeros_like(sample['mask'])
+
+        # concat normalized idx
+        sample['smoothing_strength'] = torch.ones_like(sample['mask'])*normalized_idx
+
+        return sample
+
+
+def example00():
+    url = "pipe:aws s3 cp s3://s-datasets/laion5b/laion2B-data/000000.tar -"
+    dataset = wds.WebDataset(url)
+    example = next(iter(dataset))
+    for k in example:
+        print(k, type(example[k]))
+
+    print(example["__key__"])
+    for k in ["json", "txt"]:
+        print(example[k].decode())
+
+    image = Image.open(io.BytesIO(example["jpg"]))
+    outdir = "tmp"
+    os.makedirs(outdir, exist_ok=True)
+    image.save(os.path.join(outdir, example["__key__"] + ".png"))
+
+
+    def load_example(example):
+        return {
+            "key": example["__key__"],
+            "image": Image.open(io.BytesIO(example["jpg"])),
+            "text": example["txt"].decode(),
+        }
+
+
+    for i, example in tqdm(enumerate(dataset)):
+        ex = load_example(example)
+        print(ex["image"].size, ex["text"])
+        if i >= 100:
+            break
+
+
+def example01():
+    # the first laion shards contain ~10k examples each
+    url = "pipe:aws s3 cp s3://s-datasets/laion5b/laion2B-data/{000000..000002}.tar -"
+
+    batch_size = 3
+    shuffle_buffer = 10000
+    dset = wds.WebDataset(
+            url,
+            nodesplitter=wds.shardlists.split_by_node,
+            shardshuffle=True,
+            )
+    dset = (dset
+            .shuffle(shuffle_buffer, initial=shuffle_buffer)
+            .decode('pil', handler=warn_and_continue)
+            .batched(batch_size, partial=False,
+                collation_fn=dict_collation_fn)
+            )
+
+    num_workers = 2
+    loader = wds.WebLoader(dset, batch_size=None, shuffle=False, num_workers=num_workers)
+
+    batch_sizes = list()
+    keys_per_epoch = list()
+    for epoch in range(5):
+        keys = list()
+        for batch in tqdm(loader):
+            batch_sizes.append(len(batch["__key__"]))
+            keys.append(batch["__key__"])
+
+        for bs in batch_sizes:
+            assert bs==batch_size
+        print(f"{len(batch_sizes)} batches of size {batch_size}.")
+        batch_sizes = list()
+
+        keys_per_epoch.append(keys)
+        for i_batch in [0, 1, -1]:
+            print(f"Batch {i_batch} of epoch {epoch}:")
+            print(keys[i_batch])
+        print("next epoch.")
+
+
+def example02():
+    from omegaconf import OmegaConf
+    from torch.utils.data.distributed import DistributedSampler
+    from torch.utils.data import IterableDataset
+    from torch.utils.data import DataLoader, RandomSampler, Sampler, SequentialSampler
+    from pytorch_lightning.trainer.supporters import CombinedLoader, CycleIterator
+
+    #config = OmegaConf.load("configs/stable-diffusion/txt2img-1p4B-multinode-clip-encoder-high-res-512.yaml")
+    #config = OmegaConf.load("configs/stable-diffusion/txt2img-upscale-clip-encoder-f16-1024.yaml")
+    config = OmegaConf.load("configs/stable-diffusion/txt2img-v2-clip-encoder-improved_aesthetics-256.yaml")
+    datamod = WebDataModuleFromConfig(**config["data"]["params"])
+    dataloader = datamod.train_dataloader()
+
+    for batch in dataloader:
+        print(batch.keys())
+        print(batch["jpg"].shape)
+        break
+
+
+def example03():
+    # improved aesthetics
+    tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{000000..060207}.tar -"
+    dataset = wds.WebDataset(tars)
+
+    def filter_keys(x):
+        try:
+            return ("jpg" in x) and ("txt" in x)
+        except Exception:
+            return False
+
+    def filter_size(x):
+        try:
+            return x['json']['original_width'] >= 512 and x['json']['original_height'] >= 512
+        except Exception:
+            return False
+
+    def filter_watermark(x):
+        try:
+            return x['json']['pwatermark'] < 0.5
+        except Exception:
+            return False
+
+    dataset = (dataset
+                .select(filter_keys)
+                .decode('pil', handler=wds.warn_and_continue))
+    n_save = 20
+    n_total = 0
+    n_large = 0
+    n_large_nowm = 0
+    for i, example in enumerate(dataset):
+        n_total += 1
+        if filter_size(example):
+            n_large += 1
+            if filter_watermark(example):
+                n_large_nowm += 1
+                if n_large_nowm < n_save+1:
+                    image = example["jpg"]
+                    image.save(os.path.join("tmp", f"{n_large_nowm-1:06}.png"))
+
+        if i%500 == 0:
+            print(i)
+            print(f"Large: {n_large}/{n_total} | {n_large/n_total*100:.2f}%")
+            if n_large > 0:
+                print(f"No Watermark: {n_large_nowm}/{n_large} | {n_large_nowm/n_large*100:.2f}%")
+
+
+
+def example04():
+    # improved aesthetics
+    for i_shard in range(60208)[::-1]:
+        print(i_shard)
+        tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{:06}.tar -".format(i_shard)
+        dataset = wds.WebDataset(tars)
+
+        def filter_keys(x):
+            try:
+                return ("jpg" in x) and ("txt" in x)
+            except Exception:
+                return False
+
+        def filter_size(x):
+            try:
+                return x['json']['original_width'] >= 512 and x['json']['original_height'] >= 512
+            except Exception:
+                return False
+
+        dataset = (dataset
+                    .select(filter_keys)
+                    .decode('pil', handler=wds.warn_and_continue))
+        try:
+            example = next(iter(dataset))
+        except Exception:
+            print(f"Error @ {i_shard}")
+
+
+if __name__ == "__main__":
+    #example01()
+    #example02()
+    example03()
+    #example04()