Delete One-2-3-45-master 2

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

@@ -1,35 +0,0 @@
-*.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

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

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

@@ -1,201 +0,0 @@
-                                 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

@@ -1,221 +0,0 @@
-<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

@@ -1,117 +0,0 @@
-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/download_ckpt.py

@@ -1,30 +0,0 @@
-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

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

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

@@ -1,10 +0,0 @@
-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

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

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

@@ -1,11 +0,0 @@
-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

@@ -1,199 +0,0 @@
-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

@@ -1,81 +0,0 @@
-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

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

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

@@ -1,59 +0,0 @@
-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

@@ -1,81 +0,0 @@
-"""
-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

@@ -1,101 +0,0 @@
-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

@@ -1,261 +0,0 @@
-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

@@ -1,50 +0,0 @@
-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

@@ -1,74 +0,0 @@
-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

@@ -1,54 +0,0 @@
-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

@@ -1,42 +0,0 @@
-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

@@ -1,151 +0,0 @@
-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

@@ -1,7 +0,0 @@
-[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

@@ -1,205 +0,0 @@
-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

@@ -1,154 +0,0 @@
-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

@@ -1,318 +0,0 @@
-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

@@ -1,62 +0,0 @@
-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

@@ -1,40 +0,0 @@
-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

@@ -1,253 +0,0 @@
-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

@@ -1,34 +0,0 @@
-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

@@ -1,394 +0,0 @@
-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

@@ -1,166 +0,0 @@
-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

@@ -1,537 +0,0 @@
-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()

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

@@ -1,92 +0,0 @@
-import os
-import numpy as np
-import PIL
-from PIL import Image
-from torch.utils.data import Dataset
-from torchvision import transforms
-
-
-class LSUNBase(Dataset):
-    def __init__(self,
-                 txt_file,
-                 data_root,
-                 size=None,
-                 interpolation="bicubic",
-                 flip_p=0.5
-                 ):
-        self.data_paths = txt_file
-        self.data_root = data_root
-        with open(self.data_paths, "r") as f:
-            self.image_paths = f.read().splitlines()
-        self._length = len(self.image_paths)
-        self.labels = {
-            "relative_file_path_": [l for l in self.image_paths],
-            "file_path_": [os.path.join(self.data_root, l)
-                           for l in self.image_paths],
-        }
-
-        self.size = size
-        self.interpolation = {"linear": PIL.Image.LINEAR,
-                              "bilinear": PIL.Image.BILINEAR,
-                              "bicubic": PIL.Image.BICUBIC,
-                              "lanczos": PIL.Image.LANCZOS,
-                              }[interpolation]
-        self.flip = transforms.RandomHorizontalFlip(p=flip_p)
-
-    def __len__(self):
-        return self._length
-
-    def __getitem__(self, i):
-        example = dict((k, self.labels[k][i]) for k in self.labels)
-        image = Image.open(example["file_path_"])
-        if not image.mode == "RGB":
-            image = image.convert("RGB")
-
-        # default to score-sde preprocessing
-        img = np.array(image).astype(np.uint8)
-        crop = min(img.shape[0], img.shape[1])
-        h, w, = img.shape[0], img.shape[1]
-        img = img[(h - crop) // 2:(h + crop) // 2,
-              (w - crop) // 2:(w + crop) // 2]
-
-        image = Image.fromarray(img)
-        if self.size is not None:
-            image = image.resize((self.size, self.size), resample=self.interpolation)
-
-        image = self.flip(image)
-        image = np.array(image).astype(np.uint8)
-        example["image"] = (image / 127.5 - 1.0).astype(np.float32)
-        return example
-
-
-class LSUNChurchesTrain(LSUNBase):
-    def __init__(self, **kwargs):
-        super().__init__(txt_file="data/lsun/church_outdoor_train.txt", data_root="data/lsun/churches", **kwargs)
-
-
-class LSUNChurchesValidation(LSUNBase):
-    def __init__(self, flip_p=0., **kwargs):
-        super().__init__(txt_file="data/lsun/church_outdoor_val.txt", data_root="data/lsun/churches",
-                         flip_p=flip_p, **kwargs)
-
-
-class LSUNBedroomsTrain(LSUNBase):
-    def __init__(self, **kwargs):
-        super().__init__(txt_file="data/lsun/bedrooms_train.txt", data_root="data/lsun/bedrooms", **kwargs)
-
-
-class LSUNBedroomsValidation(LSUNBase):
-    def __init__(self, flip_p=0.0, **kwargs):
-        super().__init__(txt_file="data/lsun/bedrooms_val.txt", data_root="data/lsun/bedrooms",
-                         flip_p=flip_p, **kwargs)
-
-
-class LSUNCatsTrain(LSUNBase):
-    def __init__(self, **kwargs):
-        super().__init__(txt_file="data/lsun/cat_train.txt", data_root="data/lsun/cats", **kwargs)
-
-
-class LSUNCatsValidation(LSUNBase):
-    def __init__(self, flip_p=0., **kwargs):
-        super().__init__(txt_file="data/lsun/cat_val.txt", data_root="data/lsun/cats",
-                         flip_p=flip_p, **kwargs)

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

@@ -1,165 +0,0 @@
-from torch.utils.data import Dataset
-import os
-import json
-import numpy as np
-import torch
-import imageio
-import math
-import cv2
-from torchvision import transforms
-
-def cartesian_to_spherical(xyz):
-    ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
-    xy = xyz[:,0]**2 + xyz[:,1]**2
-    z = np.sqrt(xy + xyz[:,2]**2)
-    theta = np.arctan2(np.sqrt(xy), xyz[:,2]) # for elevation angle defined from Z-axis down
-    #ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
-    azimuth = np.arctan2(xyz[:,1], xyz[:,0])
-    return np.array([theta, azimuth, z])
-
-
-def get_T(T_target, T_cond):
-    theta_cond, azimuth_cond, z_cond = cartesian_to_spherical(T_cond[None, :])
-    theta_target, azimuth_target, z_target = cartesian_to_spherical(T_target[None, :])
-    
-    d_theta = theta_target - theta_cond
-    d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
-    d_z = z_target - z_cond
-    
-    d_T = torch.tensor([d_theta.item(), math.sin(d_azimuth.item()), math.cos(d_azimuth.item()), d_z.item()])
-    return d_T
-
-def get_spherical(T_target, T_cond):
-    theta_cond, azimuth_cond, z_cond = cartesian_to_spherical(T_cond[None, :])
-    theta_target, azimuth_target, z_target = cartesian_to_spherical(T_target[None, :])
-    
-    d_theta = theta_target - theta_cond
-    d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
-    d_z = z_target - z_cond
-    
-    d_T = torch.tensor([math.degrees(d_theta.item()), math.degrees(d_azimuth.item()), d_z.item()])
-    return d_T
-
-class RTMV(Dataset):
-    def __init__(self, root_dir='datasets/RTMV/google_scanned',\
-                 first_K=64, resolution=256, load_target=False):
-        self.root_dir = root_dir
-        self.scene_list = sorted(next(os.walk(root_dir))[1])
-        self.resolution = resolution
-        self.first_K = first_K
-        self.load_target = load_target
-
-    def __len__(self):
-        return len(self.scene_list)
-
-    def __getitem__(self, idx):
-        scene_dir = os.path.join(self.root_dir, self.scene_list[idx])
-        with open(os.path.join(scene_dir, 'transforms.json'), "r") as f:
-            meta = json.load(f)
-        imgs = []
-        poses = []
-        for i_img in range(self.first_K):
-            meta_img = meta['frames'][i_img]
-
-            if i_img == 0 or self.load_target:
-                img_path = os.path.join(scene_dir, meta_img['file_path'])
-                img = imageio.imread(img_path)
-                img = cv2.resize(img, (self.resolution, self.resolution), interpolation = cv2.INTER_LINEAR)
-                imgs.append(img)
-            
-            c2w = meta_img['transform_matrix']
-            poses.append(c2w)
-            
-        imgs = (np.array(imgs) / 255.).astype(np.float32)  # (RGBA) imgs
-        imgs = torch.tensor(self.blend_rgba(imgs)).permute(0, 3, 1, 2)
-        imgs = imgs * 2 - 1. # convert to stable diffusion range
-        poses = torch.tensor(np.array(poses).astype(np.float32))
-        return imgs, poses
-                
-    def blend_rgba(self, img):
-        img = img[..., :3] * img[..., -1:] + (1. - img[..., -1:])  # blend A to RGB
-        return img
-            
-
-class GSO(Dataset):
-    def __init__(self, root_dir='datasets/GoogleScannedObjects',\
-                 split='val', first_K=5, resolution=256, load_target=False, name='render_mvs'):
-        self.root_dir = root_dir
-        with open(os.path.join(root_dir, '%s.json' % split), "r") as f:
-            self.scene_list = json.load(f)
-        self.resolution = resolution
-        self.first_K = first_K
-        self.load_target = load_target
-        self.name = name
-
-    def __len__(self):
-        return len(self.scene_list)
-
-    def __getitem__(self, idx):
-        scene_dir = os.path.join(self.root_dir, self.scene_list[idx])
-        with open(os.path.join(scene_dir, 'transforms_%s.json' % self.name), "r") as f:
-            meta = json.load(f)
-        imgs = []
-        poses = []
-        for i_img in range(self.first_K):
-            meta_img = meta['frames'][i_img]
-
-            if i_img == 0 or self.load_target:
-                img_path = os.path.join(scene_dir, meta_img['file_path'])
-                img = imageio.imread(img_path)
-                img = cv2.resize(img, (self.resolution, self.resolution), interpolation = cv2.INTER_LINEAR)
-                imgs.append(img)
-            
-            c2w = meta_img['transform_matrix']
-            poses.append(c2w)
-            
-        imgs = (np.array(imgs) / 255.).astype(np.float32)  # (RGBA) imgs
-        mask = imgs[:, :, :, -1]
-        imgs = torch.tensor(self.blend_rgba(imgs)).permute(0, 3, 1, 2)
-        imgs = imgs * 2 - 1. # convert to stable diffusion range
-        poses = torch.tensor(np.array(poses).astype(np.float32))
-        return imgs, poses
-                
-    def blend_rgba(self, img):
-        img = img[..., :3] * img[..., -1:] + (1. - img[..., -1:])  # blend A to RGB
-        return img
-             
-class WILD(Dataset):
-    def __init__(self, root_dir='data/nerf_wild',\
-                 first_K=33, resolution=256, load_target=False):
-        self.root_dir = root_dir
-        self.scene_list = sorted(next(os.walk(root_dir))[1])
-        self.resolution = resolution
-        self.first_K = first_K
-        self.load_target = load_target
-
-    def __len__(self):
-        return len(self.scene_list)
-
-    def __getitem__(self, idx):
-        scene_dir = os.path.join(self.root_dir, self.scene_list[idx])
-        with open(os.path.join(scene_dir, 'transforms_train.json'), "r") as f:
-            meta = json.load(f)
-        imgs = []
-        poses = []
-        for i_img in range(self.first_K):
-            meta_img = meta['frames'][i_img]
-
-            if i_img == 0 or self.load_target:
-                img_path = os.path.join(scene_dir, meta_img['file_path'])
-                img = imageio.imread(img_path + '.png')
-                img = cv2.resize(img, (self.resolution, self.resolution), interpolation = cv2.INTER_LINEAR)
-                imgs.append(img)
-            
-            c2w = meta_img['transform_matrix']
-            poses.append(c2w)
-            
-        imgs = (np.array(imgs) / 255.).astype(np.float32)  # (RGBA) imgs
-        imgs = torch.tensor(self.blend_rgba(imgs)).permute(0, 3, 1, 2)
-        imgs = imgs * 2 - 1. # convert to stable diffusion range
-        poses = torch.tensor(np.array(poses).astype(np.float32))
-        return imgs, poses
-                
-    def blend_rgba(self, img):
-        img = img[..., :3] * img[..., -1:] + (1. - img[..., -1:])  # blend A to RGB
-        return img

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

@@ -1,526 +0,0 @@
-from typing import Dict
-import webdataset as wds
-import numpy as np
-from omegaconf import DictConfig, ListConfig
-import torch
-from torch.utils.data import Dataset
-from pathlib import Path
-import json
-from PIL import Image
-from torchvision import transforms
-import torchvision
-from einops import rearrange
-from ldm.util import instantiate_from_config
-from datasets import load_dataset
-import pytorch_lightning as pl
-import copy
-import csv
-import cv2
-import random
-import matplotlib.pyplot as plt
-from torch.utils.data import DataLoader
-import json
-import os, sys
-import webdataset as wds
-import math
-from torch.utils.data.distributed import DistributedSampler
-
-# Some hacky things to make experimentation easier
-def make_transform_multi_folder_data(paths, caption_files=None, **kwargs):
-    ds = make_multi_folder_data(paths, caption_files, **kwargs)
-    return TransformDataset(ds)
-
-def make_nfp_data(base_path):
-    dirs = list(Path(base_path).glob("*/"))
-    print(f"Found {len(dirs)} folders")
-    print(dirs)
-    tforms = [transforms.Resize(512), transforms.CenterCrop(512)]
-    datasets = [NfpDataset(x, image_transforms=copy.copy(tforms), default_caption="A view from a train window") for x in dirs]
-    return torch.utils.data.ConcatDataset(datasets)
-
-
-class VideoDataset(Dataset):
-    def __init__(self, root_dir, image_transforms, caption_file, offset=8, n=2):
-        self.root_dir = Path(root_dir)
-        self.caption_file = caption_file
-        self.n = n
-        ext = "mp4"
-        self.paths = sorted(list(self.root_dir.rglob(f"*.{ext}")))
-        self.offset = offset
-
-        if isinstance(image_transforms, ListConfig):
-            image_transforms = [instantiate_from_config(tt) for tt in image_transforms]
-        image_transforms.extend([transforms.ToTensor(),
-                                 transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
-        image_transforms = transforms.Compose(image_transforms)
-        self.tform = image_transforms
-        with open(self.caption_file) as f:
-            reader = csv.reader(f)
-            rows = [row for row in reader]
-        self.captions = dict(rows)
-
-    def __len__(self):
-        return len(self.paths)
-
-    def __getitem__(self, index):
-        for i in range(10):
-            try:
-                return self._load_sample(index)
-            except Exception:
-                # Not really good enough but...
-                print("uh oh")
-
-    def _load_sample(self, index):
-        n = self.n
-        filename = self.paths[index]
-        min_frame = 2*self.offset + 2
-        vid = cv2.VideoCapture(str(filename))
-        max_frames = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
-        curr_frame_n = random.randint(min_frame, max_frames)
-        vid.set(cv2.CAP_PROP_POS_FRAMES,curr_frame_n)
-        _, curr_frame = vid.read()
-
-        prev_frames = []
-        for i in range(n):
-            prev_frame_n = curr_frame_n - (i+1)*self.offset
-            vid.set(cv2.CAP_PROP_POS_FRAMES,prev_frame_n)
-            _, prev_frame = vid.read()
-            prev_frame = self.tform(Image.fromarray(prev_frame[...,::-1]))
-            prev_frames.append(prev_frame)
-
-        vid.release()
-        caption = self.captions[filename.name]
-        data = {
-            "image": self.tform(Image.fromarray(curr_frame[...,::-1])),
-            "prev": torch.cat(prev_frames, dim=-1),
-            "txt": caption
-        }
-        return data
-
-# end hacky things
-
-
-def make_tranforms(image_transforms):
-    # if isinstance(image_transforms, ListConfig):
-    #     image_transforms = [instantiate_from_config(tt) for tt in image_transforms]
-    image_transforms = []
-    image_transforms.extend([transforms.ToTensor(),
-                                transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
-    image_transforms = transforms.Compose(image_transforms)
-    return image_transforms
-
-
-def make_multi_folder_data(paths, caption_files=None, **kwargs):
-    """Make a concat dataset from multiple folders
-    Don't suport captions yet
-
-    If paths is a list, that's ok, if it's a Dict interpret it as:
-    k=folder v=n_times to repeat that
-    """
-    list_of_paths = []
-    if isinstance(paths, (Dict, DictConfig)):
-        assert caption_files is None, \
-            "Caption files not yet supported for repeats"
-        for folder_path, repeats in paths.items():
-            list_of_paths.extend([folder_path]*repeats)
-        paths = list_of_paths
-
-    if caption_files is not None:
-        datasets = [FolderData(p, caption_file=c, **kwargs) for (p, c) in zip(paths, caption_files)]
-    else:
-        datasets = [FolderData(p, **kwargs) for p in paths]
-    return torch.utils.data.ConcatDataset(datasets)
-
-
-
-class NfpDataset(Dataset):
-    def __init__(self,
-        root_dir,
-        image_transforms=[],
-        ext="jpg",
-        default_caption="",
-        ) -> None:
-        """assume sequential frames and a deterministic transform"""
-
-        self.root_dir = Path(root_dir)
-        self.default_caption = default_caption
-
-        self.paths = sorted(list(self.root_dir.rglob(f"*.{ext}")))
-        self.tform = make_tranforms(image_transforms)
-
-    def __len__(self):
-        return len(self.paths) - 1
-
-
-    def __getitem__(self, index):
-        prev = self.paths[index]
-        curr = self.paths[index+1]
-        data = {}
-        data["image"] = self._load_im(curr)
-        data["prev"] = self._load_im(prev)
-        data["txt"] = self.default_caption
-        return data
-
-    def _load_im(self, filename):
-        im = Image.open(filename).convert("RGB")
-        return self.tform(im)
-
-class ObjaverseDataModuleFromConfig(pl.LightningDataModule):
-    def __init__(self, root_dir, batch_size, total_view, train=None, validation=None,
-                 test=None, num_workers=4, **kwargs):
-        super().__init__(self)
-        self.root_dir = root_dir
-        self.batch_size = batch_size
-        self.num_workers = num_workers
-        self.total_view = total_view
-
-        if train is not None:
-            dataset_config = train
-        if validation is not None:
-            dataset_config = validation
-
-        if 'image_transforms' in dataset_config:
-            image_transforms = [torchvision.transforms.Resize(dataset_config.image_transforms.size)]
-        else:
-            image_transforms = []
-        image_transforms.extend([transforms.ToTensor(),
-                                transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
-        self.image_transforms = torchvision.transforms.Compose(image_transforms)
-
-
-    def train_dataloader(self):
-        dataset = ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=False, \
-                                image_transforms=self.image_transforms)
-        sampler = DistributedSampler(dataset)
-        return wds.WebLoader(dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, sampler=sampler)
-
-    def val_dataloader(self):
-        dataset = ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=True, \
-                                image_transforms=self.image_transforms)
-        sampler = DistributedSampler(dataset)
-        return wds.WebLoader(dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
-    
-    def test_dataloader(self):
-        return wds.WebLoader(ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=self.validation),\
-                          batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
-
-
-class ObjaverseData(Dataset):
-    def __init__(self,
-        root_dir='.objaverse/hf-objaverse-v1/views',
-        image_transforms=[],
-        ext="png",
-        default_trans=torch.zeros(3),
-        postprocess=None,
-        return_paths=False,
-        total_view=4,
-        validation=False
-        ) -> None:
-        """Create a dataset from a folder of images.
-        If you pass in a root directory it will be searched for images
-        ending in ext (ext can be a list)
-        """
-        self.root_dir = Path(root_dir)
-        self.default_trans = default_trans
-        self.return_paths = return_paths
-        if isinstance(postprocess, DictConfig):
-            postprocess = instantiate_from_config(postprocess)
-        self.postprocess = postprocess
-        self.total_view = total_view
-
-        if not isinstance(ext, (tuple, list, ListConfig)):
-            ext = [ext]
-
-        with open(os.path.join(root_dir, 'valid_paths.json')) as f:
-            self.paths = json.load(f)
-            
-        total_objects = len(self.paths)
-        if validation:
-            self.paths = self.paths[math.floor(total_objects / 100. * 99.):] # used last 1% as validation
-        else:
-            self.paths = self.paths[:math.floor(total_objects / 100. * 99.)] # used first 99% as training
-        print('============= length of dataset %d =============' % len(self.paths))
-        self.tform = image_transforms
-
-    def __len__(self):
-        return len(self.paths)
-        
-    def cartesian_to_spherical(self, xyz):
-        ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
-        xy = xyz[:,0]**2 + xyz[:,1]**2
-        z = np.sqrt(xy + xyz[:,2]**2)
-        theta = np.arctan2(np.sqrt(xy), xyz[:,2]) # for elevation angle defined from Z-axis down
-        #ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
-        azimuth = np.arctan2(xyz[:,1], xyz[:,0])
-        return np.array([theta, azimuth, z])
-
-    def get_T(self, target_RT, cond_RT):
-        R, T = target_RT[:3, :3], target_RT[:, -1]
-        T_target = -R.T @ T
-
-        R, T = cond_RT[:3, :3], cond_RT[:, -1]
-        T_cond = -R.T @ T
-
-        theta_cond, azimuth_cond, z_cond = self.cartesian_to_spherical(T_cond[None, :])
-        theta_target, azimuth_target, z_target = self.cartesian_to_spherical(T_target[None, :])
-        
-        d_theta = theta_target - theta_cond
-        d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
-        d_z = z_target - z_cond
-        
-        d_T = torch.tensor([d_theta.item(), math.sin(d_azimuth.item()), math.cos(d_azimuth.item()), d_z.item()])
-        return d_T
-
-    def load_im(self, path, color):
-        '''
-        replace background pixel with random color in rendering
-        '''
-        try:
-            img = plt.imread(path)
-        except:
-            print(path)
-            sys.exit()
-        img[img[:, :, -1] == 0.] = color
-        img = Image.fromarray(np.uint8(img[:, :, :3] * 255.))
-        return img
-
-    def __getitem__(self, index):
-
-        data = {}
-        if self.paths[index][-2:] == '_1': # dirty fix for rendering dataset twice
-            total_view = 8
-        else:
-            total_view = 4
-        index_target, index_cond = random.sample(range(total_view), 2) # without replacement
-        filename = os.path.join(self.root_dir, self.paths[index])
-
-        # print(self.paths[index])
-
-        if self.return_paths:
-            data["path"] = str(filename)
-        
-        color = [1., 1., 1., 1.]
-
-        try:
-            target_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_target), color))
-            cond_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_cond), color))
-            target_RT = np.load(os.path.join(filename, '%03d.npy' % index_target))
-            cond_RT = np.load(os.path.join(filename, '%03d.npy' % index_cond))
-        except:
-            # very hacky solution, sorry about this
-            filename = os.path.join(self.root_dir, '692db5f2d3a04bb286cb977a7dba903e_1') # this one we know is valid
-            target_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_target), color))
-            cond_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_cond), color))
-            target_RT = np.load(os.path.join(filename, '%03d.npy' % index_target))
-            cond_RT = np.load(os.path.join(filename, '%03d.npy' % index_cond))
-            target_im = torch.zeros_like(target_im)
-            cond_im = torch.zeros_like(cond_im)
-
-        data["image_target"] = target_im
-        data["image_cond"] = cond_im
-        data["T"] = self.get_T(target_RT, cond_RT)
-
-        if self.postprocess is not None:
-            data = self.postprocess(data)
-
-        return data
-
-    def process_im(self, im):
-        im = im.convert("RGB")
-        return self.tform(im)
-
-class FolderData(Dataset):
-    def __init__(self,
-        root_dir,
-        caption_file=None,
-        image_transforms=[],
-        ext="jpg",
-        default_caption="",
-        postprocess=None,
-        return_paths=False,
-        ) -> None:
-        """Create a dataset from a folder of images.
-        If you pass in a root directory it will be searched for images
-        ending in ext (ext can be a list)
-        """
-        self.root_dir = Path(root_dir)
-        self.default_caption = default_caption
-        self.return_paths = return_paths
-        if isinstance(postprocess, DictConfig):
-            postprocess = instantiate_from_config(postprocess)
-        self.postprocess = postprocess
-        if caption_file is not None:
-            with open(caption_file, "rt") as f:
-                ext = Path(caption_file).suffix.lower()
-                if ext == ".json":
-                    captions = json.load(f)
-                elif ext == ".jsonl":
-                    lines = f.readlines()
-                    lines = [json.loads(x) for x in lines]
-                    captions = {x["file_name"]: x["text"].strip("\n") for x in lines}
-                else:
-                    raise ValueError(f"Unrecognised format: {ext}")
-            self.captions = captions
-        else:
-            self.captions = None
-
-        if not isinstance(ext, (tuple, list, ListConfig)):
-            ext = [ext]
-
-        # Only used if there is no caption file
-        self.paths = []
-        for e in ext:
-            self.paths.extend(sorted(list(self.root_dir.rglob(f"*.{e}"))))
-        self.tform = make_tranforms(image_transforms)
-
-    def __len__(self):
-        if self.captions is not None:
-            return len(self.captions.keys())
-        else:
-            return len(self.paths)
-
-    def __getitem__(self, index):
-        data = {}
-        if self.captions is not None:
-            chosen = list(self.captions.keys())[index]
-            caption = self.captions.get(chosen, None)
-            if caption is None:
-                caption = self.default_caption
-            filename = self.root_dir/chosen
-        else:
-            filename = self.paths[index]
-
-        if self.return_paths:
-            data["path"] = str(filename)
-
-        im = Image.open(filename).convert("RGB")
-        im = self.process_im(im)
-        data["image"] = im
-
-        if self.captions is not None:
-            data["txt"] = caption
-        else:
-            data["txt"] = self.default_caption
-
-        if self.postprocess is not None:
-            data = self.postprocess(data)
-
-        return data
-
-    def process_im(self, im):
-        im = im.convert("RGB")
-        return self.tform(im)
-import random
-
-class TransformDataset():
-    def __init__(self, ds, extra_label="sksbspic"):
-        self.ds = ds
-        self.extra_label = extra_label
-        self.transforms = {
-            "align": transforms.Resize(768),
-            "centerzoom": transforms.CenterCrop(768),
-            "randzoom": transforms.RandomCrop(768),
-        }
-
-
-    def __getitem__(self, index):
-        data = self.ds[index]
-
-        im = data['image']
-        im = im.permute(2,0,1)
-        # In case data is smaller than expected
-        im = transforms.Resize(1024)(im)
-
-        tform_name = random.choice(list(self.transforms.keys()))
-        im = self.transforms[tform_name](im)
-
-        im = im.permute(1,2,0)
-
-        data['image'] = im
-        data['txt'] = data['txt'] + f" {self.extra_label} {tform_name}"
-
-        return data
-
-    def __len__(self):
-        return len(self.ds)
-
-def hf_dataset(
-    name,
-    image_transforms=[],
-    image_column="image",
-    text_column="text",
-    split='train',
-    image_key='image',
-    caption_key='txt',
-    ):
-    """Make huggingface dataset with appropriate list of transforms applied
-    """
-    ds = load_dataset(name, split=split)
-    tform = make_tranforms(image_transforms)
-
-    assert image_column in ds.column_names, f"Didn't find column {image_column} in {ds.column_names}"
-    assert text_column in ds.column_names, f"Didn't find column {text_column} in {ds.column_names}"
-
-    def pre_process(examples):
-        processed = {}
-        processed[image_key] = [tform(im) for im in examples[image_column]]
-        processed[caption_key] = examples[text_column]
-        return processed
-
-    ds.set_transform(pre_process)
-    return ds
-
-class TextOnly(Dataset):
-    def __init__(self, captions, output_size, image_key="image", caption_key="txt", n_gpus=1):
-        """Returns only captions with dummy images"""
-        self.output_size = output_size
-        self.image_key = image_key
-        self.caption_key = caption_key
-        if isinstance(captions, Path):
-            self.captions = self._load_caption_file(captions)
-        else:
-            self.captions = captions
-
-        if n_gpus > 1:
-            # hack to make sure that all the captions appear on each gpu
-            repeated = [n_gpus*[x] for x in self.captions]
-            self.captions = []
-            [self.captions.extend(x) for x in repeated]
-
-    def __len__(self):
-        return len(self.captions)
-
-    def __getitem__(self, index):
-        dummy_im = torch.zeros(3, self.output_size, self.output_size)
-        dummy_im = rearrange(dummy_im * 2. - 1., 'c h w -> h w c')
-        return {self.image_key: dummy_im, self.caption_key: self.captions[index]}
-
-    def _load_caption_file(self, filename):
-        with open(filename, 'rt') as f:
-            captions = f.readlines()
-        return [x.strip('\n') for x in captions]
-
-
-
-import random
-import json
-class IdRetreivalDataset(FolderData):
-    def __init__(self, ret_file, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        with open(ret_file, "rt") as f:
-            self.ret = json.load(f)
-
-    def __getitem__(self, index):
-        data = super().__getitem__(index)
-        key = self.paths[index].name
-        matches = self.ret[key]
-        if len(matches) > 0:
-            retreived = random.choice(matches)
-        else:
-            retreived = key
-        filename = self.root_dir/retreived
-        im = Image.open(filename).convert("RGB")
-        im = self.process_im(im)
-        # data["match"] = im
-        data["match"] = torch.cat((data["image"], im), dim=-1)
-        return data

One-2-3-45-master 2/ldm/extras.py

@@ -1,77 +0,0 @@
-from pathlib import Path
-from omegaconf import OmegaConf
-import torch
-from ldm.util import instantiate_from_config
-import logging
-from contextlib import contextmanager
-
-from contextlib import contextmanager
-import logging
-
-@contextmanager
-def all_logging_disabled(highest_level=logging.CRITICAL):
-    """
-    A context manager that will prevent any logging messages
-    triggered during the body from being processed.
-
-    :param highest_level: the maximum logging level in use.
-      This would only need to be changed if a custom level greater than CRITICAL
-      is defined.
-
-    https://gist.github.com/simon-weber/7853144
-    """
-    # two kind-of hacks here:
-    #    * can't get the highest logging level in effect => delegate to the user
-    #    * can't get the current module-level override => use an undocumented
-    #       (but non-private!) interface
-
-    previous_level = logging.root.manager.disable
-
-    logging.disable(highest_level)
-
-    try:
-        yield
-    finally:
-        logging.disable(previous_level)
-
-def load_training_dir(train_dir, device, epoch="last"):
-    """Load a checkpoint and config from training directory"""
-    train_dir = Path(train_dir)
-    ckpt = list(train_dir.rglob(f"*{epoch}.ckpt"))
-    assert len(ckpt) == 1, f"found {len(ckpt)} matching ckpt files"
-    config = list(train_dir.rglob(f"*-project.yaml"))
-    assert len(ckpt) > 0, f"didn't find any config in {train_dir}"
-    if len(config) > 1:
-        print(f"found {len(config)} matching config files")
-        config = sorted(config)[-1]
-        print(f"selecting {config}")
-    else:
-        config = config[0]
-
-
-    config = OmegaConf.load(config)
-    return load_model_from_config(config, ckpt[0], device)
-
-def load_model_from_config(config, ckpt, device="cpu", verbose=False):
-    """Loads a model from config and a ckpt
-    if config is a path will use omegaconf to load
-    """
-    if isinstance(config, (str, Path)):
-        config = OmegaConf.load(config)
-
-    with all_logging_disabled():
-        print(f"Loading model from {ckpt}")
-        pl_sd = torch.load(ckpt, map_location="cpu")
-        global_step = pl_sd["global_step"]
-        sd = pl_sd["state_dict"]
-        model = instantiate_from_config(config.model)
-        m, u = model.load_state_dict(sd, strict=False)
-        if len(m) > 0 and verbose:
-            print("missing keys:")
-            print(m)
-        if len(u) > 0 and verbose:
-            print("unexpected keys:")
-        model.to(device)
-        model.eval()
-        model.cond_stage_model.device = device
-        return model

One-2-3-45-master 2/ldm/guidance.py

@@ -1,96 +0,0 @@
-from typing import List, Tuple
-from scipy import interpolate
-import numpy as np
-import torch
-import matplotlib.pyplot as plt
-from IPython.display import clear_output
-import abc
-
-
-class GuideModel(torch.nn.Module, abc.ABC):
-    def __init__(self) -> None:
-        super().__init__()
-
-    @abc.abstractmethod
-    def preprocess(self, x_img):
-        pass
-
-    @abc.abstractmethod
-    def compute_loss(self, inp):
-        pass
-
-
-class Guider(torch.nn.Module):
-    def __init__(self, sampler, guide_model, scale=1.0, verbose=False):
-        """Apply classifier guidance
-
-        Specify a guidance scale as either a scalar
-        Or a schedule as a list of tuples t = 0->1 and scale, e.g.
-        [(0, 10), (0.5, 20), (1, 50)]
-        """
-        super().__init__()
-        self.sampler = sampler
-        self.index = 0
-        self.show = verbose
-        self.guide_model = guide_model
-        self.history = []
-
-        if isinstance(scale, (Tuple, List)):
-            times = np.array([x[0] for x in scale])
-            values = np.array([x[1] for x in scale])
-            self.scale_schedule = {"times": times, "values": values}
-        else:
-            self.scale_schedule = float(scale)
-
-        self.ddim_timesteps = sampler.ddim_timesteps
-        self.ddpm_num_timesteps = sampler.ddpm_num_timesteps
-
-
-    def get_scales(self):
-        if isinstance(self.scale_schedule, float):
-            return len(self.ddim_timesteps)*[self.scale_schedule]
-
-        interpolater = interpolate.interp1d(self.scale_schedule["times"], self.scale_schedule["values"])
-        fractional_steps = np.array(self.ddim_timesteps)/self.ddpm_num_timesteps
-        return interpolater(fractional_steps)
-
-    def modify_score(self, model, e_t, x, t, c):
-
-        # TODO look up index by t
-        scale = self.get_scales()[self.index]
-
-        if (scale == 0):
-            return e_t
-
-        sqrt_1ma = self.sampler.ddim_sqrt_one_minus_alphas[self.index].to(x.device)
-        with torch.enable_grad():
-            x_in = x.detach().requires_grad_(True)
-            pred_x0 = model.predict_start_from_noise(x_in, t=t, noise=e_t)
-            x_img = model.first_stage_model.decode((1/0.18215)*pred_x0)
-
-            inp = self.guide_model.preprocess(x_img)
-            loss = self.guide_model.compute_loss(inp)
-            grads = torch.autograd.grad(loss.sum(), x_in)[0]
-            correction = grads * scale
-
-            if self.show:
-                clear_output(wait=True)
-                print(loss.item(), scale, correction.abs().max().item(), e_t.abs().max().item())
-                self.history.append([loss.item(), scale, correction.min().item(), correction.max().item()])
-                plt.imshow((inp[0].detach().permute(1,2,0).clamp(-1,1).cpu()+1)/2)
-                plt.axis('off')
-                plt.show()
-                plt.imshow(correction[0][0].detach().cpu())
-                plt.axis('off')
-                plt.show()
-
-
-        e_t_mod = e_t - sqrt_1ma*correction
-        if self.show:
-            fig, axs = plt.subplots(1, 3)
-            axs[0].imshow(e_t[0][0].detach().cpu(), vmin=-2, vmax=+2)
-            axs[1].imshow(e_t_mod[0][0].detach().cpu(), vmin=-2, vmax=+2)
-            axs[2].imshow(correction[0][0].detach().cpu(), vmin=-2, vmax=+2)
-            plt.show()
-        self.index += 1
-        return e_t_mod

One-2-3-45-master 2/ldm/lr_scheduler.py

@@ -1,98 +0,0 @@
-import numpy as np
-
-
-class LambdaWarmUpCosineScheduler:
-    """
-    note: use with a base_lr of 1.0
-    """
-    def __init__(self, warm_up_steps, lr_min, lr_max, lr_start, max_decay_steps, verbosity_interval=0):
-        self.lr_warm_up_steps = warm_up_steps
-        self.lr_start = lr_start
-        self.lr_min = lr_min
-        self.lr_max = lr_max
-        self.lr_max_decay_steps = max_decay_steps
-        self.last_lr = 0.
-        self.verbosity_interval = verbosity_interval
-
-    def schedule(self, n, **kwargs):
-        if self.verbosity_interval > 0:
-            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
-        if n < self.lr_warm_up_steps:
-            lr = (self.lr_max - self.lr_start) / self.lr_warm_up_steps * n + self.lr_start
-            self.last_lr = lr
-            return lr
-        else:
-            t = (n - self.lr_warm_up_steps) / (self.lr_max_decay_steps - self.lr_warm_up_steps)
-            t = min(t, 1.0)
-            lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
-                    1 + np.cos(t * np.pi))
-            self.last_lr = lr
-            return lr
-
-    def __call__(self, n, **kwargs):
-        return self.schedule(n,**kwargs)
-
-
-class LambdaWarmUpCosineScheduler2:
-    """
-    supports repeated iterations, configurable via lists
-    note: use with a base_lr of 1.0.
-    """
-    def __init__(self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0):
-        assert len(warm_up_steps) == len(f_min) == len(f_max) == len(f_start) == len(cycle_lengths)
-        self.lr_warm_up_steps = warm_up_steps
-        self.f_start = f_start
-        self.f_min = f_min
-        self.f_max = f_max
-        self.cycle_lengths = cycle_lengths
-        self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths))
-        self.last_f = 0.
-        self.verbosity_interval = verbosity_interval
-
-    def find_in_interval(self, n):
-        interval = 0
-        for cl in self.cum_cycles[1:]:
-            if n <= cl:
-                return interval
-            interval += 1
-
-    def schedule(self, n, **kwargs):
-        cycle = self.find_in_interval(n)
-        n = n - self.cum_cycles[cycle]
-        if self.verbosity_interval > 0:
-            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
-                                                       f"current cycle {cycle}")
-        if n < self.lr_warm_up_steps[cycle]:
-            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
-            self.last_f = f
-            return f
-        else:
-            t = (n - self.lr_warm_up_steps[cycle]) / (self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle])
-            t = min(t, 1.0)
-            f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (
-                    1 + np.cos(t * np.pi))
-            self.last_f = f
-            return f
-
-    def __call__(self, n, **kwargs):
-        return self.schedule(n, **kwargs)
-
-
-class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2):
-
-    def schedule(self, n, **kwargs):
-        cycle = self.find_in_interval(n)
-        n = n - self.cum_cycles[cycle]
-        if self.verbosity_interval > 0:
-            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
-                                                       f"current cycle {cycle}")
-
-        if n < self.lr_warm_up_steps[cycle]:
-            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
-            self.last_f = f
-            return f
-        else:
-            f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (self.cycle_lengths[cycle] - n) / (self.cycle_lengths[cycle])
-            self.last_f = f
-            return f
-

One-2-3-45-master 2/ldm/models/autoencoder.py

@@ -1,443 +0,0 @@
-import torch
-import pytorch_lightning as pl
-import torch.nn.functional as F
-from contextlib import contextmanager
-
-from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
-
-from ldm.modules.diffusionmodules.model import Encoder, Decoder
-from ldm.modules.distributions.distributions import DiagonalGaussianDistribution
-
-from ldm.util import instantiate_from_config
-
-
-class VQModel(pl.LightningModule):
-    def __init__(self,
-                 ddconfig,
-                 lossconfig,
-                 n_embed,
-                 embed_dim,
-                 ckpt_path=None,
-                 ignore_keys=[],
-                 image_key="image",
-                 colorize_nlabels=None,
-                 monitor=None,
-                 batch_resize_range=None,
-                 scheduler_config=None,
-                 lr_g_factor=1.0,
-                 remap=None,
-                 sane_index_shape=False, # tell vector quantizer to return indices as bhw
-                 use_ema=False
-                 ):
-        super().__init__()
-        self.embed_dim = embed_dim
-        self.n_embed = n_embed
-        self.image_key = image_key
-        self.encoder = Encoder(**ddconfig)
-        self.decoder = Decoder(**ddconfig)
-        self.loss = instantiate_from_config(lossconfig)
-        self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25,
-                                        remap=remap,
-                                        sane_index_shape=sane_index_shape)
-        self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
-        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
-        if colorize_nlabels is not None:
-            assert type(colorize_nlabels)==int
-            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
-        if monitor is not None:
-            self.monitor = monitor
-        self.batch_resize_range = batch_resize_range
-        if self.batch_resize_range is not None:
-            print(f"{self.__class__.__name__}: Using per-batch resizing in range {batch_resize_range}.")
-
-        self.use_ema = use_ema
-        if self.use_ema:
-            self.model_ema = LitEma(self)
-            print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
-
-        if ckpt_path is not None:
-            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
-        self.scheduler_config = scheduler_config
-        self.lr_g_factor = lr_g_factor
-
-    @contextmanager
-    def ema_scope(self, context=None):
-        if self.use_ema:
-            self.model_ema.store(self.parameters())
-            self.model_ema.copy_to(self)
-            if context is not None:
-                print(f"{context}: Switched to EMA weights")
-        try:
-            yield None
-        finally:
-            if self.use_ema:
-                self.model_ema.restore(self.parameters())
-                if context is not None:
-                    print(f"{context}: Restored training weights")
-
-    def init_from_ckpt(self, path, ignore_keys=list()):
-        sd = torch.load(path, map_location="cpu")["state_dict"]
-        keys = list(sd.keys())
-        for k in keys:
-            for ik in ignore_keys:
-                if k.startswith(ik):
-                    print("Deleting key {} from state_dict.".format(k))
-                    del sd[k]
-        missing, unexpected = self.load_state_dict(sd, strict=False)
-        print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
-        if len(missing) > 0:
-            print(f"Missing Keys: {missing}")
-            print(f"Unexpected Keys: {unexpected}")
-
-    def on_train_batch_end(self, *args, **kwargs):
-        if self.use_ema:
-            self.model_ema(self)
-
-    def encode(self, x):
-        h = self.encoder(x)
-        h = self.quant_conv(h)
-        quant, emb_loss, info = self.quantize(h)
-        return quant, emb_loss, info
-
-    def encode_to_prequant(self, x):
-        h = self.encoder(x)
-        h = self.quant_conv(h)
-        return h
-
-    def decode(self, quant):
-        quant = self.post_quant_conv(quant)
-        dec = self.decoder(quant)
-        return dec
-
-    def decode_code(self, code_b):
-        quant_b = self.quantize.embed_code(code_b)
-        dec = self.decode(quant_b)
-        return dec
-
-    def forward(self, input, return_pred_indices=False):
-        quant, diff, (_,_,ind) = self.encode(input)
-        dec = self.decode(quant)
-        if return_pred_indices:
-            return dec, diff, ind
-        return dec, diff
-
-    def get_input(self, batch, k):
-        x = batch[k]
-        if len(x.shape) == 3:
-            x = x[..., None]
-        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
-        if self.batch_resize_range is not None:
-            lower_size = self.batch_resize_range[0]
-            upper_size = self.batch_resize_range[1]
-            if self.global_step <= 4:
-                # do the first few batches with max size to avoid later oom
-                new_resize = upper_size
-            else:
-                new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16))
-            if new_resize != x.shape[2]:
-                x = F.interpolate(x, size=new_resize, mode="bicubic")
-            x = x.detach()
-        return x
-
-    def training_step(self, batch, batch_idx, optimizer_idx):
-        # https://github.com/pytorch/pytorch/issues/37142
-        # try not to fool the heuristics
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss, ind = self(x, return_pred_indices=True)
-
-        if optimizer_idx == 0:
-            # autoencode
-            aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
-                                            last_layer=self.get_last_layer(), split="train",
-                                            predicted_indices=ind)
-
-            self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-            return aeloss
-
-        if optimizer_idx == 1:
-            # discriminator
-            discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
-                                            last_layer=self.get_last_layer(), split="train")
-            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-            return discloss
-
-    def validation_step(self, batch, batch_idx):
-        log_dict = self._validation_step(batch, batch_idx)
-        with self.ema_scope():
-            log_dict_ema = self._validation_step(batch, batch_idx, suffix="_ema")
-        return log_dict
-
-    def _validation_step(self, batch, batch_idx, suffix=""):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss, ind = self(x, return_pred_indices=True)
-        aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0,
-                                        self.global_step,
-                                        last_layer=self.get_last_layer(),
-                                        split="val"+suffix,
-                                        predicted_indices=ind
-                                        )
-
-        discloss, log_dict_disc = self.loss(qloss, x, xrec, 1,
-                                            self.global_step,
-                                            last_layer=self.get_last_layer(),
-                                            split="val"+suffix,
-                                            predicted_indices=ind
-                                            )
-        rec_loss = log_dict_ae[f"val{suffix}/rec_loss"]
-        self.log(f"val{suffix}/rec_loss", rec_loss,
-                   prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
-        self.log(f"val{suffix}/aeloss", aeloss,
-                   prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
-        if version.parse(pl.__version__) >= version.parse('1.4.0'):
-            del log_dict_ae[f"val{suffix}/rec_loss"]
-        self.log_dict(log_dict_ae)
-        self.log_dict(log_dict_disc)
-        return self.log_dict
-
-    def configure_optimizers(self):
-        lr_d = self.learning_rate
-        lr_g = self.lr_g_factor*self.learning_rate
-        print("lr_d", lr_d)
-        print("lr_g", lr_g)
-        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
-                                  list(self.decoder.parameters())+
-                                  list(self.quantize.parameters())+
-                                  list(self.quant_conv.parameters())+
-                                  list(self.post_quant_conv.parameters()),
-                                  lr=lr_g, betas=(0.5, 0.9))
-        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
-                                    lr=lr_d, betas=(0.5, 0.9))
-
-        if self.scheduler_config is not None:
-            scheduler = instantiate_from_config(self.scheduler_config)
-
-            print("Setting up LambdaLR scheduler...")
-            scheduler = [
-                {
-                    'scheduler': LambdaLR(opt_ae, lr_lambda=scheduler.schedule),
-                    'interval': 'step',
-                    'frequency': 1
-                },
-                {
-                    'scheduler': LambdaLR(opt_disc, lr_lambda=scheduler.schedule),
-                    'interval': 'step',
-                    'frequency': 1
-                },
-            ]
-            return [opt_ae, opt_disc], scheduler
-        return [opt_ae, opt_disc], []
-
-    def get_last_layer(self):
-        return self.decoder.conv_out.weight
-
-    def log_images(self, batch, only_inputs=False, plot_ema=False, **kwargs):
-        log = dict()
-        x = self.get_input(batch, self.image_key)
-        x = x.to(self.device)
-        if only_inputs:
-            log["inputs"] = x
-            return log
-        xrec, _ = self(x)
-        if x.shape[1] > 3:
-            # colorize with random projection
-            assert xrec.shape[1] > 3
-            x = self.to_rgb(x)
-            xrec = self.to_rgb(xrec)
-        log["inputs"] = x
-        log["reconstructions"] = xrec
-        if plot_ema:
-            with self.ema_scope():
-                xrec_ema, _ = self(x)
-                if x.shape[1] > 3: xrec_ema = self.to_rgb(xrec_ema)
-                log["reconstructions_ema"] = xrec_ema
-        return log
-
-    def to_rgb(self, x):
-        assert self.image_key == "segmentation"
-        if not hasattr(self, "colorize"):
-            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
-        x = F.conv2d(x, weight=self.colorize)
-        x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
-        return x
-
-
-class VQModelInterface(VQModel):
-    def __init__(self, embed_dim, *args, **kwargs):
-        super().__init__(embed_dim=embed_dim, *args, **kwargs)
-        self.embed_dim = embed_dim
-
-    def encode(self, x):
-        h = self.encoder(x)
-        h = self.quant_conv(h)
-        return h
-
-    def decode(self, h, force_not_quantize=False):
-        # also go through quantization layer
-        if not force_not_quantize:
-            quant, emb_loss, info = self.quantize(h)
-        else:
-            quant = h
-        quant = self.post_quant_conv(quant)
-        dec = self.decoder(quant)
-        return dec
-
-
-class AutoencoderKL(pl.LightningModule):
-    def __init__(self,
-                 ddconfig,
-                 lossconfig,
-                 embed_dim,
-                 ckpt_path=None,
-                 ignore_keys=[],
-                 image_key="image",
-                 colorize_nlabels=None,
-                 monitor=None,
-                 ):
-        super().__init__()
-        self.image_key = image_key
-        self.encoder = Encoder(**ddconfig)
-        self.decoder = Decoder(**ddconfig)
-        self.loss = instantiate_from_config(lossconfig)
-        assert ddconfig["double_z"]
-        self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
-        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
-        self.embed_dim = embed_dim
-        if colorize_nlabels is not None:
-            assert type(colorize_nlabels)==int
-            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
-        if monitor is not None:
-            self.monitor = monitor
-        if ckpt_path is not None:
-            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
-
-    def init_from_ckpt(self, path, ignore_keys=list()):
-        sd = torch.load(path, map_location="cpu")["state_dict"]
-        keys = list(sd.keys())
-        for k in keys:
-            for ik in ignore_keys:
-                if k.startswith(ik):
-                    print("Deleting key {} from state_dict.".format(k))
-                    del sd[k]
-        self.load_state_dict(sd, strict=False)
-        print(f"Restored from {path}")
-
-    def encode(self, x):
-        h = self.encoder(x)
-        moments = self.quant_conv(h)
-        posterior = DiagonalGaussianDistribution(moments)
-        return posterior
-
-    def decode(self, z):
-        z = self.post_quant_conv(z)
-        dec = self.decoder(z)
-        return dec
-
-    def forward(self, input, sample_posterior=True):
-        posterior = self.encode(input)
-        if sample_posterior:
-            z = posterior.sample()
-        else:
-            z = posterior.mode()
-        dec = self.decode(z)
-        return dec, posterior
-
-    def get_input(self, batch, k):
-        x = batch[k]
-        if len(x.shape) == 3:
-            x = x[..., None]
-        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
-        return x
-
-    def training_step(self, batch, batch_idx, optimizer_idx):
-        inputs = self.get_input(batch, self.image_key)
-        reconstructions, posterior = self(inputs)
-
-        if optimizer_idx == 0:
-            # train encoder+decoder+logvar
-            aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
-                                            last_layer=self.get_last_layer(), split="train")
-            self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
-            self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False)
-            return aeloss
-
-        if optimizer_idx == 1:
-            # train the discriminator
-            discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
-                                                last_layer=self.get_last_layer(), split="train")
-
-            self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
-            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False)
-            return discloss
-
-    def validation_step(self, batch, batch_idx):
-        inputs = self.get_input(batch, self.image_key)
-        reconstructions, posterior = self(inputs)
-        aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,
-                                        last_layer=self.get_last_layer(), split="val")
-
-        discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,
-                                            last_layer=self.get_last_layer(), split="val")
-
-        self.log("val/rec_loss", log_dict_ae["val/rec_loss"])
-        self.log_dict(log_dict_ae)
-        self.log_dict(log_dict_disc)
-        return self.log_dict
-
-    def configure_optimizers(self):
-        lr = self.learning_rate
-        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
-                                  list(self.decoder.parameters())+
-                                  list(self.quant_conv.parameters())+
-                                  list(self.post_quant_conv.parameters()),
-                                  lr=lr, betas=(0.5, 0.9))
-        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
-                                    lr=lr, betas=(0.5, 0.9))
-        return [opt_ae, opt_disc], []
-
-    def get_last_layer(self):
-        return self.decoder.conv_out.weight
-
-    @torch.no_grad()
-    def log_images(self, batch, only_inputs=False, **kwargs):
-        log = dict()
-        x = self.get_input(batch, self.image_key)
-        x = x.to(self.device)
-        if not only_inputs:
-            xrec, posterior = self(x)
-            if x.shape[1] > 3:
-                # colorize with random projection
-                assert xrec.shape[1] > 3
-                x = self.to_rgb(x)
-                xrec = self.to_rgb(xrec)
-            log["samples"] = self.decode(torch.randn_like(posterior.sample()))
-            log["reconstructions"] = xrec
-        log["inputs"] = x
-        return log
-
-    def to_rgb(self, x):
-        assert self.image_key == "segmentation"
-        if not hasattr(self, "colorize"):
-            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
-        x = F.conv2d(x, weight=self.colorize)
-        x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
-        return x
-
-
-class IdentityFirstStage(torch.nn.Module):
-    def __init__(self, *args, vq_interface=False, **kwargs):
-        self.vq_interface = vq_interface  # TODO: Should be true by default but check to not break older stuff
-        super().__init__()
-
-    def encode(self, x, *args, **kwargs):
-        return x
-
-    def decode(self, x, *args, **kwargs):
-        return x
-
-    def quantize(self, x, *args, **kwargs):
-        if self.vq_interface:
-            return x, None, [None, None, None]
-        return x
-
-    def forward(self, x, *args, **kwargs):
-        return x

One-2-3-45-master 2/ldm/models/diffusion/classifier.py

@@ -1,267 +0,0 @@
-import os
-import torch
-import pytorch_lightning as pl
-from omegaconf import OmegaConf
-from torch.nn import functional as F
-from torch.optim import AdamW
-from torch.optim.lr_scheduler import LambdaLR
-from copy import deepcopy
-from einops import rearrange
-from glob import glob
-from natsort import natsorted
-
-from ldm.modules.diffusionmodules.openaimodel import EncoderUNetModel, UNetModel
-from ldm.util import log_txt_as_img, default, ismap, instantiate_from_config
-
-__models__ = {
-    'class_label': EncoderUNetModel,
-    'segmentation': UNetModel
-}
-
-
-def disabled_train(self, mode=True):
-    """Overwrite model.train with this function to make sure train/eval mode
-    does not change anymore."""
-    return self
-
-
-class NoisyLatentImageClassifier(pl.LightningModule):
-
-    def __init__(self,
-                 diffusion_path,
-                 num_classes,
-                 ckpt_path=None,
-                 pool='attention',
-                 label_key=None,
-                 diffusion_ckpt_path=None,
-                 scheduler_config=None,
-                 weight_decay=1.e-2,
-                 log_steps=10,
-                 monitor='val/loss',
-                 *args,
-                 **kwargs):
-        super().__init__(*args, **kwargs)
-        self.num_classes = num_classes
-        # get latest config of diffusion model
-        diffusion_config = natsorted(glob(os.path.join(diffusion_path, 'configs', '*-project.yaml')))[-1]
-        self.diffusion_config = OmegaConf.load(diffusion_config).model
-        self.diffusion_config.params.ckpt_path = diffusion_ckpt_path
-        self.load_diffusion()
-
-        self.monitor = monitor
-        self.numd = self.diffusion_model.first_stage_model.encoder.num_resolutions - 1
-        self.log_time_interval = self.diffusion_model.num_timesteps // log_steps
-        self.log_steps = log_steps
-
-        self.label_key = label_key if not hasattr(self.diffusion_model, 'cond_stage_key') \
-            else self.diffusion_model.cond_stage_key
-
-        assert self.label_key is not None, 'label_key neither in diffusion model nor in model.params'
-
-        if self.label_key not in __models__:
-            raise NotImplementedError()
-
-        self.load_classifier(ckpt_path, pool)
-
-        self.scheduler_config = scheduler_config
-        self.use_scheduler = self.scheduler_config is not None
-        self.weight_decay = weight_decay
-
-    def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
-        sd = torch.load(path, map_location="cpu")
-        if "state_dict" in list(sd.keys()):
-            sd = sd["state_dict"]
-        keys = list(sd.keys())
-        for k in keys:
-            for ik in ignore_keys:
-                if k.startswith(ik):
-                    print("Deleting key {} from state_dict.".format(k))
-                    del sd[k]
-        missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
-            sd, strict=False)
-        print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
-        if len(missing) > 0:
-            print(f"Missing Keys: {missing}")
-        if len(unexpected) > 0:
-            print(f"Unexpected Keys: {unexpected}")
-
-    def load_diffusion(self):
-        model = instantiate_from_config(self.diffusion_config)
-        self.diffusion_model = model.eval()
-        self.diffusion_model.train = disabled_train
-        for param in self.diffusion_model.parameters():
-            param.requires_grad = False
-
-    def load_classifier(self, ckpt_path, pool):
-        model_config = deepcopy(self.diffusion_config.params.unet_config.params)
-        model_config.in_channels = self.diffusion_config.params.unet_config.params.out_channels
-        model_config.out_channels = self.num_classes
-        if self.label_key == 'class_label':
-            model_config.pool = pool
-
-        self.model = __models__[self.label_key](**model_config)
-        if ckpt_path is not None:
-            print('#####################################################################')
-            print(f'load from ckpt "{ckpt_path}"')
-            print('#####################################################################')
-            self.init_from_ckpt(ckpt_path)
-
-    @torch.no_grad()
-    def get_x_noisy(self, x, t, noise=None):
-        noise = default(noise, lambda: torch.randn_like(x))
-        continuous_sqrt_alpha_cumprod = None
-        if self.diffusion_model.use_continuous_noise:
-            continuous_sqrt_alpha_cumprod = self.diffusion_model.sample_continuous_noise_level(x.shape[0], t + 1)
-            # todo: make sure t+1 is correct here
-
-        return self.diffusion_model.q_sample(x_start=x, t=t, noise=noise,
-                                             continuous_sqrt_alpha_cumprod=continuous_sqrt_alpha_cumprod)
-
-    def forward(self, x_noisy, t, *args, **kwargs):
-        return self.model(x_noisy, t)
-
-    @torch.no_grad()
-    def get_input(self, batch, k):
-        x = batch[k]
-        if len(x.shape) == 3:
-            x = x[..., None]
-        x = rearrange(x, 'b h w c -> b c h w')
-        x = x.to(memory_format=torch.contiguous_format).float()
-        return x
-
-    @torch.no_grad()
-    def get_conditioning(self, batch, k=None):
-        if k is None:
-            k = self.label_key
-        assert k is not None, 'Needs to provide label key'
-
-        targets = batch[k].to(self.device)
-
-        if self.label_key == 'segmentation':
-            targets = rearrange(targets, 'b h w c -> b c h w')
-            for down in range(self.numd):
-                h, w = targets.shape[-2:]
-                targets = F.interpolate(targets, size=(h // 2, w // 2), mode='nearest')
-
-            # targets = rearrange(targets,'b c h w -> b h w c')
-
-        return targets
-
-    def compute_top_k(self, logits, labels, k, reduction="mean"):
-        _, top_ks = torch.topk(logits, k, dim=1)
-        if reduction == "mean":
-            return (top_ks == labels[:, None]).float().sum(dim=-1).mean().item()
-        elif reduction == "none":
-            return (top_ks == labels[:, None]).float().sum(dim=-1)
-
-    def on_train_epoch_start(self):
-        # save some memory
-        self.diffusion_model.model.to('cpu')
-
-    @torch.no_grad()
-    def write_logs(self, loss, logits, targets):
-        log_prefix = 'train' if self.training else 'val'
-        log = {}
-        log[f"{log_prefix}/loss"] = loss.mean()
-        log[f"{log_prefix}/acc@1"] = self.compute_top_k(
-            logits, targets, k=1, reduction="mean"
-        )
-        log[f"{log_prefix}/acc@5"] = self.compute_top_k(
-            logits, targets, k=5, reduction="mean"
-        )
-
-        self.log_dict(log, prog_bar=False, logger=True, on_step=self.training, on_epoch=True)
-        self.log('loss', log[f"{log_prefix}/loss"], prog_bar=True, logger=False)
-        self.log('global_step', self.global_step, logger=False, on_epoch=False, prog_bar=True)
-        lr = self.optimizers().param_groups[0]['lr']
-        self.log('lr_abs', lr, on_step=True, logger=True, on_epoch=False, prog_bar=True)
-
-    def shared_step(self, batch, t=None):
-        x, *_ = self.diffusion_model.get_input(batch, k=self.diffusion_model.first_stage_key)
-        targets = self.get_conditioning(batch)
-        if targets.dim() == 4:
-            targets = targets.argmax(dim=1)
-        if t is None:
-            t = torch.randint(0, self.diffusion_model.num_timesteps, (x.shape[0],), device=self.device).long()
-        else:
-            t = torch.full(size=(x.shape[0],), fill_value=t, device=self.device).long()
-        x_noisy = self.get_x_noisy(x, t)
-        logits = self(x_noisy, t)
-
-        loss = F.cross_entropy(logits, targets, reduction='none')
-
-        self.write_logs(loss.detach(), logits.detach(), targets.detach())
-
-        loss = loss.mean()
-        return loss, logits, x_noisy, targets
-
-    def training_step(self, batch, batch_idx):
-        loss, *_ = self.shared_step(batch)
-        return loss
-
-    def reset_noise_accs(self):
-        self.noisy_acc = {t: {'acc@1': [], 'acc@5': []} for t in
-                          range(0, self.diffusion_model.num_timesteps, self.diffusion_model.log_every_t)}
-
-    def on_validation_start(self):
-        self.reset_noise_accs()
-
-    @torch.no_grad()
-    def validation_step(self, batch, batch_idx):
-        loss, *_ = self.shared_step(batch)
-
-        for t in self.noisy_acc:
-            _, logits, _, targets = self.shared_step(batch, t)
-            self.noisy_acc[t]['acc@1'].append(self.compute_top_k(logits, targets, k=1, reduction='mean'))
-            self.noisy_acc[t]['acc@5'].append(self.compute_top_k(logits, targets, k=5, reduction='mean'))
-
-        return loss
-
-    def configure_optimizers(self):
-        optimizer = AdamW(self.model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
-
-        if self.use_scheduler:
-            scheduler = instantiate_from_config(self.scheduler_config)
-
-            print("Setting up LambdaLR scheduler...")
-            scheduler = [
-                {
-                    'scheduler': LambdaLR(optimizer, lr_lambda=scheduler.schedule),
-                    'interval': 'step',
-                    'frequency': 1
-                }]
-            return [optimizer], scheduler
-
-        return optimizer
-
-    @torch.no_grad()
-    def log_images(self, batch, N=8, *args, **kwargs):
-        log = dict()
-        x = self.get_input(batch, self.diffusion_model.first_stage_key)
-        log['inputs'] = x
-
-        y = self.get_conditioning(batch)
-
-        if self.label_key == 'class_label':
-            y = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"])
-            log['labels'] = y
-
-        if ismap(y):
-            log['labels'] = self.diffusion_model.to_rgb(y)
-
-            for step in range(self.log_steps):
-                current_time = step * self.log_time_interval
-
-                _, logits, x_noisy, _ = self.shared_step(batch, t=current_time)
-
-                log[f'inputs@t{current_time}'] = x_noisy
-
-                pred = F.one_hot(logits.argmax(dim=1), num_classes=self.num_classes)
-                pred = rearrange(pred, 'b h w c -> b c h w')
-
-                log[f'pred@t{current_time}'] = self.diffusion_model.to_rgb(pred)
-
-        for key in log:
-            log[key] = log[key][:N]
-
-        return log

One-2-3-45-master 2/ldm/models/diffusion/ddim.py

@@ -1,326 +0,0 @@
-"""SAMPLING ONLY."""
-
-import torch
-import numpy as np
-from tqdm import tqdm
-from functools import partial
-from einops import rearrange
-
-from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor
-from ldm.models.diffusion.sampling_util import renorm_thresholding, norm_thresholding, spatial_norm_thresholding
-
-
-class DDIMSampler(object):
-    def __init__(self, model, schedule="linear", **kwargs):
-        super().__init__()
-        self.model = model
-        self.ddpm_num_timesteps = model.num_timesteps
-        self.schedule = schedule
-        self.device = model.device
-
-    def to(self, device):
-        """Same as to in torch module
-        Don't really underestand why this isn't a module in the first place"""
-        for k, v in self.__dict__.items():
-            if isinstance(v, torch.Tensor):
-                new_v = getattr(self, k).to(device)
-                setattr(self, k, new_v)
-
-
-    def register_buffer(self, name, attr, device=None):
-        if type(attr) == torch.Tensor:
-            attr = attr.to(device)
-            # if attr.device != torch.device("cuda"):
-            #     attr = attr.to(torch.device("cuda"))
-        setattr(self, name, attr)
-
-    def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
-        self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
-                                                  num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
-        alphas_cumprod = self.model.alphas_cumprod
-        assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
-        to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
-
-        self.register_buffer('betas', to_torch(self.model.betas), self.device)
-        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod), self.device)
-        self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev), self.device)
-
-        # calculations for diffusion q(x_t | x_{t-1}) and others
-        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())), self.device)
-        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())), self.device)
-        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())), self.device)
-        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())), self.device)
-        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)), self.device)
-
-        # ddim sampling parameters
-        ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
-                                                                                   ddim_timesteps=self.ddim_timesteps,
-                                                                                   eta=ddim_eta,verbose=verbose)
-        self.register_buffer('ddim_sigmas', ddim_sigmas, self.device)
-        self.register_buffer('ddim_alphas', ddim_alphas, self.device)
-        self.register_buffer('ddim_alphas_prev', ddim_alphas_prev, self.device)
-        self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas), self.device)
-        sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
-            (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
-                        1 - self.alphas_cumprod / self.alphas_cumprod_prev))
-        self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps, self.device)
-
-    @torch.no_grad()
-    def sample(self,
-               S,
-               batch_size,
-               shape,
-               conditioning=None,
-               callback=None,
-               normals_sequence=None,
-               img_callback=None,
-               quantize_x0=False,
-               eta=0.,
-               mask=None,
-               x0=None,
-               temperature=1.,
-               noise_dropout=0.,
-               score_corrector=None,
-               corrector_kwargs=None,
-               verbose=True,
-               x_T=None,
-               log_every_t=100,
-               unconditional_guidance_scale=1.,
-               unconditional_conditioning=None, # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
-               dynamic_threshold=None,
-               **kwargs
-               ):
-        if conditioning is not None:
-            if isinstance(conditioning, dict):
-                ctmp = conditioning[list(conditioning.keys())[0]]
-                while isinstance(ctmp, list): ctmp = ctmp[0]
-                cbs = ctmp.shape[0]
-                if cbs != batch_size:
-                    print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
-
-            else:
-                if conditioning.shape[0] != batch_size:
-                    print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
-
-        self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
-        # sampling
-        C, H, W = shape
-        size = (batch_size, C, H, W)
-        print(f'Data shape for DDIM sampling is {size}, eta {eta}')
-
-        samples, intermediates = self.ddim_sampling(conditioning, size,
-                                                    callback=callback,
-                                                    img_callback=img_callback,
-                                                    quantize_denoised=quantize_x0,
-                                                    mask=mask, x0=x0,
-                                                    ddim_use_original_steps=False,
-                                                    noise_dropout=noise_dropout,
-                                                    temperature=temperature,
-                                                    score_corrector=score_corrector,
-                                                    corrector_kwargs=corrector_kwargs,
-                                                    x_T=x_T,
-                                                    log_every_t=log_every_t,
-                                                    unconditional_guidance_scale=unconditional_guidance_scale,
-                                                    unconditional_conditioning=unconditional_conditioning,
-                                                    dynamic_threshold=dynamic_threshold,
-                                                    )
-        return samples, intermediates
-
-    @torch.no_grad()
-    def ddim_sampling(self, cond, shape,
-                      x_T=None, ddim_use_original_steps=False,
-                      callback=None, timesteps=None, quantize_denoised=False,
-                      mask=None, x0=None, img_callback=None, log_every_t=100,
-                      temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
-                      unconditional_guidance_scale=1., unconditional_conditioning=None, dynamic_threshold=None,
-                      t_start=-1):
-        device = self.model.betas.device
-        b = shape[0]
-        if x_T is None:
-            img = torch.randn(shape, device=device)
-        else:
-            img = x_T
-
-        if timesteps is None:
-            timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
-        elif timesteps is not None and not ddim_use_original_steps:
-            subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
-            timesteps = self.ddim_timesteps[:subset_end]
-
-        timesteps = timesteps[:t_start]
-
-        intermediates = {'x_inter': [img], 'pred_x0': [img]}
-        time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
-        total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
-        print(f"Running DDIM Sampling with {total_steps} timesteps")
-
-        iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
-
-        for i, step in enumerate(iterator):
-            index = total_steps - i - 1
-            ts = torch.full((b,), step, device=device, dtype=torch.long)
-
-            if mask is not None:
-                assert x0 is not None
-                img_orig = self.model.q_sample(x0, ts)  # TODO: deterministic forward pass?
-                img = img_orig * mask + (1. - mask) * img
-
-            outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
-                                      quantize_denoised=quantize_denoised, temperature=temperature,
-                                      noise_dropout=noise_dropout, score_corrector=score_corrector,
-                                      corrector_kwargs=corrector_kwargs,
-                                      unconditional_guidance_scale=unconditional_guidance_scale,
-                                      unconditional_conditioning=unconditional_conditioning,
-                                      dynamic_threshold=dynamic_threshold)
-            img, pred_x0 = outs
-            if callback:
-                img = callback(i, img, pred_x0)
-            if img_callback: img_callback(pred_x0, i)
-
-            if index % log_every_t == 0 or index == total_steps - 1:
-                intermediates['x_inter'].append(img)
-                intermediates['pred_x0'].append(pred_x0)
-
-        return img, intermediates
-
-    @torch.no_grad()
-    def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
-                      temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
-                      unconditional_guidance_scale=1., unconditional_conditioning=None,
-                      dynamic_threshold=None):
-        b, *_, device = *x.shape, x.device
-
-        if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
-            e_t = self.model.apply_model(x, t, c)
-        else:
-            x_in = torch.cat([x] * 2)
-            t_in = torch.cat([t] * 2)
-            if isinstance(c, dict):
-                assert isinstance(unconditional_conditioning, dict)
-                c_in = dict()
-                for k in c:
-                    if isinstance(c[k], list):
-                        c_in[k] = [torch.cat([
-                            unconditional_conditioning[k][i],
-                            c[k][i]]) for i in range(len(c[k]))]
-                    else:
-                        c_in[k] = torch.cat([
-                                unconditional_conditioning[k],
-                                c[k]])
-            else:
-                c_in = torch.cat([unconditional_conditioning, c])
-            e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
-            e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
-
-        if score_corrector is not None:
-            assert self.model.parameterization == "eps"
-            e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
-
-        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
-        alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
-        sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
-        sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
-        # select parameters corresponding to the currently considered timestep
-        a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
-        a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
-        sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
-        sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
-
-        # current prediction for x_0
-        pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
-        if quantize_denoised:
-            pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
-
-        if dynamic_threshold is not None:
-            pred_x0 = norm_thresholding(pred_x0, dynamic_threshold)
-
-        # direction pointing to x_t
-        dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
-        noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
-        if noise_dropout > 0.:
-            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
-        x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
-        return x_prev, pred_x0
-
-    @torch.no_grad()
-    def encode(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None,
-               unconditional_guidance_scale=1.0, unconditional_conditioning=None):
-        num_reference_steps = self.ddpm_num_timesteps if use_original_steps else self.ddim_timesteps.shape[0]
-
-        assert t_enc <= num_reference_steps
-        num_steps = t_enc
-
-        if use_original_steps:
-            alphas_next = self.alphas_cumprod[:num_steps]
-            alphas = self.alphas_cumprod_prev[:num_steps]
-        else:
-            alphas_next = self.ddim_alphas[:num_steps]
-            alphas = torch.tensor(self.ddim_alphas_prev[:num_steps])
-
-        x_next = x0
-        intermediates = []
-        inter_steps = []
-        for i in tqdm(range(num_steps), desc='Encoding Image'):
-            t = torch.full((x0.shape[0],), i, device=self.model.device, dtype=torch.long)
-            if unconditional_guidance_scale == 1.:
-                noise_pred = self.model.apply_model(x_next, t, c)
-            else:
-                assert unconditional_conditioning is not None
-                e_t_uncond, noise_pred = torch.chunk(
-                    self.model.apply_model(torch.cat((x_next, x_next)), torch.cat((t, t)),
-                                           torch.cat((unconditional_conditioning, c))), 2)
-                noise_pred = e_t_uncond + unconditional_guidance_scale * (noise_pred - e_t_uncond)
-
-            xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next
-            weighted_noise_pred = alphas_next[i].sqrt() * (
-                    (1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt()) * noise_pred
-            x_next = xt_weighted + weighted_noise_pred
-            if return_intermediates and i % (
-                    num_steps // return_intermediates) == 0 and i < num_steps - 1:
-                intermediates.append(x_next)
-                inter_steps.append(i)
-            elif return_intermediates and i >= num_steps - 2:
-                intermediates.append(x_next)
-                inter_steps.append(i)
-
-        out = {'x_encoded': x_next, 'intermediate_steps': inter_steps}
-        if return_intermediates:
-            out.update({'intermediates': intermediates})
-        return x_next, out
-
-    @torch.no_grad()
-    def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
-        # fast, but does not allow for exact reconstruction
-        # t serves as an index to gather the correct alphas
-        if use_original_steps:
-            sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
-            sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
-        else:
-            sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
-            sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
-
-        if noise is None:
-            noise = torch.randn_like(x0)
-        return (extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 +
-                extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise)
-
-    @torch.no_grad()
-    def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
-               use_original_steps=False):
-
-        timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
-        timesteps = timesteps[:t_start]
-
-        time_range = np.flip(timesteps)
-        total_steps = timesteps.shape[0]
-        print(f"Running DDIM Sampling with {total_steps} timesteps")
-
-        iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
-        x_dec = x_latent
-        for i, step in enumerate(iterator):
-            index = total_steps - i - 1
-            ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)
-            x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
-                                          unconditional_guidance_scale=unconditional_guidance_scale,
-                                          unconditional_conditioning=unconditional_conditioning)
-        return x_dec