├── .dockerignore ├── .gitignore ├── Dockerfile.base ├── INSTALL.md ├── LICENSE.md ├── MODELZOO.md ├── README.md ├── configs ├── projects │ ├── coco_funit │ │ ├── animal_faces │ │ │ ├── base64_bs8_class119.yaml │ │ │ └── base64_bs8_class149.yaml │ │ └── mammals │ │ │ └── base64_bs8_class305.yaml │ ├── fs_vid2vid │ │ ├── face_forensics │ │ │ └── ampO1.yaml │ │ └── youtube_dancing │ │ │ └── ampO1.yaml │ ├── funit │ │ └── animal_faces │ │ │ ├── base64_bs8_class119.yaml │ │ │ └── base64_bs8_class149.yaml │ ├── gancraft │ │ ├── demoworld.yaml │ │ ├── landscape.yaml │ │ └── survivalisland.yaml │ ├── munit │ │ ├── afhq_dog2cat │ │ │ └── ampO1.yaml │ │ ├── edges2shoes │ │ │ └── ampO1.yaml │ │ └── summer2winter_hd │ │ │ └── ampO1.yaml │ ├── pix2pixhd │ │ └── cityscapes │ │ │ ├── ampO1.yaml │ │ │ └── ampO1_2k.yaml │ ├── spade │ │ ├── cocostuff │ │ │ ├── base128_bs4.yaml │ │ │ └── base128_bs4_lmdb.yaml │ │ └── landscapes │ │ │ └── lhq.yaml │ ├── unit │ │ └── winter2summer │ │ │ └── base48_bs1.yaml │ ├── vid2vid │ │ ├── cityscapes │ │ │ └── ampO1.yaml │ │ └── dancing │ │ │ └── ampO1.yaml │ └── wc_vid2vid │ │ ├── cityscapes │ │ ├── seg_ampO1.yaml │ │ └── seg_single.yaml │ │ └── mannequin │ │ ├── hed_ampO0.yaml │ │ └── hed_single.yaml └── unit_test │ ├── coco_funit.yaml │ ├── fs_vid2vid_face.yaml │ ├── fs_vid2vid_pose.yaml │ ├── funit.yaml │ ├── gancraft.yaml │ ├── munit.yaml │ ├── munit_patch.yaml │ ├── pix2pixHD.yaml │ ├── spade.yaml │ ├── unit.yaml │ ├── vid2vid_pose.yaml │ ├── vid2vid_street.yaml │ └── wc_vid2vid.yaml ├── dataset └── unit_test │ └── raw │ ├── fs_vid2vid │ └── face │ │ ├── images │ │ └── ZzTpJYDc62E_4_ZzTpJYDc62E_4 │ │ │ ├── 00000.jpg │ │ │ ├── 00001.jpg │ │ │ ├── 00002.jpg │ │ │ ├── 00003.jpg │ │ │ └── 00004.jpg │ │ └── landmarks-dlib68 │ │ └── ZzTpJYDc62E_4_ZzTpJYDc62E_4 │ │ ├── 00000.json │ │ ├── 00001.json │ │ ├── 00002.json │ │ ├── 00003.json │ │ └── 00004.json │ ├── funit │ ├── images_content │ │ ├── cat │ │ │ └── cat_001.jpg │ │ ├── dog │ │ │ └── dog_001.jpg │ │ └── wild │ │ │ └── wild_001.jpg │ └── images_style │ │ └── images_content │ │ ├── cat │ │ └── cat_001.jpg │ │ ├── dog │ │ └── dog_001.jpg │ │ └── wild │ │ └── wild_001.jpg │ ├── munit │ ├── images_a │ │ └── dog_001.jpg │ └── images_b │ │ └── cat_001.jpg │ ├── pix2pixHD │ ├── images │ │ └── 001.jpg │ ├── instance_maps │ │ └── 001.png │ └── seg_maps │ │ └── 001.png │ ├── spade │ ├── edge_maps │ │ └── 000000581781.png │ ├── images │ │ └── 000000581781.jpg │ └── seg_maps │ │ └── 000000581781.png │ ├── vid2vid │ ├── pose │ │ ├── human_instance_maps │ │ │ └── 001_001 │ │ │ │ ├── frame000100.png │ │ │ │ ├── frame000101.png │ │ │ │ ├── frame000102.png │ │ │ │ ├── frame000103.png │ │ │ │ └── frame000104.png │ │ ├── images │ │ │ └── 001_001 │ │ │ │ ├── frame000100.jpg │ │ │ │ ├── frame000101.jpg │ │ │ │ ├── frame000102.jpg │ │ │ │ ├── frame000103.jpg │ │ │ │ └── frame000104.jpg │ │ ├── pose_maps-densepose │ │ │ └── 001_001 │ │ │ │ ├── frame000100.png │ │ │ │ ├── frame000101.png │ │ │ │ ├── frame000102.png │ │ │ │ ├── frame000103.png │ │ │ │ └── frame000104.png │ │ └── poses-openpose │ │ │ └── 001_001 │ │ │ ├── frame000100.json │ │ │ ├── frame000101.json │ │ │ ├── frame000102.json │ │ │ ├── frame000103.json │ │ │ └── frame000104.json │ └── street │ │ ├── images │ │ └── seq0001 │ │ │ ├── aachen_000001_000000_leftImg8bit.png │ │ │ ├── aachen_000001_000001_leftImg8bit.png │ │ │ ├── aachen_000001_000002_leftImg8bit.png │ │ │ └── aachen_000001_000003_leftImg8bit.png │ │ └── seg_maps │ │ └── seq0001 │ │ ├── aachen_000001_000000_leftImg8bit.png │ │ ├── aachen_000001_000001_leftImg8bit.png │ │ ├── aachen_000001_000002_leftImg8bit.png │ │ └── aachen_000001_000003_leftImg8bit.png │ └── wc_vid2vid │ └── cityscapes │ ├── depth_maps-megadepth │ └── seq0000 │ │ ├── aachen_000000_000000_leftImg8bit.png │ │ ├── aachen_000000_000001_leftImg8bit.png │ │ ├── aachen_000000_000002_leftImg8bit.png │ │ ├── aachen_000000_000003_leftImg8bit.png │ │ ├── aachen_000000_000004_leftImg8bit.png │ │ ├── aachen_000000_000005_leftImg8bit.png │ │ ├── aachen_000000_000006_leftImg8bit.png │ │ ├── aachen_000000_000007_leftImg8bit.png │ │ ├── aachen_000000_000008_leftImg8bit.png │ │ └── aachen_000000_000009_leftImg8bit.png │ ├── images │ └── seq0000 │ │ ├── aachen_000000_000000_leftImg8bit.jpg │ │ ├── aachen_000000_000001_leftImg8bit.jpg │ │ ├── aachen_000000_000002_leftImg8bit.jpg │ │ ├── aachen_000000_000003_leftImg8bit.jpg │ │ ├── aachen_000000_000004_leftImg8bit.jpg │ │ ├── aachen_000000_000005_leftImg8bit.jpg │ │ ├── aachen_000000_000006_leftImg8bit.jpg │ │ ├── aachen_000000_000007_leftImg8bit.jpg │ │ ├── aachen_000000_000008_leftImg8bit.jpg │ │ └── aachen_000000_000009_leftImg8bit.jpg │ ├── seg_maps │ └── seq0000 │ │ ├── aachen_000000_000000_leftImg8bit.png │ │ ├── aachen_000000_000001_leftImg8bit.png │ │ ├── aachen_000000_000002_leftImg8bit.png │ │ ├── aachen_000000_000003_leftImg8bit.png │ │ ├── aachen_000000_000004_leftImg8bit.png │ │ ├── aachen_000000_000005_leftImg8bit.png │ │ ├── aachen_000000_000006_leftImg8bit.png │ │ ├── aachen_000000_000007_leftImg8bit.png │ │ ├── aachen_000000_000008_leftImg8bit.png │ │ └── aachen_000000_000009_leftImg8bit.png │ └── unprojections │ └── seq0000 │ ├── aachen_000000_000000_leftImg8bit.pkl │ ├── aachen_000000_000001_leftImg8bit.pkl │ ├── aachen_000000_000002_leftImg8bit.pkl │ ├── aachen_000000_000003_leftImg8bit.pkl │ ├── aachen_000000_000004_leftImg8bit.pkl │ ├── aachen_000000_000005_leftImg8bit.pkl │ ├── aachen_000000_000006_leftImg8bit.pkl │ ├── aachen_000000_000007_leftImg8bit.pkl │ ├── aachen_000000_000008_leftImg8bit.pkl │ └── aachen_000000_000009_leftImg8bit.pkl ├── evaluate.py ├── imaginaire ├── __init__.py ├── config.py ├── datasets │ ├── __init__.py │ ├── base.py │ ├── cache.py │ ├── dummy.py │ ├── folder.py │ ├── images.py │ ├── lmdb.py │ ├── object_store.py │ ├── paired_few_shot_videos.py │ ├── paired_few_shot_videos_native.py │ ├── paired_images.py │ ├── paired_videos.py │ ├── unpaired_few_shot_images.py │ └── unpaired_images.py ├── discriminators │ ├── __init__.py │ ├── dummy.py │ ├── fpse.py │ ├── fs_vid2vid.py │ ├── funit.py │ ├── gancraft.py │ ├── mlp_multiclass.py │ ├── multires_patch.py │ ├── munit.py │ ├── residual.py │ ├── spade.py │ └── unit.py ├── evaluation │ ├── __init__.py │ ├── caption │ │ ├── __init__.py │ │ ├── clip.py │ │ ├── common.py │ │ └── r_precision.py │ ├── common.py │ ├── fid.py │ ├── kid.py │ ├── knn.py │ ├── lpips.py │ ├── msid.py │ ├── prdc.py │ ├── pretrained.py │ └── segmentation │ │ ├── __init__.py │ │ ├── celebamask_hq.py │ │ ├── cocostuff.py │ │ └── common.py ├── generators │ ├── __init__.py │ ├── coco_funit.py │ ├── dummy.py │ ├── fs_vid2vid.py │ ├── funit.py │ ├── gancraft.py │ ├── gancraft_base.py │ ├── munit.py │ ├── pix2pixHD.py │ ├── spade.py │ ├── unit.py │ ├── vid2vid.py │ └── wc_vid2vid.py ├── layers │ ├── __init__.py │ ├── activation_norm.py │ ├── conv.py │ ├── misc.py │ ├── non_local.py │ ├── nonlinearity.py │ ├── residual.py │ ├── residual_deep.py │ ├── vit.py │ └── weight_norm.py ├── losses │ ├── __init__.py │ ├── dict.py │ ├── feature_matching.py │ ├── flow.py │ ├── gan.py │ ├── info_nce.py │ ├── kl.py │ ├── perceptual.py │ └── weighted_mse.py ├── model_utils │ ├── __init__.py │ ├── fs_vid2vid.py │ ├── gancraft │ │ ├── camctl.py │ │ ├── gaugan_lbl2col.csv │ │ ├── gaugan_reduction.csv │ │ ├── id2name_gg.csv │ │ ├── layers.py │ │ ├── loss.py │ │ ├── mc_lbl_reduction.py │ │ ├── mc_reduction.csv │ │ ├── mc_utils.py │ │ ├── reduced_coco_lbls.csv │ │ └── voxlib │ │ │ ├── Makefile │ │ │ ├── __init__.py │ │ │ ├── positional_encoding.py │ │ │ ├── positional_encoding_kernel.cu │ │ │ ├── ray_voxel_intersection.cu │ │ │ ├── setup.py │ │ │ ├── sp_trilinear.py │ │ │ ├── sp_trilinear_worldcoord_kernel.cu │ │ │ ├── voxlib.cpp │ │ │ └── voxlib_common.h │ ├── label.py │ ├── pix2pixHD.py │ ├── rename_inputs.py │ └── wc_vid2vid │ │ └── render.py ├── optimizers │ ├── __init__.py │ ├── fromage.py │ └── madam.py ├── third_party │ ├── __init__.py │ ├── bias_act │ │ ├── __init__.py │ │ ├── bias_act.py │ │ ├── setup.py │ │ └── src │ │ │ ├── bias_act_cuda.cc │ │ │ ├── bias_act_cuda.h │ │ │ └── bias_act_cuda_kernel.cu │ ├── channelnorm │ │ ├── channelnorm.py │ │ ├── setup.py │ │ └── src │ │ │ ├── channelnorm_cuda.cc │ │ │ ├── channelnorm_kernel.cu │ │ │ └── channelnorm_kernel.cuh │ ├── correlation │ │ ├── correlation.py │ │ ├── setup.py │ │ └── src │ │ │ ├── correlation_cuda.cc │ │ │ ├── correlation_cuda_kernel.cu │ │ │ └── correlation_cuda_kernel.cuh │ ├── flow_net │ │ ├── __init__.py │ │ ├── flow_net.py │ │ ├── flownet2 │ │ │ ├── models.py │ │ │ ├── networks │ │ │ │ ├── __init__.py │ │ │ │ ├── flownet_c.py │ │ │ │ ├── flownet_fusion.py │ │ │ │ ├── flownet_s.py │ │ │ │ ├── flownet_sd.py │ │ │ │ └── submodules.py │ │ │ └── utils │ │ │ │ ├── __init__.py │ │ │ │ ├── flow_utils.py │ │ │ │ ├── frame_utils.py │ │ │ │ ├── param_utils.py │ │ │ │ └── tools.py │ │ └── test_flownet2.py │ ├── resample2d │ │ ├── resample2d.py │ │ ├── setup.py │ │ └── src │ │ │ ├── resample2d_cuda.cc │ │ │ ├── resample2d_kernel.cu │ │ │ └── resample2d_kernel.cuh │ └── upfirdn2d │ │ ├── __init__.py │ │ ├── setup.py │ │ ├── src │ │ ├── upfirdn2d_cuda.cc │ │ ├── upfirdn2d_cuda.h │ │ └── upfirdn2d_cuda_kernel.cu │ │ └── upfirdn2d.py ├── trainers │ ├── __init__.py │ ├── base.py │ ├── fs_vid2vid.py │ ├── funit.py │ ├── gancraft.py │ ├── munit.py │ ├── pix2pixHD.py │ ├── spade.py │ ├── unit.py │ ├── vid2vid.py │ └── wc_vid2vid.py └── utils │ ├── __init__.py │ ├── cudnn.py │ ├── data.py │ ├── dataset.py │ ├── diff_aug.py │ ├── distributed.py │ ├── gpu_affinity.py │ ├── init_weight.py │ ├── io.py │ ├── lmdb.py │ ├── logging.py │ ├── meters.py │ ├── misc.py │ ├── model_average.py │ ├── path.py │ ├── trainer.py │ └── visualization │ ├── __init__.py │ ├── common.py │ ├── face.py │ └── pose.py ├── imaginaire_logo.svg ├── inference.py ├── projects ├── coco_funit │ ├── README.md │ ├── animal_faces_content.jpg │ ├── animal_faces_output.jpg │ ├── animal_faces_style.jpg │ ├── mammals_content.png │ ├── mammals_output.jpg │ └── mammals_style.png ├── fs_vid2vid │ ├── 001.gif │ └── README.md ├── funit │ ├── README.md │ ├── content.jpg │ ├── funit.gif │ ├── output.jpg │ └── style.jpg ├── gancraft │ ├── README.md │ └── sch2vox │ │ ├── oldid2newid.json │ │ └── sch2vox.py ├── munit │ ├── README.md │ ├── example_input.jpg │ └── example_output.jpg ├── pix2pixhd │ ├── README.md │ ├── images │ │ ├── frankfurt_000000_000576.jpg │ │ ├── frankfurt_000000_011810.jpg │ │ └── frankfurt_000000_012868.jpg │ ├── instance_maps │ │ ├── frankfurt_000001_015328.png │ │ ├── frankfurt_000001_031266.png │ │ └── frankfurt_000001_044413.png │ └── seg_maps │ │ ├── frankfurt_000000_000576.png │ │ ├── frankfurt_000000_011810.png │ │ └── frankfurt_000000_012868.png ├── spade │ ├── README.md │ ├── edge_maps │ │ ├── 000000044195.png │ │ ├── 000000058384.png │ │ └── 000000072795.png │ ├── images │ │ ├── 000000044195.jpg │ │ ├── 000000058384.jpg │ │ └── 000000072795.jpg │ └── seg_maps │ │ ├── 000000044195.png │ │ ├── 000000058384.png │ │ └── 000000072795.png ├── unit │ ├── README.md │ ├── summer01.jpg │ ├── summer02.jpg │ ├── summer2winter_01.jpg │ ├── summer2winter_02.jpg │ ├── winter01.jpg │ ├── winter02.jpg │ ├── winter2summer_01.jpg │ └── winter2summer_02.jpg ├── vid2vid │ ├── README.md │ └── stuttgart_00.gif └── wc_vid2vid │ └── README.md ├── scripts ├── build_docker.sh ├── build_index.py ├── build_lmdb.py ├── build_lmdb.sh ├── download_dataset.py ├── download_test_data.py ├── install.bat ├── install.sh ├── requirements.txt ├── requirements_conda.yaml ├── start_local_docker.sh ├── test_inference.sh └── test_training.sh └── train.py /.dockerignore: -------------------------------------------------------------------------------- 1 | .git 2 | docs 3 | dataset 4 | checkpoints 5 | logs 6 | projects 7 | hub -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | checkpoints/ 2 | .idea/ 3 | logs/ 4 | results/ 5 | 6 | .vscode 7 | *.ipynb 8 | Dockerfile 9 | 10 | docs/_build/ 11 | nsys_profile* 12 | 13 | # Byte-compiled / optimized / DLL files 14 | __pycache__ 15 | *.py[cod] 16 | *$py.class 17 | *.pyc 18 | 19 | # C extensions 20 | *.so 21 | 22 | # Distribution / packaging 23 | .Python 24 | build/ 25 | develop-eggs/ 26 | dist/ 27 | downloads/ 28 | eggs/ 29 | .eggs/ 30 | lib/ 31 | lib64/ 32 | parts/ 33 | sdist/ 34 | var/ 35 | wheels/ 36 | pip-wheel-metadata/ 37 | share/python-wheels/ 38 | *.egg-info/ 39 | .installed.cfg 40 | *.egg 41 | MANIFEST 42 | 43 | # PyInstaller 44 | # Usually these files are written by a python script from a template 45 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 46 | *.manifest 47 | *.spec 48 | 49 | # Installer logs 50 | pip-log.txt 51 | pip-delete-this-directory.txt 52 | 53 | # Unit test / coverage reports 54 | htmlcov/ 55 | .tox/ 56 | .nox/ 57 | .coverage 58 | .coverage.* 59 | .cache 60 | nosetests.xml 61 | coverage.xml 62 | *.cover 63 | *.py,cover 64 | .hypothesis/ 65 | .pytest_cache/ 66 | cover/ 67 | 68 | # Translations 69 | *.mo 70 | *.pot 71 | 72 | # Django stuff: 73 | *.log 74 | local_settings.py 75 | db.sqlite3 76 | db.sqlite3-journal 77 | 78 | # Flask stuff: 79 | instance/ 80 | .webassets-cache 81 | 82 | # Scrapy stuff: 83 | .scrapy 84 | 85 | # Sphinx documentation 86 | docs/source/_build/ 87 | 88 | # PyBuilder 89 | .pybuilder/ 90 | target/ 91 | 92 | # Jupyter Notebook 93 | .ipynb_checkpoints 94 | 95 | # IPython 96 | profile_default/ 97 | ipython_config.py 98 | 99 | # pyenv 100 | # For a library or package, you might want to ignore these files since the code is 101 | # intended to run in multiple environments; otherwise, check them in: 102 | # .python-version 103 | 104 | # pipenv 105 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 106 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 107 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 108 | # install all needed dependencies. 109 | #Pipfile.lock 110 | 111 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 112 | __pypackages__/ 113 | 114 | # Celery stuff 115 | celerybeat-schedule 116 | celerybeat.pid 117 | 118 | # SageMath parsed files 119 | *.sage.py 120 | 121 | # Environments 122 | .env 123 | .venv 124 | env/ 125 | venv/ 126 | ENV/ 127 | env.bak/ 128 | venv.bak/ 129 | 130 | # Spyder project settings 131 | .spyderproject 132 | .spyproject 133 | 134 | # Rope project settings 135 | .ropeproject 136 | 137 | # mkdocs documentation 138 | /site 139 | 140 | # mypy 141 | .mypy_cache/ 142 | .dmypy.json 143 | dmypy.json 144 | 145 | # Pyre type checker 146 | .pyre/ 147 | 148 | # pytype static type analyzer 149 | .pytype/ 150 | 151 | # Cython debug symbols 152 | cython_debug/ 153 | 154 | # static files generated from Django application using `collectstatic` 155 | media 156 | static 157 | 158 | -------------------------------------------------------------------------------- /Dockerfile.base: -------------------------------------------------------------------------------- 1 | ARG DEBIAN_FRONTEND=noninteractive 2 | 3 | # Install basics 4 | RUN apt-get update && apt-get install -y --allow-downgrades --allow-change-held-packages --no-install-recommends \ 5 | build-essential \ 6 | cmake \ 7 | git \ 8 | curl \ 9 | vim \ 10 | tmux \ 11 | wget \ 12 | bzip2 \ 13 | unzip \ 14 | g++ \ 15 | ca-certificates \ 16 | ffmpeg \ 17 | libx264-dev \ 18 | imagemagick \ 19 | libnss3-dev 20 | -------------------------------------------------------------------------------- /INSTALL.md: -------------------------------------------------------------------------------- 1 | imaginaire_logo.svg 2 | 3 | # Imaginaire 4 | ### [Docs](http://imaginaire.cc/docs) | [License](LICENSE.md) | [Installation](INSTALL.md) | [Model Zoo](MODELZOO.md) 5 | 6 | # Installation Guide 7 | 8 | Our library is developed using an Ubuntu 18.04 machine. We have not yet tested our library on other operating systems. 9 | 10 | ## Prerequisite 11 | - [Anaconda3](https://www.anaconda.com/products/individual) 12 | - [cuda11.1](https://developer.nvidia.com/cuda-toolkit) 13 | - [cudnn](https://developer.nvidia.com/cudnn) 14 | 15 | We provide three different ways of installing Imaginaire. 16 | 17 | ### 1. Installation (default) 18 | Note: sudo privilege is required. 19 | ```bash 20 | git clone https://github.com/nvlabs/imaginaire 21 | cd imaginaire 22 | bash scripts/install.sh 23 | bash scripts/test_training.sh 24 | ``` 25 | If installation is not successful, error message will be prompted. 26 | 27 | ### 2. Installation with Docker 28 | We use NVIDIA docker image. We provide two ways to build the docker image. 29 | 1. Build a target docker image 30 | ```bash 31 | bash scripts/build_docker.sh 21.06 32 | ``` 33 | 34 | 2. Launch an interactive docker container and test the imaginaire repo. 35 | ```bash 36 | cd scripts 37 | bash start_local_docker.sh 21.06 38 | cd ${IMAGINAIRE_ROOT} 39 | bash scripts/test_training.sh 40 | ``` 41 | 42 | ### 3. Installation with Conda 43 | Set up the Conda environment and install packages with 44 | ```bash 45 | conda env create --file scripts/requirements_conda.yaml 46 | # install third-party libraries 47 | export CUDA_VERSION=$(nvcc --version| grep -Po "(\d+\.)+\d+" | head -1) 48 | CURRENT=$(pwd) 49 | for p in correlation channelnorm resample2d bias_act upfirdn2d; do 50 | cd imaginaire/third_party/${p}; 51 | rm -rf build dist *info; 52 | python setup.py install; 53 | cd ${CURRENT}; 54 | done 55 | ``` 56 | To activate the environment and test the repo: 57 | ```bash 58 | conda activate imaginaire 59 | bash scripts/test_training.sh 60 | ``` 61 | 62 | ## Flake8 63 | We follow the PEP8 style using flake8. To follow our practice, please do 64 | ```bash 65 | pip install flake8 66 | flake8 --install-hook git 67 | git config --bool flake8.strict true 68 | ``` 69 | We set the maximum line length to 80. To avoid error messages due to different line length, create a file `~/.config/flake8` with the following content: 70 | ``` 71 | [flake8] 72 | max-line-length = 200 73 | ``` 74 | 75 | ## Windows Installation [Out-dated] 76 | 77 | - Install [git for windows](https://gitforwindows.org/) 78 | - Install [Microsoft C++ Build Tools](https://visualstudio.microsoft.com/visual-cpp-build-tools/) 79 | - Install [Anaconda3](https://repo.anaconda.com/archive/Anaconda3-2020.02-Windows-x86_64.exe) 80 | - Install [CUDA11.1](https://developer.nvidia.com/cuda-11.1-download-archive) 81 | - Install [cudnn](https://developer.nvidia.com/cudnn) 82 | - Open an anaconda prompt. 83 | ``` 84 | cd https://github.com/NVlabs/imaginaire 85 | .\scripts\install.bat 86 | ``` 87 | 88 | Powershell 89 | ``` 90 | $env:PYTHONPATH = pwd 91 | Get-ChildItem Env:PYTHONPATH 92 | ``` 93 | 94 | -------------------------------------------------------------------------------- /configs/projects/spade/landscapes/lhq.yaml: -------------------------------------------------------------------------------- 1 | # This SPADE network has been trained on the publicly available LHQ dataset - https://github.com/universome/alis 2 | # using 183-dim COCO-stuff label space (184 including don't care label). 3 | 4 | pretrained_weight: 1NIh3_UZ6uqvzS4mJ4JVhfyYQuG9ZMmvA 5 | 6 | gen: 7 | type: imaginaire.generators.spade 8 | version: v20 9 | output_multiplier: 0.5 10 | image_channels: 3 11 | num_labels: 184 12 | style_dims: 256 13 | num_filters: 128 14 | kernel_size: 3 15 | weight_norm_type: 'spectral' 16 | use_posenc_in_input_layer: False 17 | global_adaptive_norm_type: 'sync_batch' 18 | activation_norm_params: 19 | num_filters: 128 20 | kernel_size: 5 21 | separate_projection: True 22 | activation_norm_type: 'sync_batch' 23 | style_enc: 24 | num_filters: 64 25 | kernel_size: 3 26 | 27 | data: 28 | train: 29 | augmentations: 30 | random_crop_h_w: 512, 512 31 | -------------------------------------------------------------------------------- /configs/projects/wc_vid2vid/cityscapes/seg_single.yaml: -------------------------------------------------------------------------------- 1 | image_save_iter: 5000 2 | snapshot_save_epoch: 5 3 | max_epoch: 400 4 | logging_iter: 100 5 | trainer: 6 | distributed_data_parallel: apex 7 | type: imaginaire.trainers.spade 8 | amp_config: 9 | enabled: True 10 | model_average_config: 11 | enabled: False 12 | beta: 0.9999 13 | start_iteration: 1000 14 | num_batch_norm_estimation_iterations: 30 15 | 16 | fid_type: image 17 | gan_mode: hinge 18 | gan_relativistic: False 19 | perceptual_loss: 20 | mode: 'vgg19' 21 | layers: ['relu_1_1', 'relu_2_1', 'relu_3_1', 'relu_4_1', 'relu_5_1'] 22 | weights: [0.03125, 0.0625, 0.125, 0.25, 1.0] 23 | fp16: True 24 | loss_weight: 25 | gan: 1.0 26 | perceptual: 10.0 27 | feature_matching: 10.0 28 | kl: 0.05 29 | init: 30 | type: xavier 31 | gain: 0.02 32 | gen_opt: 33 | type: adam 34 | lr: 0.0001 35 | adam_beta1: 0. 36 | adam_beta2: 0.999 37 | lr_policy: 38 | iteration_mode: False 39 | type: step 40 | step_size: 100 41 | gamma: 0.01 42 | dis_opt: 43 | type: adam 44 | lr: 0.0004 45 | adam_beta1: 0. 46 | adam_beta2: 0.999 47 | lr_policy: 48 | iteration_mode: False 49 | type: step 50 | step_size: 100 51 | gamma: 0.01 52 | gen: 53 | type: imaginaire.generators.spade 54 | version: v20 55 | style_dims: 256 56 | num_filters: 128 57 | kernel_size: 3 58 | weight_norm_type: 'spectral' 59 | global_adaptive_norm_type: 'sync_batch' 60 | activation_norm_params: 61 | num_filters: 128 62 | kernel_size: 5 63 | separate_projection: False 64 | weight_norm_type: 'none' 65 | activation_norm_type: 'sync_batch' 66 | style_enc: 67 | num_filters: 64 68 | kernel_size: 3 69 | dis: 70 | type: imaginaire.discriminators.spade 71 | kernel_size: 3 72 | num_filters: 128 73 | max_num_filters: 512 74 | num_discriminators: 2 75 | num_layers: 5 76 | activation_norm_type: 'none' 77 | weight_norm_type: 'spectral' 78 | 79 | # Data options. 80 | data: 81 | # Name of this dataset. 82 | name: 'cityscapes' 83 | # Which dataloader to use? 84 | type: imaginaire.datasets.paired_images 85 | # How many data loading workers per GPU? 86 | num_workers: 8 87 | 88 | input_types: 89 | - images: 90 | ext: jpg 91 | num_channels: 3 92 | normalize: True 93 | - seg_maps: 94 | ext: png 95 | num_channels: 1 96 | is_mask: True 97 | normalize: False 98 | 99 | full_data_ops: imaginaire.model_utils.label::make_one_hot, imaginaire.model_utils.label::concat_labels 100 | use_dont_care: False 101 | one_hot_num_classes: 102 | seg_maps: 35 103 | input_labels: 104 | - seg_maps 105 | 106 | # Which lmdb contains the ground truth image. 107 | input_image: 108 | - images 109 | 110 | # Train dataset details. 111 | train: 112 | # Data augmentations to be performed in given order. 113 | augmentations: 114 | resize_h_w: 512, 1024 115 | random_crop_h_w: 512, 512 116 | 117 | # Validation dataset details. 118 | val: 119 | # If resize_h_w is not given, then it is assumed to be same as crop_h_w. 120 | augmentations: 121 | resize_h_w: 512, 1024 122 | horizontal_flip: False 123 | -------------------------------------------------------------------------------- /configs/unit_test/pix2pixHD.yaml: -------------------------------------------------------------------------------- 1 | # How often do you want to log the training stats. 2 | logging_iter: 2 3 | # Number of training epochs. 4 | max_iter: 2 5 | # Whether to benchmark speed or not. 6 | speed_benchmark: True 7 | # trainer options 8 | trainer: 9 | type: imaginaire.trainers.pix2pixHD 10 | amp_config: 11 | enabled: True 12 | perceptual_loss: 13 | mode: 'vgg19' 14 | layers: ['relu_1_1', 'relu_2_1', 'relu_3_1', 'relu_4_1', 'relu_5_1'] 15 | weights: [0.03125, 0.0625, 0.125, 0.25, 1.0] 16 | gan_mode: hinge 17 | gan_relativistic: False 18 | loss_weight: 19 | gan: 1.0 20 | feature_matching: 10.0 21 | perceptual: 10.0 22 | init: 23 | type: xavier 24 | gain: 0.02 25 | 26 | # model options 27 | gen: 28 | type: imaginaire.generators.pix2pixHD 29 | global_generator: 30 | num_filters: 32 31 | num_downsamples: 4 32 | num_res_blocks: 9 33 | local_enhancer: 34 | num_enhancers: 0 35 | num_res_blocks: 3 36 | weight_norm_type: spectral 37 | activation_norm_type: instance 38 | padding_mode: reflect 39 | dis: 40 | type: imaginaire.discriminators.multires_patch 41 | num_filters: 32 42 | max_num_filters: 512 43 | num_discriminators: 2 44 | num_layers: 3 45 | weight_norm_type: spectral 46 | activation_norm_type: instance 47 | 48 | # optimization option 49 | gen_opt: 50 | type: adam 51 | lr: 0.0001 52 | adam_beta1: 0. 53 | adam_beta2: 0.999 54 | lr_policy: 55 | iteration_mode: False 56 | type: step 57 | step_size: 20 58 | gamma: 0.9 59 | dis_opt: 60 | type: adam 61 | lr: 0.0004 62 | adam_beta1: 0. 63 | adam_beta2: 0.999 64 | lr_policy: 65 | iteration_mode: False 66 | type: step 67 | step_size: 20 68 | gamma: 0.9 69 | 70 | # Data options. 71 | data: 72 | # Name of this dataset. 73 | name: 'cityscapes' 74 | # Which dataloader to use? 75 | type: imaginaire.datasets.paired_images 76 | # How many data loading workers per GPU? 77 | num_workers: 4 78 | input_types: 79 | - images: 80 | ext: jpg 81 | num_channels: 3 82 | normalize: True 83 | - seg_maps: 84 | ext: png 85 | num_channels: 1 86 | is_mask: True 87 | normalize: False 88 | - instance_maps: 89 | ext: png 90 | num_channels: 1 91 | is_mask: True 92 | normalize: False 93 | 94 | full_data_ops: imaginaire.model_utils.label::make_one_hot, imaginaire.model_utils.label::concat_labels 95 | use_dont_care: False 96 | one_hot_num_classes: 97 | seg_maps: 35 98 | input_labels: 99 | - seg_maps 100 | - instance_maps 101 | 102 | input_image: 103 | - images 104 | 105 | # Train dataset details. 106 | train: 107 | # Input LMDBs. 108 | roots: 109 | - dataset/unit_test/raw/pix2pixHD 110 | # Batch size per GPU. 111 | batch_size: 1 112 | # Data augmentations to be performed in given order. 113 | augmentations: 114 | resize_h_w: 256, 512 115 | # Horizontal flip? 116 | horizontal_flip: False 117 | 118 | # Validation dataset details. 119 | val: 120 | # Input LMDBs. 121 | roots: 122 | - dataset/unit_test/raw/pix2pixHD 123 | # Batch size per GPU. 124 | batch_size: 1 125 | # If resize_h_w is not given, then it is assumed to be same as crop_h_w. 126 | augmentations: 127 | resize_h_w: 256, 512 128 | horizontal_flip: False 129 | -------------------------------------------------------------------------------- /configs/unit_test/unit.yaml: -------------------------------------------------------------------------------- 1 | # How often do you want to log the training stats. 2 | logging_iter: 2 3 | # Number of training epochs. 4 | max_iter: 2 5 | # Whether to benchmark speed or not. 6 | speed_benchmark: True 7 | trainer: 8 | type: imaginaire.trainers.unit 9 | model_average_config: 10 | enabled: True 11 | gan_mode: hinge 12 | perceptual_mode: vgg19 13 | perceptual_layers: 'relu_4_1' 14 | loss_weight: 15 | gan: 1 16 | image_recon: 1 17 | cycle_recon: 1 18 | perceptual: 0 19 | init: 20 | type: normal 21 | gain: 0.02 22 | # Optimization option. 23 | gen_opt: 24 | type: adam 25 | lr: 0.0001 26 | adam_beta1: 0. 27 | adam_beta2: 0.999 28 | lr_policy: 29 | iteration_mode: False 30 | type: step 31 | step_size: 100 32 | gamma: 0.5 33 | dis_opt: 34 | type: adam 35 | lr: 0.0004 36 | adam_beta1: 0. 37 | adam_beta2: 0.999 38 | lr_policy: 39 | iteration_mode: False 40 | type: step 41 | step_size: 100 42 | gamma: 0.5 43 | gen: 44 | type: imaginaire.generators.unit 45 | num_filters: 24 46 | max_num_filters: 64 47 | num_res_blocks: 4 48 | num_downsamples_content: 2 49 | content_norm_type: instance 50 | decoder_norm_type: instance 51 | weight_norm_type: none 52 | pre_act: True 53 | dis: 54 | type: imaginaire.discriminators.unit 55 | num_filters: 24 56 | max_num_filters: 64 57 | num_discriminators: 3 58 | num_layers: 4 59 | kernel_size: 4 60 | weight_norm_type: spectral 61 | 62 | # Data options. 63 | data: 64 | # Name of this dataset. 65 | name: 'test' 66 | # Which dataloader to use? 67 | type: 'imaginaire.datasets.unpaired_images' 68 | # How many data loading workers per GPU? 69 | num_workers: 8 70 | input_types: 71 | - images_a: 72 | # If not specified, is None by default. 73 | ext: jpg 74 | # If not specified, is None by default. 75 | num_channels: 3 76 | # If not specified, is None by default. 77 | normalize: True 78 | # If not specified, is None by default. 79 | pre_aug_ops: None 80 | # If not specified, is None by default. 81 | post_aug_ops: None 82 | - images_b: 83 | # If not specified, is None by default. 84 | ext: jpg 85 | # If not specified, is None by default. 86 | num_channels: 3 87 | # If not specified, is None by default. 88 | normalize: True 89 | # If not specified, is None by default. 90 | pre_aug_ops: None 91 | # If not specified, is None by default. 92 | post_aug_ops: None 93 | 94 | # Train dataset details. 95 | train: 96 | # Input LMDBs. 97 | roots: 98 | - dataset/unit_test/raw/munit 99 | # Batch size per GPU. 100 | batch_size: 1 101 | # Data augmentations to be performed in given order. 102 | augmentations: 103 | # First resize all inputs to this size. 104 | resize_smallest_side: 512 105 | # Horizontal flip? 106 | horizontal_flip: True 107 | # Crop size. 108 | random_crop_h_w: 512, 512 109 | 110 | # Val dataset details. 111 | val: 112 | # Input LMDBS. 113 | roots: 114 | - dataset/unit_test/raw/munit 115 | # Batch size per GPU. 116 | batch_size: 1 117 | # If resize_h_w is not given, then it is assumed to be same as crop_h_w. 118 | augmentations: 119 | resize_smallest_side: 512 120 | center_crop_h_w: 512, 512 -------------------------------------------------------------------------------- /dataset/unit_test/raw/fs_vid2vid/face/images/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00000.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/fs_vid2vid/face/images/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00000.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/fs_vid2vid/face/images/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00001.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/fs_vid2vid/face/images/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00001.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/fs_vid2vid/face/images/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00002.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/fs_vid2vid/face/images/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00002.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/fs_vid2vid/face/images/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00003.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/fs_vid2vid/face/images/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00003.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/fs_vid2vid/face/images/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00004.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/fs_vid2vid/face/images/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00004.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/fs_vid2vid/face/landmarks-dlib68/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00000.json: -------------------------------------------------------------------------------- 1 | [[292, 136], [293, 153], [296, 169], [299, 186], [304, 201], [311, 215], [322, 226], [337, 233], [353, 234], [369, 233], [383, 226], [394, 215], [403, 201], [407, 185], [410, 169], [413, 152], [415, 135], [304, 117], [312, 111], [322, 109], [332, 111], [342, 116], [370, 115], [378, 110], [388, 108], [398, 110], [404, 117], [355, 130], [355, 141], [355, 151], [355, 162], [346, 171], [350, 173], [355, 174], [359, 173], [364, 171], [316, 134], [323, 130], [331, 130], [338, 134], [330, 136], [323, 137], [372, 134], [378, 130], [386, 130], [392, 135], [386, 137], [378, 137], [334, 196], [341, 190], [349, 187], [355, 188], [361, 186], [368, 189], [374, 196], [368, 202], [361, 205], [355, 206], [348, 205], [341, 203], [338, 196], [349, 196], [355, 196], [361, 196], [371, 196], [360, 195], [355, 196], [349, 195]] -------------------------------------------------------------------------------- /dataset/unit_test/raw/fs_vid2vid/face/landmarks-dlib68/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00001.json: -------------------------------------------------------------------------------- 1 | [[289, 138], [290, 154], [293, 170], [297, 186], [302, 201], [310, 215], [321, 225], [336, 232], [352, 233], [368, 231], [382, 224], [393, 213], [401, 199], [405, 184], [408, 168], [411, 152], [412, 135], [301, 116], [308, 109], [318, 106], [329, 108], [339, 114], [366, 112], [375, 107], [385, 105], [394, 108], [401, 115], [351, 128], [351, 138], [351, 148], [351, 158], [342, 168], [347, 169], [351, 171], [356, 169], [361, 167], [313, 132], [320, 127], [328, 127], [335, 132], [328, 134], [320, 135], [368, 132], [375, 127], [383, 127], [389, 132], [383, 135], [375, 134], [332, 194], [339, 186], [346, 182], [352, 184], [358, 182], [365, 185], [371, 193], [365, 198], [358, 202], [352, 202], [346, 202], [339, 200], [336, 194], [346, 192], [352, 192], [358, 191], [368, 192], [358, 191], [352, 192], [346, 192]] -------------------------------------------------------------------------------- /dataset/unit_test/raw/fs_vid2vid/face/landmarks-dlib68/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00002.json: -------------------------------------------------------------------------------- 1 | [[287, 136], [288, 152], [291, 169], [295, 185], [300, 200], [307, 213], [318, 224], [333, 230], [349, 230], [366, 229], [381, 222], [392, 211], [399, 198], [403, 183], [406, 167], [409, 151], [410, 134], [298, 115], [305, 108], [315, 105], [326, 107], [336, 111], [362, 110], [372, 105], [382, 103], [392, 106], [399, 114], [348, 125], [348, 135], [348, 144], [348, 154], [339, 165], [343, 167], [348, 168], [353, 166], [357, 164], [310, 131], [317, 125], [325, 125], [332, 130], [325, 133], [317, 133], [365, 130], [372, 124], [380, 124], [386, 130], [380, 132], [372, 132], [329, 192], [337, 184], [344, 181], [349, 182], [354, 180], [361, 183], [368, 191], [361, 196], [355, 199], [349, 199], [344, 199], [337, 198], [334, 192], [344, 190], [349, 189], [354, 189], [364, 190], [354, 189], [349, 190], [344, 190]] -------------------------------------------------------------------------------- /dataset/unit_test/raw/fs_vid2vid/face/landmarks-dlib68/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00003.json: -------------------------------------------------------------------------------- 1 | [[286, 136], [286, 152], [290, 168], [293, 184], [298, 199], [304, 212], [315, 222], [330, 228], [346, 229], [363, 228], [378, 221], [390, 211], [397, 198], [402, 183], [404, 167], [407, 150], [407, 134], [296, 113], [302, 106], [312, 103], [323, 104], [333, 110], [358, 109], [367, 103], [377, 101], [388, 104], [395, 112], [345, 123], [345, 133], [345, 142], [344, 152], [336, 163], [340, 164], [345, 166], [350, 164], [355, 162], [308, 130], [314, 126], [321, 125], [329, 129], [322, 131], [314, 131], [362, 128], [369, 125], [376, 125], [383, 129], [376, 131], [369, 130], [325, 189], [332, 182], [340, 178], [346, 179], [352, 178], [360, 181], [367, 189], [360, 195], [353, 198], [346, 198], [340, 198], [332, 196], [330, 189], [340, 187], [346, 187], [352, 187], [363, 188], [352, 187], [346, 188], [340, 187]] -------------------------------------------------------------------------------- /dataset/unit_test/raw/fs_vid2vid/face/landmarks-dlib68/ZzTpJYDc62E_4_ZzTpJYDc62E_4/00004.json: -------------------------------------------------------------------------------- 1 | [[282, 133], [284, 149], [287, 166], [291, 181], [295, 196], [303, 210], [314, 220], [329, 226], [345, 228], [362, 226], [377, 220], [388, 209], [396, 196], [400, 181], [402, 165], [405, 148], [406, 131], [293, 113], [299, 105], [309, 102], [320, 103], [330, 108], [355, 107], [364, 101], [375, 99], [385, 103], [393, 110], [342, 123], [342, 131], [342, 140], [342, 148], [333, 161], [338, 162], [343, 163], [348, 162], [353, 160], [305, 132], [311, 130], [318, 129], [326, 130], [318, 132], [312, 133], [360, 129], [367, 128], [374, 128], [380, 130], [373, 131], [367, 131], [323, 189], [330, 181], [337, 177], [343, 178], [349, 176], [356, 179], [364, 187], [358, 194], [350, 197], [344, 197], [338, 197], [330, 196], [328, 188], [337, 186], [343, 185], [349, 185], [360, 187], [350, 187], [344, 187], [338, 187]] -------------------------------------------------------------------------------- /dataset/unit_test/raw/funit/images_content/cat/cat_001.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/funit/images_content/cat/cat_001.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/funit/images_content/dog/dog_001.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/funit/images_content/dog/dog_001.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/funit/images_content/wild/wild_001.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/funit/images_content/wild/wild_001.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/funit/images_style/images_content/cat/cat_001.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/funit/images_style/images_content/cat/cat_001.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/funit/images_style/images_content/dog/dog_001.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/funit/images_style/images_content/dog/dog_001.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/funit/images_style/images_content/wild/wild_001.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/funit/images_style/images_content/wild/wild_001.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/munit/images_a/dog_001.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/munit/images_a/dog_001.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/munit/images_b/cat_001.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/munit/images_b/cat_001.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/pix2pixHD/images/001.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/pix2pixHD/images/001.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/pix2pixHD/instance_maps/001.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/pix2pixHD/instance_maps/001.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/pix2pixHD/seg_maps/001.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/pix2pixHD/seg_maps/001.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/spade/edge_maps/000000581781.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/spade/edge_maps/000000581781.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/spade/images/000000581781.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/spade/images/000000581781.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/spade/seg_maps/000000581781.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/spade/seg_maps/000000581781.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/human_instance_maps/001_001/frame000100.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/human_instance_maps/001_001/frame000100.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/human_instance_maps/001_001/frame000101.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/human_instance_maps/001_001/frame000101.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/human_instance_maps/001_001/frame000102.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/human_instance_maps/001_001/frame000102.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/human_instance_maps/001_001/frame000103.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/human_instance_maps/001_001/frame000103.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/human_instance_maps/001_001/frame000104.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/human_instance_maps/001_001/frame000104.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/images/001_001/frame000100.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/images/001_001/frame000100.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/images/001_001/frame000101.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/images/001_001/frame000101.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/images/001_001/frame000102.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/images/001_001/frame000102.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/images/001_001/frame000103.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/images/001_001/frame000103.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/images/001_001/frame000104.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/images/001_001/frame000104.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/pose_maps-densepose/001_001/frame000100.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/pose_maps-densepose/001_001/frame000100.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/pose_maps-densepose/001_001/frame000101.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/pose_maps-densepose/001_001/frame000101.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/pose_maps-densepose/001_001/frame000102.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/pose_maps-densepose/001_001/frame000102.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/pose_maps-densepose/001_001/frame000103.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/pose_maps-densepose/001_001/frame000103.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/pose_maps-densepose/001_001/frame000104.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/pose/pose_maps-densepose/001_001/frame000104.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/poses-openpose/001_001/frame000100.json: -------------------------------------------------------------------------------- 1 | {"version":1.2,"people":[{"pose_keypoints_2d":[989.392,409.487,0.896952,974.779,480.162,0.948245,924.642,474.288,0.827813,924.587,550.792,0.877172,1001.31,550.909,0.781895,1024.79,483.14,0.875447,1042.36,556.728,0.884154,971.685,536.236,0.758101,977.641,647.922,0.800718,948.149,647.959,0.743706,974.786,786.304,0.838945,1004.19,921.723,0.780684,1012.95,645.152,0.729135,1001.36,789.223,0.832865,974.772,904.027,0.839149,974.542,400.772,0.882267,995.354,391.975,0.909277,951.025,421.41,0.868128,0,0,0,960.002,951.2,0.774895,977.752,951.125,0.759855,962.983,907,0.574491,1015.86,957.063,0.716831,998.103,954.12,0.817665,1018.83,930.623,0.65345],"face_keypoints_2d":[949.508,412.241,0.548972,952.678,420.023,0.623743,956.713,428.381,0.66788,961.324,435.298,0.623044,968.529,440.197,0.620277,977.464,442.503,0.658985,986.399,442.791,0.692881,993.892,442.503,0.739515,1000.52,439.909,0.792257,1004.56,435.586,0.716717,1005.71,430.398,0.736368,1007.15,424.057,0.701478,1007.15,417.429,0.689222,1005.71,411.088,0.618149,1004.56,405.036,0.64413,1001.96,398.983,0.62905,998.503,393.795,0.535485,957.001,400.712,0.802538,959.307,395.813,0.889584,963.63,391.778,0.893319,968.529,390.913,0.836494,974.294,391.201,0.794647,982.94,387.743,0.719777,984.669,386.59,0.774043,986.975,384.861,0.893816,989.569,384.284,0.791669,992.163,384.573,0.675626,982.94,396.677,0.808389,986.399,400.424,0.897604,989.281,403.306,0.934042,992.451,406.477,0.770227,986.975,414.547,0.89801,989.857,414.258,0.868223,993.316,413.394,0.828126,994.18,412.241,0.823823,995.91,410.223,0.776254,964.206,405.612,0.717644,968.241,403.595,0.643539,971.123,402.442,0.728754,974.582,402.153,0.848513,972.276,403.306,0.783209,968.241,405.324,0.730009,987.551,397.254,0.708198,989.281,396.101,0.695611,991.298,395.236,0.613121,993.604,394.66,0.553855,991.586,396.389,0.705006,989.569,396.677,0.676177,986.975,426.363,0.74681,989.281,422.04,0.852127,991.875,419.158,0.864877,994.468,418.87,0.794511,996.198,417.14,0.760127,998.503,419.158,0.6951,1000.23,421.464,0.639184,1000.52,423.769,0.693368,998.792,425.499,0.628336,997.062,426.075,0.675825,994.468,426.651,0.709053,991.298,427.228,0.742071,988.992,425.787,0.780266,993.316,421.752,0.840678,995.621,421.464,0.778341,996.774,421.175,0.662196,998.792,421.464,0.634825,997.927,421.752,0.664832,996.198,422.905,0.704388,993.604,423.769,0.81792,970.547,403.306,0.582712,989.857,396.101,0.625781],"hand_left_keypoints_2d":[945.061,518.675,0.0210557,939.662,525.574,0.0358395,945.961,533.073,0.0357281,959.459,528.874,0.0437127,967.558,526.174,0.097172,946.261,535.473,0.0331476,958.859,537.572,0.0455975,964.559,533.373,0.0665946,969.358,527.974,0.0850925,949.561,536.372,0.02788,957.36,543.571,0.0455821,959.459,541.172,0.033063,959.759,547.171,0.0236449,947.161,528.274,0.0164021,951.66,540.872,0.0217454,953.46,546.271,0.0213297,960.359,547.471,0.0227669,945.061,525.274,0.0145665,942.962,535.173,0.0119985,959.159,528.574,0.0139527,959.759,528.274,0.00995805],"hand_right_keypoints_2d":[1047.9,562.519,0.0153641,1038.51,570.337,0.0186044,1036.32,571.901,0.0524353,1030.38,528.118,0.0812377,1022.88,532.809,0.0517582,1043.52,541.253,0.0375972,1032.88,526.867,0.0867068,1025.07,529.369,0.129042,1020.06,532.497,0.118747,1040.39,537.5,0.0442179,1033.82,530.62,0.0748278,1026.63,530.307,0.161309,1020.69,531.871,0.112858,1039.14,537.188,0.0421561,1034.76,531.558,0.0741886,1027.26,530.307,0.0992351,1021.31,531.246,0.0997959,1035.39,572.214,0.0257112,1034.14,525.929,0.0309488,1027.57,529.369,0.0424813,1023.19,530.307,0.0410194],"pose_keypoints_3d":[],"face_keypoints_3d":[],"hand_left_keypoints_3d":[],"hand_right_keypoints_3d":[]}]} -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/poses-openpose/001_001/frame000101.json: -------------------------------------------------------------------------------- 1 | {"version":1.2,"people":[{"pose_keypoints_2d":[983.593,409.533,0.909504,974.817,480.124,0.942257,924.699,468.519,0.83783,921.732,547.983,0.845838,1001.22,550.821,0.757409,1024.81,483.154,0.872246,1042.43,556.722,0.895325,971.787,536.184,0.758654,977.598,647.943,0.807591,945.32,647.974,0.740126,974.804,786.274,0.83036,1004.18,921.703,0.778573,1010.07,647.925,0.732271,1001.37,789.226,0.836636,974.77,903.981,0.838684,971.688,403.711,0.867032,995.275,391.997,0.897872,948.277,424.326,0.85864,1004.17,412.432,0.168901,960.006,951.202,0.772192,977.756,951.126,0.757975,962.975,906.944,0.584187,1015.86,957.064,0.71538,998.099,954.125,0.814262,1018.82,930.603,0.646473],"face_keypoints_2d":[948.272,412.622,0.536211,951.579,420.139,0.632717,955.488,427.655,0.615579,959.396,434.87,0.575063,966.01,440.282,0.558402,975.03,442.988,0.571638,984.35,443.589,0.6822,992.167,442.988,0.701227,999.081,440.282,0.779864,1003.59,436.674,0.777412,1005.7,430.661,0.69496,1006.6,424.949,0.759219,1006.6,419.237,0.707776,1005.7,412.622,0.649352,1003.89,406.309,0.632992,1001.79,400.597,0.603727,999.081,395.185,0.489519,956.089,400.897,0.80567,957.893,396.087,0.874479,961.501,393.381,0.86415,966.612,391.277,0.838317,972.324,391.577,0.82416,981.644,388.571,0.769958,984.049,386.466,0.80863,986.755,384.963,0.821012,989.461,384.362,0.803169,992.467,385.865,0.623562,981.644,397.891,0.815331,984.65,400.897,0.947593,987.356,404.204,0.835289,990.363,407.812,0.729717,985.252,415.028,0.90315,988.258,415.028,0.897554,990.964,414.126,0.836473,992.167,412.923,0.870311,993.369,411.721,0.737557,963.004,406.309,0.707074,966.311,405.106,0.640279,970.219,403.002,0.771004,972.925,403.002,0.922759,970.219,404.505,0.778103,966.912,405.708,0.723072,987.056,398.191,0.827817,988.859,396.087,0.797608,990.964,395.786,0.785111,993.67,395.486,0.749769,991.565,396.388,0.805391,989.461,397.891,0.80208,985.552,427.354,0.745661,987.356,422.844,0.811571,989.761,419.838,0.909815,992.467,418.635,0.827026,994.572,417.733,0.790228,996.977,419.838,0.748395,998.781,422.544,0.742806,998.781,424.949,0.707668,996.977,426.753,0.722937,995.173,427.354,0.739336,992.167,427.955,0.789001,989.461,428.256,0.757337,987.356,426.753,0.785328,991.565,422.544,0.793873,993.67,421.942,0.784511,994.872,421.942,0.708331,996.977,422.544,0.752995,995.774,422.844,0.704062,994.271,423.746,0.714366,991.866,424.648,0.824082,968.415,404.806,0.563037,989.461,396.388,0.808571],"hand_left_keypoints_2d":[940.377,513.513,0.0276548,939.778,523.408,0.0319507,947.274,530.605,0.0318811,959.569,528.206,0.0435013,968.266,526.407,0.102764,945.175,535.703,0.0413921,957.77,537.503,0.0464342,964.067,533.004,0.0814467,969.165,527.907,0.12324,945.175,534.504,0.0582758,952.672,539.302,0.0538213,958.67,539.602,0.0438802,959.869,547.099,0.0252485,940.977,532.405,0.0448116,949.374,537.203,0.0390725,951.473,535.703,0.0373433,961.369,547.099,0.0258886,940.677,524.308,0.0306871,947.874,527.907,0.0215856,957.17,528.206,0.0364523,958.97,528.506,0.0182037],"hand_right_keypoints_2d":[1049.49,570.561,0.0160748,1038.8,572.831,0.0189794,1035.23,571.858,0.0488461,1031.01,529.064,0.123601,1023.23,532.63,0.0977182,1042.69,537.817,0.0308391,1034.58,526.794,0.0918298,1025.83,529.388,0.16401,1020.32,532.63,0.166035,1039.77,537.817,0.0295079,1035.55,526.47,0.0518523,1028.1,529.712,0.103495,1022.91,531.657,0.0617105,1040.09,537.493,0.0302868,1036.85,527.443,0.0440225,1030.04,531.009,0.0664725,1022.26,532.63,0.0483061,1034.26,571.534,0.0272904,1035.88,525.173,0.0250872,1029.72,531.657,0.0318258,1021.29,532.954,0.0238266],"pose_keypoints_3d":[],"face_keypoints_3d":[],"hand_left_keypoints_3d":[],"hand_right_keypoints_3d":[]}]} -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/poses-openpose/001_001/frame000102.json: -------------------------------------------------------------------------------- 1 | {"version":1.2,"people":[{"pose_keypoints_2d":[980.723,409.584,0.923551,974.82,480.104,0.94031,924.71,468.489,0.840958,921.696,550.725,0.847989,1001.17,550.875,0.765961,1024.84,483.134,0.875031,1042.45,556.723,0.902057,971.823,533.342,0.762227,977.591,645.16,0.807438,945.295,647.947,0.737741,974.806,786.331,0.832231,1004.15,918.85,0.77572,1010.07,645.143,0.732196,1001.35,789.279,0.833163,974.757,903.989,0.83238,968.836,406.58,0.87376,992.422,392.009,0.883024,948.278,427.2,0.863675,1004.14,409.68,0.256046,959.996,951.198,0.765278,977.758,951.128,0.748775,962.964,906.951,0.578804,1015.87,957.056,0.719445,998.103,954.118,0.81762,1018.81,930.583,0.634611],"face_keypoints_2d":[948.172,414.049,0.596571,951.888,421.788,0.703195,954.984,429.219,0.699369,959.627,435.72,0.670827,966.439,441.293,0.614643,973.559,443.77,0.596394,981.918,444.389,0.733024,989.348,444.079,0.778787,996.779,441.912,0.846287,1001.73,438.816,0.823227,1004.52,432.624,0.758021,1006.38,427.051,0.785761,1006.69,420.55,0.703968,1006.07,414.358,0.672499,1004.21,407.237,0.644881,1001.73,401.046,0.585102,999.255,394.544,0.473633,954.364,402.594,0.806752,956.222,397.64,0.824628,959.937,394.544,0.901772,965.2,392.377,0.849476,971.392,392.377,0.831615,979.751,389.591,0.791782,982.228,387.423,0.828944,985.633,385.256,0.808417,988.729,384.947,0.827573,991.516,384.947,0.65186,979.751,398.259,0.828706,982.847,401.974,0.875043,986.253,404.761,0.887909,989.039,407.857,0.733987,984.085,415.287,0.853852,986.872,414.977,0.91204,989.348,414.668,0.914767,991.206,414.049,0.843085,991.825,412.191,0.828338,962.104,407.237,0.8548,964.581,405.07,0.778109,968.296,403.832,0.752342,971.702,403.832,0.870839,969.225,405.07,0.829837,965.51,406.928,0.822325,985.633,397.64,0.7969,986.872,397.021,0.840801,989.348,395.783,0.774285,991.825,395.163,0.680592,989.658,397.021,0.80512,987.491,397.33,0.807909,982.847,429.219,0.748995,985.014,424.265,0.856739,988.42,421.169,0.828027,991.516,420.55,0.806823,993.683,419.621,0.805539,996.469,421.479,0.794984,998.636,423.956,0.83354,997.398,426.742,0.787633,995.54,428.599,0.780489,993.373,429.219,0.847612,989.968,429.838,0.832268,986.562,429.838,0.845906,984.705,427.361,0.786626,989.348,424.884,0.83488,991.825,424.265,0.846298,993.992,424.265,0.859517,996.779,424.265,0.824377,993.992,424.265,0.858854,991.825,424.884,0.839713,989.348,425.813,0.837965,966.748,405.07,0.683511,988.42,397.021,0.784041],"hand_left_keypoints_2d":[940.024,514.405,0.0329577,941.541,523.806,0.0277824,945.786,525.019,0.0243327,972.776,538.969,0.0290321,967.621,526.535,0.078634,945.483,535.936,0.0608439,957.613,534.723,0.0583183,963.982,531.388,0.0988171,969.44,526.839,0.216014,944.877,534.42,0.0883373,952.458,537.756,0.0975707,959.736,538.969,0.0559232,968.834,526.839,0.0425614,940.024,532.601,0.088387,949.425,535.936,0.101295,952.155,536.543,0.0599895,966.711,528.052,0.0400319,940.631,525.019,0.0435125,946.999,527.142,0.0397179,958.22,528.355,0.0494711,961.556,528.052,0.0291901],"hand_right_keypoints_2d":[1050.07,567.769,0.0145988,1046.5,536.348,0.0187861,1035.81,572.304,0.0426545,1032.25,529.221,0.102354,1029.01,526.306,0.0551356,1043.59,538.615,0.0485236,1036.14,528.249,0.123396,1027.07,529.221,0.187019,1020.91,533.432,0.158446,1040.35,536.672,0.0330553,1035.81,528.897,0.101567,1027.39,530.841,0.192002,1021.24,533.432,0.162759,1039.38,536.672,0.0314099,1035.49,528.897,0.104451,1027.07,531.165,0.150516,1020.91,533.108,0.171756,1034.52,571.98,0.0292887,1033.87,526.954,0.0769301,1030.63,531.165,0.0718004,1025.45,531.489,0.0635766],"pose_keypoints_3d":[],"face_keypoints_3d":[],"hand_left_keypoints_3d":[],"hand_right_keypoints_3d":[]}]} -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/poses-openpose/001_001/frame000103.json: -------------------------------------------------------------------------------- 1 | {"version":1.2,"people":[{"pose_keypoints_2d":[980.644,409.583,0.936084,974.832,480.114,0.941322,924.727,471.37,0.843195,918.841,550.729,0.845186,1001.19,550.861,0.76846,1024.85,483.121,0.867901,1042.44,556.781,0.889278,971.743,539.063,0.742016,977.589,647.933,0.810152,945.297,647.978,0.739443,974.821,786.394,0.835345,1004.16,921.704,0.78003,1010.06,645.161,0.735747,1001.33,789.259,0.835515,974.775,903.991,0.83266,968.777,406.606,0.8758,992.332,392.034,0.870146,948.279,427.233,0.868259,1001.38,409.688,0.281189,960.019,951.199,0.764472,977.774,951.125,0.749095,962.99,906.961,0.577598,1015.88,957.054,0.719746,998.111,954.12,0.817392,1018.78,930.611,0.633987],"face_keypoints_2d":[947.228,415.39,0.580392,951.495,422.4,0.646118,954.848,429.41,0.683006,959.42,436.42,0.675944,964.906,441.297,0.646067,971.611,444.04,0.620627,979.84,445.564,0.649198,988.069,444.65,0.685536,995.079,443.43,0.713015,1000.57,439.468,0.80582,1003.61,434.287,0.805466,1005.75,428.801,0.758408,1006.05,422.095,0.778763,1005.44,415.39,0.71185,1003,409.295,0.671675,999.956,403.199,0.628261,995.689,397.713,0.530766,953.934,402.589,0.830502,955.458,398.018,0.839954,959.42,394.97,0.948078,964.601,393.141,0.87742,970.087,393.446,0.824508,978.926,390.093,0.82805,981.059,387.96,0.862723,983.802,386.741,0.843294,986.545,385.826,0.815093,989.593,386.741,0.660954,978.926,399.237,0.881025,981.669,402.589,0.916185,984.412,405.942,0.827127,987.155,409.295,0.761256,983.193,416.305,0.804123,985.631,416,0.85561,988.374,415.085,0.883012,989.898,414.476,0.853606,991.117,412.952,0.822712,961.248,408.075,0.751958,963.991,405.942,0.706838,967.344,404.418,0.768592,971.002,404.418,0.902538,968.563,405.942,0.778673,964.296,407.466,0.795561,983.802,399.237,0.823723,985.936,397.408,0.840284,987.765,396.494,0.696753,990.203,396.494,0.646277,988.374,397.713,0.791936,986.241,398.627,0.79558,980.755,429.41,0.781772,983.498,424.534,0.867927,986.545,421.791,0.852767,990.508,421.791,0.810541,992.641,419.962,0.743281,995.689,421.791,0.846958,998.127,424.229,0.880114,996.299,427.277,0.776913,994.47,429.41,0.784259,991.422,430.934,0.79967,988.679,431.544,0.915255,984.412,431.544,0.809892,982.583,428.496,0.740496,988.069,426.667,0.88748,990.812,425.753,0.813026,993.251,424.534,0.836313,995.994,424.534,0.803134,993.251,424.839,0.838777,990.812,426.362,0.839833,988.069,426.667,0.900897,966.43,405.637,0.595544,986.241,397.408,0.805411],"hand_left_keypoints_2d":[940.095,513.906,0.0202261,940.392,521.927,0.017354,954.949,530.839,0.0303942,973.071,539.157,0.0397996,967.426,526.68,0.0299051,945.74,536.186,0.0665135,957.623,534.701,0.0626674,963.267,530.245,0.113342,969.209,527.274,0.203654,944.848,534.107,0.0888368,952.275,537.078,0.101922,961.188,529.948,0.0650305,967.723,527.571,0.0729568,940.095,530.245,0.0667782,949.602,536.484,0.111518,959.702,529.651,0.0758699,967.723,526.977,0.0698343,938.61,526.977,0.0488883,947.522,529.057,0.0265845,949.305,526.977,0.043188,960.594,527.274,0.0351838],"hand_right_keypoints_2d":[1049.85,569.03,0.0155539,1055.22,530.093,0.0241942,1038.1,525.394,0.0599738,1032.39,528.079,0.139321,1022.66,533.114,0.106517,1042.46,536.471,0.0313604,1035.41,526.401,0.0716546,1026.69,529.758,0.131257,1019.64,533.45,0.141782,1039.78,523.044,0.0199173,1035.75,525.394,0.0376324,1027.02,529.758,0.0691273,1021.32,532.443,0.0559435,1040.45,537.813,0.0184415,1037.09,525.394,0.0395165,1027.36,529.758,0.0615685,1021.65,533.114,0.0542473,1035.41,573.058,0.0308379,1035.08,526.737,0.0411084,1029.71,531.436,0.0525679,1026.69,530.764,0.0359305],"pose_keypoints_3d":[],"face_keypoints_3d":[],"hand_left_keypoints_3d":[],"hand_right_keypoints_3d":[]}]} -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/pose/poses-openpose/001_001/frame000104.json: -------------------------------------------------------------------------------- 1 | {"version":1.2,"people":[{"pose_keypoints_2d":[980.607,409.642,0.944162,977.532,480.125,0.938141,927.512,471.376,0.85049,921.693,550.752,0.85741,1001.3,553.778,0.782964,1024.9,483.144,0.860912,1045.17,556.786,0.891419,971.759,536.218,0.770378,977.612,645.143,0.808065,945.327,645.165,0.739919,974.78,783.44,0.816984,1004.13,918.798,0.770604,1010.08,645.124,0.733768,1001.34,786.449,0.827347,974.793,901.185,0.828969,965.935,406.681,0.885148,989.462,394.858,0.851514,948.247,427.229,0.860073,1001.37,409.644,0.257524,960.003,951.172,0.75794,977.765,951.11,0.744206,965.834,904.13,0.577398,1015.87,957.05,0.713311,998.097,954.107,0.807773,1018.79,930.552,0.621312],"face_keypoints_2d":[947.42,416.852,0.610676,951.63,424.07,0.657689,954.637,430.385,0.652729,959.148,437.001,0.687548,965.463,441.812,0.661215,972.079,444.519,0.641893,980.199,446.023,0.698524,987.717,445.12,0.710606,994.934,443.617,0.770893,1000.05,439.407,0.834197,1003.96,434.294,0.790065,1005.76,428.581,0.720958,1006.66,422.265,0.781775,1005.46,415.65,0.720732,1003.35,409.936,0.65543,1000.05,403.32,0.646652,996.438,398.208,0.534121,952.232,404.523,0.76062,954.337,399.711,0.860441,958.547,395.802,0.894774,963.659,393.697,0.897921,968.771,393.697,0.865055,978.394,390.99,0.825312,980.5,388.885,0.858197,982.905,387.382,0.845624,985.612,386.479,0.863658,988.619,386.479,0.689323,978.094,400.012,0.850204,980.5,403.32,0.973015,982.905,406.628,0.848327,985.311,410.237,0.801126,981.101,417.153,0.845374,984.108,417.153,0.845936,986.815,416.552,0.838845,988.318,415.349,0.882174,990.123,413.845,0.827396,960.952,408.432,0.780962,963.659,406.327,0.741708,966.365,405.124,0.802901,969.974,405.124,0.849439,967.568,406.327,0.778574,963.96,408.131,0.833285,983.206,398.809,0.908407,985.311,397.606,0.889696,987.416,396.403,0.816713,989.521,396.102,0.771447,987.717,397.907,0.879613,985.612,398.508,0.94984,979.597,429.784,0.809516,982.905,424.972,0.906378,985.913,422.265,0.846709,989.521,422.265,0.827961,991.927,420.461,0.769564,994.934,421.965,0.807718,997.34,424.371,0.805225,995.836,427.378,0.752713,993.731,429.483,0.784498,991.025,430.686,0.783153,988.018,431.889,0.894811,983.807,431.889,0.829434,981.101,429.182,0.816597,987.717,427.077,0.867946,990.123,425.874,0.808731,992.528,424.972,0.873918,995.235,424.671,0.77147,992.528,424.972,0.881364,990.123,426.476,0.838466,987.717,427.077,0.894881,965.764,406.026,0.645151,985.612,397.907,0.871693],"hand_left_keypoints_2d":[943.804,526.634,0.0245631,941.939,528.188,0.0383901,948.154,535.646,0.0472801,975.19,535.024,0.0397656,969.286,526.323,0.0781313,969.596,528.499,0.0594958,968.664,546.522,0.0569026,970.839,533.471,0.0540028,970.839,535.024,0.0582226,964.003,536.268,0.0422166,968.353,545.279,0.0379151,971.15,552.116,0.0268802,969.907,526.013,0.0231355,964.003,542.172,0.0259248,964.313,544.658,0.0222914,961.517,544.969,0.0170297,961.517,545.901,0.0165331,963.381,543.104,0.0145252,948.776,537.821,0.0160402,961.517,545.901,0.0139166,961.517,546.833,0.0142382],"hand_right_keypoints_2d":[1050.3,572.635,0.0202862,1054.52,531.397,0.0203782,1035.03,573.284,0.0553058,1034.06,528.15,0.117276,1023.35,532.696,0.0634688,1038.93,571.011,0.0309456,1036.01,527.501,0.056417,1026.27,530.098,0.0884901,1020.1,532.696,0.0711614,1038.93,571.336,0.0161528,1036.01,526.851,0.0250459,1026.92,529.773,0.0492533,1021.07,532.371,0.0363066,1034.39,571.011,0.0201017,1035.36,527.176,0.0196526,1026.27,530.098,0.0409866,1021.72,532.696,0.0392283,1034.71,572.635,0.0414154,1035.03,527.176,0.031508,1030.81,530.423,0.0430554,1021.4,533.021,0.0238244],"pose_keypoints_3d":[],"face_keypoints_3d":[],"hand_left_keypoints_3d":[],"hand_right_keypoints_3d":[]}]} -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/street/images/seq0001/aachen_000001_000000_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/street/images/seq0001/aachen_000001_000000_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/street/images/seq0001/aachen_000001_000001_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/street/images/seq0001/aachen_000001_000001_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/street/images/seq0001/aachen_000001_000002_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/street/images/seq0001/aachen_000001_000002_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/street/images/seq0001/aachen_000001_000003_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/street/images/seq0001/aachen_000001_000003_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/street/seg_maps/seq0001/aachen_000001_000000_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/street/seg_maps/seq0001/aachen_000001_000000_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/street/seg_maps/seq0001/aachen_000001_000001_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/street/seg_maps/seq0001/aachen_000001_000001_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/street/seg_maps/seq0001/aachen_000001_000002_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/street/seg_maps/seq0001/aachen_000001_000002_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/vid2vid/street/seg_maps/seq0001/aachen_000001_000003_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/vid2vid/street/seg_maps/seq0001/aachen_000001_000003_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000000_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000000_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000001_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000001_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000002_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000002_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000003_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000003_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000004_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000004_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000005_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000005_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000006_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000006_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000007_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000007_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000008_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000008_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000009_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/depth_maps-megadepth/seq0000/aachen_000000_000009_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000000_leftImg8bit.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000000_leftImg8bit.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000001_leftImg8bit.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000001_leftImg8bit.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000002_leftImg8bit.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000002_leftImg8bit.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000003_leftImg8bit.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000003_leftImg8bit.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000004_leftImg8bit.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000004_leftImg8bit.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000005_leftImg8bit.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000005_leftImg8bit.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000006_leftImg8bit.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000006_leftImg8bit.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000007_leftImg8bit.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000007_leftImg8bit.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000008_leftImg8bit.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000008_leftImg8bit.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000009_leftImg8bit.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/images/seq0000/aachen_000000_000009_leftImg8bit.jpg -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000000_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000000_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000001_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000001_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000002_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000002_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000003_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000003_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000004_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000004_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000005_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000005_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000006_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000006_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000007_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000007_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000008_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000008_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000009_leftImg8bit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/seg_maps/seq0000/aachen_000000_000009_leftImg8bit.png -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000000_leftImg8bit.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000000_leftImg8bit.pkl -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000001_leftImg8bit.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000001_leftImg8bit.pkl -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000002_leftImg8bit.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000002_leftImg8bit.pkl -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000003_leftImg8bit.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000003_leftImg8bit.pkl -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000004_leftImg8bit.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000004_leftImg8bit.pkl -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000005_leftImg8bit.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000005_leftImg8bit.pkl -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000006_leftImg8bit.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000006_leftImg8bit.pkl -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000007_leftImg8bit.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000007_leftImg8bit.pkl -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000008_leftImg8bit.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000008_leftImg8bit.pkl -------------------------------------------------------------------------------- /dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000009_leftImg8bit.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/dataset/unit_test/raw/wc_vid2vid/cityscapes/unprojections/seq0000/aachen_000000_000009_leftImg8bit.pkl -------------------------------------------------------------------------------- /imaginaire/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | -------------------------------------------------------------------------------- /imaginaire/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | -------------------------------------------------------------------------------- /imaginaire/datasets/cache.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import diskcache 6 | 7 | """ 8 | INFO: 9 | Cache objects are thread-safe and may be shared between threads. 10 | Two Cache objects may also reference the same directory from separate 11 | threads or processes. In this way, they are also process-safe and support 12 | cross-process communication. 13 | """ 14 | 15 | 16 | class Cache(object): 17 | r"""This creates an on disk cache, which saves files as bytes. 18 | Args: 19 | root (str): Path to the cache dir. 20 | size_MB (float): Size of cache in MB. 21 | """ 22 | 23 | def __init__(self, root, size_GB): 24 | self.root = root 25 | self.size_limit_B = size_GB * 1024 * 1024 * 1024 26 | self.cache = diskcache.Cache(root, size_limit=self.size_limit_B) 27 | print('Created cache of max size %d GB at %s' % 28 | (size_GB, self.cache.directory)) 29 | 30 | def read(self, key): 31 | if key in self.cache: 32 | return self.cache[key] 33 | return False 34 | 35 | def write(self, key, value): 36 | try: 37 | self.cache[key] = value 38 | except Exception as e: # noqa 39 | print(e) 40 | return False 41 | -------------------------------------------------------------------------------- /imaginaire/datasets/dummy.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch 6 | 7 | 8 | class Dataset(torch.utils.data.Dataset): 9 | r"""Dummy dataset, returns nothing.""" 10 | 11 | def __init__(self, cfg, is_inference=False, is_test=False): 12 | super(Dataset, self).__init__() 13 | 14 | def __getitem__(self, index): 15 | return {} 16 | 17 | def __len__(self): 18 | return 65535 19 | -------------------------------------------------------------------------------- /imaginaire/datasets/folder.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import os 6 | 7 | import cv2 8 | import numpy as np 9 | import torch.utils.data as data 10 | from PIL import Image 11 | 12 | from imaginaire.utils.data import IMG_EXTENSIONS, HDR_IMG_EXTENSIONS 13 | import imageio 14 | 15 | 16 | class FolderDataset(data.Dataset): 17 | r"""This deals with opening, and reading from an Folder dataset. 18 | 19 | Args: 20 | root (str): Path to the folder. 21 | metadata (dict): Containing extensions. 22 | """ 23 | 24 | def __init__(self, root, metadata): 25 | self.root = os.path.expanduser(root) 26 | self.extensions = metadata 27 | 28 | print('Folder at %s opened.' % (root)) 29 | 30 | def getitem_by_path(self, path, data_type): 31 | r"""Load data item stored for key = path. 32 | 33 | Args: 34 | path (str): Key into Folder dataset. 35 | data_type (str): Key into self.extensions e.g. data/data_segmaps/... 36 | Returns: 37 | img (PIL.Image) or buf (str): Contents of file for this key. 38 | """ 39 | # Figure out decoding params. 40 | ext = self.extensions[data_type] 41 | is_image = False 42 | is_hdr = False 43 | if ext in IMG_EXTENSIONS: 44 | is_image = True 45 | if 'tif' in ext: 46 | dtype, mode = np.uint16, -1 47 | elif 'JPEG' in ext or 'JPG' in ext \ 48 | or 'jpeg' in ext or 'jpg' in ext: 49 | dtype, mode = np.uint8, 3 50 | else: 51 | dtype, mode = np.uint8, -1 52 | elif ext in HDR_IMG_EXTENSIONS: 53 | is_hdr = True 54 | else: 55 | is_image = False 56 | 57 | # Get value from key. 58 | filepath = os.path.join(self.root, path.decode() + '.' + ext) 59 | assert os.path.exists(filepath), '%s does not exist' % (filepath) 60 | with open(filepath, 'rb') as f: 61 | buf = f.read() 62 | 63 | # Decode and return. 64 | if is_image: 65 | try: 66 | img = cv2.imdecode(np.fromstring(buf, dtype=dtype), mode) 67 | except Exception: 68 | print(path) 69 | # BGR to RGB if 3 channels. 70 | if img.ndim == 3 and img.shape[-1] == 3: 71 | img = img[:, :, ::-1] 72 | img = Image.fromarray(img) 73 | return img 74 | elif is_hdr: 75 | try: 76 | imageio.plugins.freeimage.download() 77 | img = imageio.imread(buf) 78 | except Exception: 79 | print(path) 80 | return img # Return a numpy array 81 | else: 82 | return buf 83 | 84 | def __len__(self): 85 | r"""Return number of keys in Folder dataset.""" 86 | return self.length 87 | -------------------------------------------------------------------------------- /imaginaire/datasets/lmdb.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import json 6 | import os 7 | 8 | import cv2 9 | import lmdb 10 | import numpy as np 11 | import torch.utils.data as data 12 | from PIL import Image 13 | 14 | from imaginaire.utils.data import IMG_EXTENSIONS, HDR_IMG_EXTENSIONS 15 | from imaginaire.utils.distributed import master_only_print as print 16 | import imageio 17 | 18 | 19 | class LMDBDataset(data.Dataset): 20 | r"""This deals with opening, and reading from an LMDB dataset. 21 | Args: 22 | root (str): Path to the LMDB file. 23 | """ 24 | 25 | def __init__(self, root): 26 | self.root = os.path.expanduser(root) 27 | self.env = lmdb.open(root, max_readers=126, readonly=True, lock=False, 28 | readahead=False, meminit=False) 29 | with self.env.begin(write=False) as txn: 30 | self.length = txn.stat()['entries'] 31 | 32 | # Read metadata. 33 | with open(os.path.join(self.root, '..', 'metadata.json')) as fin: 34 | self.extensions = json.load(fin) 35 | 36 | print('LMDB file at %s opened.' % (root)) 37 | 38 | def getitem_by_path(self, path, data_type): 39 | r"""Load data item stored for key = path. 40 | 41 | Args: 42 | path (str): Key into LMDB dataset. 43 | data_type (str): Key into self.extensions e.g. data/data_segmaps/... 44 | Returns: 45 | img (PIL.Image) or buf (str): Contents of LMDB value for this key. 46 | """ 47 | # Figure out decoding params. 48 | ext = self.extensions[data_type] 49 | is_image = False 50 | is_hdr = False 51 | if ext in IMG_EXTENSIONS: 52 | is_image = True 53 | if 'tif' in ext: 54 | dtype, mode = np.uint16, -1 55 | elif 'JPEG' in ext or 'JPG' in ext \ 56 | or 'jpeg' in ext or 'jpg' in ext: 57 | dtype, mode = np.uint8, 3 58 | else: 59 | dtype, mode = np.uint8, -1 60 | elif ext in HDR_IMG_EXTENSIONS: 61 | is_hdr = True 62 | else: 63 | is_image = False 64 | 65 | # Get value from key. 66 | with self.env.begin(write=False) as txn: 67 | buf = txn.get(path) 68 | 69 | # Decode and return. 70 | if is_image: 71 | try: 72 | img = cv2.imdecode(np.fromstring(buf, dtype=dtype), mode) 73 | except Exception: 74 | print(path) 75 | # BGR to RGB if 3 channels. 76 | if img.ndim == 3 and img.shape[-1] == 3: 77 | img = img[:, :, ::-1] 78 | img = Image.fromarray(img) 79 | return img 80 | elif is_hdr: 81 | try: 82 | imageio.plugins.freeimage.download() 83 | img = imageio.imread(buf) 84 | except Exception: 85 | print(path) 86 | return img # Return a numpy array 87 | else: 88 | return buf 89 | 90 | def __len__(self): 91 | r"""Return number of keys in LMDB dataset.""" 92 | return self.length 93 | -------------------------------------------------------------------------------- /imaginaire/datasets/paired_images.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | 6 | from imaginaire.datasets.paired_videos import Dataset as VideoDataset 7 | 8 | 9 | class Dataset(VideoDataset): 10 | r"""Paired image dataset for use in pix2pixHD, SPADE. 11 | 12 | Args: 13 | cfg (Config): Loaded config object. 14 | is_inference (bool): In train or inference mode? 15 | """ 16 | 17 | def __init__(self, cfg, is_inference=False, is_test=False): 18 | self.paired = True 19 | super(Dataset, self).__init__(cfg, is_inference, 20 | sequence_length=1, 21 | is_test=is_test) 22 | self.is_video_dataset = False 23 | 24 | def _create_mapping(self): 25 | r"""Creates mapping from idx to key in LMDB. 26 | 27 | Returns: 28 | (tuple): 29 | - self.mapping (list): List mapping idx to key. 30 | - self.epoch_length (int): Number of samples in an epoch. 31 | """ 32 | idx_to_key = [] 33 | for lmdb_idx, sequence_list in enumerate(self.sequence_lists): 34 | for sequence_name, filenames in sequence_list.items(): 35 | for filename in filenames: 36 | idx_to_key.append({ 37 | 'lmdb_root': self.lmdb_roots[lmdb_idx], 38 | 'lmdb_idx': lmdb_idx, 39 | 'sequence_name': sequence_name, 40 | 'filenames': [filename], 41 | }) 42 | self.mapping = idx_to_key 43 | self.epoch_length = len(self.mapping) 44 | return self.mapping, self.epoch_length 45 | 46 | def _sample_keys(self, index): 47 | r"""Gets files to load for this sample. 48 | 49 | Args: 50 | index (int): Index in [0, len(dataset)]. 51 | Returns: 52 | key (dict): 53 | - lmdb_idx (int): Chosen LMDB dataset root. 54 | - sequence_name (str): Chosen sequence in chosen dataset. 55 | - filenames (list of str): Chosen filenames in chosen sequence. 56 | """ 57 | assert self.sequence_length == 1, \ 58 | 'Image dataset can only have sequence length = 1, not %d' % ( 59 | self.sequence_length) 60 | return self.mapping[index] 61 | 62 | def set_sequence_length(self, sequence_length): 63 | r"""Set the length of sequence you want as output from dataloader. 64 | Ignore this as this is an image loader. 65 | 66 | Args: 67 | sequence_length (int): Length of output sequences. 68 | """ 69 | pass 70 | 71 | def set_inference_sequence_idx(self, index): 72 | r"""Get frames from this sequence during inference. 73 | Overriden from super as this is not applicable for images. 74 | 75 | Args: 76 | index (int): Index of inference sequence. 77 | """ 78 | raise RuntimeError('Image dataset does not have sequences.') 79 | 80 | def num_inference_sequences(self): 81 | r"""Number of sequences available for inference. 82 | Overriden from super as this is not applicable for images. 83 | 84 | Returns: 85 | (int) 86 | """ 87 | raise RuntimeError('Image dataset does not have sequences.') 88 | -------------------------------------------------------------------------------- /imaginaire/discriminators/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | -------------------------------------------------------------------------------- /imaginaire/discriminators/dummy.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch.nn as nn 6 | 7 | from imaginaire.layers import LinearBlock 8 | 9 | 10 | class Discriminator(nn.Module): 11 | """Dummy Discriminator constructor. 12 | 13 | Args: 14 | dis_cfg (obj): Discriminator definition part of the yaml config file. 15 | data_cfg (obj): Data definition part of the yaml config file 16 | """ 17 | 18 | def __init__(self, dis_cfg, data_cfg): 19 | super(Discriminator, self).__init__() 20 | self.dummy_layer = LinearBlock(1, 1) 21 | pass 22 | 23 | def forward(self, data): 24 | """Dummy discriminator forward. 25 | 26 | Args: 27 | data (dict): 28 | """ 29 | return 30 | -------------------------------------------------------------------------------- /imaginaire/discriminators/mlp_multiclass.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import functools 6 | 7 | import numpy as np 8 | import torch.nn as nn 9 | 10 | from imaginaire.layers import LinearBlock 11 | 12 | 13 | class Discriminator(nn.Module): 14 | r"""Multi-layer Perceptron Classifier constructor. 15 | 16 | Args: 17 | dis_cfg (obj): Discriminator definition part of the yaml config file. 18 | data_cfg (obj): Data definition part of the yaml config file 19 | """ 20 | 21 | def __init__(self, dis_cfg, data_cfg): 22 | super(Discriminator, self).__init__() 23 | num_input_channels = dis_cfg.input_dims 24 | num_labels = dis_cfg.num_labels 25 | num_layers = getattr(dis_cfg, 'num_layers', 5) 26 | num_filters = getattr(dis_cfg, 'num_filters', 512) 27 | activation_norm_type = getattr(dis_cfg, 28 | 'activation_norm_type', 29 | 'batch_norm') 30 | nonlinearity = getattr(dis_cfg, 'nonlinearity', 'leakyrelu') 31 | base_linear_block = \ 32 | functools.partial(LinearBlock, 33 | activation_norm_type=activation_norm_type, 34 | nonlinearity=nonlinearity, 35 | order='CNA') 36 | dropout_ratio = 0.1 37 | layers = [base_linear_block(num_input_channels, num_filters), 38 | nn.Dropout(dropout_ratio)] 39 | for n in range(num_layers): 40 | dropout_ratio *= 1.5 41 | dropout_ratio = np.min([dropout_ratio, 0.5]) 42 | layers += [base_linear_block(num_filters, num_filters), 43 | nn.Dropout(dropout_ratio)] 44 | layers += [LinearBlock(num_filters, num_labels)] 45 | self.model = nn.Sequential(*layers) 46 | 47 | def forward(self, data): 48 | r"""Patch Discriminator forward. 49 | 50 | Args: 51 | data (dict): 52 | - data (N x -1 tensor): We will reshape the tensor to this format. 53 | Returns: 54 | (dict): 55 | - results (N x C tensor): Output scores before softmax. 56 | """ 57 | input_x = data['data'] 58 | bs = input_x.size()[0] 59 | input_x = input_x.view(bs, -1) 60 | pre_softmax_scores = self.model(input_x) 61 | outputs = dict() 62 | outputs['results'] = pre_softmax_scores 63 | return outputs 64 | -------------------------------------------------------------------------------- /imaginaire/evaluation/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | from .fid import compute_fid, compute_fid_data 6 | from .kid import compute_kid, compute_kid_data 7 | from .prdc import compute_prdc 8 | from .common import compute_all_metrics, compute_all_metrics_data 9 | 10 | __all__ = ['compute_fid', 'compute_fid_data', 'compute_kid', 'compute_kid_data', 11 | 'compute_prdc', 'compute_all_metrics', 'compute_all_metrics_data'] 12 | -------------------------------------------------------------------------------- /imaginaire/evaluation/caption/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | from .r_precision import get_r_precision 6 | from .common import get_image_encoder 7 | 8 | __all__ = ['get_image_encoder', 'get_r_precision'] 9 | -------------------------------------------------------------------------------- /imaginaire/evaluation/caption/common.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import os 6 | 7 | import boto3 8 | import torch 9 | from torch import nn, distributed as dist 10 | from torch.nn import functional as F 11 | from torch.distributed import barrier 12 | 13 | from imaginaire.utils.distributed import is_local_master 14 | from .clip import build_model 15 | from ...utils.io import download_file_from_google_drive 16 | 17 | 18 | def get_image_encoder(aws_credentials=None): 19 | if dist.is_initialized() and not is_local_master(): 20 | # Make sure only the first process in distributed training downloads the model, and the others use the cache. 21 | barrier() 22 | 23 | # Load the CLIP image encoder. 24 | print("Loading CLIP image encoder.") 25 | model_path = os.path.join(torch.hub.get_dir(), 'checkpoints', 'ViT-B-32.pt') 26 | if not os.path.exists(model_path): 27 | if aws_credentials is not None: 28 | s3 = boto3.client('s3', **aws_credentials) 29 | s3.download_file('lpi-poe', 'model_zoo/ViT-B-32.pt', model_path) 30 | else: 31 | download_file_from_google_drive("1Ri5APYM34A_IjG4F3Admutsf2oUwDjfW", model_path) 32 | model = torch.load(model_path, map_location='cpu') 33 | 34 | if dist.is_initialized() and is_local_master(): 35 | # Make sure only the first process in distributed training downloads the model, and the others use the cache. 36 | barrier() 37 | 38 | encoder = build_model(model).cuda() 39 | return ImageEncoder(encoder) 40 | 41 | 42 | class ImageEncoder(nn.Module): 43 | def __init__(self, encoder): 44 | super().__init__() 45 | self.model = encoder 46 | self.image_size = self.model.visual.input_resolution 47 | self.mean = torch.tensor([0.48145466, 0.4578275, 0.40821073], device="cuda") 48 | self.std = torch.tensor([0.26862954, 0.26130258, 0.27577711], device="cuda") 49 | 50 | @torch.no_grad() 51 | def forward(self, data, fake_images, align_corners=True): 52 | images = 0.5 * (1 + fake_images) 53 | images = F.interpolate(images, (self.image_size, self.image_size), mode='bicubic', align_corners=align_corners) 54 | images.clamp_(0, 1) 55 | images = (images - self.mean[None, :, None, None]) / (self.std[None, :, None, None]) 56 | image_code = self.model.encode_image(images) 57 | return torch.cat((image_code, data['captions-clip']), dim=1) 58 | -------------------------------------------------------------------------------- /imaginaire/evaluation/caption/r_precision.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | # flake8: noqa 6 | 7 | import torch 8 | import torch.nn.functional as F 9 | 10 | 11 | def get_r_precision(image_text_code, eps=1e-5): 12 | all_image_code, all_text_code = torch.chunk(image_text_code, 2, dim=1) 13 | P_rates = [] 14 | num_samples = len(all_image_code) 15 | assert num_samples >= 100 16 | for i in range(0, num_samples, 100): 17 | if i + 100 <= num_samples: 18 | cur_image_code = all_image_code[i:i + 100] 19 | cur_text_code = all_text_code[i:i + 100] 20 | cur_image_code = F.normalize(cur_image_code, dim=1, eps=eps) 21 | cur_text_code = F.normalize(cur_text_code, dim=1, eps=eps) 22 | cosine_similarities = cur_image_code @ cur_text_code.T 23 | top1_indices = torch.topk(cosine_similarities, dim=1, k=1)[1][:, 0] 24 | P_rate = torch.sum(top1_indices == torch.arange(100, device=top1_indices.device)).item() 25 | P_rates.append(P_rate) 26 | A_precision = sum(P_rates) * 1.0 / len(P_rates) 27 | return {"caption_rprec": A_precision} 28 | -------------------------------------------------------------------------------- /imaginaire/evaluation/knn.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch 6 | 7 | from imaginaire.evaluation.common import compute_nn 8 | 9 | 10 | def _get_1nn_acc(data_x, data_y, k=1): 11 | device = data_x.device 12 | n0 = data_x.size(0) 13 | n1 = data_y.size(0) 14 | data_all = torch.cat((data_x, data_y), dim=0) 15 | val, idx = compute_nn(data_all, k) 16 | label = torch.cat((torch.ones(n0, device=device), 17 | torch.zeros(n1, device=device))) 18 | 19 | count = torch.zeros(n0 + n1, device=device) 20 | for i in range(0, k): 21 | count = count + label.index_select(0, idx[:, i]) 22 | pred = torch.ge(count, (float(k) / 2) * 23 | torch.ones(n0 + n1, device=device)).float() 24 | 25 | tp = (pred * label).sum() 26 | fp = (pred * (1 - label)).sum() 27 | fn = ((1 - pred) * label).sum() 28 | tn = ((1 - pred) * (1 - label)).sum() 29 | acc_r = (tp / (tp + fn)).item() 30 | acc_f = (tn / (tn + fp)).item() 31 | acc = torch.eq(label, pred).float().mean().item() 32 | 33 | return {'1NN_acc': acc, 34 | '1NN_acc_real': acc_r, 35 | '1NN_acc_fake': acc_f} 36 | -------------------------------------------------------------------------------- /imaginaire/evaluation/segmentation/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | from .common import get_segmentation_hist_model, get_miou 6 | 7 | __all__ = ['get_segmentation_hist_model', 'get_miou'] 8 | -------------------------------------------------------------------------------- /imaginaire/evaluation/segmentation/cocostuff.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | from torch import nn 6 | from torch.nn import functional as F 7 | import torch.hub 8 | 9 | 10 | class DeepLabV2(nn.Module): 11 | def __init__(self, n_classes=182, image_size=512, use_dont_care=True): 12 | super(DeepLabV2, self).__init__() 13 | self.model = torch.hub.load( 14 | "kazuto1011/deeplab-pytorch", "deeplabv2_resnet101", 15 | pretrained=False, n_classes=182 16 | ) 17 | state_dict = torch.hub.load_state_dict_from_url( 18 | 'https://github.com/kazuto1011/deeplab-pytorch/releases/download/' 19 | 'v1.0/deeplabv2_resnet101_msc-cocostuff164k-100000.pth', 20 | map_location="cpu" 21 | ) 22 | self.model.load_state_dict(state_dict) 23 | 24 | self.image_size = image_size 25 | # self.mean = torch.tensor([122.675, 116.669, 104.008], device="cuda") 26 | self.mean = torch.tensor([104.008, 116.669, 122.675], device="cuda") 27 | self.n_classes = n_classes 28 | self.use_dont_care = use_dont_care 29 | 30 | def forward(self, images, align_corners=True): 31 | scale = self.image_size / max(images.shape[2:]) 32 | images = F.interpolate( 33 | images, scale_factor=scale, mode='bilinear', 34 | align_corners=align_corners 35 | ) 36 | images = 255 * 0.5 * (images + 1) # (-1, 1) -> (0, 255) 37 | images = images.flip(1) # RGB to BGR 38 | images -= self.mean[None, :, None, None] 39 | _, _, H, W = images.shape 40 | 41 | logits = self.model(images) 42 | logits = F.interpolate( 43 | logits, size=(H, W), mode="bilinear", 44 | align_corners=align_corners 45 | ) 46 | probs = F.softmax(logits, dim=1) 47 | pred = torch.argmax(probs, dim=1) 48 | return pred 49 | -------------------------------------------------------------------------------- /imaginaire/generators/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | -------------------------------------------------------------------------------- /imaginaire/generators/dummy.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch.nn as nn 6 | 7 | from imaginaire.layers import LinearBlock 8 | 9 | 10 | class Generator(nn.Module): 11 | r"""Dummy generator. 12 | 13 | Args: 14 | gen_cfg (obj): Generator definition part of the yaml config file. 15 | data_cfg (obj): Data definition part of the yaml config file. 16 | """ 17 | 18 | def __init__(self, gen_cfg, data_cfg): 19 | super(Generator, self).__init__() 20 | self.dummy_layer = LinearBlock(1, 1) 21 | pass 22 | 23 | def forward(self, data): 24 | r"""Dummy Generator forward. 25 | 26 | Args: 27 | data (dict): 28 | """ 29 | return 30 | -------------------------------------------------------------------------------- /imaginaire/layers/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | from .conv import LinearBlock, Conv1dBlock, Conv2dBlock, Conv3dBlock, \ 6 | HyperConv2dBlock, MultiOutConv2dBlock, \ 7 | PartialConv2dBlock, PartialConv3dBlock 8 | from .residual import ResLinearBlock, Res1dBlock, Res2dBlock, Res3dBlock, \ 9 | HyperRes2dBlock, MultiOutRes2dBlock, UpRes2dBlock, DownRes2dBlock, \ 10 | PartialRes2dBlock, PartialRes3dBlock 11 | from .non_local import NonLocal2dBlock 12 | 13 | __all__ = ['Conv1dBlock', 'Conv2dBlock', 'Conv3dBlock', 'LinearBlock', 14 | 'HyperConv2dBlock', 'MultiOutConv2dBlock', 15 | 'PartialConv2dBlock', 'PartialConv3dBlock', 16 | 'Res1dBlock', 'Res2dBlock', 'Res3dBlock', 17 | 'UpRes2dBlock', 'DownRes2dBlock', 18 | 'ResLinearBlock', 'HyperRes2dBlock', 'MultiOutRes2dBlock', 19 | 'PartialRes2dBlock', 'PartialRes3dBlock', 20 | 'NonLocal2dBlock'] 21 | 22 | try: 23 | from .repvgg import RepVGG1dBlock, RepVGG2dBlock, RepVGG3dBlock 24 | from .attn import MultiheadAttention 25 | __all__.extend(['RepVGG1dBlock', 'RepVGG2dBlock', 'RepVGG3dBlock']) 26 | except: # noqa 27 | pass 28 | -------------------------------------------------------------------------------- /imaginaire/layers/misc.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch 6 | from torch import nn 7 | 8 | 9 | class ApplyNoise(nn.Module): 10 | r"""Add Gaussian noise to the input tensor.""" 11 | 12 | def __init__(self): 13 | super().__init__() 14 | # scale of the noise 15 | self.scale = nn.Parameter(torch.zeros(1)) 16 | self.conditional = True 17 | 18 | def forward(self, x, *_args, noise=None, **_kwargs): 19 | r""" 20 | 21 | Args: 22 | x (tensor): Input tensor. 23 | noise (tensor, optional, default=``None``) : Noise tensor to be 24 | added to the input. 25 | """ 26 | if noise is None: 27 | sz = x.size() 28 | noise = x.new_empty(sz[0], 1, *sz[2:]).normal_() 29 | 30 | return x + self.scale * noise 31 | 32 | 33 | class PartialSequential(nn.Sequential): 34 | r"""Sequential block for partial convolutions.""" 35 | def __init__(self, *modules): 36 | super(PartialSequential, self).__init__(*modules) 37 | 38 | def forward(self, x): 39 | r""" 40 | 41 | Args: 42 | x (tensor): Input tensor. 43 | """ 44 | act = x[:, :-1] 45 | mask = x[:, -1].unsqueeze(1) 46 | for module in self: 47 | act, mask = module(act, mask_in=mask) 48 | return act 49 | 50 | 51 | class ConstantInput(nn.Module): 52 | def __init__(self, channel, size=4): 53 | super().__init__() 54 | if isinstance(size, int): 55 | h, w = size, size 56 | else: 57 | h, w = size 58 | self.input = nn.Parameter(torch.randn(1, channel, h, w)) 59 | 60 | def forward(self): 61 | return self.input 62 | -------------------------------------------------------------------------------- /imaginaire/layers/non_local.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | from functools import partial 6 | 7 | import torch 8 | import torch.nn as nn 9 | 10 | from imaginaire.layers import Conv2dBlock 11 | 12 | 13 | class NonLocal2dBlock(nn.Module): 14 | r"""Self attention Layer 15 | 16 | Args: 17 | in_channels (int): Number of channels in the input tensor. 18 | scale (bool, optional, default=True): If ``True``, scale the 19 | output by a learnable parameter. 20 | clamp (bool, optional, default=``False``): If ``True``, clamp the 21 | scaling parameter to (-1, 1). 22 | weight_norm_type (str, optional, default='none'): 23 | Type of weight normalization. 24 | ``'none'``, ``'spectral'``, ``'weight'``. 25 | weight_norm_params (obj, optional, default=None): 26 | Parameters of weight normalization. 27 | If not ``None``, weight_norm_params.__dict__ will be used as 28 | keyword arguments when initializing weight normalization. 29 | bias (bool, optional, default=True): If ``True``, adds bias in the 30 | convolutional blocks. 31 | """ 32 | 33 | def __init__(self, 34 | in_channels, 35 | scale=True, 36 | clamp=False, 37 | weight_norm_type='none', 38 | weight_norm_params=None, 39 | bias=True): 40 | super(NonLocal2dBlock, self).__init__() 41 | self.clamp = clamp 42 | self.gamma = nn.Parameter(torch.zeros(1)) if scale else 1.0 43 | self.in_channels = in_channels 44 | base_conv2d_block = partial(Conv2dBlock, 45 | kernel_size=1, 46 | stride=1, 47 | padding=0, 48 | weight_norm_type=weight_norm_type, 49 | weight_norm_params=weight_norm_params, 50 | bias=bias) 51 | self.theta = base_conv2d_block(in_channels, in_channels // 8) 52 | self.phi = base_conv2d_block(in_channels, in_channels // 8) 53 | self.g = base_conv2d_block(in_channels, in_channels // 2) 54 | self.out_conv = base_conv2d_block(in_channels // 2, in_channels) 55 | self.softmax = nn.Softmax(dim=-1) 56 | self.max_pool = nn.MaxPool2d(2) 57 | 58 | def forward(self, x): 59 | r""" 60 | 61 | Args: 62 | x (tensor) : input feature maps (B X C X W X H) 63 | Returns: 64 | (tuple): 65 | - out (tensor) : self attention value + input feature 66 | - attention (tensor): B x N x N (N is Width*Height) 67 | """ 68 | n, c, h, w = x.size() 69 | theta = self.theta(x).view(n, -1, h * w).permute(0, 2, 1) 70 | 71 | phi = self.phi(x) 72 | phi = self.max_pool(phi).view(n, -1, h * w // 4) 73 | 74 | energy = torch.bmm(theta, phi) 75 | attention = self.softmax(energy) 76 | 77 | g = self.g(x) 78 | g = self.max_pool(g).view(n, -1, h * w // 4) 79 | 80 | out = torch.bmm(g, attention.permute(0, 2, 1)) 81 | out = out.view(n, c // 2, h, w) 82 | out = self.out_conv(out) 83 | 84 | if self.clamp: 85 | out = self.gamma.clamp(-1, 1) * out + x 86 | else: 87 | out = self.gamma * out + x 88 | return out 89 | -------------------------------------------------------------------------------- /imaginaire/layers/nonlinearity.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch 6 | from torch import nn 7 | import torch.nn.functional as F 8 | 9 | from imaginaire.third_party.bias_act.bias_act import FusedNonlinearity 10 | 11 | 12 | class ScaledLeakyReLU(nn.Module): 13 | def __init__(self, negative_slope=0.2, scale=2 ** 0.5, inplace=False): 14 | super().__init__() 15 | 16 | self.negative_slope = negative_slope 17 | self.scale = scale 18 | self.inplace = inplace 19 | 20 | def forward(self, x): 21 | return F.leaky_relu(x, self.negative_slope, inplace=self.inplace) * self.scale 22 | # return _fused_scaled_leakyrelu(x, self.negative_slope, self.inplace, self.scale) 23 | 24 | 25 | # @torch.jit.script 26 | # def _fused_scaled_leakyrelu(x: torch.Tensor, negative_slope: float, inplace: bool, scale: float): 27 | # return F.leaky_relu(x, negative_slope, inplace=inplace) * scale 28 | 29 | 30 | def get_nonlinearity_layer(nonlinearity_type, inplace, **kwargs): 31 | r"""Return a nonlinearity layer. 32 | 33 | Args: 34 | nonlinearity_type (str): 35 | Type of nonlinear activation function. 36 | ``'none'``, ``'relu'``, ``'leakyrelu'``, ``'prelu'``, 37 | ``'tanh'`` , ``'sigmoid'`` or ``'softmax'``. 38 | inplace (bool): If ``True``, set ``inplace=True`` when initializing 39 | the nonlinearity layer. 40 | """ 41 | if nonlinearity_type.startswith('fused'): 42 | nonlinearity = FusedNonlinearity(nonlinearity=nonlinearity_type[6:], **kwargs) 43 | elif nonlinearity_type == 'relu': 44 | nonlinearity = nn.ReLU(inplace=inplace) 45 | elif nonlinearity_type == 'leakyrelu': 46 | nonlinearity = nn.LeakyReLU(0.2, inplace=inplace) 47 | elif nonlinearity_type == 'scaled_leakyrelu': 48 | nonlinearity = ScaledLeakyReLU(0.2, inplace=inplace) 49 | import imaginaire.config 50 | if imaginaire.config.USE_JIT: 51 | nonlinearity = torch.jit.script(nonlinearity) 52 | elif nonlinearity_type == 'prelu': 53 | nonlinearity = nn.PReLU() 54 | elif nonlinearity_type == 'tanh': 55 | nonlinearity = nn.Tanh() 56 | elif nonlinearity_type == 'sigmoid': 57 | nonlinearity = nn.Sigmoid() 58 | elif nonlinearity_type.startswith('softmax'): 59 | dim = nonlinearity_type.split(',')[1] if ',' in nonlinearity_type else 1 60 | nonlinearity = nn.Softmax(dim=int(dim)) 61 | elif nonlinearity_type == 'none' or nonlinearity_type == '': 62 | nonlinearity = None 63 | else: 64 | raise ValueError('Nonlinearity %s is not recognized' % nonlinearity_type) 65 | return nonlinearity 66 | -------------------------------------------------------------------------------- /imaginaire/losses/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | from .gan import GANLoss 6 | from .perceptual import PerceptualLoss 7 | from .feature_matching import FeatureMatchingLoss 8 | from .kl import GaussianKLLoss 9 | from .flow import MaskedL1Loss, FlowLoss 10 | from .dict import DictLoss 11 | from .weighted_mse import WeightedMSELoss 12 | 13 | __all__ = ['GANLoss', 'PerceptualLoss', 'FeatureMatchingLoss', 'GaussianKLLoss', 14 | 'MaskedL1Loss', 'FlowLoss', 'DictLoss', 15 | 'WeightedMSELoss'] 16 | 17 | try: 18 | from .gradient_penalty import GradientPenaltyLoss 19 | __all__.extend(['GradientPenaltyLoss']) 20 | except: # noqa 21 | pass 22 | -------------------------------------------------------------------------------- /imaginaire/losses/dict.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch.nn as nn 6 | 7 | 8 | class DictLoss(nn.Module): 9 | def __init__(self, criterion='l1'): 10 | super(DictLoss, self).__init__() 11 | if criterion == 'l1': 12 | self.criterion = nn.L1Loss() 13 | elif criterion == 'l2' or criterion == 'mse': 14 | self.criterion = nn.MSELoss() 15 | else: 16 | raise ValueError('Criterion %s is not recognized' % criterion) 17 | 18 | def forward(self, fake, real): 19 | """Return the target vector for the l1/l2 loss computation. 20 | 21 | Args: 22 | fake (dict, list or tuple): Discriminator features of fake images. 23 | real (dict, list or tuple): Discriminator features of real images. 24 | Returns: 25 | loss (tensor): Loss value. 26 | """ 27 | loss = 0 28 | if type(fake) == dict: 29 | for key in fake.keys(): 30 | loss += self.criterion(fake[key], real[key].detach()) 31 | elif type(fake) == list or type(fake) == tuple: 32 | for f, r in zip(fake, real): 33 | loss += self.criterion(f, r.detach()) 34 | else: 35 | loss += self.criterion(fake, real.detach()) 36 | return loss 37 | -------------------------------------------------------------------------------- /imaginaire/losses/feature_matching.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch.nn as nn 6 | 7 | 8 | class FeatureMatchingLoss(nn.Module): 9 | r"""Compute feature matching loss""" 10 | def __init__(self, criterion='l1'): 11 | super(FeatureMatchingLoss, self).__init__() 12 | if criterion == 'l1': 13 | self.criterion = nn.L1Loss() 14 | elif criterion == 'l2' or criterion == 'mse': 15 | self.criterion = nn.MSELoss() 16 | else: 17 | raise ValueError('Criterion %s is not recognized' % criterion) 18 | 19 | def forward(self, fake_features, real_features): 20 | r"""Return the target vector for the binary cross entropy loss 21 | computation. 22 | 23 | Args: 24 | fake_features (list of lists): Discriminator features of fake images. 25 | real_features (list of lists): Discriminator features of real images. 26 | 27 | Returns: 28 | (tensor): Loss value. 29 | """ 30 | num_d = len(fake_features) 31 | dis_weight = 1.0 / num_d 32 | loss = fake_features[0][0].new_tensor(0) 33 | for i in range(num_d): 34 | for j in range(len(fake_features[i])): 35 | tmp_loss = self.criterion(fake_features[i][j], 36 | real_features[i][j].detach()) 37 | loss += dis_weight * tmp_loss 38 | return loss 39 | -------------------------------------------------------------------------------- /imaginaire/losses/info_nce.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import math 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torch.distributed as dist 10 | 11 | from imaginaire.utils.distributed import get_world_size, get_rank, \ 12 | dist_all_reduce_tensor 13 | 14 | 15 | class GatherLayer(torch.autograd.Function): 16 | @staticmethod 17 | def forward(ctx, input): 18 | ctx.save_for_backward(input) 19 | output = [torch.zeros_like(input) for _ in range(dist.get_world_size())] 20 | dist.all_gather(output, input) 21 | return tuple(output) 22 | 23 | @staticmethod 24 | def backward(ctx, *grads): 25 | input, = ctx.saved_tensors 26 | grad_out = torch.zeros_like(input) 27 | all_grads = torch.stack(grads) 28 | all_grads = dist_all_reduce_tensor(all_grads, reduce='sum') 29 | grad_out[:] = all_grads[get_rank()] 30 | return grad_out 31 | 32 | 33 | class InfoNCELoss(nn.Module): 34 | def __init__(self, 35 | temperature=0.07, 36 | gather_distributed=True, 37 | learn_temperature=True, 38 | single_direction=False, 39 | flatten=True): 40 | super(InfoNCELoss, self).__init__() 41 | self.logit_scale = nn.Parameter(torch.tensor([math.log(1/temperature)])) 42 | self.logit_scale.requires_grad = learn_temperature 43 | self.gather_distributed = gather_distributed 44 | self.single_direction = single_direction 45 | self.flatten = flatten 46 | 47 | def forward(self, features_a, features_b, gather_distributed=None, eps=1e-8): 48 | if gather_distributed is None: 49 | gather_distributed = self.gather_distributed 50 | 51 | if features_a is None or features_b is None: 52 | return torch.tensor(0, device='cuda'), torch.tensor(0, device='cuda') 53 | 54 | bs_a, bs_b = features_a.size(0), features_b.size(0) 55 | if self.flatten: 56 | features_a, features_b = features_a.reshape(bs_a, -1), features_b.reshape(bs_b, -1) 57 | else: 58 | features_a = features_a.reshape(bs_a, features_a.size(1), -1).mean(-1) 59 | features_b = features_b.reshape(bs_b, features_b.size(1), -1).mean(-1) 60 | 61 | # Temperature clipping. 62 | self.logit_scale.data = torch.clamp(self.logit_scale.data, 0, 4.6052) 63 | 64 | # normalized features 65 | features_a = features_a / (features_a.norm(dim=1, keepdim=True) + eps) 66 | features_b = features_b / (features_b.norm(dim=1, keepdim=True) + eps) 67 | 68 | loss_a = self._forward_single_direction(features_a, features_b, gather_distributed) 69 | if self.single_direction: 70 | return loss_a 71 | else: 72 | loss_b = self._forward_single_direction(features_b, features_a, gather_distributed) 73 | return loss_a + loss_b 74 | 75 | def _forward_single_direction( 76 | self, features_a, features_b, gather_distributed): 77 | bs_a = features_a.shape[0] 78 | logit_scale = self.logit_scale.exp() 79 | if get_world_size() > 1 and gather_distributed: 80 | gather_features_b = torch.cat(GatherLayer.apply(features_b)) 81 | gather_labels_a = torch.arange(bs_a, device='cuda') + get_rank() * bs_a 82 | logits_a = logit_scale * features_a @ gather_features_b.t() 83 | else: 84 | gather_labels_a = torch.arange(bs_a, device='cuda') 85 | logits_a = logit_scale * features_a @ features_b.t() 86 | loss_a = F.cross_entropy(logits_a, gather_labels_a) 87 | return loss_a 88 | -------------------------------------------------------------------------------- /imaginaire/losses/kl.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch 6 | import torch.nn as nn 7 | 8 | 9 | class GaussianKLLoss(nn.Module): 10 | r"""Compute KL loss in VAE for Gaussian distributions""" 11 | def __init__(self): 12 | super(GaussianKLLoss, self).__init__() 13 | 14 | def forward(self, mu, logvar=None): 15 | r"""Compute loss 16 | 17 | Args: 18 | mu (tensor): mean 19 | logvar (tensor): logarithm of variance 20 | """ 21 | if logvar is None: 22 | logvar = torch.zeros_like(mu) 23 | return -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) 24 | -------------------------------------------------------------------------------- /imaginaire/losses/weighted_mse.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch 6 | import torch.nn as nn 7 | 8 | 9 | class WeightedMSELoss(nn.Module): 10 | r"""Compute Weighted MSE loss""" 11 | def __init__(self, reduction='mean'): 12 | super(WeightedMSELoss, self).__init__() 13 | self.reduction = reduction 14 | 15 | def forward(self, input, target, weight): 16 | r"""Return weighted MSE Loss. 17 | Args: 18 | input (tensor): 19 | target (tensor): 20 | weight (tensor): 21 | Returns: 22 | (tensor): Loss value. 23 | """ 24 | if self.reduction == 'mean': 25 | loss = torch.mean(weight * (input - target) ** 2) 26 | else: 27 | loss = torch.sum(weight * (input - target) ** 2) 28 | return loss 29 | -------------------------------------------------------------------------------- /imaginaire/model_utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | -------------------------------------------------------------------------------- /imaginaire/model_utils/gancraft/mc_lbl_reduction.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import os 6 | import csv 7 | 8 | 9 | class ReducedLabelMapper: 10 | def __init__(self): 11 | this_path = os.path.dirname(os.path.abspath(__file__)) 12 | print('[ReducedLabelMapper] Loading from {}'.format(this_path)) 13 | 14 | # Load Minecraft LUT 15 | mcid2rdlbl_lut = {} 16 | mcid2mclbl_lut = {} 17 | with open(os.path.join(this_path, 'mc_reduction.csv'), newline='') as csvfile: 18 | csvreader = csv.reader(csvfile, delimiter=',') 19 | for row in csvreader: 20 | mcid = int(row[0]) 21 | mcid2rdlbl_lut[mcid] = row[3] 22 | mcid2mclbl_lut[mcid] = row[1] 23 | 24 | # Load reduced label set 25 | reduced_lbls = [] 26 | rdlbl2rdid = {} 27 | with open(os.path.join(this_path, 'reduced_coco_lbls.csv'), newline='') as csvfile: 28 | csvreader = csv.reader(csvfile, delimiter=',') 29 | for idx, row in enumerate(csvreader): 30 | rdlbl2rdid[row[0]] = idx 31 | reduced_lbls.append(row[0]) 32 | print(['{}: {}'.format(rdid, rdlbl) for rdid, rdlbl in enumerate(reduced_lbls)]) 33 | # The first label should always be 'ignore' 34 | assert reduced_lbls[0] == 'ignore' 35 | 36 | # Generate Minecraft ID to Reduced ID LUT 37 | mcid2rdid_lut = [] 38 | for mcid in range(len(mcid2rdlbl_lut)): 39 | rdlbl = mcid2rdlbl_lut[mcid] 40 | if rdlbl == '': 41 | rdlbl = 'ignore' 42 | rdid = rdlbl2rdid[rdlbl] 43 | mcid2rdid_lut.append(rdid) 44 | 45 | # ================= coco part ================== 46 | gg_label_list = [] 47 | gglbl2ggid = {} 48 | with open(os.path.join(this_path, 'gaugan_lbl2col.csv'), newline='') as csvfile: 49 | csvreader = csv.reader(csvfile, delimiter=',') 50 | for idx, row in enumerate(csvreader): 51 | gg_label_list.append(row[0]) 52 | gglbl2ggid[row[0]] = idx 53 | 54 | # Load coco -> reduced mapping table 55 | gglbl2rdid = {} 56 | with open(os.path.join(this_path, 'gaugan_reduction.csv'), newline='') as csvfile: 57 | csvreader = csv.reader(csvfile, delimiter=',') 58 | for idx, row in enumerate(csvreader): 59 | gglbl = row[0] 60 | target_rdlbl = row[1] 61 | ggid = gglbl2ggid[gglbl] 62 | target_rdid = rdlbl2rdid[target_rdlbl] 63 | gglbl2rdid[ggid] = target_rdid 64 | ggid2rdid = [gglbl2rdid[i] for i in range(len(gglbl2rdid))] 65 | 66 | print('[ReducedLabelMapper] #Reduced Labels: {}'.format(len(reduced_lbls))) 67 | 68 | self.mcid2rdid_lut = mcid2rdid_lut 69 | self.ggid2rdid = ggid2rdid 70 | self.reduced_lbls = reduced_lbls 71 | 72 | self.ignore_id = rdlbl2rdid['ignore'] 73 | self.dirt_id = rdlbl2rdid['dirt'] 74 | self.water_id = rdlbl2rdid['water'] 75 | 76 | self.gglbl2ggid = gglbl2ggid 77 | 78 | def gglbl2ggid(self, gglbl): 79 | return self.gglbl2ggid[gglbl] 80 | 81 | 82 | if __name__ == '__main__': 83 | mapper = ReducedLabelMapper() 84 | -------------------------------------------------------------------------------- /imaginaire/model_utils/gancraft/reduced_coco_lbls.csv: -------------------------------------------------------------------------------- 1 | ignore 2 | sky 3 | tree 4 | dirt 5 | flower 6 | grass 7 | gravel 8 | water 9 | rock 10 | stone 11 | sand 12 | snow -------------------------------------------------------------------------------- /imaginaire/model_utils/gancraft/voxlib/Makefile: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | 6 | all: 7 | python setup.py build_ext --inplace 8 | python setup.py install 9 | 10 | clean: 11 | rm -rf *.o *.a *.so test build 12 | -------------------------------------------------------------------------------- /imaginaire/model_utils/gancraft/voxlib/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | from .positional_encoding import positional_encoding 6 | from .sp_trilinear import sparse_trilinear_interp_worldcoord 7 | from voxlib import ray_voxel_intersection_perspective 8 | -------------------------------------------------------------------------------- /imaginaire/model_utils/gancraft/voxlib/positional_encoding.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch 6 | from torch.autograd import Function 7 | import voxlib 8 | 9 | # Cheatsheet: 10 | # mark_dirty() must be used to mark any input that is modified inplace by the forward function. 11 | # mark_non_differentiable() 12 | 13 | 14 | class PositionalEncodingFunction(Function): 15 | @staticmethod 16 | def forward(ctx, in_feature, pe_degrees, dim, incl_orig): 17 | out_feature = voxlib.positional_encoding(in_feature, pe_degrees, dim, incl_orig) 18 | 19 | ctx.save_for_backward(out_feature) 20 | ctx.pe_degrees = pe_degrees 21 | ctx.dim = dim 22 | ctx.incl_orig = incl_orig 23 | 24 | return out_feature 25 | 26 | @staticmethod 27 | def backward(ctx, out_feature_grad): 28 | out_feature, = ctx.saved_tensors 29 | 30 | # torch::Tensor positional_encoding_backward(const torch::Tensor& out_feature_grad, 31 | # const torch::Tensor& out_feature, int ndegrees, int dim, bool incl_orig) { 32 | in_feature_grad = voxlib.positional_encoding_backward( 33 | out_feature_grad, out_feature, ctx.pe_degrees, ctx.dim, ctx.incl_orig) 34 | 35 | return in_feature_grad, None, None, None 36 | 37 | 38 | def positional_encoding(in_feature, pe_degrees, dim=-1, incl_orig=False): 39 | return PositionalEncodingFunction.apply(in_feature, pe_degrees, dim, incl_orig) 40 | 41 | # input: N, C 42 | # output: N, pe_degrees*C 43 | 44 | 45 | def positional_encoding_pt(pts, pe_degrees, dim=-1, incl_orig=False): 46 | import numpy as np 47 | pe_stor = [] 48 | for i in range(pe_degrees): 49 | pe_stor.append(torch.sin(pts * np.pi * 2 ** i)) 50 | pe_stor.append(torch.cos(pts * np.pi * 2 ** i)) 51 | if incl_orig: 52 | pe_stor.append(pts) 53 | pe = torch.cat(pe_stor, dim=dim) 54 | return pe 55 | 56 | 57 | if __name__ == '__main__': 58 | x = torch.rand(384, 512, 5, 48).cuda() * 1024 59 | y = positional_encoding_pt(x, 4, incl_orig=True) 60 | y2 = positional_encoding(x, 4, incl_orig=True) 61 | 62 | print(torch.abs(y - y2)) 63 | print(torch.allclose(y, y2, rtol=1e-05, atol=1e-05)) 64 | -------------------------------------------------------------------------------- /imaginaire/model_utils/gancraft/voxlib/setup.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | from setuptools import setup 6 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 7 | 8 | cxx_args = ['-fopenmp'] 9 | nvcc_args = [] 10 | 11 | setup( 12 | name='voxrender', 13 | ext_modules=[ 14 | CUDAExtension('voxlib', [ 15 | 'voxlib.cpp', 16 | 'ray_voxel_intersection.cu', 17 | 'sp_trilinear_worldcoord_kernel.cu', 18 | 'positional_encoding_kernel.cu' 19 | ], 20 | extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args} 21 | ) 22 | ], 23 | cmdclass={ 24 | 'build_ext': BuildExtension 25 | }) 26 | -------------------------------------------------------------------------------- /imaginaire/model_utils/gancraft/voxlib/sp_trilinear.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | from torch.autograd import Function 6 | import voxlib 7 | 8 | """ 9 | It takes world coordinate as input instead of block-local coordinate. Corner IDs are looked up on-the-fly to 10 | save memory. 11 | """ 12 | 13 | 14 | class SparseTrilinearWorldCoordFunction(Function): 15 | @staticmethod 16 | def forward(ctx, in_feature, corner_lut_t, in_worldcoord, ign_zero): 17 | 18 | out_feature = voxlib.sp_trilinear_worldcoord(in_feature, corner_lut_t, in_worldcoord, ign_zero, -1) 19 | ctx.ign_zero = ign_zero 20 | ctx.save_for_backward(in_feature, corner_lut_t, in_worldcoord) 21 | 22 | return out_feature 23 | 24 | @staticmethod 25 | def backward(ctx, out_feature_grad): 26 | in_feature, corner_lut_t, in_worldcoord = ctx.saved_tensors 27 | 28 | assert ctx.needs_input_grad[2] is False 29 | in_feature_grad, = voxlib.sp_trilinear_worldcoord_backward( 30 | out_feature_grad, in_feature, corner_lut_t, in_worldcoord, ctx.ign_zero, False) 31 | return in_feature_grad, None, None, None, None 32 | 33 | 34 | def sparse_trilinear_interp_worldcoord(in_feature, corner_lut_t, in_worldcoord, ign_zero=False): 35 | return SparseTrilinearWorldCoordFunction.apply(in_feature, corner_lut_t, in_worldcoord, ign_zero) 36 | -------------------------------------------------------------------------------- /imaginaire/model_utils/gancraft/voxlib/voxlib.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // This work is made available under the Nvidia Source Code License-NC. 4 | // To view a copy of this license, check out LICENSE.md 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | // Fast voxel traversal along rays 11 | std::vector ray_voxel_intersection_perspective_cuda(const torch::Tensor& in_voxel, const torch::Tensor& cam_ori, const torch::Tensor& cam_dir, const torch::Tensor& cam_up, float cam_f, const std::vector& cam_c, const std::vector& img_dims, int max_samples); 12 | 13 | 14 | // World Coordinate Sparse Trilinear Interpolation 15 | torch::Tensor sp_trilinear_worldcoord_cuda(const torch::Tensor& in_feature, const torch::Tensor& in_corner_lut, const torch::Tensor& in_worldcoord, bool ign_zero, int channel_pos); 16 | 17 | std::vector sp_trilinear_worldcoord_backward_cuda(const torch::Tensor& out_feature_grad , const torch::Tensor& in_feature, const torch::Tensor& in_corner_lut, const torch::Tensor& in_worldcoord, bool ign_zero, bool need_coord_grad); 18 | 19 | // Fast & Memory Efficient Positional Encoding 20 | torch::Tensor positional_encoding_cuda(const torch::Tensor& in_feature, int ndegrees, int dim, bool incl_orig); 21 | 22 | torch::Tensor positional_encoding_backward_cuda(const torch::Tensor& out_feature_grad, const torch::Tensor& out_feature, int ndegrees, int dim, bool incl_orig); 23 | 24 | 25 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 26 | m.def("ray_voxel_intersection_perspective", &ray_voxel_intersection_perspective_cuda, "Ray-voxel intersections given perspective camera parameters (CUDA)"); 27 | m.def("sp_trilinear_worldcoord", &sp_trilinear_worldcoord_cuda, "Sparse Trilinear interpolation, world coordinate [forward] (CUDA)"); 28 | m.def("sp_trilinear_worldcoord_backward", &sp_trilinear_worldcoord_backward_cuda, "Sparse Trilinear interpolation, world coordinate [backward] (CUDA)"); 29 | m.def("positional_encoding", &positional_encoding_cuda, "Fused Positional Encoding [forward] (CUDA)"); 30 | m.def("positional_encoding_backward", &positional_encoding_backward_cuda, "Fused Positional Encoding [backward] (CUDA)"); 31 | } -------------------------------------------------------------------------------- /imaginaire/model_utils/gancraft/voxlib/voxlib_common.h: -------------------------------------------------------------------------------- 1 | // Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // This work is made available under the Nvidia Source Code License-NC. 4 | // To view a copy of this license, check out LICENSE.md 5 | #ifndef VOXLIB_COMMON_H 6 | #define VOXLIB_COMMON_H 7 | 8 | #define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor") 9 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 10 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 11 | #define CHECK_CPU(x) TORCH_CHECK(x.device().is_cpu(), #x " must be a CPU tensor") 12 | 13 | #include 14 | #include 15 | // CUDA vector math functions 16 | __host__ __device__ __forceinline__ int floor_div(int a, int b) { 17 | int c = a / b; 18 | 19 | if (c * b > a) { 20 | c--; 21 | } 22 | 23 | return c; 24 | } 25 | 26 | template 27 | __host__ __forceinline__ void cross(scalar_t* r, const scalar_t* a, const scalar_t* b) { 28 | r[0] = a[1]*b[2] - a[2]*b[1]; 29 | r[1] = a[2]*b[0] - a[0]*b[2]; 30 | r[2] = a[0]*b[1] - a[1]*b[0]; 31 | } 32 | 33 | __device__ __host__ __forceinline__ float dot(const float* a, const float* b) { 34 | return a[0] * b[0] + a[1] * b[1] + a[2] * b[2]; 35 | } 36 | 37 | template 38 | __device__ __host__ __forceinline__ void copyarr(scalar_t* r, const scalar_t* a) { 39 | #pragma unroll 40 | for (int i=0; i 48 | __device__ __host__ __forceinline__ void normalize(scalar_t* a) { 49 | scalar_t vec_len=0.0f; 50 | #pragma unroll 51 | for (int i=0; i 63 | __device__ __host__ __forceinline__ void normalize(scalar_t* r, const scalar_t* a) { 64 | scalar_t vec_len=0.0f; 65 | #pragma unroll 66 | for (int i=0; i 0.0 and d_p_norm > 0.0: 39 | p.data.add_(-group['lr'], d_p * (p_norm / d_p_norm)) 40 | else: 41 | p.data.add_(-group['lr'], d_p) 42 | p.data /= math.sqrt(1 + group['lr'] ** 2) 43 | 44 | return loss 45 | -------------------------------------------------------------------------------- /imaginaire/optimizers/madam.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch 6 | from torch.optim.optimizer import Optimizer, required 7 | 8 | 9 | class Madam(Optimizer): 10 | r"""MADAM optimizer implementation (https://arxiv.org/abs/2006.14560)""" 11 | def __init__(self, params, lr=required, scale=3.0, 12 | g_bound=None, momentum=0): 13 | self.scale = scale 14 | self.g_bound = g_bound 15 | defaults = dict(lr=lr, momentum=momentum) 16 | super(Madam, self).__init__(params, defaults) 17 | 18 | def step(self, closure=None): 19 | r"""Performs a single optimization step. 20 | 21 | Args: 22 | closure (callable, optional): A closure that reevaluates the model 23 | and returns the loss. 24 | """ 25 | loss = None 26 | if closure is not None: 27 | loss = closure() 28 | 29 | for group in self.param_groups: 30 | for p in group['params']: 31 | if p.grad is None: 32 | continue 33 | 34 | state = self.state[p] 35 | if len(state) == 0: 36 | state['max'] = self.scale * (p * p).mean().sqrt().item() 37 | state['step'] = 0 38 | state['exp_avg_sq'] = torch.zeros_like(p) 39 | 40 | state['step'] += 1 41 | bias_correction = 1 - 0.999 ** state['step'] 42 | state['exp_avg_sq'] = 0.999 * state[ 43 | 'exp_avg_sq'] + 0.001 * p.grad.data ** 2 44 | g_normed = \ 45 | p.grad.data / (state['exp_avg_sq'] / bias_correction).sqrt() 46 | g_normed[torch.isnan(g_normed)] = 0 47 | if self.g_bound is not None: 48 | g_normed.clamp_(-self.g_bound, self.g_bound) 49 | 50 | p.data *= torch.exp( 51 | -group['lr'] * g_normed * torch.sign(p.data)) 52 | p.data.clamp_(-state['max'], state['max']) 53 | 54 | return loss 55 | -------------------------------------------------------------------------------- /imaginaire/third_party/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/imaginaire/third_party/__init__.py -------------------------------------------------------------------------------- /imaginaire/third_party/bias_act/__init__.py: -------------------------------------------------------------------------------- 1 | from .bias_act import FusedNonlinearity 2 | 3 | __all__ = ['FusedNonlinearity'] 4 | -------------------------------------------------------------------------------- /imaginaire/third_party/bias_act/setup.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | from setuptools import setup 3 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 4 | import os 5 | 6 | 7 | cuda_version = os.getenv('CUDA_VERSION') 8 | print('CUDA_VERSION: {}'.format(cuda_version)) 9 | 10 | nvcc_args = list() 11 | # nvcc_args.append('-gencode') 12 | # nvcc_args.append('arch=compute_50,code=sm_50') 13 | # nvcc_args.append('-gencode') 14 | # nvcc_args.append('arch=compute_52,code=sm_52') 15 | # nvcc_args.append('-gencode') 16 | # nvcc_args.append('arch=compute_60,code=sm_60') 17 | # nvcc_args.append('-gencode') 18 | # nvcc_args.append('arch=compute_61,code=sm_61') 19 | nvcc_args.append('-gencode') 20 | nvcc_args.append('arch=compute_70,code=sm_70') 21 | nvcc_args.append('-gencode') 22 | nvcc_args.append('arch=compute_75,code=sm_75') 23 | if cuda_version is not None: 24 | if cuda_version >= '11.0': 25 | nvcc_args.append('-gencode') 26 | nvcc_args.append('arch=compute_80,code=sm_80') 27 | nvcc_args.append('-Xcompiler') 28 | nvcc_args.append('-Wall') 29 | nvcc_args.append('-std=c++14') 30 | 31 | setup( 32 | name='bias_act_cuda', 33 | py_modules=['bias_act'], 34 | ext_modules=[ 35 | CUDAExtension('bias_act_cuda', [ 36 | './src/bias_act_cuda.cc', 37 | './src/bias_act_cuda_kernel.cu' 38 | ], extra_compile_args={'cxx': ['-Wall', '-std=c++14'], 39 | 'nvcc': nvcc_args}) 40 | ], 41 | cmdclass={ 42 | 'build_ext': BuildExtension 43 | }) 44 | -------------------------------------------------------------------------------- /imaginaire/third_party/bias_act/src/bias_act_cuda.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | // 3 | // NVIDIA CORPORATION and its licensors retain all intellectual property 4 | // and proprietary rights in and to this software, related documentation 5 | // and any modifications thereto. Any use, reproduction, disclosure or 6 | // distribution of this software and related documentation without an express 7 | // license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | 9 | //------------------------------------------------------------------------ 10 | // CUDA kernel parameters. 11 | 12 | struct bias_act_kernel_params 13 | { 14 | const void* x; // [sizeX] 15 | const void* b; // [sizeB] or NULL 16 | const void* xref; // [sizeX] or NULL 17 | const void* yref; // [sizeX] or NULL 18 | const void* dy; // [sizeX] or NULL 19 | void* y; // [sizeX] 20 | 21 | int grad; 22 | int act; 23 | float alpha; 24 | float gain; 25 | float clamp; 26 | 27 | int sizeX; 28 | int sizeB; 29 | int stepB; 30 | int loopX; 31 | }; 32 | 33 | //------------------------------------------------------------------------ 34 | // CUDA kernel selection. 35 | 36 | template void* choose_bias_act_kernel(const bias_act_kernel_params& p); 37 | 38 | //------------------------------------------------------------------------ 39 | -------------------------------------------------------------------------------- /imaginaire/third_party/channelnorm/channelnorm.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | from torch.autograd import Function, Variable 3 | from torch.nn.modules.module import Module 4 | import channelnorm_cuda 5 | 6 | 7 | class ChannelNormFunction(Function): 8 | @staticmethod 9 | def forward(ctx, input1, norm_deg=2): 10 | assert input1.is_contiguous() 11 | b, _, h, w = input1.size() 12 | output = input1.new(b, 1, h, w).zero_() 13 | 14 | channelnorm_cuda.forward(input1, output, norm_deg) 15 | ctx.save_for_backward(input1, output) 16 | ctx.norm_deg = norm_deg 17 | 18 | return output 19 | 20 | @staticmethod 21 | def backward(ctx, grad_output): 22 | input1, output = ctx.saved_tensors 23 | 24 | grad_input1 = Variable(input1.new(input1.size()).zero_()) 25 | 26 | channelnorm_cuda.backward(input1, output, grad_output.data, 27 | grad_input1.data, ctx.norm_deg) 28 | 29 | return grad_input1, None 30 | 31 | 32 | class ChannelNorm(Module): 33 | 34 | def __init__(self, norm_deg=2): 35 | super(ChannelNorm, self).__init__() 36 | self.norm_deg = norm_deg 37 | 38 | def forward(self, input1): 39 | return ChannelNormFunction.apply(input1, self.norm_deg) 40 | -------------------------------------------------------------------------------- /imaginaire/third_party/channelnorm/setup.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | from setuptools import setup 3 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 4 | import os 5 | 6 | 7 | cuda_version = os.getenv('CUDA_VERSION') 8 | print('CUDA_VERSION: {}'.format(cuda_version)) 9 | 10 | nvcc_args = list() 11 | # nvcc_args.append('-gencode') 12 | # nvcc_args.append('arch=compute_50,code=sm_50') 13 | # nvcc_args.append('-gencode') 14 | # nvcc_args.append('arch=compute_52,code=sm_52') 15 | # nvcc_args.append('-gencode') 16 | # nvcc_args.append('arch=compute_60,code=sm_60') 17 | # nvcc_args.append('-gencode') 18 | # nvcc_args.append('arch=compute_61,code=sm_61') 19 | nvcc_args.append('-gencode') 20 | nvcc_args.append('arch=compute_70,code=sm_70') 21 | nvcc_args.append('-gencode') 22 | nvcc_args.append('arch=compute_75,code=sm_75') 23 | if cuda_version is not None: 24 | if cuda_version >= '11.0': 25 | nvcc_args.append('-gencode') 26 | nvcc_args.append('arch=compute_80,code=sm_80') 27 | nvcc_args.append('-Xcompiler') 28 | nvcc_args.append('-Wall') 29 | nvcc_args.append('-std=c++14') 30 | 31 | setup( 32 | name='channelnorm_cuda', 33 | py_modules=['channelnorm'], 34 | ext_modules=[ 35 | CUDAExtension('channelnorm_cuda', [ 36 | './src/channelnorm_cuda.cc', 37 | './src/channelnorm_kernel.cu' 38 | ], extra_compile_args={'cxx': ['-Wall', '-std=c++14'], 39 | 'nvcc': nvcc_args}) 40 | ], 41 | cmdclass={ 42 | 'build_ext': BuildExtension 43 | }) 44 | -------------------------------------------------------------------------------- /imaginaire/third_party/channelnorm/src/channelnorm_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "channelnorm_kernel.cuh" 5 | 6 | int channelnorm_cuda_forward( 7 | at::Tensor& input1, 8 | at::Tensor& output, 9 | int norm_deg) { 10 | 11 | channelnorm_kernel_forward(input1, output, norm_deg); 12 | return 1; 13 | } 14 | 15 | 16 | int channelnorm_cuda_backward( 17 | at::Tensor& input1, 18 | at::Tensor& output, 19 | at::Tensor& gradOutput, 20 | at::Tensor& gradInput1, 21 | int norm_deg) { 22 | 23 | channelnorm_kernel_backward(input1, output, gradOutput, gradInput1, norm_deg); 24 | return 1; 25 | } 26 | 27 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 28 | m.def("forward", &channelnorm_cuda_forward, "Channel norm forward (CUDA)"); 29 | m.def("backward", &channelnorm_cuda_backward, "Channel norm backward (CUDA)"); 30 | } 31 | 32 | -------------------------------------------------------------------------------- /imaginaire/third_party/channelnorm/src/channelnorm_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | void channelnorm_kernel_forward( 6 | at::Tensor& input1, 7 | at::Tensor& output, 8 | int norm_deg); 9 | 10 | 11 | void channelnorm_kernel_backward( 12 | at::Tensor& input1, 13 | at::Tensor& output, 14 | at::Tensor& gradOutput, 15 | at::Tensor& gradInput1, 16 | int norm_deg); 17 | -------------------------------------------------------------------------------- /imaginaire/third_party/correlation/correlation.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | import torch 3 | from torch.nn.modules.module import Module 4 | from torch.autograd import Function 5 | import correlation_cuda 6 | 7 | 8 | class CorrelationFunction(Function): 9 | 10 | @staticmethod 11 | def forward(ctx, 12 | pad_size, 13 | kernel_size, 14 | max_displacement, 15 | stride1, 16 | stride2, 17 | corr_multiply, 18 | input1, 19 | input2): 20 | ctx.save_for_backward(input1, input2) 21 | ctx.pad_size = pad_size 22 | ctx.kernel_size = kernel_size 23 | ctx.max_displacement = max_displacement 24 | ctx.stride1 = stride1 25 | ctx.stride2 = stride2 26 | ctx.corr_multiply = corr_multiply 27 | 28 | with torch.cuda.device_of(input1): 29 | rbot1 = input1.new() 30 | rbot2 = input2.new() 31 | output = input1.new() 32 | 33 | correlation_cuda.forward( 34 | input1, 35 | input2, 36 | rbot1, 37 | rbot2, 38 | output, 39 | ctx.pad_size, 40 | ctx.kernel_size, 41 | ctx.max_displacement, 42 | ctx.stride1, 43 | ctx.stride2, 44 | ctx.corr_multiply) 45 | 46 | return output 47 | 48 | @staticmethod 49 | def backward(ctx, grad_output): 50 | input1, input2 = ctx.saved_tensors 51 | 52 | with torch.cuda.device_of(input1): 53 | rbot1 = input1.new() 54 | rbot2 = input2.new() 55 | 56 | grad_input1 = input1.new() 57 | grad_input2 = input2.new() 58 | 59 | correlation_cuda.backward( 60 | input1, 61 | input2, 62 | rbot1, 63 | rbot2, 64 | grad_output, 65 | grad_input1, 66 | grad_input2, 67 | ctx.pad_size, 68 | ctx.kernel_size, 69 | ctx.max_displacement, 70 | ctx.stride1, 71 | ctx.stride2, 72 | ctx.corr_multiply) 73 | 74 | return grad_input1, grad_input2 75 | 76 | class Correlation(Module): 77 | def __init__( 78 | self, 79 | pad_size=0, 80 | kernel_size=0, 81 | max_displacement=0, 82 | stride1=1, 83 | stride2=2, 84 | corr_multiply=1): 85 | super(Correlation, self).__init__() 86 | self.pad_size = pad_size 87 | self.kernel_size = kernel_size 88 | self.max_displacement = max_displacement 89 | self.stride1 = stride1 90 | self.stride2 = stride2 91 | self.corr_multiply = corr_multiply 92 | 93 | def forward(self, input1, input2): 94 | 95 | result = CorrelationFunction.apply( 96 | self.pad_size, 97 | self.kernel_size, 98 | self.max_displacement, 99 | self.stride1, 100 | self.stride2, 101 | self.corr_multiply, 102 | input1, 103 | input2) 104 | 105 | return result 106 | -------------------------------------------------------------------------------- /imaginaire/third_party/correlation/setup.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | from setuptools import setup 3 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 4 | import os 5 | 6 | 7 | cuda_version = os.getenv('CUDA_VERSION') 8 | print('CUDA_VERSION: {}'.format(cuda_version)) 9 | 10 | nvcc_args = list() 11 | # nvcc_args.append('-gencode') 12 | # nvcc_args.append('arch=compute_50,code=sm_50') 13 | # nvcc_args.append('-gencode') 14 | # nvcc_args.append('arch=compute_52,code=sm_52') 15 | # nvcc_args.append('-gencode') 16 | # nvcc_args.append('arch=compute_60,code=sm_60') 17 | # nvcc_args.append('-gencode') 18 | # nvcc_args.append('arch=compute_61,code=sm_61') 19 | nvcc_args.append('-gencode') 20 | nvcc_args.append('arch=compute_70,code=sm_70') 21 | nvcc_args.append('-gencode') 22 | nvcc_args.append('arch=compute_75,code=sm_75') 23 | if cuda_version is not None: 24 | if cuda_version >= '11.0': 25 | nvcc_args.append('-gencode') 26 | nvcc_args.append('arch=compute_80,code=sm_80') 27 | nvcc_args.append('-Xcompiler') 28 | nvcc_args.append('-Wall') 29 | nvcc_args.append('-std=c++14') 30 | 31 | setup( 32 | name='correlation_cuda', 33 | py_modules=['correlation'], 34 | ext_modules=[ 35 | CUDAExtension('correlation_cuda', [ 36 | './src/correlation_cuda.cc', 37 | './src/correlation_cuda_kernel.cu' 38 | ], extra_compile_args={'cxx': ['-Wall', '-std=c++14'], 39 | 'nvcc': nvcc_args}) 40 | ], 41 | cmdclass={ 42 | 'build_ext': BuildExtension 43 | }) 44 | -------------------------------------------------------------------------------- /imaginaire/third_party/correlation/src/correlation_cuda_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | int correlation_forward_cuda_kernel(at::Tensor& output, 8 | int ob, 9 | int oc, 10 | int oh, 11 | int ow, 12 | int osb, 13 | int osc, 14 | int osh, 15 | int osw, 16 | 17 | at::Tensor& input1, 18 | int ic, 19 | int ih, 20 | int iw, 21 | int isb, 22 | int isc, 23 | int ish, 24 | int isw, 25 | 26 | at::Tensor& input2, 27 | int gc, 28 | int gsb, 29 | int gsc, 30 | int gsh, 31 | int gsw, 32 | 33 | at::Tensor& rInput1, 34 | at::Tensor& rInput2, 35 | int pad_size, 36 | int kernel_size, 37 | int max_displacement, 38 | int stride1, 39 | int stride2, 40 | int corr_type_multiply, 41 | cudaStream_t stream); 42 | 43 | 44 | int correlation_backward_cuda_kernel( 45 | at::Tensor& gradOutput, 46 | int gob, 47 | int goc, 48 | int goh, 49 | int gow, 50 | int gosb, 51 | int gosc, 52 | int gosh, 53 | int gosw, 54 | 55 | at::Tensor& input1, 56 | int ic, 57 | int ih, 58 | int iw, 59 | int isb, 60 | int isc, 61 | int ish, 62 | int isw, 63 | 64 | at::Tensor& input2, 65 | int gsb, 66 | int gsc, 67 | int gsh, 68 | int gsw, 69 | 70 | at::Tensor& gradInput1, 71 | int gisb, 72 | int gisc, 73 | int gish, 74 | int gisw, 75 | 76 | at::Tensor& gradInput2, 77 | int ggc, 78 | int ggsb, 79 | int ggsc, 80 | int ggsh, 81 | int ggsw, 82 | 83 | at::Tensor& rInput1, 84 | at::Tensor& rInput2, 85 | int pad_size, 86 | int kernel_size, 87 | int max_displacement, 88 | int stride1, 89 | int stride2, 90 | int corr_type_multiply, 91 | cudaStream_t stream); 92 | -------------------------------------------------------------------------------- /imaginaire/third_party/flow_net/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/imaginaire/third_party/flow_net/__init__.py -------------------------------------------------------------------------------- /imaginaire/third_party/flow_net/flownet2/networks/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | # The file is duplicated from https://github.com/NVIDIA/flownet2-pytorch 6 | # with some modifications. 7 | -------------------------------------------------------------------------------- /imaginaire/third_party/flow_net/flownet2/networks/flownet_fusion.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | # The file is duplicated from https://github.com/NVIDIA/flownet2-pytorch 6 | # with some modifications. 7 | from torch.nn import init 8 | import torch 9 | import torch.nn as nn 10 | from .submodules import conv, i_conv, predict_flow, deconv 11 | 12 | 13 | class FlowNetFusion(nn.Module): 14 | r"""FlowNet2 Fusion module. Check out the FlowNet2 paper for more details 15 | https://arxiv.org/abs/1612.01925 16 | 17 | Args: 18 | args (obj): Network initialization arguments 19 | use_batch_norm (bool): Use batch norm or not. Default is true. 20 | """ 21 | def __init__(self, args, use_batch_norm=True): 22 | super(FlowNetFusion, self).__init__() 23 | 24 | self.use_batch_norm = use_batch_norm 25 | self.conv0 = conv(self.use_batch_norm, 11, 64) 26 | self.conv1 = conv(self.use_batch_norm, 64, 64, stride=2) 27 | self.conv1_1 = conv(self.use_batch_norm, 64, 128) 28 | self.conv2 = conv(self.use_batch_norm, 128, 128, stride=2) 29 | self.conv2_1 = conv(self.use_batch_norm, 128, 128) 30 | 31 | self.deconv1 = deconv(128, 32) 32 | self.deconv0 = deconv(162, 16) 33 | 34 | self.inter_conv1 = i_conv(self.use_batch_norm, 162, 32) 35 | self.inter_conv0 = i_conv(self.use_batch_norm, 82, 16) 36 | 37 | self.predict_flow2 = predict_flow(128) 38 | self.predict_flow1 = predict_flow(32) 39 | self.predict_flow0 = predict_flow(16) 40 | 41 | self.upsampled_flow2_to_1 = nn.ConvTranspose2d(2, 2, 4, 2, 1) 42 | self.upsampled_flow1_to_0 = nn.ConvTranspose2d(2, 2, 4, 2, 1) 43 | 44 | for m in self.modules(): 45 | if isinstance(m, nn.Conv2d): 46 | if m.bias is not None: 47 | init.uniform_(m.bias) 48 | init.xavier_uniform_(m.weight) 49 | 50 | if isinstance(m, nn.ConvTranspose2d): 51 | if m.bias is not None: 52 | init.uniform_(m.bias) 53 | init.xavier_uniform_(m.weight) 54 | # init_deconv_bilinear(m.weight) 55 | 56 | def forward(self, x): 57 | r""" 58 | 59 | Args: 60 | x (tensor): Input tensors of concatenated images. 61 | Returns: 62 | flow2 (tensor): Output flow tensors. 63 | """ 64 | out_conv0 = self.conv0(x) 65 | out_conv1 = self.conv1_1(self.conv1(out_conv0)) 66 | out_conv2 = self.conv2_1(self.conv2(out_conv1)) 67 | 68 | flow2 = self.predict_flow2(out_conv2) 69 | flow2_up = self.upsampled_flow2_to_1(flow2) 70 | out_deconv1 = self.deconv1(out_conv2) 71 | 72 | concat1 = torch.cat((out_conv1, out_deconv1, flow2_up), 1) 73 | out_interconv1 = self.inter_conv1(concat1) 74 | flow1 = self.predict_flow1(out_interconv1) 75 | flow1_up = self.upsampled_flow1_to_0(flow1) 76 | out_deconv0 = self.deconv0(concat1) 77 | 78 | concat0 = torch.cat((out_conv0, out_deconv0, flow1_up), 1) 79 | out_interconv0 = self.inter_conv0(concat0) 80 | flow0 = self.predict_flow0(out_interconv0) 81 | 82 | return flow0 83 | -------------------------------------------------------------------------------- /imaginaire/third_party/flow_net/flownet2/networks/submodules.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | # The file is duplicated from https://github.com/NVIDIA/flownet2-pytorch 6 | # with some modifications. 7 | import torch 8 | import torch.nn as nn 9 | import numpy as np 10 | 11 | 12 | def conv(use_batch_norm, in_planes, out_planes, kernel_size=3, stride=1): 13 | if use_batch_norm: 14 | return nn.Sequential( 15 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, 16 | stride=stride, padding=(kernel_size - 1) // 2, 17 | bias=False), 18 | nn.BatchNorm2d(out_planes), 19 | nn.LeakyReLU(0.1, inplace=True) 20 | ) 21 | else: 22 | return nn.Sequential( 23 | nn.Conv2d( 24 | in_planes, 25 | out_planes, 26 | kernel_size=kernel_size, 27 | stride=stride, 28 | padding=( 29 | kernel_size - 1) // 2, 30 | bias=True), 31 | nn.LeakyReLU( 32 | 0.1, 33 | inplace=True)) 34 | 35 | 36 | def i_conv(use_batch_norm, in_planes, out_planes, kernel_size=3, stride=1, 37 | bias=True): 38 | if use_batch_norm: 39 | return nn.Sequential( 40 | nn.Conv2d( 41 | in_planes, 42 | out_planes, 43 | kernel_size=kernel_size, 44 | stride=stride, 45 | padding=( 46 | kernel_size - 1) // 2, 47 | bias=bias), 48 | nn.BatchNorm2d(out_planes), 49 | ) 50 | else: 51 | return nn.Sequential( 52 | nn.Conv2d( 53 | in_planes, 54 | out_planes, 55 | kernel_size=kernel_size, 56 | stride=stride, 57 | padding=( 58 | kernel_size - 59 | 1) // 60 | 2, 61 | bias=bias), 62 | ) 63 | 64 | 65 | def predict_flow(in_planes): 66 | return nn.Conv2d(in_planes, 2, kernel_size=3, stride=1, padding=1, 67 | bias=True) 68 | 69 | 70 | def deconv(in_planes, out_planes): 71 | return nn.Sequential( 72 | nn.ConvTranspose2d(in_planes, out_planes, kernel_size=4, stride=2, 73 | padding=1, bias=True), 74 | nn.LeakyReLU(0.1, inplace=True) 75 | ) 76 | 77 | 78 | class tofp16(nn.Module): 79 | def __init__(self): 80 | super(tofp16, self).__init__() 81 | 82 | def forward(self, input): 83 | return input.half() 84 | 85 | 86 | class tofp32(nn.Module): 87 | def __init__(self): 88 | super(tofp32, self).__init__() 89 | 90 | def forward(self, input): 91 | return input.float() 92 | 93 | 94 | def init_deconv_bilinear(weight): 95 | f_shape = weight.size() 96 | heigh, width = f_shape[-2], f_shape[-1] 97 | f = np.ceil(width / 2.0) 98 | c = (2 * f - 1 - f % 2) / (2.0 * f) 99 | bilinear = np.zeros([heigh, width]) 100 | for x in range(width): 101 | for y in range(heigh): 102 | value = (1 - abs(x / f - c)) * (1 - abs(y / f - c)) 103 | bilinear[x, y] = value 104 | weight.data.fill_(0.) 105 | for i in range(f_shape[0]): 106 | for j in range(f_shape[1]): 107 | weight.data[i, j, :, :] = torch.from_numpy(bilinear) 108 | 109 | 110 | def save_grad(grads, name): 111 | def hook(grad): 112 | grads[name] = grad 113 | return hook 114 | -------------------------------------------------------------------------------- /imaginaire/third_party/flow_net/flownet2/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | -------------------------------------------------------------------------------- /imaginaire/third_party/flow_net/flownet2/utils/frame_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import numpy as np 6 | from os.path import splitext 7 | from scipy.misc import imread 8 | from . import flow_utils 9 | 10 | 11 | def read_gen(file_name): 12 | ext = splitext(file_name)[-1] 13 | if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg': 14 | im = imread(file_name) 15 | if im.shape[2] > 3: 16 | return im[:, :, :3] 17 | else: 18 | return im 19 | elif ext == '.bin' or ext == '.raw': 20 | return np.load(file_name) 21 | elif ext == '.flo': 22 | return flow_utils.readFlow(file_name).astype(np.float32) 23 | return [] 24 | -------------------------------------------------------------------------------- /imaginaire/third_party/flow_net/test_flownet2.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa E402 2 | import sys 3 | sys.path.append('../../../tests/third_party/flow_net') 4 | from PIL import Image 5 | from torchvision.transforms.functional import to_tensor 6 | from imaginaire.third_party.flow_net.flow_net import FlowNet 7 | import wget 8 | import os 9 | 10 | im1_fn = 'frame_0010.png' 11 | im2_fn = 'frame_0011.png' 12 | im1_url = 'https://github.com/NVlabs/PWC-Net/raw/master/PyTorch/data/' + im1_fn 13 | im2_url = 'https://github.com/NVlabs/PWC-Net/raw/master/PyTorch/data/' + im2_fn 14 | if not os.path.exists(im1_fn): 15 | wget.download(im1_url, out=im1_fn) 16 | if not os.path.exists(im2_fn): 17 | wget.download(im2_url, out=im2_fn) 18 | img1 = to_tensor(Image.open(im1_fn).convert('RGB')).unsqueeze(0).cuda() 19 | img2 = to_tensor(Image.open(im2_fn).convert('RGB')).unsqueeze(0).cuda() 20 | # Image range between -1 and 1 21 | img1 = 2 * img1 - 1 22 | img2 = 2 * img2 - 1 23 | flow_net = FlowNet() 24 | flow = flow_net.cuda().eval() 25 | flo, conf = flow(img1, img2) 26 | print(flo.size()) 27 | print(conf.size()) 28 | -------------------------------------------------------------------------------- /imaginaire/third_party/resample2d/resample2d.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | from torch.nn.modules.module import Module 3 | from torch.autograd import Function, Variable 4 | from torch.cuda.amp import autocast 5 | import resample2d_cuda 6 | 7 | 8 | class Resample2dFunction(Function): 9 | 10 | @staticmethod 11 | # def forward(ctx, input1, input2, kernel_size=1, bilinear=True): 12 | def forward(ctx, input1, input2, kernel_size=1): 13 | assert input1.is_contiguous() 14 | assert input2.is_contiguous() 15 | 16 | ctx.save_for_backward(input1, input2) 17 | ctx.kernel_size = kernel_size 18 | ctx.bilinear = True 19 | 20 | _, d, _, _ = input1.size() 21 | b, _, h, w = input2.size() 22 | output = input1.new(b, d, h, w).zero_() 23 | 24 | resample2d_cuda.forward(input1, input2, output, kernel_size) 25 | 26 | return output 27 | 28 | @staticmethod 29 | def backward(ctx, grad_output): 30 | grad_output = grad_output.contiguous() 31 | assert grad_output.is_contiguous() 32 | 33 | input1, input2 = ctx.saved_tensors 34 | 35 | grad_input1 = Variable(input1.new(input1.size()).zero_()) 36 | grad_input2 = Variable(input1.new(input2.size()).zero_()) 37 | 38 | # resample2d_cuda.backward(input1, input2, grad_output.data, 39 | # grad_input1.data, grad_input2.data, 40 | # ctx.kernel_size, ctx.bilinear) 41 | resample2d_cuda.backward(input1, input2, grad_output.data, 42 | grad_input1.data, grad_input2.data, 43 | ctx.kernel_size) 44 | 45 | return grad_input1, grad_input2, None, None 46 | 47 | 48 | class Resample2d(Module): 49 | 50 | def __init__(self, kernel_size=1, bilinear=True): 51 | super(Resample2d, self).__init__() 52 | self.kernel_size = kernel_size 53 | self.bilinear = bilinear 54 | 55 | @autocast(False) 56 | def forward(self, input1, input2): 57 | input1, input2 = input1.float(), input2.float() 58 | input1_c = input1.contiguous() 59 | # return Resample2dFunction.apply( 60 | # input1_c, input2, self.kernel_size, self.bilinear) 61 | return Resample2dFunction.apply( 62 | input1_c, input2, self.kernel_size) -------------------------------------------------------------------------------- /imaginaire/third_party/resample2d/setup.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | from setuptools import setup 3 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 4 | import os 5 | 6 | 7 | cuda_version = os.getenv('CUDA_VERSION') 8 | print('CUDA_VERSION: {}'.format(cuda_version)) 9 | 10 | nvcc_args = list() 11 | # nvcc_args.append('-gencode') 12 | # nvcc_args.append('arch=compute_50,code=sm_50') 13 | # nvcc_args.append('-gencode') 14 | # nvcc_args.append('arch=compute_52,code=sm_52') 15 | # nvcc_args.append('-gencode') 16 | # nvcc_args.append('arch=compute_60,code=sm_60') 17 | # nvcc_args.append('-gencode') 18 | # nvcc_args.append('arch=compute_61,code=sm_61') 19 | nvcc_args.append('-gencode') 20 | nvcc_args.append('arch=compute_70,code=sm_70') 21 | nvcc_args.append('-gencode') 22 | nvcc_args.append('arch=compute_75,code=sm_75') 23 | if cuda_version is not None: 24 | if cuda_version >= '11.0': 25 | nvcc_args.append('-gencode') 26 | nvcc_args.append('arch=compute_80,code=sm_80') 27 | nvcc_args.append('-Xcompiler') 28 | nvcc_args.append('-Wall') 29 | nvcc_args.append('-std=c++14') 30 | 31 | setup( 32 | name='resample2d_cuda', 33 | py_modules=['resample2d'], 34 | ext_modules=[ 35 | CUDAExtension('resample2d_cuda', [ 36 | './src/resample2d_cuda.cc', 37 | './src/resample2d_kernel.cu' 38 | ], extra_compile_args={'cxx': ['-Wall', '-std=c++14'], 39 | 'nvcc': nvcc_args}) 40 | ], 41 | cmdclass={ 42 | 'build_ext': BuildExtension 43 | }) 44 | -------------------------------------------------------------------------------- /imaginaire/third_party/resample2d/src/resample2d_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "resample2d_kernel.cuh" 5 | 6 | int resample2d_cuda_forward( 7 | at::Tensor& input1, 8 | at::Tensor& input2, 9 | at::Tensor& output, 10 | int kernel_size/*, bool bilinear*/) { 11 | resample2d_kernel_forward(input1, input2, output, kernel_size/*, 12 | bilinear*/); 13 | return 1; 14 | } 15 | 16 | int resample2d_cuda_backward( 17 | at::Tensor& input1, 18 | at::Tensor& input2, 19 | at::Tensor& gradOutput, 20 | at::Tensor& gradInput1, 21 | at::Tensor& gradInput2, 22 | int kernel_size/*, bool bilinear*/) { 23 | resample2d_kernel_backward(input1, input2, gradOutput, gradInput1, 24 | gradInput2, kernel_size/*, bilinear*/); 25 | return 1; 26 | } 27 | 28 | 29 | 30 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 31 | m.def("forward", &resample2d_cuda_forward, "Resample2D forward (CUDA)"); 32 | m.def("backward", &resample2d_cuda_backward, "Resample2D backward (CUDA)"); 33 | } 34 | 35 | -------------------------------------------------------------------------------- /imaginaire/third_party/resample2d/src/resample2d_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | void resample2d_kernel_forward( 6 | at::Tensor& input1, 7 | at::Tensor& input2, 8 | at::Tensor& output, 9 | int kernel_size/*, 10 | bool bilinear*/); 11 | 12 | void resample2d_kernel_backward( 13 | at::Tensor& input1, 14 | at::Tensor& input2, 15 | at::Tensor& gradOutput, 16 | at::Tensor& gradInput1, 17 | at::Tensor& gradInput2, 18 | int kernel_size/*, 19 | bool bilinear*/); -------------------------------------------------------------------------------- /imaginaire/third_party/upfirdn2d/__init__.py: -------------------------------------------------------------------------------- 1 | from .upfirdn2d import BlurUpsample, BlurDownsample, Blur 2 | 3 | __all__ = ['BlurUpsample', 'BlurDownsample', 'Blur'] 4 | -------------------------------------------------------------------------------- /imaginaire/third_party/upfirdn2d/setup.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | from setuptools import setup 3 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 4 | import os 5 | 6 | 7 | cuda_version = os.getenv('CUDA_VERSION') 8 | print('CUDA_VERSION: {}'.format(cuda_version)) 9 | 10 | nvcc_args = list() 11 | # nvcc_args.append('-gencode') 12 | # nvcc_args.append('arch=compute_50,code=sm_50') 13 | # nvcc_args.append('-gencode') 14 | # nvcc_args.append('arch=compute_52,code=sm_52') 15 | # nvcc_args.append('-gencode') 16 | # nvcc_args.append('arch=compute_60,code=sm_60') 17 | # nvcc_args.append('-gencode') 18 | # nvcc_args.append('arch=compute_61,code=sm_61') 19 | nvcc_args.append('-gencode') 20 | nvcc_args.append('arch=compute_70,code=sm_70') 21 | nvcc_args.append('-gencode') 22 | nvcc_args.append('arch=compute_75,code=sm_75') 23 | if cuda_version is not None: 24 | if cuda_version >= '11.0': 25 | nvcc_args.append('-gencode') 26 | nvcc_args.append('arch=compute_80,code=sm_80') 27 | nvcc_args.append('-Xcompiler') 28 | nvcc_args.append('-Wall') 29 | nvcc_args.append('-std=c++14') 30 | 31 | setup( 32 | name='upfirdn2d_cuda', 33 | py_modules=['upfirdn2d'], 34 | ext_modules=[ 35 | CUDAExtension('upfirdn2d_cuda', [ 36 | './src/upfirdn2d_cuda.cc', 37 | './src/upfirdn2d_cuda_kernel.cu' 38 | ], extra_compile_args={'cxx': ['-Wall', '-std=c++14'], 39 | 'nvcc': nvcc_args}) 40 | ], 41 | cmdclass={ 42 | 'build_ext': BuildExtension 43 | }) 44 | -------------------------------------------------------------------------------- /imaginaire/third_party/upfirdn2d/src/upfirdn2d_cuda.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | // 3 | // NVIDIA CORPORATION and its licensors retain all intellectual property 4 | // and proprietary rights in and to this software, related documentation 5 | // and any modifications thereto. Any use, reproduction, disclosure or 6 | // distribution of this software and related documentation without an express 7 | // license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | 9 | #include 10 | 11 | //------------------------------------------------------------------------ 12 | // CUDA kernel parameters. 13 | 14 | struct upfirdn2d_kernel_params 15 | { 16 | const void* x; 17 | const float* f; 18 | void* y; 19 | 20 | int2 up; 21 | int2 down; 22 | int2 pad0; 23 | int flip; 24 | float gain; 25 | 26 | int4 inSize; // [width, height, channel, batch] 27 | int4 inStride; 28 | int2 filterSize; // [width, height] 29 | int2 filterStride; 30 | int4 outSize; // [width, height, channel, batch] 31 | int4 outStride; 32 | int sizeMinor; 33 | int sizeMajor; 34 | 35 | int loopMinor; 36 | int loopMajor; 37 | int loopX; 38 | int launchMinor; 39 | int launchMajor; 40 | }; 41 | 42 | //------------------------------------------------------------------------ 43 | // CUDA kernel specialization. 44 | 45 | struct upfirdn2d_kernel_spec 46 | { 47 | void* kernel; 48 | int tileOutW; 49 | int tileOutH; 50 | int loopMinor; 51 | int loopX; 52 | }; 53 | 54 | //------------------------------------------------------------------------ 55 | // CUDA kernel selection. 56 | 57 | template upfirdn2d_kernel_spec choose_upfirdn2d_kernel(const upfirdn2d_kernel_params& p); 58 | 59 | //------------------------------------------------------------------------ 60 | -------------------------------------------------------------------------------- /imaginaire/trainers/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | -------------------------------------------------------------------------------- /imaginaire/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | -------------------------------------------------------------------------------- /imaginaire/utils/cudnn.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch.backends.cudnn as cudnn 6 | 7 | from imaginaire.utils.distributed import master_only_print as print 8 | 9 | 10 | def init_cudnn(deterministic, benchmark): 11 | r"""Initialize the cudnn module. The two things to consider is whether to 12 | use cudnn benchmark and whether to use cudnn deterministic. If cudnn 13 | benchmark is set, then the cudnn deterministic is automatically false. 14 | 15 | Args: 16 | deterministic (bool): Whether to use cudnn deterministic. 17 | benchmark (bool): Whether to use cudnn benchmark. 18 | """ 19 | cudnn.deterministic = deterministic 20 | cudnn.benchmark = benchmark 21 | print('cudnn benchmark: {}'.format(benchmark)) 22 | print('cudnn deterministic: {}'.format(deterministic)) 23 | -------------------------------------------------------------------------------- /imaginaire/utils/distributed.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import functools 6 | import ctypes 7 | 8 | import torch 9 | import torch.distributed as dist 10 | 11 | 12 | def init_dist(local_rank, backend='nccl', **kwargs): 13 | r"""Initialize distributed training""" 14 | if dist.is_available(): 15 | if dist.is_initialized(): 16 | return torch.cuda.current_device() 17 | torch.cuda.set_device(local_rank) 18 | dist.init_process_group(backend=backend, init_method='env://', **kwargs) 19 | 20 | # Increase the L2 fetch granularity for faster speed. 21 | _libcudart = ctypes.CDLL('libcudart.so') 22 | # Set device limit on the current device 23 | # cudaLimitMaxL2FetchGranularity = 0x05 24 | pValue = ctypes.cast((ctypes.c_int * 1)(), ctypes.POINTER(ctypes.c_int)) 25 | _libcudart.cudaDeviceSetLimit(ctypes.c_int(0x05), ctypes.c_int(128)) 26 | _libcudart.cudaDeviceGetLimit(pValue, ctypes.c_int(0x05)) 27 | # assert pValue.contents.value == 128 28 | 29 | 30 | def get_rank(): 31 | r"""Get rank of the thread.""" 32 | rank = 0 33 | if dist.is_available(): 34 | if dist.is_initialized(): 35 | rank = dist.get_rank() 36 | return rank 37 | 38 | 39 | def get_world_size(): 40 | r"""Get world size. How many GPUs are available in this job.""" 41 | world_size = 1 42 | if dist.is_available(): 43 | if dist.is_initialized(): 44 | world_size = dist.get_world_size() 45 | return world_size 46 | 47 | 48 | def master_only(func): 49 | r"""Apply this function only to the master GPU.""" 50 | @functools.wraps(func) 51 | def wrapper(*args, **kwargs): 52 | r"""Simple function wrapper for the master function""" 53 | if get_rank() == 0: 54 | return func(*args, **kwargs) 55 | else: 56 | return None 57 | return wrapper 58 | 59 | 60 | def is_master(): 61 | r"""check if current process is the master""" 62 | return get_rank() == 0 63 | 64 | 65 | def is_local_master(): 66 | return torch.cuda.current_device() == 0 67 | 68 | 69 | @master_only 70 | def master_only_print(*args): 71 | r"""master-only print""" 72 | print(*args) 73 | 74 | 75 | def dist_reduce_tensor(tensor, rank=0, reduce='mean'): 76 | r""" Reduce to rank 0 """ 77 | world_size = get_world_size() 78 | if world_size < 2: 79 | return tensor 80 | with torch.no_grad(): 81 | dist.reduce(tensor, dst=rank) 82 | if get_rank() == rank: 83 | if reduce == 'mean': 84 | tensor /= world_size 85 | elif reduce == 'sum': 86 | pass 87 | else: 88 | raise NotImplementedError 89 | return tensor 90 | 91 | 92 | def dist_all_reduce_tensor(tensor, reduce='mean'): 93 | r""" Reduce to all ranks """ 94 | world_size = get_world_size() 95 | if world_size < 2: 96 | return tensor 97 | with torch.no_grad(): 98 | dist.all_reduce(tensor) 99 | if reduce == 'mean': 100 | tensor /= world_size 101 | elif reduce == 'sum': 102 | pass 103 | else: 104 | raise NotImplementedError 105 | return tensor 106 | 107 | 108 | def dist_all_gather_tensor(tensor): 109 | r""" gather to all ranks """ 110 | world_size = get_world_size() 111 | if world_size < 2: 112 | return [tensor] 113 | tensor_list = [ 114 | torch.ones_like(tensor) for _ in range(dist.get_world_size())] 115 | with torch.no_grad(): 116 | dist.all_gather(tensor_list, tensor) 117 | return tensor_list 118 | -------------------------------------------------------------------------------- /imaginaire/utils/gpu_affinity.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import math 6 | import os 7 | import pynvml 8 | 9 | pynvml.nvmlInit() 10 | 11 | 12 | def systemGetDriverVersion(): 13 | r"""Get Driver Version""" 14 | return pynvml.nvmlSystemGetDriverVersion() 15 | 16 | 17 | def deviceGetCount(): 18 | r"""Get number of devices""" 19 | return pynvml.nvmlDeviceGetCount() 20 | 21 | 22 | class device(object): 23 | r"""Device used for nvml.""" 24 | _nvml_affinity_elements = math.ceil(os.cpu_count() / 64) 25 | 26 | def __init__(self, device_idx): 27 | super().__init__() 28 | self.handle = pynvml.nvmlDeviceGetHandleByIndex(device_idx) 29 | 30 | def getName(self): 31 | r"""Get obect name""" 32 | return pynvml.nvmlDeviceGetName(self.handle) 33 | 34 | def getCpuAffinity(self): 35 | r"""Get CPU affinity""" 36 | affinity_string = '' 37 | for j in pynvml.nvmlDeviceGetCpuAffinity( 38 | self.handle, device._nvml_affinity_elements): 39 | # assume nvml returns list of 64 bit ints 40 | affinity_string = '{:064b}'.format(j) + affinity_string 41 | affinity_list = [int(x) for x in affinity_string] 42 | affinity_list.reverse() # so core 0 is in 0th element of list 43 | 44 | return [i for i, e in enumerate(affinity_list) if e != 0] 45 | 46 | 47 | def set_affinity(gpu_id=None): 48 | r"""Set GPU affinity 49 | 50 | Args: 51 | gpu_id (int): Which gpu device. 52 | """ 53 | if gpu_id is None: 54 | gpu_id = int(os.getenv('LOCAL_RANK', 0)) 55 | 56 | dev = device(gpu_id) 57 | os.sched_setaffinity(0, dev.getCpuAffinity()) 58 | 59 | # list of ints 60 | # representing the logical cores this process is now affinitied with 61 | return os.sched_getaffinity(0) 62 | -------------------------------------------------------------------------------- /imaginaire/utils/init_weight.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import torch 6 | from torch.nn import init 7 | 8 | 9 | def weights_init(init_type='normal', gain=0.02, bias=None): 10 | r"""Initialize weights in the network. 11 | 12 | Args: 13 | init_type (str): The name of the initialization scheme. 14 | gain (float): The parameter that is required for the initialization 15 | scheme. 16 | bias (object): If not ``None``, specifies the initialization parameter 17 | for bias. 18 | 19 | Returns: 20 | (obj): init function to be applied. 21 | """ 22 | 23 | def init_func(m): 24 | r"""Init function 25 | 26 | Args: 27 | m: module to be weight initialized. 28 | """ 29 | class_name = m.__class__.__name__ 30 | if hasattr(m, 'weight') and ( 31 | class_name.find('Conv') != -1 or 32 | class_name.find('Linear') != -1 or 33 | class_name.find('Embedding') != -1): 34 | lr_mul = getattr(m, 'lr_mul', 1.) 35 | gain_final = gain / lr_mul 36 | if init_type == 'normal': 37 | init.normal_(m.weight.data, 0.0, gain_final) 38 | elif init_type == 'xavier': 39 | init.xavier_normal_(m.weight.data, gain=gain_final) 40 | elif init_type == 'xavier_uniform': 41 | init.xavier_uniform_(m.weight.data, gain=gain_final) 42 | elif init_type == 'kaiming': 43 | init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') 44 | with torch.no_grad(): 45 | m.weight.data *= gain_final 46 | elif init_type == 'kaiming_linear': 47 | init.kaiming_normal_( 48 | m.weight.data, a=0, mode='fan_in', nonlinearity='linear' 49 | ) 50 | with torch.no_grad(): 51 | m.weight.data *= gain_final 52 | elif init_type == 'orthogonal': 53 | init.orthogonal_(m.weight.data, gain=gain_final) 54 | elif init_type == 'none': 55 | pass 56 | # m.reset_parameters() 57 | else: 58 | raise NotImplementedError( 59 | 'initialization method [%s] is ' 60 | 'not implemented' % init_type) 61 | if hasattr(m, 'bias') and m.bias is not None: 62 | if init_type == 'none': 63 | pass 64 | elif bias is not None: 65 | bias_type = getattr(bias, 'type', 'normal') 66 | if bias_type == 'normal': 67 | bias_gain = getattr(bias, 'gain', 0.5) 68 | init.normal_(m.bias.data, 0.0, bias_gain) 69 | else: 70 | raise NotImplementedError( 71 | 'initialization method [%s] is ' 72 | 'not implemented' % bias_type) 73 | else: 74 | init.constant_(m.bias.data, 0.0) 75 | return init_func 76 | 77 | 78 | def weights_rescale(): 79 | def init_func(m): 80 | if hasattr(m, 'init_gain'): 81 | for name, p in m.named_parameters(): 82 | if 'output_scale' not in name: 83 | p.data.mul_(m.init_gain) 84 | return init_func 85 | -------------------------------------------------------------------------------- /imaginaire/utils/logging.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | import datetime 6 | import os 7 | 8 | from imaginaire.utils.distributed import master_only 9 | from imaginaire.utils.distributed import master_only_print as print 10 | from imaginaire.utils.meters import set_summary_writer 11 | 12 | 13 | def get_date_uid(): 14 | """Generate a unique id based on date. 15 | Returns: 16 | str: Return uid string, e.g. '20171122171307111552'. 17 | """ 18 | return str(datetime.datetime.now().strftime("%Y_%m%d_%H%M_%S")) 19 | 20 | 21 | def init_logging(config_path, logdir): 22 | r"""Create log directory for storing checkpoints and output images. 23 | 24 | Args: 25 | config_path (str): Path to the configuration file. 26 | logdir (str): Log directory name 27 | Returns: 28 | str: Return log dir 29 | """ 30 | config_file = os.path.basename(config_path) 31 | root_dir = 'logs' 32 | date_uid = get_date_uid() 33 | # example: logs/2019_0125_1047_58_spade_cocostuff 34 | log_file = '_'.join([date_uid, os.path.splitext(config_file)[0]]) 35 | if logdir is None: 36 | logdir = os.path.join(root_dir, log_file) 37 | return date_uid, logdir 38 | 39 | 40 | @master_only 41 | def make_logging_dir(logdir): 42 | r"""Create the logging directory 43 | 44 | Args: 45 | logdir (str): Log directory name 46 | """ 47 | print('Make folder {}'.format(logdir)) 48 | os.makedirs(logdir, exist_ok=True) 49 | tensorboard_dir = os.path.join(logdir, 'tensorboard') 50 | os.makedirs(tensorboard_dir, exist_ok=True) 51 | set_summary_writer(tensorboard_dir) 52 | -------------------------------------------------------------------------------- /imaginaire/utils/path.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | """Utils to deal with directories and paths.""" 6 | 7 | import glob 8 | import os 9 | 10 | 11 | def get_immediate_subdirectories(input_dir): 12 | """List dirs immediately under input_dir. 13 | 14 | Args: 15 | input_dir (str): Directory to list children of. 16 | Returns: 17 | (list): List of directory paths relative to input_dir. 18 | """ 19 | return sorted([name for name in os.listdir(input_dir) 20 | if os.path.isdir(os.path.join(input_dir, name))]) 21 | 22 | 23 | def get_recursive_subdirectories(input_dir, ext): 24 | """List dirs recursively under input_dir. 25 | 26 | Args: 27 | input_dir (str): Directory to list children of. 28 | ext (str): Extension of files expected in this directory. 29 | Returns: 30 | (list): List of directory paths relative to input_dir. 31 | """ 32 | lines = glob.glob('%s/**/*.%s' % (input_dir, ext), recursive=True) 33 | dirpaths = [os.path.dirname(item) for item in lines] 34 | dirpaths = [os.path.relpath(item, input_dir) for item in dirpaths] 35 | dirpaths = sorted(list(set(dirpaths))) 36 | return dirpaths 37 | -------------------------------------------------------------------------------- /imaginaire/utils/visualization/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, check out LICENSE.md 5 | from .common import tensor2im, tensor2flow, tensor2label, tensor2pilimage 6 | from .common import save_tensor_image 7 | 8 | __all__ = ['tensor2im', 'tensor2flow', 'tensor2label', 'tensor2pilimage', 9 | 'save_tensor_image'] 10 | -------------------------------------------------------------------------------- /projects/coco_funit/animal_faces_content.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/coco_funit/animal_faces_content.jpg -------------------------------------------------------------------------------- /projects/coco_funit/animal_faces_output.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/coco_funit/animal_faces_output.jpg -------------------------------------------------------------------------------- /projects/coco_funit/animal_faces_style.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/coco_funit/animal_faces_style.jpg -------------------------------------------------------------------------------- /projects/coco_funit/mammals_content.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/coco_funit/mammals_content.png -------------------------------------------------------------------------------- /projects/coco_funit/mammals_output.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/coco_funit/mammals_output.jpg -------------------------------------------------------------------------------- /projects/coco_funit/mammals_style.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/coco_funit/mammals_style.png -------------------------------------------------------------------------------- /projects/fs_vid2vid/001.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/fs_vid2vid/001.gif -------------------------------------------------------------------------------- /projects/funit/content.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/funit/content.jpg -------------------------------------------------------------------------------- /projects/funit/funit.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/funit/funit.gif -------------------------------------------------------------------------------- /projects/funit/output.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/funit/output.jpg -------------------------------------------------------------------------------- /projects/funit/style.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/funit/style.jpg -------------------------------------------------------------------------------- /projects/gancraft/sch2vox/oldid2newid.json: -------------------------------------------------------------------------------- 1 | {"0": 0, "1": 1, "2": 8, "3": 9, "4": 12, "5": 13, "6": 19, "7": 25, "8": 26, "9": 26, "10": 27, "11": 27, "12": 28, "13": 30, "14": 31, "15": 32, "16": 33, "17": 34, "18": 58, "19": 64, "20": 66, "21": 67, "22": 68, "23": 69, "24": 70, "25": 73, "26": 74, "27": 90, "28": 91, "29": 92, "30": 93, "31": 94, "32": 96, "33": 99, "34": 100, "35": 101, "37": 118, "38": 119, "39": 131, "40": 132, "41": 133, "42": 134, "43": 450, "44": 450, "45": 135, "46": 136, "47": 137, "48": 138, "49": 139, "50": 141, "51": 142, "52": 143, "53": 144, "54": 145, "55": 146, "56": 147, "57": 148, "58": 149, "59": 150, "60": 151, "61": 152, "62": 152, "63": 153, "64": 159, "65": 160, "66": 161, "67": 162, "68": 153, "69": 169, "70": 170, "71": 171, "72": 172, "73": 178, "74": 178, "75": 180, "76": 179, "77": 181, "78": 182, "79": 183, "80": 184, "81": 185, "82": 186, "83": 187, "84": 188, "85": 189, "86": 195, "87": 191, "88": 192, "89": 193, "90": 194, "91": 196, "92": 197, "93": 198, "94": 198, "95": 199, "96": 215, "97": 225, "98": 221, "99": 221, "100": 221, "101": 234, "102": 235, "103": 236, "104": 239, "105": 240, "106": 241, "107": 242, "108": 243, "109": 244, "110": 245, "111": 246, "112": 247, "113": 248, "114": 249, "115": 250, "116": 251, "117": 252, "118": 253, "119": 254, "120": 255, "121": 256, "122": 257, "123": 258, "124": 258, "125": 444, "126": 444, "127": 259, "128": 260, "129": 261, "130": 262, "131": 263, "132": 263, "133": 265, "134": 266, "135": 267, "136": 268, "137": 269, "138": 270, "139": 271, "140": 273, "141": 298, "142": 299, "143": 300, "144": 306, "145": 318, "146": 321, "147": 322, "148": 323, "149": 324, "150": 324, "151": 325, "152": 326, "153": 327, "154": 328, "155": 329, "156": 332, "157": 333, "158": 334, "159": 335, "160": 351, "161": 62, "162": 38, "163": 367, "164": 368, "165": 369, "166": 370, "167": 371, "168": 372, "169": 381, "170": 382, "171": 383, "172": 399, "173": 400, "174": 401, "175": 402, "176": 408, "177": 424, "178": 325, "179": 440, "180": 443, "181": 450, "182": 460, "183": 467, "184": 468, "185": 469, "186": 471, "187": 470, "188": 472, "189": 473, "190": 474, "191": 476, "192": 475, "193": 477, "194": 478, "195": 479, "196": 480, "197": 481, "198": 482, "199": 483, "200": 484, "201": 485, "202": 486, "203": 487, "204": 462, "205": 462, "206": 488, "207": 489, "208": 490, "209": 491, "210": 492, "211": 493, "212": 494, "213": 495, "214": 496, "215": 497, "216": 498, "217": 499, "218": 500, "219": 502, "220": 503, "221": 504, "222": 505, "223": 506, "224": 507, "225": 508, "226": 509, "227": 510, "228": 511, "229": 512, "230": 513, "231": 514, "232": 515, "233": 516, "234": 517, "235": 518, "236": 519, "237": 520, "238": 521, "239": 522, "240": 523, "241": 524, "242": 525, "243": 526, "244": 527, "245": 528, "246": 529, "247": 530, "248": 531, "249": 532, "250": 533, "251": 534, "252": 550, "255": 673} -------------------------------------------------------------------------------- /projects/gancraft/sch2vox/sch2vox.py: -------------------------------------------------------------------------------- 1 | # share: outside-ok 2 | import numpy as np 3 | import nbtlib 4 | import json 5 | import argparse 6 | 7 | 8 | def main(args): 9 | with open('oldid2newid.json', 'r') as f: 10 | oldid2newid_lut = json.load(f) 11 | oldid2newid_lut = {int(k): v for k, v in oldid2newid_lut.items()} 12 | oldid2newid_lut[36] = 0 13 | oldid2newid_lut[253] = 0 14 | oldid2newid_lut[254] = 0 15 | oldid2newid_lut = [oldid2newid_lut[i] for i in range(256)] 16 | oldid2newid_lut = np.array(oldid2newid_lut) 17 | 18 | sch = nbtlib.load(args.sch) 19 | blocks = np.array(sch['Schematic']['Blocks'], dtype=np.int32) 20 | width = int(sch['Schematic']['Width']) # X 21 | height = int(sch['Schematic']['Height']) # Y 22 | length = int(sch['Schematic']['Length']) # Z 23 | print(width, height, length) 24 | 25 | blocks = blocks.reshape(height, length, width) 26 | 27 | blocks_newid = oldid2newid_lut[blocks] 28 | # Treat grass below sea level as seagrass (Pre-1.13 MCs don't have seagrass). 29 | # Assuming sea level at height==62 30 | seagrass_mask = blocks_newid == 8 31 | seagrass_mask[args.sea_level+1:] = 0 32 | blocks_newid[seagrass_mask] = 97 # water 33 | 34 | np.save(args.npy, blocks_newid.astype(np.int16)) 35 | print('Done') 36 | 37 | 38 | if __name__ == '__main__': 39 | parser = argparse.ArgumentParser(description='Comverting MCEdit Schematic file to plain Numpy array') 40 | parser.add_argument('sch', 41 | help='Path to the source schematic file.') 42 | parser.add_argument('npy', 43 | help='Path to the destination npy file.') 44 | parser.add_argument('--sea_level', type=int, default=62, 45 | help='Sea level of the provided schematic file. Default to 62.') 46 | 47 | args = parser.parse_args() 48 | main(args) 49 | -------------------------------------------------------------------------------- /projects/munit/example_input.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/munit/example_input.jpg -------------------------------------------------------------------------------- /projects/munit/example_output.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/munit/example_output.jpg -------------------------------------------------------------------------------- /projects/pix2pixhd/images/frankfurt_000000_000576.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/pix2pixhd/images/frankfurt_000000_000576.jpg -------------------------------------------------------------------------------- /projects/pix2pixhd/images/frankfurt_000000_011810.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/pix2pixhd/images/frankfurt_000000_011810.jpg -------------------------------------------------------------------------------- /projects/pix2pixhd/images/frankfurt_000000_012868.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/pix2pixhd/images/frankfurt_000000_012868.jpg -------------------------------------------------------------------------------- /projects/pix2pixhd/instance_maps/frankfurt_000001_015328.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/pix2pixhd/instance_maps/frankfurt_000001_015328.png -------------------------------------------------------------------------------- /projects/pix2pixhd/instance_maps/frankfurt_000001_031266.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/pix2pixhd/instance_maps/frankfurt_000001_031266.png -------------------------------------------------------------------------------- /projects/pix2pixhd/instance_maps/frankfurt_000001_044413.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/pix2pixhd/instance_maps/frankfurt_000001_044413.png -------------------------------------------------------------------------------- /projects/pix2pixhd/seg_maps/frankfurt_000000_000576.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/pix2pixhd/seg_maps/frankfurt_000000_000576.png -------------------------------------------------------------------------------- /projects/pix2pixhd/seg_maps/frankfurt_000000_011810.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/pix2pixhd/seg_maps/frankfurt_000000_011810.png -------------------------------------------------------------------------------- /projects/pix2pixhd/seg_maps/frankfurt_000000_012868.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/pix2pixhd/seg_maps/frankfurt_000000_012868.png -------------------------------------------------------------------------------- /projects/spade/edge_maps/000000044195.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/spade/edge_maps/000000044195.png -------------------------------------------------------------------------------- /projects/spade/edge_maps/000000058384.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/spade/edge_maps/000000058384.png -------------------------------------------------------------------------------- /projects/spade/edge_maps/000000072795.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/spade/edge_maps/000000072795.png -------------------------------------------------------------------------------- /projects/spade/images/000000044195.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/spade/images/000000044195.jpg -------------------------------------------------------------------------------- /projects/spade/images/000000058384.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/spade/images/000000058384.jpg -------------------------------------------------------------------------------- /projects/spade/images/000000072795.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/spade/images/000000072795.jpg -------------------------------------------------------------------------------- /projects/spade/seg_maps/000000044195.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/spade/seg_maps/000000044195.png -------------------------------------------------------------------------------- /projects/spade/seg_maps/000000058384.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/spade/seg_maps/000000058384.png -------------------------------------------------------------------------------- /projects/spade/seg_maps/000000072795.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/spade/seg_maps/000000072795.png -------------------------------------------------------------------------------- /projects/unit/summer01.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/unit/summer01.jpg -------------------------------------------------------------------------------- /projects/unit/summer02.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/unit/summer02.jpg -------------------------------------------------------------------------------- /projects/unit/summer2winter_01.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/unit/summer2winter_01.jpg -------------------------------------------------------------------------------- /projects/unit/summer2winter_02.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/unit/summer2winter_02.jpg -------------------------------------------------------------------------------- /projects/unit/winter01.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/unit/winter01.jpg -------------------------------------------------------------------------------- /projects/unit/winter02.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/unit/winter02.jpg -------------------------------------------------------------------------------- /projects/unit/winter2summer_01.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/unit/winter2summer_01.jpg -------------------------------------------------------------------------------- /projects/unit/winter2summer_02.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/unit/winter2summer_02.jpg -------------------------------------------------------------------------------- /projects/vid2vid/stuttgart_00.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/imaginaire/c6f74845c699c58975fd12b778c375b72eb00e8d/projects/vid2vid/stuttgart_00.gif -------------------------------------------------------------------------------- /scripts/build_docker.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | key=${1} 4 | 5 | rm -rf Dockerfile 6 | echo "FROM nvcr.io/nvidia/pytorch:${key}-py3" > Dockerfile 7 | input="Dockerfile.base" 8 | 9 | while IFS= read -r line 10 | do 11 | echo "$line" >> Dockerfile 12 | done < "$input" 13 | 14 | input="scripts/requirements.txt" 15 | while IFS= read -r line 16 | do 17 | echo "RUN pip install $line" >> Dockerfile 18 | done < "$input" 19 | 20 | 21 | for p in correlation channelnorm resample2d bias_act upfirdn2d; do 22 | echo "COPY imaginaire/third_party/$p $p" >> Dockerfile 23 | echo "RUN cd $p && rm -rf build dist *-info && python setup.py install" >> Dockerfile 24 | done 25 | 26 | # Compile GANcraft libraries. 27 | echo "COPY imaginaire/model_utils/gancraft/voxlib gancraft/voxlib" >> Dockerfile 28 | echo "RUN cd gancraft/voxlib && make" >> Dockerfile 29 | 30 | docker build -t nvcr.io/nvidian/lpr-imagine/imaginaire:${key}-py3 . 31 | -------------------------------------------------------------------------------- /scripts/build_index.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import os 4 | import sys 5 | 6 | sys.path.append('.') 7 | from imaginaire.utils.lmdb import create_metadata # noqa: E402 8 | from imaginaire.config import Config # noqa: E402 9 | 10 | 11 | def parse_args(): 12 | r"""Parse user input arguments""" 13 | parser = argparse.ArgumentParser(description='Folder -> LMDB conversion') 14 | parser.add_argument('--data_root', type=str, required=True, 15 | help='Input data location.') 16 | parser.add_argument('--output_root', type=str, default='', 17 | help='Input data location.') 18 | parser.add_argument('--config', type=str, required=True, 19 | help='Config with label info.') 20 | parser.add_argument('--paired', default=False, action='store_true', 21 | help='Is the input data paired?') 22 | parser.add_argument('--input_list', type=str, default='', 23 | help='list of images that will be used.') 24 | args = parser.parse_args() 25 | return args 26 | 27 | 28 | def main(): 29 | r""" Build lmdb for training/testing. 30 | Usage: 31 | python scripts/build_index.py \ 32 | --data_root /mnt/bigdata01/datasets/test_image \ 33 | --output_root /mnt/bigdata01/datasets/test_image/lmdb_0/ \ 34 | --overwrite 35 | """ 36 | args = parse_args() 37 | if args.output_root == '': 38 | args.output_root = args.data_root 39 | cfg = Config(args.config) 40 | 41 | all_filenames, extensions = \ 42 | create_metadata( 43 | data_root=args.data_root, 44 | cfg=cfg, 45 | paired=args.paired, 46 | input_list=args.input_list) 47 | 48 | os.makedirs(args.output_root, exist_ok=True) 49 | 50 | if args.paired: 51 | base = args.data_root.split('/')[-1] 52 | new_all_filenames = dict() 53 | for key in all_filenames.keys(): 54 | new_all_filenames['{}/{}'.format(base, key)] = all_filenames[key] 55 | all_filenames = new_all_filenames.copy() 56 | 57 | # Output list of all filenames. 58 | with open(args.output_root + '/all_filenames.json', 'w') as fout: 59 | json.dump(all_filenames, fout, indent=4) 60 | 61 | # Output metadata. 62 | with open(args.output_root + '/metadata.json', 'w') as fout: 63 | json.dump(extensions, fout, indent=4) 64 | 65 | 66 | if __name__ == "__main__": 67 | main() 68 | -------------------------------------------------------------------------------- /scripts/build_lmdb.sh: -------------------------------------------------------------------------------- 1 | MODEL=$1 2 | DATASET=$2 3 | 4 | for SPLIT in test train; do 5 | RAW=dataset/${DATASET}_raw/${SPLIT} 6 | LMDB=dataset/${DATASET}/${SPLIT} 7 | echo ${LMDB} 8 | python scripts/build_lmdb.py --config configs/projects/${MODEL}/${DATASET}/ampO1.yaml --data_root ${RAW} --output_root ${LMDB} --overwrite 9 | done -------------------------------------------------------------------------------- /scripts/download_dataset.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import tarfile 4 | import sys 5 | 6 | sys.path.append('.') 7 | from imaginaire.utils.io import download_file_from_google_drive # noqa: E402 8 | 9 | 10 | def parse_args(): 11 | parser = argparse.ArgumentParser(description='Download and process dataset') 12 | parser.add_argument('--dataset', help='Name of the dataset.', required=True, 13 | choices=['afhq_dog2cat', 14 | 'animal_faces']) 15 | parser.add_argument('--data_dir', default='./dataset', 16 | help='Directory to save all datasets.') 17 | args = parser.parse_args() 18 | return args 19 | 20 | 21 | def main(): 22 | args = parse_args() 23 | if args.dataset == 'afhq_dog2cat': 24 | url = '1XaiwS0eRctqm-JEDezOBy4TXriAQgc4_' 25 | elif args.dataset == 'animal_faces': 26 | url = '1ftr1xWm0VakGlLUWi7-hdAt9W37luQOA' 27 | else: 28 | raise ValueError('Invalid dataset {}.'.format(args.dataset)) 29 | 30 | # Create the dataset directory. 31 | if not os.path.exists(args.data_dir): 32 | os.makedirs(args.data_dir) 33 | 34 | # Download the compressed dataset. 35 | folder_path = os.path.join(args.data_dir, args.dataset + '_raw') 36 | compressed_path = folder_path + '.tar.gz' 37 | if not os.path.exists(compressed_path) and not os.path.exists(folder_path): 38 | print("Downloading the dataset {}.".format(args.dataset)) 39 | download_file_from_google_drive(url, compressed_path) 40 | 41 | # Extract the dataset. 42 | if not os.path.exists(folder_path): 43 | print("Extracting the dataset {}.".format(args.dataset)) 44 | with tarfile.open(compressed_path) as tar: 45 | tar.extractall(folder_path) 46 | 47 | 48 | if __name__ == "__main__": 49 | main() 50 | -------------------------------------------------------------------------------- /scripts/download_test_data.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import sys 4 | import tarfile 5 | sys.path.append('.') 6 | from imaginaire.utils.io import download_file_from_google_drive # noqa: E402 7 | 8 | URLS = { 9 | 'pix2pixhd': '1Xg9m184zkuG8H0LHdBtSzt2VbMi3SWwR', 10 | 'spade': '1ESm-gHWu_aMHnKF42qkGc8qf1SBECsgf', 11 | 'funit': '1a-EE_6RsYPUoKxEl5oXrpRmKYUltqaD-', 12 | 'coco_funit': '1JYVYB0Q1VStDLOb0SBJbN1vkaf6KrGDh', 13 | 'unit': '17BbwnCG7qF7FI-t9VkORv2XCKqlrY1CO', 14 | 'munit': '1VPgHGuQfmm1N1Vh56wr34wtAwaXzjXtH', 15 | 'vid2vid': '1SHvGPMq-55GDUQ0Ac2Ng0eyG5xCPeKhc', 16 | 'fs_vid2vid': '1fTj0HHjzcitgsSeG5O_aWMF8yvCQUQkN', 17 | 'wc_vid2vid/cityscapes': '1KKzrTHfbpBY9xtLqK8e3QvX8psSdrFcD', 18 | 'wc_vid2vid/mannequin': '1mafZf9KJrwUGGI1kBTvwgehHSqP5iaA0', 19 | 'gancraft': '1m6q7ZtYJjxFL0SQ_WzMbvoLZxXmI5_vJ', 20 | } 21 | 22 | 23 | def parse_args(): 24 | parser = argparse.ArgumentParser(description='Download test data.') 25 | parser.add_argument('--model_name', required=True, 26 | help='Name of the model.') 27 | args = parser.parse_args() 28 | return args 29 | 30 | 31 | def main(): 32 | args = parse_args() 33 | test_data_dir = 'projects/' + args.model_name + '/test_data' 34 | print(test_data_dir) 35 | assert args.model_name in URLS, 'No sample test data available' 36 | url = URLS[args.model_name] 37 | 38 | if os.path.exists(test_data_dir): 39 | print('Test data exists at', test_data_dir) 40 | compressed_path = test_data_dir + '.tar.gz' 41 | # Extract the dataset. 42 | print('Extracting test data to', test_data_dir) 43 | with tarfile.open(compressed_path) as tar: 44 | tar.extractall(path=test_data_dir) 45 | else: 46 | os.makedirs(test_data_dir, exist_ok=True) 47 | # Download the compressed dataset. 48 | compressed_path = test_data_dir + '.tar.gz' 49 | if not os.path.exists(compressed_path): 50 | print('Downloading test data to', compressed_path) 51 | download_file_from_google_drive(url, compressed_path) 52 | 53 | # Extract the dataset. 54 | print('Extracting test data to', test_data_dir) 55 | with tarfile.open(compressed_path) as tar: 56 | tar.extractall(path=test_data_dir) 57 | 58 | 59 | if __name__ == "__main__": 60 | main() 61 | -------------------------------------------------------------------------------- /scripts/install.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | FOR /F "tokens=*" %%g IN ('nvcc --version') do (set ver=%%g) 3 | 4 | echo %ver% 5 | set CUDA_VERSION=%ver:~11,4% 6 | echo %CUDA_VERSION% 7 | 8 | pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio===0.9.0 -f https://download.pytorch.org/whl/torch_stable.html 9 | 10 | pip install --upgrade -r scripts/requirements.txt 11 | 12 | echo %cd% 13 | set curr_directory=%cd% 14 | echo %curr_directory% 15 | 16 | for %%p in (correlation channelnorm resample2d bias_act upfirdn2d) do ( 17 | cd %curr_directory% 18 | cd imaginaire\third_party\%%p 19 | rmdir /s /q build dist *info 20 | python setup.py install 21 | cd %curr_directory% 22 | ) 23 | 24 | 25 | 26 | -------------------------------------------------------------------------------- /scripts/install.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | CURRENT=$(pwd) 3 | 4 | # Check CUDA_VERSION 5 | export CUDA_VERSION=$(nvcc --version| grep -Po "(\d+\.)+\d+" | head -1) 6 | 7 | sudo apt-get update && sudo apt-get install -y --allow-downgrades --allow-change-held-packages --no-install-recommends \ 8 | build-essential \ 9 | git \ 10 | curl \ 11 | vim \ 12 | tmux \ 13 | wget \ 14 | bzip2 \ 15 | unzip \ 16 | g++ \ 17 | ca-certificates \ 18 | ffmpeg \ 19 | libx264-dev \ 20 | imagemagick 21 | 22 | pip install --upgrade -r scripts/requirements.txt 23 | 24 | for p in correlation channelnorm resample2d bias_act upfirdn2d; do 25 | cd imaginaire/third_party/${p}; 26 | rm -rf build dist *info; 27 | python setup.py install; 28 | cd ${CURRENT}; 29 | done 30 | 31 | for p in gancraft/voxlib; do 32 | cd imaginaire/model_utils/${p}; 33 | make all 34 | cd ${CURRENT}; 35 | done 36 | -------------------------------------------------------------------------------- /scripts/requirements.txt: -------------------------------------------------------------------------------- 1 | cmake 2 | pynvml 3 | Pillow>=8.3.2 4 | scipy==1.3.3 5 | scikit-image 6 | tqdm==4.35.0 7 | wget 8 | cython 9 | lmdb 10 | av 11 | opencv-python 12 | opencv-contrib-python 13 | imutils 14 | imageio-ffmpeg 15 | qimage2ndarray 16 | albumentations 17 | requests==2.25.1 18 | nvidia-ml-py3==7.352.0 19 | pyglet 20 | timm 21 | diskcache 22 | boto3 23 | awscli_plugin_endpoint 24 | awscli 25 | rsa 26 | wandb 27 | tensorboard 28 | lpips 29 | face-alignment 30 | dlib 31 | clean-fid 32 | -------------------------------------------------------------------------------- /scripts/requirements_conda.yaml: -------------------------------------------------------------------------------- 1 | name: imaginaire 2 | channels: 3 | - conda-forge 4 | - pytorch 5 | dependencies: 6 | - python=3.8 7 | - pytorch=1.9 8 | - torchvision 9 | - cudatoolkit-dev 10 | - pip 11 | - numpy 12 | - scipy 13 | - ipython 14 | - ipdb 15 | - tqdm 16 | - matplotlib 17 | - jupyterlab 18 | - cmake 19 | - cython 20 | - scikit-image 21 | - scikit-video 22 | - scikit-learn 23 | - pynvml 24 | - ninja 25 | - python-lmdb 26 | - diskcache 27 | - opencv 28 | - imutils 29 | - imageio-ffmpeg 30 | - qimage2ndarray 31 | - albumentations 32 | - pyglet 33 | - timm 34 | - boto3 35 | - rsa 36 | - wandb 37 | - tensorboard 38 | - dlib 39 | - pip: 40 | - av 41 | - nvidia-ml-py3 42 | - awscli 43 | - awscli_plugin_endpoint 44 | - lpips 45 | - face-alignment 46 | - clean-fid 47 | -------------------------------------------------------------------------------- /scripts/start_local_docker.sh: -------------------------------------------------------------------------------- 1 | docker run \ 2 | --gpus all \ 3 | --shm-size 32g \ 4 | --ipc=host \ 5 | -it \ 6 | -v /mnt:/mnt \ 7 | -v ~/:/home \ 8 | nvcr.io/nvidian/lpr-imagine/imaginaire:${1}-py3 \ 9 | /bin/bash 10 | 11 | -------------------------------------------------------------------------------- /scripts/test_training.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | function output { 3 | eval ${cmd} 4 | RESULT=$? 5 | if [ $RESULT -eq 0 ]; then 6 | echo -e "\e[1;32m ${cmd} [Success] \e[0m" 7 | else 8 | echo -e "\e[1;31m ${cmd} [Failure] \e[0m" 9 | exit 1 10 | fi 11 | } 12 | 13 | LOG="/tmp/unit_test.log" 14 | BASE_CMD="python -m torch.distributed.launch --nproc_per_node=1 train.py " 15 | 16 | CONFIG=configs/unit_test/spade.yaml 17 | if test -f "$CONFIG"; then 18 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 19 | output 20 | fi 21 | 22 | CONFIG=configs/unit_test/pix2pixHD.yaml 23 | if test -f "$CONFIG"; then 24 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 25 | output 26 | fi 27 | 28 | CONFIG=configs/unit_test/munit.yaml 29 | if test -f "$CONFIG"; then 30 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 31 | output 32 | fi 33 | 34 | CONFIG=configs/unit_test/munit_patch.yaml 35 | if test -f "$CONFIG"; then 36 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 37 | output 38 | fi 39 | 40 | CONFIG=configs/unit_test/unit.yaml 41 | if test -f "$CONFIG"; then 42 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 43 | output 44 | fi 45 | 46 | CONFIG=configs/unit_test/funit.yaml 47 | if test -f "$CONFIG"; then 48 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 49 | output 50 | fi 51 | 52 | CONFIG=configs/unit_test/coco_funit.yaml 53 | if test -f "$CONFIG"; then 54 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 55 | output 56 | fi 57 | 58 | CONFIG=configs/unit_test/biggan.yaml 59 | if test -f "$CONFIG"; then 60 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 61 | output 62 | fi 63 | 64 | CONFIG=configs/unit_test/stylegan.yaml 65 | if test -f "$CONFIG"; then 66 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 67 | output 68 | fi 69 | 70 | CONFIG=configs/unit_test/vid2vid_street.yaml 71 | if test -f "$CONFIG"; then 72 | cmd="python train.py --single_gpu --config $CONFIG >> ${LOG} " 73 | output 74 | fi 75 | 76 | CONFIG=configs/unit_test/vid2vid_pose.yaml 77 | if test -f "$CONFIG"; then 78 | cmd="python train.py --single_gpu --config $CONFIG >> ${LOG} " 79 | output 80 | fi 81 | 82 | CONFIG=configs/unit_test/fs_vid2vid_face.yaml 83 | if test -f "$CONFIG"; then 84 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 85 | output 86 | fi 87 | 88 | CONFIG=configs/unit_test/fs_vid2vid_pose.yaml 89 | if test -f "$CONFIG"; then 90 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 91 | output 92 | fi 93 | 94 | CONFIG=configs/unit_test/wc_vid2vid.yaml 95 | if test -f "$CONFIG"; then 96 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 97 | output 98 | fi 99 | 100 | CONFIG=configs/unit_test/first_order_motion.yaml 101 | if test -f "$CONFIG"; then 102 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 103 | output 104 | fi 105 | 106 | CONFIG=configs/unit_test/face_vid2vid.yaml 107 | if test -f "$CONFIG"; then 108 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 109 | output 110 | fi 111 | 112 | CONFIG=configs/unit_test/gancraft.yaml 113 | if test -f "$CONFIG"; then 114 | python scripts/download_test_data.py --model_name gancraft; 115 | cmd="${BASE_CMD} --config $CONFIG >> ${LOG} " 116 | output 117 | rm -rf projects/gancraft/test_data* 118 | fi --------------------------------------------------------------------------------