upload files

LICENSE.md

@@ -0,0 +1,14 @@
+# S-Lab License 1.0
+
+Copyright 2024 S-Lab
+
+Redistribution and use for non-commercial purpose in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.\
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+4. In the event that redistribution and/or use for commercial purpose in source or binary forms, with or without modification is required, please contact the contributor(s) of the work.
+
+
+---
+For the commercial use of the code, please consult Prof. Chen Change Loy (ccloy@ntu.edu.sg)

README.md

@@ -0,0 +1,208 @@
+# FRESCO - Official PyTorch Implementation
+
+
+**FRESCO: Spatial-Temporal Correspondence for Zero-Shot Video Translation**<br>
+[Shuai Yang](https://williamyang1991.github.io/), [Yifan Zhou](https://zhouyifan.net/), [Ziwei Liu](https://liuziwei7.github.io/) and [Chen Change Loy](https://www.mmlab-ntu.com/person/ccloy/)<br>
+in CVPR 2024 <br>
+[**Project Page**](https://www.mmlab-ntu.com/project/fresco/) | [**Paper**](https://arxiv.org/abs/2403.12962) | [**Supplementary Video**](https://youtu.be/jLnGx5H-wLw) | [**Input Data and Video Results**](https://drive.google.com/file/d/12BFx3hp8_jp9m0EmKpw-cus2SABPQx2Q/view?usp=sharing) <br>
+
+**Abstract:** *The remarkable efficacy of text-to-image diffusion models has motivated extensive exploration of their potential application in video domains.
+Zero-shot methods seek to extend image diffusion models to videos without necessitating model training.
+Recent methods mainly focus on incorporating inter-frame correspondence into attention mechanisms. However, the soft constraint imposed on determining where to attend to valid features can sometimes be insufficient, resulting in temporal inconsistency.
+In this paper, we introduce FRESCO, intra-frame correspondence alongside inter-frame correspondence to establish a more robust spatial-temporal constraint. This enhancement ensures a more consistent transformation of semantically similar content across frames. Beyond mere attention guidance, our approach involves an explicit update of features to achieve high spatial-temporal consistency with the input video, significantly improving the visual coherence of the resulting translated videos.
+Extensive experiments demonstrate the effectiveness of our proposed framework in producing high-quality, coherent videos, marking a notable improvement over existing zero-shot methods.*
+
+**Features**:<br>
+- **Temporal consistency**: use intra-and inter-frame constraint with better consistency and coverage than optical flow alone.
+    - Compared with our previous work [Rerender-A-Video](https://github.com/williamyang1991/Rerender_A_Video), FRESCO is more robust to large and quick motion.
+- **Zero-shot**: no training or fine-tuning required.
+- **Flexibility**: compatible with off-the-shelf models (e.g., [ControlNet](https://github.com/lllyasviel/ControlNet), [LoRA](https://civitai.com/)) for customized translation.
+
+https://github.com/williamyang1991/FRESCO/assets/18130694/aad358af-4d27-4f18-b069-89a1abd94d38
+
+
+## Updates
+- [03/2023] Paper is released.
+- [03/2023] Code is released.
+- [03/2024] This website is created.
+
+### TODO
+- [ ] Integrate into Diffusers
+- [ ] Add Huggingface web demo
+- [x] ~~Add webUI.~~
+- [x] ~~Update readme~~
+- [x] ~~Upload paper to arXiv, release related material~~
+
+## Installation
+
+1. Clone the repository. 
+
+```shell
+git clone https://github.com/williamyang1991/FRESCO.git
+cd FRESCO
+```
+
+2. You can simply set up the environment with pip based on [requirements.txt](https://github.com/williamyang1991/FRESCO/blob/main/requirements.txt)
+    - Create a conda environment and install torch >= 2.0.0. Here is an example script to install torch 2.0.0 + CUDA 11.8 :
+    ```
+    conda create --name diffusers python==3.8.5
+    conda activate diffusers
+    pip install torch==2.0.0 torchvision==0.15.1 --index-url https://download.pytorch.org/whl/cu118
+    ```
+    - Run `pip install -r requirements.txt` in an environment where torch is installed.
+    - We have tested on torch 2.0.0/2.1.0 and diffusers 0.19.3
+    - If you use new versions of diffusers, you need to modify [my_forward()](https://github.com/williamyang1991/FRESCO/blob/fb991262615665de88f7a8f2cc903d9539e1b234/src/diffusion_hacked.py#L496)
+
+3. Run the installation script. The required models will be downloaded in `./model`, `./src/ControlNet/annotator` and `./src/ebsynth/deps/ebsynth/bin`.
+    - Requires access to huggingface.co
+
+```shell
+python install.py
+```
+
+4. You can run the demo with `run_fresco.py`
+
+```shell
+python run_fresco.py ./config/config_music.yaml
+```
+
+5. For issues with Ebsynth, please refer to [issues](https://github.com/williamyang1991/Rerender_A_Video#issues)
+
+
+## (1) Inference
+
+### WebUI (recommended)
+
+```
+python webUI.py
+```
+The Gradio app also allows you to flexibly change the inference options. Just try it for more details. 
+
+Upload your video, input the prompt, select the model and seed, and hit:
+- **Run Key Frames**: detect keyframes, translate all keyframes.
+- **Run Propagation**: propagate the keyframes to other frames for full video translation
+- **Run All**: **Run Key Frames** and **Run Propagation**
+
+Select the model:
+- **Base model**: base Stable Diffusion model (SD 1.5)
+    - Stable Diffusion 1.5: official model
+    - [rev-Animated](https://huggingface.co/stablediffusionapi/rev-animated): a semi-realistic (2.5D) model
+    - [realistic-Vision](https://huggingface.co/SG161222/Realistic_Vision_V2.0): a photo-realistic model
+    - [flat2d-animerge](https://huggingface.co/stablediffusionapi/flat-2d-animerge): a cartoon model
+    - You can add other models on huggingface.co by modifying this [line](https://github.com/williamyang1991/FRESCO/blob/1afcca9c7b1bc1ac68254f900be9bd768fbb6988/webUI.py#L362) 
+   
+![overview](https://github.com/williamyang1991/FRESCO/assets/18130694/6ce5d54e-b020-4e43-95e7-72ab1783f482)
+
+We provide abundant advanced options to play with
+
+</details>
+
+<details id="option1">
+<summary> <b>Advanced options for single frame processing</b></summary>
+
+1. **Frame resolution**: resize the short side of the video to 512.
+2. ControlNet related:
+   - **ControlNet strength**: how well the output matches the input control edges
+   - **Control type**: HED edge, Canny edge, Depth map
+   - **Canny low/high threshold**: low values for more edge details
+3. SDEdit related:
+   - **Denoising strength**: repaint degree (low value to make the output look more like the original video)
+   - **Preserve color**: preserve the color of the original video
+4. SD related:
+   - **Steps**: denoising step
+   - **CFG scale**: how well the output matches the prompt
+   - **Added prompt/Negative prompt**: supplementary prompts
+5. FreeU related:
+   - **FreeU first/second-stage backbone factor**: =1 do nothing; >1 enhance output color and details
+   - **FreeU first/second-stage skip factor**: =1 do nothing; <1 enhance output color and details
+
+</details>
+
+<details id="option2">
+<summary> <b>Advanced options for FRESCO constraints</b></summary>
+
+1. Keyframe related
+   - **Number of frames**: Total frames to be translated
+   - **Number of frames in a batch**: To avoid out-of-memory, use small batch size
+   - **Min keyframe interval (s_min)**: The keyframes will be detected at least every s_min frames
+   - **Max keyframe interval (s_max)**: The keyframes will be detected at most every s_max frames
+2. FRESCO constraints
+   - FRESCO-guided Attention:
+     - **spatial-guided attention**: Check to enable spatial-guided attention
+     - **cross-frame attention**: Check to enable efficient cross-frame attention
+     - **temporal-guided attention**: Check to enable temporal-guided attention
+   - FRESCO-guided optimization:
+     - **spatial-guided optimization**: Check to enable spatial-guided optimization
+     - **temporal-guided optimization**: Check to enable temporal-guided optimization
+3. **Background smoothing**: Check to enable background smoothing (best for static background)
+   
+</details>
+
+<details id="option3">
+<summary> <b>Advanced options for the full video translation</b></summary>
+
+1. **Gradient blending**: apply Poisson Blending to reduce ghosting artifacts. May slow the process and increase flickers.
+2. **Number of parallel processes**: multiprocessing to speed up the process. Large value (4) is recommended.
+</details>
+
+![option](https://github.com/williamyang1991/FRESCO/assets/18130694/72600758-1dff-4b7c-8f3f-65ee3909f8f6)
+
+### Command Line
+
+We provide a flexible script `run_fresco.py` to run our method.
+
+Set the options via a config file. For example,
+```shell
+python run_fresco.py ./config/config_music.yaml
+```
+We provide some examples of the config in `config` directory.
+Most options in the config is the same as those in WebUI.
+Please check the explanations in the WebUI section.
+
+We provide a separate Ebsynth python script `video_blend.py` with the temporal blending algorithm introduced in
+[Stylizing Video by Example](https://dcgi.fel.cvut.cz/home/sykorad/ebsynth.html) for interpolating style between key frames.
+It can work on your own stylized key frames independently of our FRESCO algorithm.
+For the details, please refer to our previous work [Rerender-A-Video](https://github.com/williamyang1991/Rerender_A_Video/tree/main?tab=readme-ov-file#our-ebsynth-implementation)
+
+## (2) Results
+
+### Key frame translation
+
+<table class="center">
+<tr>
+  <td><img src="https://github.com/williamyang1991/FRESCO/assets/18130694/e8d5776a-37c5-49ae-8ab4-15669df6f572" raw=true></td>
+  <td><img src="https://github.com/williamyang1991/FRESCO/assets/18130694/8a792af6-555c-4e82-ac1e-5c2e1ee35fdb" raw=true></td>
+  <td><img src="https://github.com/williamyang1991/FRESCO/assets/18130694/10f9a964-85ac-4433-84c5-1611a6c2c434" raw=true></td>
+  <td><img src="https://github.com/williamyang1991/FRESCO/assets/18130694/0ec0fbf9-90dd-4d8b-964d-945b5f6687c2" raw=true></td>
+</tr>
+<tr>
+  <td width=26.5% align="center">a red car turns in the winter</td>
+  <td width=26.5% align="center">an African American boxer wearing black boxing gloves punches towards the camera, cartoon style</td>
+  <td width=26.5% align="center">a cartoon spiderman in black suit, black shoes and white gloves is dancing</td>
+  <td width=20.5% align="center">a beautiful woman holding her glasses in CG style</td>
+</tr>
+</table>
+
+
+### Full video translation
+
+https://github.com/williamyang1991/FRESCO/assets/18130694/bf8bfb82-5cb7-4b2f-8169-cf8dbf408b54
+
+## Citation
+
+If you find this work useful for your research, please consider citing our paper:
+
+```bibtex
+@inproceedings{yang2024fresco,
+ title = {FRESCO: Spatial-Temporal Correspondence for Zero-Shot Video Translation},
+ author = {Yang, Shuai and Zhou, Yifan and Liu, Ziwei and and Loy, Chen Change},
+ booktitle = {CVPR},
+ year = {2024},
+}
+```
+
+## Acknowledgments
+
+The code is mainly developed based on [Rerender-A-Video](https://github.com/williamyang1991/Rerender_A_Video), [ControlNet](https://github.com/lllyasviel/ControlNet), [Stable Diffusion](https://github.com/Stability-AI/stablediffusion), [GMFlow](https://github.com/haofeixu/gmflow) and [Ebsynth](https://github.com/jamriska/ebsynth).
+
+

config/config_boxer.yaml

@@ -0,0 +1,27 @@
+# data
+file_path: './data/boxer-punching-towards-camera.mp4'
+save_path: './output/boxer-punching-towards-camera/'
+mininterv: 2 # for keyframe selection
+maxinterv: 2 # for keyframe selection
+
+# diffusion
+seed: 0
+prompt: 'An African American boxer wearing black boxing gloves punches towards the camera, cartoon style'
+sd_path: 'stablediffusionapi/flat-2d-animerge' 
+use_controlnet: True
+controlnet_type: 'depth'  # 'hed', 'canny'
+cond_scale: 0.7
+use_freeu: False
+
+# video-to-video translation 
+batch_size: 8
+num_inference_steps: 20
+num_warmup_steps: 5
+end_opt_step: 15
+run_ebsynth: False
+max_process: 4
+
+# supporting model
+gmflow_path: './model/gmflow_sintel-0c07dcb3.pth'
+sod_path: './model/epoch_resnet.pth'
+use_salinecy: True

config/config_carturn.yaml

@@ -0,0 +1,30 @@
+# data
+file_path: './data/car-turn.mp4'
+save_path: './output/car-turn/'
+mininterv: 5 # for keyframe selection
+maxinterv: 5 # for keyframe selection
+
+# diffusion
+seed: 0
+prompt: 'a red car turns in the winter'
+# sd_path: 'runwayml/stable-diffusion-v1-5' 
+# sd_path: 'stablediffusionapi/rev-animated' 
+# sd_path: 'stablediffusionapi/flat-2d-animerge' 
+sd_path: 'SG161222/Realistic_Vision_V2.0' 
+use_controlnet: True
+controlnet_type: 'hed'  # 'depth', 'canny'
+cond_scale: 0.7
+use_freeu: False
+
+# video-to-video translation 
+batch_size: 8
+num_inference_steps: 20
+num_warmup_steps: 5
+end_opt_step: 15
+run_ebsynth: False
+max_process: 4
+
+# supporting model
+gmflow_path: './model/gmflow_sintel-0c07dcb3.pth'
+sod_path: './model/epoch_resnet.pth'
+use_salinecy: True

config/config_dog.yaml

@@ -0,0 +1,27 @@
+# data
+file_path: './data/dog.mp4'
+save_path: './output/dog/'
+mininterv: 10 # for keyframe selection
+maxinterv: 30 # for keyframe selection
+
+# diffusion
+seed: 0
+prompt: 'greetings from a fox by shaking front paws'
+sd_path: 'SG161222/Realistic_Vision_V2.0' 
+use_controlnet: True
+controlnet_type: 'hed'  # 'depth', 'canny'
+cond_scale: 1.0
+use_freeu: False
+
+# video-to-video translation 
+batch_size: 8
+num_inference_steps: 20
+num_warmup_steps: 8
+end_opt_step: 15
+run_ebsynth: False
+max_process: 4
+
+# supporting model
+gmflow_path: './model/gmflow_sintel-0c07dcb3.pth'
+sod_path: './model/epoch_resnet.pth'
+use_salinecy: True

config/config_music.yaml

@@ -0,0 +1,27 @@
+# data
+file_path: './data/music.mp4'
+save_path: './output/music/'
+mininterv: 10 # for keyframe selection
+maxinterv: 30 # for keyframe selection
+
+# diffusion
+seed: 0
+prompt: 'A beautiful woman with headphones listening to music in CG cyberpunk style, neon, closed eyes, colorful'
+sd_path: 'stablediffusionapi/rev-animated' 
+use_controlnet: True
+controlnet_type: 'hed'  # 'depth', 'canny'
+cond_scale: 1.0
+use_freeu: False
+
+# video-to-video translation 
+batch_size: 8
+num_inference_steps: 20
+num_warmup_steps: 3
+end_opt_step: 15
+run_ebsynth: False
+max_process: 4
+
+# supporting model
+gmflow_path: './model/gmflow_sintel-0c07dcb3.pth'
+sod_path: './model/epoch_resnet.pth'
+use_salinecy: True

install.py

@@ -0,0 +1,95 @@
+import os
+import platform
+
+import requests
+
+
+def build_ebsynth():
+    if os.path.exists('src/ebsynth/deps/ebsynth/bin/ebsynth'):
+        print('Ebsynth has been built.')
+        return
+
+    os_str = platform.system()
+
+    if os_str == 'Windows':
+        print('Build Ebsynth Windows 64 bit.',
+              'If you want to build for 32 bit, please modify install.py.')
+        cmd = '.\\build-win64-cpu+cuda.bat'
+        exe_file = 'src/ebsynth/deps/ebsynth/bin/ebsynth.exe'
+    elif os_str == 'Linux':
+        cmd = 'bash build-linux-cpu+cuda.sh'
+        exe_file = 'src/ebsynth/deps/ebsynth/bin/ebsynth'
+    elif os_str == 'Darwin':
+        cmd = 'sh build-macos-cpu_only.sh'
+        exe_file = 'src/ebsynth/deps/ebsynth/bin/ebsynth.app'
+    else:
+        print('Cannot recognize OS. Ebsynth installation stopped.')
+        return
+
+    os.chdir('src/ebsynth/deps/ebsynth')
+    print(cmd)
+    os.system(cmd)
+    os.chdir('../../../..')
+    if os.path.exists(exe_file):
+        print('Ebsynth installed successfully.')
+    else:
+        print('Failed to install Ebsynth.')
+
+
+def download(url, dir, name=None):
+    os.makedirs(dir, exist_ok=True)
+    if name is None:
+        name = url.split('/')[-1]
+    path = os.path.join(dir, name)
+    if not os.path.exists(path):
+        print(f'Install {name} ...')
+        open(path, 'wb').write(requests.get(url).content)
+        print('Install successfully.')
+
+
+def download_gmflow_ckpt():
+    url = ('https://huggingface.co/PKUWilliamYang/Rerender/'
+           'resolve/main/models/gmflow_sintel-0c07dcb3.pth')
+    download(url, 'model')
+
+
+def download_egnet_ckpt():
+    url = ('https://huggingface.co/PKUWilliamYang/Rerender/'
+           'resolve/main/models/epoch_resnet.pth')
+    download(url, 'model')
+
+def download_hed_ckpt():
+    url = ('https://huggingface.co/lllyasviel/Annotators/'
+           'resolve/main/ControlNetHED.pth')
+    download(url, 'src/ControlNet/annotator/ckpts')
+
+def download_depth_ckpt():
+    url = ('https://huggingface.co/lllyasviel/ControlNet/'
+           'resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt')
+    download(url, 'src/ControlNet/annotator/ckpts')
+
+def download_ebsynth_ckpt():
+    os_str = platform.system()
+    if os_str == 'Linux':
+        url = ('https://huggingface.co/PKUWilliamYang/Rerender/'
+               'resolve/main/models/ebsynth')
+        download(url, 'src/ebsynth/deps/ebsynth/bin')
+    elif os_str == 'Windows':
+        url = ('https://huggingface.co/PKUWilliamYang/Rerender/'
+               'resolve/main/models/ebsynth.exe')
+        download(url, 'src/ebsynth/deps/ebsynth/bin')  
+        url = ('https://huggingface.co/PKUWilliamYang/Rerender/'
+               'resolve/main/models/ebsynth_cpu.dll')
+        download(url, 'src/ebsynth/deps/ebsynth/bin')
+        url = ('https://huggingface.co/PKUWilliamYang/Rerender/'
+               'resolve/main/models/ebsynth_cpu.exe')
+        download(url, 'src/ebsynth/deps/ebsynth/bin')
+    else:
+        print('No available compiled Ebsynth.')
+
+#build_ebsynth()
+download_ebsynth_ckpt()
+download_gmflow_ckpt()
+download_egnet_ckpt()
+download_hed_ckpt()
+download_depth_ckpt()

requirements.txt

@@ -0,0 +1,11 @@
+diffusers[torch]==0.19.3
+transformers
+opencv-python
+einops
+matplotlib
+timm
+av
+basicsr==1.4.2
+numba==0.57.0
+imageio-ffmpeg
+gradio==3.44.4

run_fresco.py

@@ -0,0 +1,318 @@
+import os
+#os.environ['CUDA_VISIBLE_DEVICES'] = "6"
+
+# In China, set this to use huggingface
+# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
+
+import cv2
+import io
+import gc
+import yaml
+import argparse
+import torch
+import torchvision
+import diffusers
+from diffusers import StableDiffusionPipeline, AutoencoderKL, DDPMScheduler, ControlNetModel
+
+from src.utils import *
+from src.keyframe_selection import get_keyframe_ind
+from src.diffusion_hacked import apply_FRESCO_attn, apply_FRESCO_opt, disable_FRESCO_opt
+from src.diffusion_hacked import get_flow_and_interframe_paras, get_intraframe_paras
+from src.pipe_FRESCO import inference
+
+def get_models(config):
+    print('\n' + '=' * 100)
+    print('creating models...')
+    import sys
+    sys.path.append("./src/ebsynth/deps/gmflow/")
+    sys.path.append("./src/EGNet/")
+    sys.path.append("./src/ControlNet/")
+    
+    from gmflow.gmflow import GMFlow
+    from model import build_model
+    from annotator.hed import HEDdetector
+    from annotator.canny import CannyDetector
+    from annotator.midas import MidasDetector
+
+    # optical flow
+    flow_model = GMFlow(feature_channels=128,
+                   num_scales=1,
+                   upsample_factor=8,
+                   num_head=1,
+                   attention_type='swin',
+                   ffn_dim_expansion=4,
+                   num_transformer_layers=6,
+                   ).to('cuda')
+    
+    checkpoint = torch.load(config['gmflow_path'], map_location=lambda storage, loc: storage)
+    weights = checkpoint['model'] if 'model' in checkpoint else checkpoint
+    flow_model.load_state_dict(weights, strict=False)
+    flow_model.eval() 
+    print('create optical flow estimation model successfully!')
+    
+    # saliency detection
+    sod_model = build_model('resnet')
+    sod_model.load_state_dict(torch.load(config['sod_path']))
+    sod_model.to("cuda").eval()
+    print('create saliency detection model successfully!')
+    
+    # controlnet
+    if config['controlnet_type'] not in ['hed', 'depth', 'canny']:
+        print('unsupported control type, set to hed')
+        config['controlnet_type'] = 'hed'
+    controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-"+config['controlnet_type'], 
+                                                 torch_dtype=torch.float16)
+    controlnet.to("cuda") 
+    if config['controlnet_type'] == 'depth':
+        detector = MidasDetector()
+    elif config['controlnet_type'] == 'canny':
+        detector = CannyDetector()
+    else:
+        detector = HEDdetector()
+    print('create controlnet model-' + config['controlnet_type'] + ' successfully!')
+    
+    # diffusion model
+    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)
+    pipe = StableDiffusionPipeline.from_pretrained(config['sd_path'], vae=vae, torch_dtype=torch.float16)
+    pipe.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
+    #noise_scheduler = DDPMScheduler.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="scheduler")
+    pipe.to("cuda")
+    pipe.scheduler.set_timesteps(config['num_inference_steps'], device=pipe._execution_device)
+    
+    if config['use_freeu']:
+        from src.free_lunch_utils import apply_freeu
+        apply_freeu(pipe, b1=1.2, b2=1.5, s1=1.0, s2=1.0)
+
+    frescoProc = apply_FRESCO_attn(pipe)
+    frescoProc.controller.disable_controller()
+    apply_FRESCO_opt(pipe)
+    print('create diffusion model ' + config['sd_path'] + ' successfully!')
+    
+    for param in flow_model.parameters():
+        param.requires_grad = False    
+    for param in sod_model.parameters():
+        param.requires_grad = False
+    for param in controlnet.parameters():
+        param.requires_grad = False
+    for param in pipe.unet.parameters():
+        param.requires_grad = False
+    
+    return pipe, frescoProc, controlnet, detector, flow_model, sod_model
+
+def apply_control(x, detector, config):
+    if config['controlnet_type'] == 'depth':
+        detected_map, _ = detector(x)
+    elif config['controlnet_type'] == 'canny':
+        detected_map = detector(x, 50, 100)
+    else:
+        detected_map = detector(x)
+    return detected_map
+
+def run_keyframe_translation(config):
+    pipe, frescoProc, controlnet, detector, flow_model, sod_model = get_models(config)
+    device = pipe._execution_device
+    guidance_scale = 7.5
+    do_classifier_free_guidance = guidance_scale > 1
+    assert(do_classifier_free_guidance)
+    timesteps = pipe.scheduler.timesteps
+    cond_scale = [config['cond_scale']] * config['num_inference_steps']
+    dilate = Dilate(device=device)
+    
+    base_prompt = config['prompt']
+    if 'Realistic' in config['sd_path'] or 'realistic' in config['sd_path']:
+        a_prompt = ', RAW photo, subject, (high detailed skin:1.2), 8k uhd, dslr, soft lighting, high quality, film grain, Fujifilm XT3, '
+        n_prompt = '(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, mutated hands and fingers:1.4), (deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, amputation'
+    else:
+        a_prompt = ', best quality, extremely detailed, '
+        n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing finger, extra digit, fewer digits, cropped, worst quality, low quality'    
+
+    print('\n' + '=' * 100)
+    print('key frame selection for \"%s\"...'%(config['file_path']))
+    
+    video_cap = cv2.VideoCapture(config['file_path'])
+    frame_num = int(video_cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    
+    # you can set extra_prompts for individual keyframe
+    # for example, extra_prompts[38] = ', closed eyes' to specify the person frame38 closes the eyes
+    extra_prompts = [''] * frame_num
+    
+    keys = get_keyframe_ind(config['file_path'], frame_num, config['mininterv'], config['maxinterv'])
+    
+    os.makedirs(config['save_path'], exist_ok=True)
+    os.makedirs(config['save_path']+'keys', exist_ok=True)
+    os.makedirs(config['save_path']+'video', exist_ok=True)
+    
+    sublists = [keys[i:i+config['batch_size']-2] for i in range(2, len(keys), config['batch_size']-2)]
+    sublists[0].insert(0, keys[0])
+    sublists[0].insert(1, keys[1])
+    if len(sublists) > 1 and len(sublists[-1]) < 3:
+        add_num = 3 - len(sublists[-1])
+        sublists[-1] = sublists[-2][-add_num:] + sublists[-1]
+        sublists[-2] = sublists[-2][:-add_num]
+
+    if not sublists[-2]:
+        del sublists[-2]
+        
+    print('processing %d batches:\nkeyframe indexes'%(len(sublists)), sublists)    
+
+    print('\n' + '=' * 100)
+    print('video to video translation...')
+    
+    batch_ind = 0
+    propagation_mode = batch_ind > 0
+    imgs = []
+    record_latents = []
+    video_cap = cv2.VideoCapture(config['file_path'])
+    for i in range(frame_num):
+        # prepare a batch of frame based on sublists
+        success, frame = video_cap.read()
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        img = resize_image(frame, 512)
+        H, W, C = img.shape
+        Image.fromarray(img).save(os.path.join(config['save_path'], 'video/%04d.png'%(i)))
+        if i not in sublists[batch_ind]:
+            continue
+            
+        imgs += [img]
+        if i != sublists[batch_ind][-1]:
+            continue
+        
+        print('processing batch [%d/%d] with %d frames'%(batch_ind+1, len(sublists), len(sublists[batch_ind])))
+        
+        # prepare input
+        batch_size = len(imgs)
+        n_prompts = [n_prompt] * len(imgs)
+        prompts = [base_prompt + a_prompt + extra_prompts[ind] for ind in sublists[batch_ind]]
+        if propagation_mode: # restore the extra_prompts from previous batch
+            assert len(imgs) == len(sublists[batch_ind]) + 2
+            prompts = ref_prompt + prompts
+        
+        prompt_embeds = pipe._encode_prompt(
+            prompts,
+            device,
+            1,
+            do_classifier_free_guidance,
+            n_prompts,
+        ) 
+            
+        imgs_torch = torch.cat([numpy2tensor(img) for img in imgs], dim=0)
+        edges = torch.cat([numpy2tensor(apply_control(img, detector, config)[:, :, None]) for img in imgs], dim=0)
+        edges = edges.repeat(1,3,1,1).cuda() * 0.5 + 0.5
+        if do_classifier_free_guidance:
+            edges = torch.cat([edges.to(pipe.unet.dtype)] * 2)
+            
+        if config['use_salinecy']:
+            saliency = get_saliency(imgs, sod_model, dilate) 
+        else:
+            saliency = None
+        
+        # prepare parameters for inter-frame and intra-frame consistency
+        flows, occs, attn_mask, interattn_paras = get_flow_and_interframe_paras(flow_model, imgs)
+        correlation_matrix = get_intraframe_paras(pipe, imgs_torch, frescoProc, 
+                            prompt_embeds, seed = config['seed'])
+    
+        '''
+        Flexible settings for attention:
+        * Turn off FRESCO-guided attention: frescoProc.controller.disable_controller() 
+        Then you can turn on one specific attention submodule
+        * Turn on Cross-frame attention: frescoProc.controller.enable_cfattn(attn_mask) 
+        * Turn on Spatial-guided attention: frescoProc.controller.enable_intraattn() 
+        * Turn on Temporal-guided attention: frescoProc.controller.enable_interattn(interattn_paras)
+    
+        Flexible settings for optimization:
+        * Turn off Spatial-guided optimization: set optimize_temporal = False in apply_FRESCO_opt()
+        * Turn off Temporal-guided optimization: set correlation_matrix = [] in apply_FRESCO_opt()
+        * Turn off FRESCO-guided optimization: disable_FRESCO_opt(pipe)
+    
+        Flexible settings for background smoothing:
+        * Turn off background smoothing: set saliency = None in apply_FRESCO_opt()
+        '''    
+        # Turn on all FRESCO support
+        frescoProc.controller.enable_controller(interattn_paras=interattn_paras, attn_mask=attn_mask)
+        apply_FRESCO_opt(pipe, steps = timesteps[:config['end_opt_step']],
+                         flows = flows, occs = occs, correlation_matrix=correlation_matrix, 
+                         saliency=saliency, optimize_temporal = True)
+        
+        gc.collect()
+        torch.cuda.empty_cache()   
+        
+        # run!
+        latents = inference(pipe, controlnet, frescoProc, 
+                  imgs_torch, prompt_embeds, edges, timesteps,
+                  cond_scale, config['num_inference_steps'], config['num_warmup_steps'], 
+                  do_classifier_free_guidance, config['seed'], guidance_scale, config['use_controlnet'],         
+                  record_latents, propagation_mode,
+                  flows = flows, occs = occs, saliency=saliency, repeat_noise=True)
+
+        gc.collect()
+        torch.cuda.empty_cache()
+        
+        with torch.no_grad():
+            image = pipe.vae.decode(latents / pipe.vae.config.scaling_factor, return_dict=False)[0]
+            image = torch.clamp(image, -1 , 1)
+            save_imgs = tensor2numpy(image)
+            bias = 2 if propagation_mode else 0
+            for ind, num in enumerate(sublists[batch_ind]):
+                Image.fromarray(save_imgs[ind+bias]).save(os.path.join(config['save_path'], 'keys/%04d.png'%(num)))
+                
+        gc.collect()
+        torch.cuda.empty_cache()
+        
+        batch_ind += 1
+        # current batch uses the last frame of the previous batch as ref
+        ref_prompt= [prompts[0], prompts[-1]]
+        imgs = [imgs[0], imgs[-1]]
+        propagation_mode = batch_ind > 0
+        if batch_ind == len(sublists):
+            gc.collect()
+            torch.cuda.empty_cache()
+            break    
+    return keys
+
+def run_full_video_translation(config, keys):
+    print('\n' + '=' * 100)
+    if not config['run_ebsynth']:
+        print('to translate full video with ebsynth, install ebsynth and run:')
+    else:
+        print('translating full video with:')
+        
+    video_cap = cv2.VideoCapture(config['file_path'])    
+    fps = int(video_cap.get(cv2.CAP_PROP_FPS))
+    o_video = os.path.join(config['save_path'], 'blend.mp4')
+    max_process = config['max_process']
+    save_path = config['save_path']
+    key_ind = io.StringIO()
+    for k in keys:
+        print('%d'%(k), end=' ', file=key_ind)
+    cmd = (
+        f'python video_blend.py {save_path} --key keys '
+        f'--key_ind {key_ind.getvalue()} --output {o_video} --fps {fps} '
+        f'--n_proc {max_process} -ps')
+    
+    print('\n```')
+    print(cmd)
+    print('```')
+    
+    if config['run_ebsynth']:
+        os.system(cmd)
+    
+    print('\n' + '=' * 100)
+    print('Done')    
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('config_path', type=str, 
+                        default='./config/config_carturn.yaml',
+                        help='The configuration file.')
+    opt = parser.parse_args()
+
+    print('=' * 100)
+    print('loading configuration...')
+    with open(opt.config_path, "r") as f:
+        config = yaml.safe_load(f)
+        
+    for name, value in sorted(config.items()):
+        print('%s: %s' % (str(name), str(value)))  
+
+    keys = run_keyframe_translation(config)
+    run_full_video_translation(config, keys)

src/ControlNet/annotator/canny/__init__.py

@@ -0,0 +1,6 @@
+import cv2
+
+
+class CannyDetector:
+    def __call__(self, img, low_threshold, high_threshold):
+        return cv2.Canny(img, low_threshold, high_threshold)

src/ControlNet/annotator/ckpts/ckpts.txt

@@ -0,0 +1 @@
+Weights here.

src/ControlNet/annotator/hed/__init__.py

@@ -0,0 +1,96 @@
+# This is an improved version and model of HED edge detection with Apache License, Version 2.0.
+# Please use this implementation in your products
+# This implementation may produce slightly different results from Saining Xie's official implementations,
+# but it generates smoother edges and is more suitable for ControlNet as well as other image-to-image translations.
+# Different from official models and other implementations, this is an RGB-input model (rather than BGR)
+# and in this way it works better for gradio's RGB protocol
+
+import os
+import cv2
+import torch
+import numpy as np
+
+from einops import rearrange
+from annotator.util import annotator_ckpts_path
+
+
+class DoubleConvBlock(torch.nn.Module):
+    def __init__(self, input_channel, output_channel, layer_number):
+        super().__init__()
+        self.convs = torch.nn.Sequential()
+        self.convs.append(torch.nn.Conv2d(in_channels=input_channel, out_channels=output_channel, kernel_size=(3, 3), stride=(1, 1), padding=1))
+        for i in range(1, layer_number):
+            self.convs.append(torch.nn.Conv2d(in_channels=output_channel, out_channels=output_channel, kernel_size=(3, 3), stride=(1, 1), padding=1))
+        self.projection = torch.nn.Conv2d(in_channels=output_channel, out_channels=1, kernel_size=(1, 1), stride=(1, 1), padding=0)
+
+    def __call__(self, x, down_sampling=False):
+        h = x
+        if down_sampling:
+            h = torch.nn.functional.max_pool2d(h, kernel_size=(2, 2), stride=(2, 2))
+        for conv in self.convs:
+            h = conv(h)
+            h = torch.nn.functional.relu(h)
+        return h, self.projection(h)
+
+
+class ControlNetHED_Apache2(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.norm = torch.nn.Parameter(torch.zeros(size=(1, 3, 1, 1)))
+        self.block1 = DoubleConvBlock(input_channel=3, output_channel=64, layer_number=2)
+        self.block2 = DoubleConvBlock(input_channel=64, output_channel=128, layer_number=2)
+        self.block3 = DoubleConvBlock(input_channel=128, output_channel=256, layer_number=3)
+        self.block4 = DoubleConvBlock(input_channel=256, output_channel=512, layer_number=3)
+        self.block5 = DoubleConvBlock(input_channel=512, output_channel=512, layer_number=3)
+
+    def __call__(self, x):
+        h = x - self.norm
+        h, projection1 = self.block1(h)
+        h, projection2 = self.block2(h, down_sampling=True)
+        h, projection3 = self.block3(h, down_sampling=True)
+        h, projection4 = self.block4(h, down_sampling=True)
+        h, projection5 = self.block5(h, down_sampling=True)
+        return projection1, projection2, projection3, projection4, projection5
+
+
+class HEDdetector:
+    def __init__(self):
+        remote_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetHED.pth"
+        modelpath = os.path.join(annotator_ckpts_path, "ControlNetHED.pth")
+        if not os.path.exists(modelpath):
+            from basicsr.utils.download_util import load_file_from_url
+            load_file_from_url(remote_model_path, model_dir=annotator_ckpts_path)
+        self.netNetwork = ControlNetHED_Apache2().float().cuda().eval()
+        self.netNetwork.load_state_dict(torch.load(modelpath))
+
+    def __call__(self, input_image):
+        assert input_image.ndim == 3
+        H, W, C = input_image.shape
+        with torch.no_grad():
+            image_hed = torch.from_numpy(input_image.copy()).float().cuda()
+            image_hed = rearrange(image_hed, 'h w c -> 1 c h w')
+            edges = self.netNetwork(image_hed)
+            edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges]
+            edges = [cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR) for e in edges]
+            edges = np.stack(edges, axis=2)
+            edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64)))
+            edge = (edge * 255.0).clip(0, 255).astype(np.uint8)
+            return edge
+
+
+def nms(x, t, s):
+    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
+
+    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
+    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
+    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
+    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
+
+    y = np.zeros_like(x)
+
+    for f in [f1, f2, f3, f4]:
+        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
+
+    z = np.zeros_like(y, dtype=np.uint8)
+    z[y > t] = 255
+    return z

src/ControlNet/annotator/midas/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Intel ISL (Intel Intelligent Systems Lab)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

src/ControlNet/annotator/midas/__init__.py

@@ -0,0 +1,42 @@
+# Midas Depth Estimation
+# From https://github.com/isl-org/MiDaS
+# MIT LICENSE
+
+import cv2
+import numpy as np
+import torch
+
+from einops import rearrange
+from .api import MiDaSInference
+
+
+class MidasDetector:
+    def __init__(self):
+        self.model = MiDaSInference(model_type="dpt_hybrid").cuda()
+
+    def __call__(self, input_image, a=np.pi * 2.0, bg_th=0.1):
+        assert input_image.ndim == 3
+        image_depth = input_image
+        with torch.no_grad():
+            image_depth = torch.from_numpy(image_depth).float().cuda()
+            image_depth = image_depth / 127.5 - 1.0
+            image_depth = rearrange(image_depth, 'h w c -> 1 c h w')
+            depth = self.model(image_depth)[0]
+
+            depth_pt = depth.clone()
+            depth_pt -= torch.min(depth_pt)
+            depth_pt /= torch.max(depth_pt)
+            depth_pt = depth_pt.cpu().numpy()
+            depth_image = (depth_pt * 255.0).clip(0, 255).astype(np.uint8)
+
+            depth_np = depth.cpu().numpy()
+            x = cv2.Sobel(depth_np, cv2.CV_32F, 1, 0, ksize=3)
+            y = cv2.Sobel(depth_np, cv2.CV_32F, 0, 1, ksize=3)
+            z = np.ones_like(x) * a
+            x[depth_pt < bg_th] = 0
+            y[depth_pt < bg_th] = 0
+            normal = np.stack([x, y, z], axis=2)
+            normal /= np.sum(normal ** 2.0, axis=2, keepdims=True) ** 0.5
+            normal_image = (normal * 127.5 + 127.5).clip(0, 255).astype(np.uint8)
+
+            return depth_image, normal_image

src/ControlNet/annotator/midas/api.py

@@ -0,0 +1,169 @@
+# based on https://github.com/isl-org/MiDaS
+
+import cv2
+import os
+import torch
+import torch.nn as nn
+from torchvision.transforms import Compose
+
+from .midas.dpt_depth import DPTDepthModel
+from .midas.midas_net import MidasNet
+from .midas.midas_net_custom import MidasNet_small
+from .midas.transforms import Resize, NormalizeImage, PrepareForNet
+from annotator.util import annotator_ckpts_path
+
+
+ISL_PATHS = {
+    "dpt_large": os.path.join(annotator_ckpts_path, "dpt_large-midas-2f21e586.pt"),
+    "dpt_hybrid": os.path.join(annotator_ckpts_path, "dpt_hybrid-midas-501f0c75.pt"),
+    "midas_v21": "",
+    "midas_v21_small": "",
+}
+
+remote_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"
+
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+
+def load_midas_transform(model_type):
+    # https://github.com/isl-org/MiDaS/blob/master/run.py
+    # load transform only
+    if model_type == "dpt_large":  # DPT-Large
+        net_w, net_h = 384, 384
+        resize_mode = "minimal"
+        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    elif model_type == "dpt_hybrid":  # DPT-Hybrid
+        net_w, net_h = 384, 384
+        resize_mode = "minimal"
+        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    elif model_type == "midas_v21":
+        net_w, net_h = 384, 384
+        resize_mode = "upper_bound"
+        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    elif model_type == "midas_v21_small":
+        net_w, net_h = 256, 256
+        resize_mode = "upper_bound"
+        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    else:
+        assert False, f"model_type '{model_type}' not implemented, use: --model_type large"
+
+    transform = Compose(
+        [
+            Resize(
+                net_w,
+                net_h,
+                resize_target=None,
+                keep_aspect_ratio=True,
+                ensure_multiple_of=32,
+                resize_method=resize_mode,
+                image_interpolation_method=cv2.INTER_CUBIC,
+            ),
+            normalization,
+            PrepareForNet(),
+        ]
+    )
+
+    return transform
+
+
+def load_model(model_type):
+    # https://github.com/isl-org/MiDaS/blob/master/run.py
+    # load network
+    model_path = ISL_PATHS[model_type]
+    if model_type == "dpt_large":  # DPT-Large
+        model = DPTDepthModel(
+            path=model_path,
+            backbone="vitl16_384",
+            non_negative=True,
+        )
+        net_w, net_h = 384, 384
+        resize_mode = "minimal"
+        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    elif model_type == "dpt_hybrid":  # DPT-Hybrid
+        if not os.path.exists(model_path):
+            from basicsr.utils.download_util import load_file_from_url
+            load_file_from_url(remote_model_path, model_dir=annotator_ckpts_path)
+
+        model = DPTDepthModel(
+            path=model_path,
+            backbone="vitb_rn50_384",
+            non_negative=True,
+        )
+        net_w, net_h = 384, 384
+        resize_mode = "minimal"
+        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    elif model_type == "midas_v21":
+        model = MidasNet(model_path, non_negative=True)
+        net_w, net_h = 384, 384
+        resize_mode = "upper_bound"
+        normalization = NormalizeImage(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+
+    elif model_type == "midas_v21_small":
+        model = MidasNet_small(model_path, features=64, backbone="efficientnet_lite3", exportable=True,
+                               non_negative=True, blocks={'expand': True})
+        net_w, net_h = 256, 256
+        resize_mode = "upper_bound"
+        normalization = NormalizeImage(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+
+    else:
+        print(f"model_type '{model_type}' not implemented, use: --model_type large")
+        assert False
+
+    transform = Compose(
+        [
+            Resize(
+                net_w,
+                net_h,
+                resize_target=None,
+                keep_aspect_ratio=True,
+                ensure_multiple_of=32,
+                resize_method=resize_mode,
+                image_interpolation_method=cv2.INTER_CUBIC,
+            ),
+            normalization,
+            PrepareForNet(),
+        ]
+    )
+
+    return model.eval(), transform
+
+
+class MiDaSInference(nn.Module):
+    MODEL_TYPES_TORCH_HUB = [
+        "DPT_Large",
+        "DPT_Hybrid",
+        "MiDaS_small"
+    ]
+    MODEL_TYPES_ISL = [
+        "dpt_large",
+        "dpt_hybrid",
+        "midas_v21",
+        "midas_v21_small",
+    ]
+
+    def __init__(self, model_type):
+        super().__init__()
+        assert (model_type in self.MODEL_TYPES_ISL)
+        model, _ = load_model(model_type)
+        self.model = model
+        self.model.train = disabled_train
+
+    def forward(self, x):
+        with torch.no_grad():
+            prediction = self.model(x)
+        return prediction
+

src/ControlNet/annotator/midas/midas/base_model.py

@@ -0,0 +1,16 @@
+import torch
+
+
+class BaseModel(torch.nn.Module):
+    def load(self, path):
+        """Load model from file.
+
+        Args:
+            path (str): file path
+        """
+        parameters = torch.load(path, map_location=torch.device('cpu'))
+
+        if "optimizer" in parameters:
+            parameters = parameters["model"]
+
+        self.load_state_dict(parameters)

src/ControlNet/annotator/midas/midas/blocks.py

@@ -0,0 +1,342 @@
+import torch
+import torch.nn as nn
+
+from .vit import (
+    _make_pretrained_vitb_rn50_384,
+    _make_pretrained_vitl16_384,
+    _make_pretrained_vitb16_384,
+    forward_vit,
+)
+
+def _make_encoder(backbone, features, use_pretrained, groups=1, expand=False, exportable=True, hooks=None, use_vit_only=False, use_readout="ignore",):
+    if backbone == "vitl16_384":
+        pretrained = _make_pretrained_vitl16_384(
+            use_pretrained, hooks=hooks, use_readout=use_readout
+        )
+        scratch = _make_scratch(
+            [256, 512, 1024, 1024], features, groups=groups, expand=expand
+        )  # ViT-L/16 - 85.0% Top1 (backbone)
+    elif backbone == "vitb_rn50_384":
+        pretrained = _make_pretrained_vitb_rn50_384(
+            use_pretrained,
+            hooks=hooks,
+            use_vit_only=use_vit_only,
+            use_readout=use_readout,
+        )
+        scratch = _make_scratch(
+            [256, 512, 768, 768], features, groups=groups, expand=expand
+        )  # ViT-H/16 - 85.0% Top1 (backbone)
+    elif backbone == "vitb16_384":
+        pretrained = _make_pretrained_vitb16_384(
+            use_pretrained, hooks=hooks, use_readout=use_readout
+        )
+        scratch = _make_scratch(
+            [96, 192, 384, 768], features, groups=groups, expand=expand
+        )  # ViT-B/16 - 84.6% Top1 (backbone)
+    elif backbone == "resnext101_wsl":
+        pretrained = _make_pretrained_resnext101_wsl(use_pretrained)
+        scratch = _make_scratch([256, 512, 1024, 2048], features, groups=groups, expand=expand)     # efficientnet_lite3  
+    elif backbone == "efficientnet_lite3":
+        pretrained = _make_pretrained_efficientnet_lite3(use_pretrained, exportable=exportable)
+        scratch = _make_scratch([32, 48, 136, 384], features, groups=groups, expand=expand)  # efficientnet_lite3     
+    else:
+        print(f"Backbone '{backbone}' not implemented")
+        assert False
+        
+    return pretrained, scratch
+
+
+def _make_scratch(in_shape, out_shape, groups=1, expand=False):
+    scratch = nn.Module()
+
+    out_shape1 = out_shape
+    out_shape2 = out_shape
+    out_shape3 = out_shape
+    out_shape4 = out_shape
+    if expand==True:
+        out_shape1 = out_shape
+        out_shape2 = out_shape*2
+        out_shape3 = out_shape*4
+        out_shape4 = out_shape*8
+
+    scratch.layer1_rn = nn.Conv2d(
+        in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+    )
+    scratch.layer2_rn = nn.Conv2d(
+        in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+    )
+    scratch.layer3_rn = nn.Conv2d(
+        in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+    )
+    scratch.layer4_rn = nn.Conv2d(
+        in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+    )
+
+    return scratch
+
+
+def _make_pretrained_efficientnet_lite3(use_pretrained, exportable=False):
+    efficientnet = torch.hub.load(
+        "rwightman/gen-efficientnet-pytorch",
+        "tf_efficientnet_lite3",
+        pretrained=use_pretrained,
+        exportable=exportable
+    )
+    return _make_efficientnet_backbone(efficientnet)
+
+
+def _make_efficientnet_backbone(effnet):
+    pretrained = nn.Module()
+
+    pretrained.layer1 = nn.Sequential(
+        effnet.conv_stem, effnet.bn1, effnet.act1, *effnet.blocks[0:2]
+    )
+    pretrained.layer2 = nn.Sequential(*effnet.blocks[2:3])
+    pretrained.layer3 = nn.Sequential(*effnet.blocks[3:5])
+    pretrained.layer4 = nn.Sequential(*effnet.blocks[5:9])
+
+    return pretrained
+    
+
+def _make_resnet_backbone(resnet):
+    pretrained = nn.Module()
+    pretrained.layer1 = nn.Sequential(
+        resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1
+    )
+
+    pretrained.layer2 = resnet.layer2
+    pretrained.layer3 = resnet.layer3
+    pretrained.layer4 = resnet.layer4
+
+    return pretrained
+
+
+def _make_pretrained_resnext101_wsl(use_pretrained):
+    resnet = torch.hub.load("facebookresearch/WSL-Images", "resnext101_32x8d_wsl")
+    return _make_resnet_backbone(resnet)
+
+
+
+class Interpolate(nn.Module):
+    """Interpolation module.
+    """
+
+    def __init__(self, scale_factor, mode, align_corners=False):
+        """Init.
+
+        Args:
+            scale_factor (float): scaling
+            mode (str): interpolation mode
+        """
+        super(Interpolate, self).__init__()
+
+        self.interp = nn.functional.interpolate
+        self.scale_factor = scale_factor
+        self.mode = mode
+        self.align_corners = align_corners
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x (tensor): input
+
+        Returns:
+            tensor: interpolated data
+        """
+
+        x = self.interp(
+            x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners
+        )
+
+        return x
+
+
+class ResidualConvUnit(nn.Module):
+    """Residual convolution module.
+    """
+
+    def __init__(self, features):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(
+            features, features, kernel_size=3, stride=1, padding=1, bias=True
+        )
+
+        self.conv2 = nn.Conv2d(
+            features, features, kernel_size=3, stride=1, padding=1, bias=True
+        )
+
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x (tensor): input
+
+        Returns:
+            tensor: output
+        """
+        out = self.relu(x)
+        out = self.conv1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+
+        return out + x
+
+
+class FeatureFusionBlock(nn.Module):
+    """Feature fusion block.
+    """
+
+    def __init__(self, features):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super(FeatureFusionBlock, self).__init__()
+
+        self.resConfUnit1 = ResidualConvUnit(features)
+        self.resConfUnit2 = ResidualConvUnit(features)
+
+    def forward(self, *xs):
+        """Forward pass.
+
+        Returns:
+            tensor: output
+        """
+        output = xs[0]
+
+        if len(xs) == 2:
+            output += self.resConfUnit1(xs[1])
+
+        output = self.resConfUnit2(output)
+
+        output = nn.functional.interpolate(
+            output, scale_factor=2, mode="bilinear", align_corners=True
+        )
+
+        return output
+
+
+
+
+class ResidualConvUnit_custom(nn.Module):
+    """Residual convolution module.
+    """
+
+    def __init__(self, features, activation, bn):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super().__init__()
+
+        self.bn = bn
+
+        self.groups=1
+
+        self.conv1 = nn.Conv2d(
+            features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
+        )
+        
+        self.conv2 = nn.Conv2d(
+            features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
+        )
+
+        if self.bn==True:
+            self.bn1 = nn.BatchNorm2d(features)
+            self.bn2 = nn.BatchNorm2d(features)
+
+        self.activation = activation
+
+        self.skip_add = nn.quantized.FloatFunctional()
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x (tensor): input
+
+        Returns:
+            tensor: output
+        """
+        
+        out = self.activation(x)
+        out = self.conv1(out)
+        if self.bn==True:
+            out = self.bn1(out)
+       
+        out = self.activation(out)
+        out = self.conv2(out)
+        if self.bn==True:
+            out = self.bn2(out)
+
+        if self.groups > 1:
+            out = self.conv_merge(out)
+
+        return self.skip_add.add(out, x)
+
+        # return out + x
+
+
+class FeatureFusionBlock_custom(nn.Module):
+    """Feature fusion block.
+    """
+
+    def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super(FeatureFusionBlock_custom, self).__init__()
+
+        self.deconv = deconv
+        self.align_corners = align_corners
+
+        self.groups=1
+
+        self.expand = expand
+        out_features = features
+        if self.expand==True:
+            out_features = features//2
+        
+        self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
+
+        self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn)
+        self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn)
+        
+        self.skip_add = nn.quantized.FloatFunctional()
+
+    def forward(self, *xs):
+        """Forward pass.
+
+        Returns:
+            tensor: output
+        """
+        output = xs[0]
+
+        if len(xs) == 2:
+            res = self.resConfUnit1(xs[1])
+            output = self.skip_add.add(output, res)
+            # output += res
+
+        output = self.resConfUnit2(output)
+
+        output = nn.functional.interpolate(
+            output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
+        )
+
+        output = self.out_conv(output)
+
+        return output
+

src/ControlNet/annotator/midas/midas/dpt_depth.py

@@ -0,0 +1,109 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .base_model import BaseModel
+from .blocks import (
+    FeatureFusionBlock,
+    FeatureFusionBlock_custom,
+    Interpolate,
+    _make_encoder,
+    forward_vit,
+)
+
+
+def _make_fusion_block(features, use_bn):
+    return FeatureFusionBlock_custom(
+        features,
+        nn.ReLU(False),
+        deconv=False,
+        bn=use_bn,
+        expand=False,
+        align_corners=True,
+    )
+
+
+class DPT(BaseModel):
+    def __init__(
+        self,
+        head,
+        features=256,
+        backbone="vitb_rn50_384",
+        readout="project",
+        channels_last=False,
+        use_bn=False,
+    ):
+
+        super(DPT, self).__init__()
+
+        self.channels_last = channels_last
+
+        hooks = {
+            "vitb_rn50_384": [0, 1, 8, 11],
+            "vitb16_384": [2, 5, 8, 11],
+            "vitl16_384": [5, 11, 17, 23],
+        }
+
+        # Instantiate backbone and reassemble blocks
+        self.pretrained, self.scratch = _make_encoder(
+            backbone,
+            features,
+            False, # Set to true of you want to train from scratch, uses ImageNet weights
+            groups=1,
+            expand=False,
+            exportable=False,
+            hooks=hooks[backbone],
+            use_readout=readout,
+        )
+
+        self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
+
+        self.scratch.output_conv = head
+
+
+    def forward(self, x):
+        if self.channels_last == True:
+            x.contiguous(memory_format=torch.channels_last)
+
+        layer_1, layer_2, layer_3, layer_4 = forward_vit(self.pretrained, x)
+
+        layer_1_rn = self.scratch.layer1_rn(layer_1)
+        layer_2_rn = self.scratch.layer2_rn(layer_2)
+        layer_3_rn = self.scratch.layer3_rn(layer_3)
+        layer_4_rn = self.scratch.layer4_rn(layer_4)
+
+        path_4 = self.scratch.refinenet4(layer_4_rn)
+        path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
+        path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
+        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
+
+        out = self.scratch.output_conv(path_1)
+
+        return out
+
+
+class DPTDepthModel(DPT):
+    def __init__(self, path=None, non_negative=True, **kwargs):
+        features = kwargs["features"] if "features" in kwargs else 256
+
+        head = nn.Sequential(
+            nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1),
+            Interpolate(scale_factor=2, mode="bilinear", align_corners=True),
+            nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(True),
+            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
+            nn.ReLU(True) if non_negative else nn.Identity(),
+            nn.Identity(),
+        )
+
+        super().__init__(head, **kwargs)
+
+        if path is not None:
+           self.load(path)
+
+    def forward(self, x):
+        return super().forward(x).squeeze(dim=1)
+

src/ControlNet/annotator/midas/midas/midas_net.py

@@ -0,0 +1,76 @@
+"""MidashNet: Network for monocular depth estimation trained by mixing several datasets.
+This file contains code that is adapted from
+https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py
+"""
+import torch
+import torch.nn as nn
+
+from .base_model import BaseModel
+from .blocks import FeatureFusionBlock, Interpolate, _make_encoder
+
+
+class MidasNet(BaseModel):
+    """Network for monocular depth estimation.
+    """
+
+    def __init__(self, path=None, features=256, non_negative=True):
+        """Init.
+
+        Args:
+            path (str, optional): Path to saved model. Defaults to None.
+            features (int, optional): Number of features. Defaults to 256.
+            backbone (str, optional): Backbone network for encoder. Defaults to resnet50
+        """
+        print("Loading weights: ", path)
+
+        super(MidasNet, self).__init__()
+
+        use_pretrained = False if path is None else True
+
+        self.pretrained, self.scratch = _make_encoder(backbone="resnext101_wsl", features=features, use_pretrained=use_pretrained)
+
+        self.scratch.refinenet4 = FeatureFusionBlock(features)
+        self.scratch.refinenet3 = FeatureFusionBlock(features)
+        self.scratch.refinenet2 = FeatureFusionBlock(features)
+        self.scratch.refinenet1 = FeatureFusionBlock(features)
+
+        self.scratch.output_conv = nn.Sequential(
+            nn.Conv2d(features, 128, kernel_size=3, stride=1, padding=1),
+            Interpolate(scale_factor=2, mode="bilinear"),
+            nn.Conv2d(128, 32, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(True),
+            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
+            nn.ReLU(True) if non_negative else nn.Identity(),
+        )
+
+        if path:
+            self.load(path)
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x (tensor): input data (image)
+
+        Returns:
+            tensor: depth
+        """
+
+        layer_1 = self.pretrained.layer1(x)
+        layer_2 = self.pretrained.layer2(layer_1)
+        layer_3 = self.pretrained.layer3(layer_2)
+        layer_4 = self.pretrained.layer4(layer_3)
+
+        layer_1_rn = self.scratch.layer1_rn(layer_1)
+        layer_2_rn = self.scratch.layer2_rn(layer_2)
+        layer_3_rn = self.scratch.layer3_rn(layer_3)
+        layer_4_rn = self.scratch.layer4_rn(layer_4)
+
+        path_4 = self.scratch.refinenet4(layer_4_rn)
+        path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
+        path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
+        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
+
+        out = self.scratch.output_conv(path_1)
+
+        return torch.squeeze(out, dim=1)

src/ControlNet/annotator/midas/midas/midas_net_custom.py

@@ -0,0 +1,128 @@
+"""MidashNet: Network for monocular depth estimation trained by mixing several datasets.
+This file contains code that is adapted from
+https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py
+"""
+import torch
+import torch.nn as nn
+
+from .base_model import BaseModel
+from .blocks import FeatureFusionBlock, FeatureFusionBlock_custom, Interpolate, _make_encoder
+
+
+class MidasNet_small(BaseModel):
+    """Network for monocular depth estimation.
+    """
+
+    def __init__(self, path=None, features=64, backbone="efficientnet_lite3", non_negative=True, exportable=True, channels_last=False, align_corners=True,
+        blocks={'expand': True}):
+        """Init.
+
+        Args:
+            path (str, optional): Path to saved model. Defaults to None.
+            features (int, optional): Number of features. Defaults to 256.
+            backbone (str, optional): Backbone network for encoder. Defaults to resnet50
+        """
+        print("Loading weights: ", path)
+
+        super(MidasNet_small, self).__init__()
+
+        use_pretrained = False if path else True
+                
+        self.channels_last = channels_last
+        self.blocks = blocks
+        self.backbone = backbone
+
+        self.groups = 1
+
+        features1=features
+        features2=features
+        features3=features
+        features4=features
+        self.expand = False
+        if "expand" in self.blocks and self.blocks['expand'] == True:
+            self.expand = True
+            features1=features
+            features2=features*2
+            features3=features*4
+            features4=features*8
+
+        self.pretrained, self.scratch = _make_encoder(self.backbone, features, use_pretrained, groups=self.groups, expand=self.expand, exportable=exportable)
+  
+        self.scratch.activation = nn.ReLU(False)    
+
+        self.scratch.refinenet4 = FeatureFusionBlock_custom(features4, self.scratch.activation, deconv=False, bn=False, expand=self.expand, align_corners=align_corners)
+        self.scratch.refinenet3 = FeatureFusionBlock_custom(features3, self.scratch.activation, deconv=False, bn=False, expand=self.expand, align_corners=align_corners)
+        self.scratch.refinenet2 = FeatureFusionBlock_custom(features2, self.scratch.activation, deconv=False, bn=False, expand=self.expand, align_corners=align_corners)
+        self.scratch.refinenet1 = FeatureFusionBlock_custom(features1, self.scratch.activation, deconv=False, bn=False, align_corners=align_corners)
+
+        
+        self.scratch.output_conv = nn.Sequential(
+            nn.Conv2d(features, features//2, kernel_size=3, stride=1, padding=1, groups=self.groups),
+            Interpolate(scale_factor=2, mode="bilinear"),
+            nn.Conv2d(features//2, 32, kernel_size=3, stride=1, padding=1),
+            self.scratch.activation,
+            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
+            nn.ReLU(True) if non_negative else nn.Identity(),
+            nn.Identity(),
+        )
+        
+        if path:
+            self.load(path)
+
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x (tensor): input data (image)
+
+        Returns:
+            tensor: depth
+        """
+        if self.channels_last==True:
+            print("self.channels_last = ", self.channels_last)
+            x.contiguous(memory_format=torch.channels_last)
+
+
+        layer_1 = self.pretrained.layer1(x)
+        layer_2 = self.pretrained.layer2(layer_1)
+        layer_3 = self.pretrained.layer3(layer_2)
+        layer_4 = self.pretrained.layer4(layer_3)
+        
+        layer_1_rn = self.scratch.layer1_rn(layer_1)
+        layer_2_rn = self.scratch.layer2_rn(layer_2)
+        layer_3_rn = self.scratch.layer3_rn(layer_3)
+        layer_4_rn = self.scratch.layer4_rn(layer_4)
+
+
+        path_4 = self.scratch.refinenet4(layer_4_rn)
+        path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
+        path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
+        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
+        
+        out = self.scratch.output_conv(path_1)
+
+        return torch.squeeze(out, dim=1)
+
+
+
+def fuse_model(m):
+    prev_previous_type = nn.Identity()
+    prev_previous_name = ''
+    previous_type = nn.Identity()
+    previous_name = ''
+    for name, module in m.named_modules():
+        if prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d and type(module) == nn.ReLU:
+            # print("FUSED ", prev_previous_name, previous_name, name)
+            torch.quantization.fuse_modules(m, [prev_previous_name, previous_name, name], inplace=True)
+        elif prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d:
+            # print("FUSED ", prev_previous_name, previous_name)
+            torch.quantization.fuse_modules(m, [prev_previous_name, previous_name], inplace=True)
+        # elif previous_type == nn.Conv2d and type(module) == nn.ReLU:
+        #    print("FUSED ", previous_name, name)
+        #    torch.quantization.fuse_modules(m, [previous_name, name], inplace=True)
+
+        prev_previous_type = previous_type
+        prev_previous_name = previous_name
+        previous_type = type(module)
+        previous_name = name

src/ControlNet/annotator/midas/midas/transforms.py

@@ -0,0 +1,234 @@
+import numpy as np
+import cv2
+import math
+
+
+def apply_min_size(sample, size, image_interpolation_method=cv2.INTER_AREA):
+    """Rezise the sample to ensure the given size. Keeps aspect ratio.
+
+    Args:
+        sample (dict): sample
+        size (tuple): image size
+
+    Returns:
+        tuple: new size
+    """
+    shape = list(sample["disparity"].shape)
+
+    if shape[0] >= size[0] and shape[1] >= size[1]:
+        return sample
+
+    scale = [0, 0]
+    scale[0] = size[0] / shape[0]
+    scale[1] = size[1] / shape[1]
+
+    scale = max(scale)
+
+    shape[0] = math.ceil(scale * shape[0])
+    shape[1] = math.ceil(scale * shape[1])
+
+    # resize
+    sample["image"] = cv2.resize(
+        sample["image"], tuple(shape[::-1]), interpolation=image_interpolation_method
+    )
+
+    sample["disparity"] = cv2.resize(
+        sample["disparity"], tuple(shape[::-1]), interpolation=cv2.INTER_NEAREST
+    )
+    sample["mask"] = cv2.resize(
+        sample["mask"].astype(np.float32),
+        tuple(shape[::-1]),
+        interpolation=cv2.INTER_NEAREST,
+    )
+    sample["mask"] = sample["mask"].astype(bool)
+
+    return tuple(shape)
+
+
+class Resize(object):
+    """Resize sample to given size (width, height).
+    """
+
+    def __init__(
+        self,
+        width,
+        height,
+        resize_target=True,
+        keep_aspect_ratio=False,
+        ensure_multiple_of=1,
+        resize_method="lower_bound",
+        image_interpolation_method=cv2.INTER_AREA,
+    ):
+        """Init.
+
+        Args:
+            width (int): desired output width
+            height (int): desired output height
+            resize_target (bool, optional):
+                True: Resize the full sample (image, mask, target).
+                False: Resize image only.
+                Defaults to True.
+            keep_aspect_ratio (bool, optional):
+                True: Keep the aspect ratio of the input sample.
+                Output sample might not have the given width and height, and
+                resize behaviour depends on the parameter 'resize_method'.
+                Defaults to False.
+            ensure_multiple_of (int, optional):
+                Output width and height is constrained to be multiple of this parameter.
+                Defaults to 1.
+            resize_method (str, optional):
+                "lower_bound": Output will be at least as large as the given size.
+                "upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.)
+                "minimal": Scale as least as possible.  (Output size might be smaller than given size.)
+                Defaults to "lower_bound".
+        """
+        self.__width = width
+        self.__height = height
+
+        self.__resize_target = resize_target
+        self.__keep_aspect_ratio = keep_aspect_ratio
+        self.__multiple_of = ensure_multiple_of
+        self.__resize_method = resize_method
+        self.__image_interpolation_method = image_interpolation_method
+
+    def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
+        y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        if max_val is not None and y > max_val:
+            y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        if y < min_val:
+            y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        return y
+
+    def get_size(self, width, height):
+        # determine new height and width
+        scale_height = self.__height / height
+        scale_width = self.__width / width
+
+        if self.__keep_aspect_ratio:
+            if self.__resize_method == "lower_bound":
+                # scale such that output size is lower bound
+                if scale_width > scale_height:
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            elif self.__resize_method == "upper_bound":
+                # scale such that output size is upper bound
+                if scale_width < scale_height:
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            elif self.__resize_method == "minimal":
+                # scale as least as possbile
+                if abs(1 - scale_width) < abs(1 - scale_height):
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            else:
+                raise ValueError(
+                    f"resize_method {self.__resize_method} not implemented"
+                )
+
+        if self.__resize_method == "lower_bound":
+            new_height = self.constrain_to_multiple_of(
+                scale_height * height, min_val=self.__height
+            )
+            new_width = self.constrain_to_multiple_of(
+                scale_width * width, min_val=self.__width
+            )
+        elif self.__resize_method == "upper_bound":
+            new_height = self.constrain_to_multiple_of(
+                scale_height * height, max_val=self.__height
+            )
+            new_width = self.constrain_to_multiple_of(
+                scale_width * width, max_val=self.__width
+            )
+        elif self.__resize_method == "minimal":
+            new_height = self.constrain_to_multiple_of(scale_height * height)
+            new_width = self.constrain_to_multiple_of(scale_width * width)
+        else:
+            raise ValueError(f"resize_method {self.__resize_method} not implemented")
+
+        return (new_width, new_height)
+
+    def __call__(self, sample):
+        width, height = self.get_size(
+            sample["image"].shape[1], sample["image"].shape[0]
+        )
+
+        # resize sample
+        sample["image"] = cv2.resize(
+            sample["image"],
+            (width, height),
+            interpolation=self.__image_interpolation_method,
+        )
+
+        if self.__resize_target:
+            if "disparity" in sample:
+                sample["disparity"] = cv2.resize(
+                    sample["disparity"],
+                    (width, height),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+
+            if "depth" in sample:
+                sample["depth"] = cv2.resize(
+                    sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST
+                )
+
+            sample["mask"] = cv2.resize(
+                sample["mask"].astype(np.float32),
+                (width, height),
+                interpolation=cv2.INTER_NEAREST,
+            )
+            sample["mask"] = sample["mask"].astype(bool)
+
+        return sample
+
+
+class NormalizeImage(object):
+    """Normlize image by given mean and std.
+    """
+
+    def __init__(self, mean, std):
+        self.__mean = mean
+        self.__std = std
+
+    def __call__(self, sample):
+        sample["image"] = (sample["image"] - self.__mean) / self.__std
+
+        return sample
+
+
+class PrepareForNet(object):
+    """Prepare sample for usage as network input.
+    """
+
+    def __init__(self):
+        pass
+
+    def __call__(self, sample):
+        image = np.transpose(sample["image"], (2, 0, 1))
+        sample["image"] = np.ascontiguousarray(image).astype(np.float32)
+
+        if "mask" in sample:
+            sample["mask"] = sample["mask"].astype(np.float32)
+            sample["mask"] = np.ascontiguousarray(sample["mask"])
+
+        if "disparity" in sample:
+            disparity = sample["disparity"].astype(np.float32)
+            sample["disparity"] = np.ascontiguousarray(disparity)
+
+        if "depth" in sample:
+            depth = sample["depth"].astype(np.float32)
+            sample["depth"] = np.ascontiguousarray(depth)
+
+        return sample

src/ControlNet/annotator/midas/midas/vit.py

@@ -0,0 +1,491 @@
+import torch
+import torch.nn as nn
+import timm
+import types
+import math
+import torch.nn.functional as F
+
+
+class Slice(nn.Module):
+    def __init__(self, start_index=1):
+        super(Slice, self).__init__()
+        self.start_index = start_index
+
+    def forward(self, x):
+        return x[:, self.start_index :]
+
+
+class AddReadout(nn.Module):
+    def __init__(self, start_index=1):
+        super(AddReadout, self).__init__()
+        self.start_index = start_index
+
+    def forward(self, x):
+        if self.start_index == 2:
+            readout = (x[:, 0] + x[:, 1]) / 2
+        else:
+            readout = x[:, 0]
+        return x[:, self.start_index :] + readout.unsqueeze(1)
+
+
+class ProjectReadout(nn.Module):
+    def __init__(self, in_features, start_index=1):
+        super(ProjectReadout, self).__init__()
+        self.start_index = start_index
+
+        self.project = nn.Sequential(nn.Linear(2 * in_features, in_features), nn.GELU())
+
+    def forward(self, x):
+        readout = x[:, 0].unsqueeze(1).expand_as(x[:, self.start_index :])
+        features = torch.cat((x[:, self.start_index :], readout), -1)
+
+        return self.project(features)
+
+
+class Transpose(nn.Module):
+    def __init__(self, dim0, dim1):
+        super(Transpose, self).__init__()
+        self.dim0 = dim0
+        self.dim1 = dim1
+
+    def forward(self, x):
+        x = x.transpose(self.dim0, self.dim1)
+        return x
+
+
+def forward_vit(pretrained, x):
+    b, c, h, w = x.shape
+
+    glob = pretrained.model.forward_flex(x)
+
+    layer_1 = pretrained.activations["1"]
+    layer_2 = pretrained.activations["2"]
+    layer_3 = pretrained.activations["3"]
+    layer_4 = pretrained.activations["4"]
+
+    layer_1 = pretrained.act_postprocess1[0:2](layer_1)
+    layer_2 = pretrained.act_postprocess2[0:2](layer_2)
+    layer_3 = pretrained.act_postprocess3[0:2](layer_3)
+    layer_4 = pretrained.act_postprocess4[0:2](layer_4)
+
+    unflatten = nn.Sequential(
+        nn.Unflatten(
+            2,
+            torch.Size(
+                [
+                    h // pretrained.model.patch_size[1],
+                    w // pretrained.model.patch_size[0],
+                ]
+            ),
+        )
+    )
+
+    if layer_1.ndim == 3:
+        layer_1 = unflatten(layer_1)
+    if layer_2.ndim == 3:
+        layer_2 = unflatten(layer_2)
+    if layer_3.ndim == 3:
+        layer_3 = unflatten(layer_3)
+    if layer_4.ndim == 3:
+        layer_4 = unflatten(layer_4)
+
+    layer_1 = pretrained.act_postprocess1[3 : len(pretrained.act_postprocess1)](layer_1)
+    layer_2 = pretrained.act_postprocess2[3 : len(pretrained.act_postprocess2)](layer_2)
+    layer_3 = pretrained.act_postprocess3[3 : len(pretrained.act_postprocess3)](layer_3)
+    layer_4 = pretrained.act_postprocess4[3 : len(pretrained.act_postprocess4)](layer_4)
+
+    return layer_1, layer_2, layer_3, layer_4
+
+
+def _resize_pos_embed(self, posemb, gs_h, gs_w):
+    posemb_tok, posemb_grid = (
+        posemb[:, : self.start_index],
+        posemb[0, self.start_index :],
+    )
+
+    gs_old = int(math.sqrt(len(posemb_grid)))
+
+    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
+    posemb_grid = F.interpolate(posemb_grid, size=(gs_h, gs_w), mode="bilinear")
+    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_h * gs_w, -1)
+
+    posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
+
+    return posemb
+
+
+def forward_flex(self, x):
+    b, c, h, w = x.shape
+
+    pos_embed = self._resize_pos_embed(
+        self.pos_embed, h // self.patch_size[1], w // self.patch_size[0]
+    )
+
+    B = x.shape[0]
+
+    if hasattr(self.patch_embed, "backbone"):
+        x = self.patch_embed.backbone(x)
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+
+    x = self.patch_embed.proj(x).flatten(2).transpose(1, 2)
+
+    if getattr(self, "dist_token", None) is not None:
+        cls_tokens = self.cls_token.expand(
+            B, -1, -1
+        )  # stole cls_tokens impl from Phil Wang, thanks
+        dist_token = self.dist_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, dist_token, x), dim=1)
+    else:
+        cls_tokens = self.cls_token.expand(
+            B, -1, -1
+        )  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+
+    x = x + pos_embed
+    x = self.pos_drop(x)
+
+    for blk in self.blocks:
+        x = blk(x)
+
+    x = self.norm(x)
+
+    return x
+
+
+activations = {}
+
+
+def get_activation(name):
+    def hook(model, input, output):
+        activations[name] = output
+
+    return hook
+
+
+def get_readout_oper(vit_features, features, use_readout, start_index=1):
+    if use_readout == "ignore":
+        readout_oper = [Slice(start_index)] * len(features)
+    elif use_readout == "add":
+        readout_oper = [AddReadout(start_index)] * len(features)
+    elif use_readout == "project":
+        readout_oper = [
+            ProjectReadout(vit_features, start_index) for out_feat in features
+        ]
+    else:
+        assert (
+            False
+        ), "wrong operation for readout token, use_readout can be 'ignore', 'add', or 'project'"
+
+    return readout_oper
+
+
+def _make_vit_b16_backbone(
+    model,
+    features=[96, 192, 384, 768],
+    size=[384, 384],
+    hooks=[2, 5, 8, 11],
+    vit_features=768,
+    use_readout="ignore",
+    start_index=1,
+):
+    pretrained = nn.Module()
+
+    pretrained.model = model
+    pretrained.model.blocks[hooks[0]].register_forward_hook(get_activation("1"))
+    pretrained.model.blocks[hooks[1]].register_forward_hook(get_activation("2"))
+    pretrained.model.blocks[hooks[2]].register_forward_hook(get_activation("3"))
+    pretrained.model.blocks[hooks[3]].register_forward_hook(get_activation("4"))
+
+    pretrained.activations = activations
+
+    readout_oper = get_readout_oper(vit_features, features, use_readout, start_index)
+
+    # 32, 48, 136, 384
+    pretrained.act_postprocess1 = nn.Sequential(
+        readout_oper[0],
+        Transpose(1, 2),
+        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+        nn.Conv2d(
+            in_channels=vit_features,
+            out_channels=features[0],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        ),
+        nn.ConvTranspose2d(
+            in_channels=features[0],
+            out_channels=features[0],
+            kernel_size=4,
+            stride=4,
+            padding=0,
+            bias=True,
+            dilation=1,
+            groups=1,
+        ),
+    )
+
+    pretrained.act_postprocess2 = nn.Sequential(
+        readout_oper[1],
+        Transpose(1, 2),
+        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+        nn.Conv2d(
+            in_channels=vit_features,
+            out_channels=features[1],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        ),
+        nn.ConvTranspose2d(
+            in_channels=features[1],
+            out_channels=features[1],
+            kernel_size=2,
+            stride=2,
+            padding=0,
+            bias=True,
+            dilation=1,
+            groups=1,
+        ),
+    )
+
+    pretrained.act_postprocess3 = nn.Sequential(
+        readout_oper[2],
+        Transpose(1, 2),
+        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+        nn.Conv2d(
+            in_channels=vit_features,
+            out_channels=features[2],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        ),
+    )
+
+    pretrained.act_postprocess4 = nn.Sequential(
+        readout_oper[3],
+        Transpose(1, 2),
+        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+        nn.Conv2d(
+            in_channels=vit_features,
+            out_channels=features[3],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        ),
+        nn.Conv2d(
+            in_channels=features[3],
+            out_channels=features[3],
+            kernel_size=3,
+            stride=2,
+            padding=1,
+        ),
+    )
+
+    pretrained.model.start_index = start_index
+    pretrained.model.patch_size = [16, 16]
+
+    # We inject this function into the VisionTransformer instances so that
+    # we can use it with interpolated position embeddings without modifying the library source.
+    pretrained.model.forward_flex = types.MethodType(forward_flex, pretrained.model)
+    pretrained.model._resize_pos_embed = types.MethodType(
+        _resize_pos_embed, pretrained.model
+    )
+
+    return pretrained
+
+
+def _make_pretrained_vitl16_384(pretrained, use_readout="ignore", hooks=None):
+    model = timm.create_model("vit_large_patch16_384", pretrained=pretrained)
+
+    hooks = [5, 11, 17, 23] if hooks == None else hooks
+    return _make_vit_b16_backbone(
+        model,
+        features=[256, 512, 1024, 1024],
+        hooks=hooks,
+        vit_features=1024,
+        use_readout=use_readout,
+    )
+
+
+def _make_pretrained_vitb16_384(pretrained, use_readout="ignore", hooks=None):
+    model = timm.create_model("vit_base_patch16_384", pretrained=pretrained)
+
+    hooks = [2, 5, 8, 11] if hooks == None else hooks
+    return _make_vit_b16_backbone(
+        model, features=[96, 192, 384, 768], hooks=hooks, use_readout=use_readout
+    )
+
+
+def _make_pretrained_deitb16_384(pretrained, use_readout="ignore", hooks=None):
+    model = timm.create_model("vit_deit_base_patch16_384", pretrained=pretrained)
+
+    hooks = [2, 5, 8, 11] if hooks == None else hooks
+    return _make_vit_b16_backbone(
+        model, features=[96, 192, 384, 768], hooks=hooks, use_readout=use_readout
+    )
+
+
+def _make_pretrained_deitb16_distil_384(pretrained, use_readout="ignore", hooks=None):
+    model = timm.create_model(
+        "vit_deit_base_distilled_patch16_384", pretrained=pretrained
+    )
+
+    hooks = [2, 5, 8, 11] if hooks == None else hooks
+    return _make_vit_b16_backbone(
+        model,
+        features=[96, 192, 384, 768],
+        hooks=hooks,
+        use_readout=use_readout,
+        start_index=2,
+    )
+
+
+def _make_vit_b_rn50_backbone(
+    model,
+    features=[256, 512, 768, 768],
+    size=[384, 384],
+    hooks=[0, 1, 8, 11],
+    vit_features=768,
+    use_vit_only=False,
+    use_readout="ignore",
+    start_index=1,
+):
+    pretrained = nn.Module()
+
+    pretrained.model = model
+
+    if use_vit_only == True:
+        pretrained.model.blocks[hooks[0]].register_forward_hook(get_activation("1"))
+        pretrained.model.blocks[hooks[1]].register_forward_hook(get_activation("2"))
+    else:
+        pretrained.model.patch_embed.backbone.stages[0].register_forward_hook(
+            get_activation("1")
+        )
+        pretrained.model.patch_embed.backbone.stages[1].register_forward_hook(
+            get_activation("2")
+        )
+
+    pretrained.model.blocks[hooks[2]].register_forward_hook(get_activation("3"))
+    pretrained.model.blocks[hooks[3]].register_forward_hook(get_activation("4"))
+
+    pretrained.activations = activations
+
+    readout_oper = get_readout_oper(vit_features, features, use_readout, start_index)
+
+    if use_vit_only == True:
+        pretrained.act_postprocess1 = nn.Sequential(
+            readout_oper[0],
+            Transpose(1, 2),
+            nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+            nn.Conv2d(
+                in_channels=vit_features,
+                out_channels=features[0],
+                kernel_size=1,
+                stride=1,
+                padding=0,
+            ),
+            nn.ConvTranspose2d(
+                in_channels=features[0],
+                out_channels=features[0],
+                kernel_size=4,
+                stride=4,
+                padding=0,
+                bias=True,
+                dilation=1,
+                groups=1,
+            ),
+        )
+
+        pretrained.act_postprocess2 = nn.Sequential(
+            readout_oper[1],
+            Transpose(1, 2),
+            nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+            nn.Conv2d(
+                in_channels=vit_features,
+                out_channels=features[1],
+                kernel_size=1,
+                stride=1,
+                padding=0,
+            ),
+            nn.ConvTranspose2d(
+                in_channels=features[1],
+                out_channels=features[1],
+                kernel_size=2,
+                stride=2,
+                padding=0,
+                bias=True,
+                dilation=1,
+                groups=1,
+            ),
+        )
+    else:
+        pretrained.act_postprocess1 = nn.Sequential(
+            nn.Identity(), nn.Identity(), nn.Identity()
+        )
+        pretrained.act_postprocess2 = nn.Sequential(
+            nn.Identity(), nn.Identity(), nn.Identity()
+        )
+
+    pretrained.act_postprocess3 = nn.Sequential(
+        readout_oper[2],
+        Transpose(1, 2),
+        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+        nn.Conv2d(
+            in_channels=vit_features,
+            out_channels=features[2],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        ),
+    )
+
+    pretrained.act_postprocess4 = nn.Sequential(
+        readout_oper[3],
+        Transpose(1, 2),
+        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+        nn.Conv2d(
+            in_channels=vit_features,
+            out_channels=features[3],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        ),
+        nn.Conv2d(
+            in_channels=features[3],
+            out_channels=features[3],
+            kernel_size=3,
+            stride=2,
+            padding=1,
+        ),
+    )
+
+    pretrained.model.start_index = start_index
+    pretrained.model.patch_size = [16, 16]
+
+    # We inject this function into the VisionTransformer instances so that
+    # we can use it with interpolated position embeddings without modifying the library source.
+    pretrained.model.forward_flex = types.MethodType(forward_flex, pretrained.model)
+
+    # We inject this function into the VisionTransformer instances so that
+    # we can use it with interpolated position embeddings without modifying the library source.
+    pretrained.model._resize_pos_embed = types.MethodType(
+        _resize_pos_embed, pretrained.model
+    )
+
+    return pretrained
+
+
+def _make_pretrained_vitb_rn50_384(
+    pretrained, use_readout="ignore", hooks=None, use_vit_only=False
+):
+    model = timm.create_model("vit_base_resnet50_384", pretrained=pretrained)
+
+    hooks = [0, 1, 8, 11] if hooks == None else hooks
+    return _make_vit_b_rn50_backbone(
+        model,
+        features=[256, 512, 768, 768],
+        size=[384, 384],
+        hooks=hooks,
+        use_vit_only=use_vit_only,
+        use_readout=use_readout,
+    )

src/ControlNet/annotator/midas/utils.py

@@ -0,0 +1,189 @@
+"""Utils for monoDepth."""
+import sys
+import re
+import numpy as np
+import cv2
+import torch
+
+
+def read_pfm(path):
+    """Read pfm file.
+
+    Args:
+        path (str): path to file
+
+    Returns:
+        tuple: (data, scale)
+    """
+    with open(path, "rb") as file:
+
+        color = None
+        width = None
+        height = None
+        scale = None
+        endian = None
+
+        header = file.readline().rstrip()
+        if header.decode("ascii") == "PF":
+            color = True
+        elif header.decode("ascii") == "Pf":
+            color = False
+        else:
+            raise Exception("Not a PFM file: " + path)
+
+        dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("ascii"))
+        if dim_match:
+            width, height = list(map(int, dim_match.groups()))
+        else:
+            raise Exception("Malformed PFM header.")
+
+        scale = float(file.readline().decode("ascii").rstrip())
+        if scale < 0:
+            # little-endian
+            endian = "<"
+            scale = -scale
+        else:
+            # big-endian
+            endian = ">"
+
+        data = np.fromfile(file, endian + "f")
+        shape = (height, width, 3) if color else (height, width)
+
+        data = np.reshape(data, shape)
+        data = np.flipud(data)
+
+        return data, scale
+
+
+def write_pfm(path, image, scale=1):
+    """Write pfm file.
+
+    Args:
+        path (str): pathto file
+        image (array): data
+        scale (int, optional): Scale. Defaults to 1.
+    """
+
+    with open(path, "wb") as file:
+        color = None
+
+        if image.dtype.name != "float32":
+            raise Exception("Image dtype must be float32.")
+
+        image = np.flipud(image)
+
+        if len(image.shape) == 3 and image.shape[2] == 3:  # color image
+            color = True
+        elif (
+            len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1
+        ):  # greyscale
+            color = False
+        else:
+            raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")
+
+        file.write("PF\n" if color else "Pf\n".encode())
+        file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))
+
+        endian = image.dtype.byteorder
+
+        if endian == "<" or endian == "=" and sys.byteorder == "little":
+            scale = -scale
+
+        file.write("%f\n".encode() % scale)
+
+        image.tofile(file)
+
+
+def read_image(path):
+    """Read image and output RGB image (0-1).
+
+    Args:
+        path (str): path to file
+
+    Returns:
+        array: RGB image (0-1)
+    """
+    img = cv2.imread(path)
+
+    if img.ndim == 2:
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
+
+    return img
+
+
+def resize_image(img):
+    """Resize image and make it fit for network.
+
+    Args:
+        img (array): image
+
+    Returns:
+        tensor: data ready for network
+    """
+    height_orig = img.shape[0]
+    width_orig = img.shape[1]
+
+    if width_orig > height_orig:
+        scale = width_orig / 384
+    else:
+        scale = height_orig / 384
+
+    height = (np.ceil(height_orig / scale / 32) * 32).astype(int)
+    width = (np.ceil(width_orig / scale / 32) * 32).astype(int)
+
+    img_resized = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)
+
+    img_resized = (
+        torch.from_numpy(np.transpose(img_resized, (2, 0, 1))).contiguous().float()
+    )
+    img_resized = img_resized.unsqueeze(0)
+
+    return img_resized
+
+
+def resize_depth(depth, width, height):
+    """Resize depth map and bring to CPU (numpy).
+
+    Args:
+        depth (tensor): depth
+        width (int): image width
+        height (int): image height
+
+    Returns:
+        array: processed depth
+    """
+    depth = torch.squeeze(depth[0, :, :, :]).to("cpu")
+
+    depth_resized = cv2.resize(
+        depth.numpy(), (width, height), interpolation=cv2.INTER_CUBIC
+    )
+
+    return depth_resized
+
+def write_depth(path, depth, bits=1):
+    """Write depth map to pfm and png file.
+
+    Args:
+        path (str): filepath without extension
+        depth (array): depth
+    """
+    write_pfm(path + ".pfm", depth.astype(np.float32))
+
+    depth_min = depth.min()
+    depth_max = depth.max()
+
+    max_val = (2**(8*bits))-1
+
+    if depth_max - depth_min > np.finfo("float").eps:
+        out = max_val * (depth - depth_min) / (depth_max - depth_min)
+    else:
+        out = np.zeros(depth.shape, dtype=depth.type)
+
+    if bits == 1:
+        cv2.imwrite(path + ".png", out.astype("uint8"))
+    elif bits == 2:
+        cv2.imwrite(path + ".png", out.astype("uint16"))
+
+    return

src/ControlNet/annotator/mlsd/LICENSE

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

src/ControlNet/annotator/mlsd/__init__.py

@@ -0,0 +1,43 @@
+# MLSD Line Detection
+# From https://github.com/navervision/mlsd
+# Apache-2.0 license
+
+import cv2
+import numpy as np
+import torch
+import os
+
+from einops import rearrange
+from .models.mbv2_mlsd_tiny import MobileV2_MLSD_Tiny
+from .models.mbv2_mlsd_large import MobileV2_MLSD_Large
+from .utils import pred_lines
+
+from annotator.util import annotator_ckpts_path
+
+
+remote_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/mlsd_large_512_fp32.pth"
+
+
+class MLSDdetector:
+    def __init__(self):
+        model_path = os.path.join(annotator_ckpts_path, "mlsd_large_512_fp32.pth")
+        if not os.path.exists(model_path):
+            from basicsr.utils.download_util import load_file_from_url
+            load_file_from_url(remote_model_path, model_dir=annotator_ckpts_path)
+        model = MobileV2_MLSD_Large()
+        model.load_state_dict(torch.load(model_path), strict=True)
+        self.model = model.cuda().eval()
+
+    def __call__(self, input_image, thr_v, thr_d):
+        assert input_image.ndim == 3
+        img = input_image
+        img_output = np.zeros_like(img)
+        try:
+            with torch.no_grad():
+                lines = pred_lines(img, self.model, [img.shape[0], img.shape[1]], thr_v, thr_d)
+                for line in lines:
+                    x_start, y_start, x_end, y_end = [int(val) for val in line]
+                    cv2.line(img_output, (x_start, y_start), (x_end, y_end), [255, 255, 255], 1)
+        except Exception as e:
+            pass
+        return img_output[:, :, 0]

src/ControlNet/annotator/mlsd/models/mbv2_mlsd_large.py

@@ -0,0 +1,292 @@
+import os
+import sys
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+from  torch.nn import  functional as F
+
+
+class BlockTypeA(nn.Module):
+    def __init__(self, in_c1, in_c2, out_c1, out_c2, upscale = True):
+        super(BlockTypeA, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c2, out_c2, kernel_size=1),
+            nn.BatchNorm2d(out_c2),
+            nn.ReLU(inplace=True)
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c1, out_c1, kernel_size=1),
+            nn.BatchNorm2d(out_c1),
+            nn.ReLU(inplace=True)
+        )
+        self.upscale = upscale
+
+    def forward(self, a, b):
+        b = self.conv1(b)
+        a = self.conv2(a)
+        if self.upscale:
+             b = F.interpolate(b, scale_factor=2.0, mode='bilinear', align_corners=True)
+        return torch.cat((a, b), dim=1)
+
+
+class BlockTypeB(nn.Module):
+    def __init__(self, in_c, out_c):
+        super(BlockTypeB, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_c),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        x = self.conv1(x) + x
+        x = self.conv2(x)
+        return x
+
+class BlockTypeC(nn.Module):
+    def __init__(self, in_c, out_c):
+        super(BlockTypeC, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=5, dilation=5),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv3 = nn.Conv2d(in_c, out_c, kernel_size=1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        return x
+
+def _make_divisible(v, divisor, min_value=None):
+    """
+    This function is taken from the original tf repo.
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_value:
+    :return:
+    """
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class ConvBNReLU(nn.Sequential):
+    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
+        self.channel_pad = out_planes - in_planes
+        self.stride = stride
+        #padding = (kernel_size - 1) // 2
+
+        # TFLite uses slightly different padding than PyTorch
+        if stride == 2:
+            padding = 0
+        else:
+            padding = (kernel_size - 1) // 2
+
+        super(ConvBNReLU, self).__init__(
+            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
+            nn.BatchNorm2d(out_planes),
+            nn.ReLU6(inplace=True)
+        )
+        self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
+
+
+    def forward(self, x):
+        # TFLite uses  different padding
+        if self.stride == 2:
+            x = F.pad(x, (0, 1, 0, 1), "constant", 0)
+            #print(x.shape)
+
+        for module in self:
+            if not isinstance(module, nn.MaxPool2d):
+                x = module(x)
+        return x
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = int(round(inp * expand_ratio))
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        layers = []
+        if expand_ratio != 1:
+            # pw
+            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
+        layers.extend([
+            # dw
+            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
+            # pw-linear
+            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(oup),
+        ])
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, pretrained=True):
+        """
+        MobileNet V2 main class
+        Args:
+            num_classes (int): Number of classes
+            width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
+            inverted_residual_setting: Network structure
+            round_nearest (int): Round the number of channels in each layer to be a multiple of this number
+            Set to 1 to turn off rounding
+            block: Module specifying inverted residual building block for mobilenet
+        """
+        super(MobileNetV2, self).__init__()
+
+        block = InvertedResidual
+        input_channel = 32
+        last_channel = 1280
+        width_mult = 1.0
+        round_nearest = 8
+
+        inverted_residual_setting = [
+            # t, c, n, s
+            [1, 16, 1, 1],
+            [6, 24, 2, 2],
+            [6, 32, 3, 2],
+            [6, 64, 4, 2],
+            [6, 96, 3, 1],
+            #[6, 160, 3, 2],
+            #[6, 320, 1, 1],
+        ]
+
+        # only check the first element, assuming user knows t,c,n,s are required
+        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
+            raise ValueError("inverted_residual_setting should be non-empty "
+                             "or a 4-element list, got {}".format(inverted_residual_setting))
+
+        # building first layer
+        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
+        self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
+        features = [ConvBNReLU(4, input_channel, stride=2)]
+        # building inverted residual blocks
+        for t, c, n, s in inverted_residual_setting:
+            output_channel = _make_divisible(c * width_mult, round_nearest)
+            for i in range(n):
+                stride = s if i == 0 else 1
+                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
+                input_channel = output_channel
+
+        self.features = nn.Sequential(*features)
+        self.fpn_selected = [1, 3, 6, 10, 13]
+        # weight initialization
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.zeros_(m.bias)
+        if pretrained:
+           self._load_pretrained_model()
+
+    def _forward_impl(self, x):
+        # This exists since TorchScript doesn't support inheritance, so the superclass method
+        # (this one) needs to have a name other than `forward` that can be accessed in a subclass
+        fpn_features = []
+        for i, f in enumerate(self.features):
+            if i > self.fpn_selected[-1]:
+                break
+            x = f(x)
+            if i in self.fpn_selected:
+                fpn_features.append(x)
+
+        c1, c2, c3, c4, c5 = fpn_features
+        return c1, c2, c3, c4, c5
+
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+    def _load_pretrained_model(self):
+        pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/mobilenet_v2-b0353104.pth')
+        model_dict = {}
+        state_dict = self.state_dict()
+        for k, v in pretrain_dict.items():
+            if k in state_dict:
+                model_dict[k] = v
+        state_dict.update(model_dict)
+        self.load_state_dict(state_dict)
+
+
+class MobileV2_MLSD_Large(nn.Module):
+    def __init__(self):
+        super(MobileV2_MLSD_Large, self).__init__()
+
+        self.backbone = MobileNetV2(pretrained=False)
+        ## A, B
+        self.block15 = BlockTypeA(in_c1= 64, in_c2= 96,
+                                  out_c1= 64, out_c2=64,
+                                  upscale=False)
+        self.block16 = BlockTypeB(128, 64)
+
+        ## A, B
+        self.block17 = BlockTypeA(in_c1 = 32,  in_c2 = 64,
+                                  out_c1= 64,  out_c2= 64)
+        self.block18 = BlockTypeB(128, 64)
+
+        ## A, B
+        self.block19 = BlockTypeA(in_c1=24, in_c2=64,
+                                  out_c1=64, out_c2=64)
+        self.block20 = BlockTypeB(128, 64)
+
+        ## A, B, C
+        self.block21 = BlockTypeA(in_c1=16, in_c2=64,
+                                  out_c1=64, out_c2=64)
+        self.block22 = BlockTypeB(128, 64)
+
+        self.block23 = BlockTypeC(64, 16)
+
+    def forward(self, x):
+        c1, c2, c3, c4, c5 = self.backbone(x)
+
+        x = self.block15(c4, c5)
+        x = self.block16(x)
+
+        x = self.block17(c3, x)
+        x = self.block18(x)
+
+        x = self.block19(c2, x)
+        x = self.block20(x)
+
+        x = self.block21(c1, x)
+        x = self.block22(x)
+        x = self.block23(x)
+        x = x[:, 7:, :, :]
+
+        return x

src/ControlNet/annotator/mlsd/models/mbv2_mlsd_tiny.py

@@ -0,0 +1,275 @@
+import os
+import sys
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+from  torch.nn import  functional as F
+
+
+class BlockTypeA(nn.Module):
+    def __init__(self, in_c1, in_c2, out_c1, out_c2, upscale = True):
+        super(BlockTypeA, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c2, out_c2, kernel_size=1),
+            nn.BatchNorm2d(out_c2),
+            nn.ReLU(inplace=True)
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c1, out_c1, kernel_size=1),
+            nn.BatchNorm2d(out_c1),
+            nn.ReLU(inplace=True)
+        )
+        self.upscale = upscale
+
+    def forward(self, a, b):
+        b = self.conv1(b)
+        a = self.conv2(a)
+        b = F.interpolate(b, scale_factor=2.0, mode='bilinear', align_corners=True)
+        return torch.cat((a, b), dim=1)
+
+
+class BlockTypeB(nn.Module):
+    def __init__(self, in_c, out_c):
+        super(BlockTypeB, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_c),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        x = self.conv1(x) + x
+        x = self.conv2(x)
+        return x
+
+class BlockTypeC(nn.Module):
+    def __init__(self, in_c, out_c):
+        super(BlockTypeC, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=5, dilation=5),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv3 = nn.Conv2d(in_c, out_c, kernel_size=1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        return x
+
+def _make_divisible(v, divisor, min_value=None):
+    """
+    This function is taken from the original tf repo.
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_value:
+    :return:
+    """
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class ConvBNReLU(nn.Sequential):
+    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
+        self.channel_pad = out_planes - in_planes
+        self.stride = stride
+        #padding = (kernel_size - 1) // 2
+
+        # TFLite uses slightly different padding than PyTorch
+        if stride == 2:
+            padding = 0
+        else:
+            padding = (kernel_size - 1) // 2
+
+        super(ConvBNReLU, self).__init__(
+            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
+            nn.BatchNorm2d(out_planes),
+            nn.ReLU6(inplace=True)
+        )
+        self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
+
+
+    def forward(self, x):
+        # TFLite uses  different padding
+        if self.stride == 2:
+            x = F.pad(x, (0, 1, 0, 1), "constant", 0)
+            #print(x.shape)
+
+        for module in self:
+            if not isinstance(module, nn.MaxPool2d):
+                x = module(x)
+        return x
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = int(round(inp * expand_ratio))
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        layers = []
+        if expand_ratio != 1:
+            # pw
+            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
+        layers.extend([
+            # dw
+            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
+            # pw-linear
+            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(oup),
+        ])
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, pretrained=True):
+        """
+        MobileNet V2 main class
+        Args:
+            num_classes (int): Number of classes
+            width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
+            inverted_residual_setting: Network structure
+            round_nearest (int): Round the number of channels in each layer to be a multiple of this number
+            Set to 1 to turn off rounding
+            block: Module specifying inverted residual building block for mobilenet
+        """
+        super(MobileNetV2, self).__init__()
+
+        block = InvertedResidual
+        input_channel = 32
+        last_channel = 1280
+        width_mult = 1.0
+        round_nearest = 8
+
+        inverted_residual_setting = [
+            # t, c, n, s
+            [1, 16, 1, 1],
+            [6, 24, 2, 2],
+            [6, 32, 3, 2],
+            [6, 64, 4, 2],
+            #[6, 96, 3, 1],
+            #[6, 160, 3, 2],
+            #[6, 320, 1, 1],
+        ]
+
+        # only check the first element, assuming user knows t,c,n,s are required
+        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
+            raise ValueError("inverted_residual_setting should be non-empty "
+                             "or a 4-element list, got {}".format(inverted_residual_setting))
+
+        # building first layer
+        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
+        self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
+        features = [ConvBNReLU(4, input_channel, stride=2)]
+        # building inverted residual blocks
+        for t, c, n, s in inverted_residual_setting:
+            output_channel = _make_divisible(c * width_mult, round_nearest)
+            for i in range(n):
+                stride = s if i == 0 else 1
+                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
+                input_channel = output_channel
+        self.features = nn.Sequential(*features)
+
+        self.fpn_selected = [3, 6, 10]
+        # weight initialization
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.zeros_(m.bias)
+
+        #if pretrained:
+        #    self._load_pretrained_model()
+
+    def _forward_impl(self, x):
+        # This exists since TorchScript doesn't support inheritance, so the superclass method
+        # (this one) needs to have a name other than `forward` that can be accessed in a subclass
+        fpn_features = []
+        for i, f in enumerate(self.features):
+            if i > self.fpn_selected[-1]:
+                break
+            x = f(x)
+            if i in self.fpn_selected:
+                fpn_features.append(x)
+
+        c2, c3, c4 = fpn_features
+        return c2, c3, c4
+
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+    def _load_pretrained_model(self):
+        pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/mobilenet_v2-b0353104.pth')
+        model_dict = {}
+        state_dict = self.state_dict()
+        for k, v in pretrain_dict.items():
+            if k in state_dict:
+                model_dict[k] = v
+        state_dict.update(model_dict)
+        self.load_state_dict(state_dict)
+
+
+class MobileV2_MLSD_Tiny(nn.Module):
+    def __init__(self):
+        super(MobileV2_MLSD_Tiny, self).__init__()
+
+        self.backbone = MobileNetV2(pretrained=True)
+
+        self.block12 = BlockTypeA(in_c1= 32, in_c2= 64,
+                                  out_c1= 64, out_c2=64)
+        self.block13 = BlockTypeB(128, 64)
+
+        self.block14 = BlockTypeA(in_c1 = 24,  in_c2 = 64,
+                                  out_c1= 32,  out_c2= 32)
+        self.block15 = BlockTypeB(64, 64)
+
+        self.block16 = BlockTypeC(64, 16)
+
+    def forward(self, x):
+        c2, c3, c4 = self.backbone(x)
+
+        x = self.block12(c3, c4)
+        x = self.block13(x)
+        x = self.block14(c2, x)
+        x = self.block15(x)
+        x = self.block16(x)
+        x = x[:, 7:, :, :]
+        #print(x.shape)
+        x = F.interpolate(x, scale_factor=2.0, mode='bilinear', align_corners=True)
+
+        return x

src/ControlNet/annotator/mlsd/utils.py

@@ -0,0 +1,580 @@
+'''
+modified by  lihaoweicv
+pytorch version
+'''
+
+'''
+M-LSD
+Copyright 2021-present NAVER Corp.
+Apache License v2.0
+'''
+
+import os
+import numpy as np
+import cv2
+import torch
+from  torch.nn import  functional as F
+
+
+def deccode_output_score_and_ptss(tpMap, topk_n = 200, ksize = 5):
+    '''
+    tpMap:
+    center: tpMap[1, 0, :, :]
+    displacement: tpMap[1, 1:5, :, :]
+    '''
+    b, c, h, w = tpMap.shape
+    assert  b==1, 'only support bsize==1'
+    displacement = tpMap[:, 1:5, :, :][0]
+    center = tpMap[:, 0, :, :]
+    heat = torch.sigmoid(center)
+    hmax = F.max_pool2d( heat, (ksize, ksize), stride=1, padding=(ksize-1)//2)
+    keep = (hmax == heat).float()
+    heat = heat * keep
+    heat = heat.reshape(-1, )
+
+    scores, indices = torch.topk(heat, topk_n, dim=-1, largest=True)
+    yy = torch.floor_divide(indices, w).unsqueeze(-1)
+    xx = torch.fmod(indices, w).unsqueeze(-1)
+    ptss = torch.cat((yy, xx),dim=-1)
+
+    ptss   = ptss.detach().cpu().numpy()
+    scores = scores.detach().cpu().numpy()
+    displacement = displacement.detach().cpu().numpy()
+    displacement = displacement.transpose((1,2,0))
+    return  ptss, scores, displacement
+
+
+def pred_lines(image, model,
+               input_shape=[512, 512],
+               score_thr=0.10,
+               dist_thr=20.0):
+    h, w, _ = image.shape
+    h_ratio, w_ratio = [h / input_shape[0], w / input_shape[1]]
+
+    resized_image = np.concatenate([cv2.resize(image, (input_shape[1], input_shape[0]), interpolation=cv2.INTER_AREA),
+                                    np.ones([input_shape[0], input_shape[1], 1])], axis=-1)
+
+    resized_image = resized_image.transpose((2,0,1))
+    batch_image = np.expand_dims(resized_image, axis=0).astype('float32')
+    batch_image = (batch_image / 127.5) - 1.0
+
+    batch_image = torch.from_numpy(batch_image).float().cuda()
+    outputs = model(batch_image)
+    pts, pts_score, vmap = deccode_output_score_and_ptss(outputs, 200, 3)
+    start = vmap[:, :, :2]
+    end = vmap[:, :, 2:]
+    dist_map = np.sqrt(np.sum((start - end) ** 2, axis=-1))
+
+    segments_list = []
+    for center, score in zip(pts, pts_score):
+        y, x = center
+        distance = dist_map[y, x]
+        if score > score_thr and distance > dist_thr:
+            disp_x_start, disp_y_start, disp_x_end, disp_y_end = vmap[y, x, :]
+            x_start = x + disp_x_start
+            y_start = y + disp_y_start
+            x_end = x + disp_x_end
+            y_end = y + disp_y_end
+            segments_list.append([x_start, y_start, x_end, y_end])
+
+    lines = 2 * np.array(segments_list)  # 256 > 512
+    lines[:, 0] = lines[:, 0] * w_ratio
+    lines[:, 1] = lines[:, 1] * h_ratio
+    lines[:, 2] = lines[:, 2] * w_ratio
+    lines[:, 3] = lines[:, 3] * h_ratio
+
+    return lines
+
+
+def pred_squares(image,
+                 model,
+                 input_shape=[512, 512],
+                 params={'score': 0.06,
+                         'outside_ratio': 0.28,
+                         'inside_ratio': 0.45,
+                         'w_overlap': 0.0,
+                         'w_degree': 1.95,
+                         'w_length': 0.0,
+                         'w_area': 1.86,
+                         'w_center': 0.14}):
+    '''
+    shape = [height, width]
+    '''
+    h, w, _ = image.shape
+    original_shape = [h, w]
+
+    resized_image = np.concatenate([cv2.resize(image, (input_shape[0], input_shape[1]), interpolation=cv2.INTER_AREA),
+                                    np.ones([input_shape[0], input_shape[1], 1])], axis=-1)
+    resized_image = resized_image.transpose((2, 0, 1))
+    batch_image = np.expand_dims(resized_image, axis=0).astype('float32')
+    batch_image = (batch_image / 127.5) - 1.0
+
+    batch_image = torch.from_numpy(batch_image).float().cuda()
+    outputs = model(batch_image)
+
+    pts, pts_score, vmap = deccode_output_score_and_ptss(outputs, 200, 3)
+    start = vmap[:, :, :2]  # (x, y)
+    end = vmap[:, :, 2:]  # (x, y)
+    dist_map = np.sqrt(np.sum((start - end) ** 2, axis=-1))
+
+    junc_list = []
+    segments_list = []
+    for junc, score in zip(pts, pts_score):
+        y, x = junc
+        distance = dist_map[y, x]
+        if score > params['score'] and distance > 20.0:
+            junc_list.append([x, y])
+            disp_x_start, disp_y_start, disp_x_end, disp_y_end = vmap[y, x, :]
+            d_arrow = 1.0
+            x_start = x + d_arrow * disp_x_start
+            y_start = y + d_arrow * disp_y_start
+            x_end = x + d_arrow * disp_x_end
+            y_end = y + d_arrow * disp_y_end
+            segments_list.append([x_start, y_start, x_end, y_end])
+
+    segments = np.array(segments_list)
+
+    ####### post processing for squares
+    # 1. get unique lines
+    point = np.array([[0, 0]])
+    point = point[0]
+    start = segments[:, :2]
+    end = segments[:, 2:]
+    diff = start - end
+    a = diff[:, 1]
+    b = -diff[:, 0]
+    c = a * start[:, 0] + b * start[:, 1]
+
+    d = np.abs(a * point[0] + b * point[1] - c) / np.sqrt(a ** 2 + b ** 2 + 1e-10)
+    theta = np.arctan2(diff[:, 0], diff[:, 1]) * 180 / np.pi
+    theta[theta < 0.0] += 180
+    hough = np.concatenate([d[:, None], theta[:, None]], axis=-1)
+
+    d_quant = 1
+    theta_quant = 2
+    hough[:, 0] //= d_quant
+    hough[:, 1] //= theta_quant
+    _, indices, counts = np.unique(hough, axis=0, return_index=True, return_counts=True)
+
+    acc_map = np.zeros([512 // d_quant + 1, 360 // theta_quant + 1], dtype='float32')
+    idx_map = np.zeros([512 // d_quant + 1, 360 // theta_quant + 1], dtype='int32') - 1
+    yx_indices = hough[indices, :].astype('int32')
+    acc_map[yx_indices[:, 0], yx_indices[:, 1]] = counts
+    idx_map[yx_indices[:, 0], yx_indices[:, 1]] = indices
+
+    acc_map_np = acc_map
+    # acc_map = acc_map[None, :, :, None]
+    #
+    # ### fast suppression using tensorflow op
+    # acc_map = tf.constant(acc_map, dtype=tf.float32)
+    # max_acc_map = tf.keras.layers.MaxPool2D(pool_size=(5, 5), strides=1, padding='same')(acc_map)
+    # acc_map = acc_map * tf.cast(tf.math.equal(acc_map, max_acc_map), tf.float32)
+    # flatten_acc_map = tf.reshape(acc_map, [1, -1])
+    # topk_values, topk_indices = tf.math.top_k(flatten_acc_map, k=len(pts))
+    # _, h, w, _ = acc_map.shape
+    # y = tf.expand_dims(topk_indices // w, axis=-1)
+    # x = tf.expand_dims(topk_indices % w, axis=-1)
+    # yx = tf.concat([y, x], axis=-1)
+
+    ### fast suppression using pytorch op
+    acc_map = torch.from_numpy(acc_map_np).unsqueeze(0).unsqueeze(0)
+    _,_, h, w = acc_map.shape
+    max_acc_map = F.max_pool2d(acc_map,kernel_size=5, stride=1, padding=2)
+    acc_map = acc_map * ( (acc_map == max_acc_map).float() )
+    flatten_acc_map = acc_map.reshape([-1, ])
+
+    scores, indices = torch.topk(flatten_acc_map, len(pts), dim=-1, largest=True)
+    yy = torch.div(indices, w, rounding_mode='floor').unsqueeze(-1)
+    xx = torch.fmod(indices, w).unsqueeze(-1)
+    yx = torch.cat((yy, xx), dim=-1)
+
+    yx = yx.detach().cpu().numpy()
+
+    topk_values = scores.detach().cpu().numpy()
+    indices = idx_map[yx[:, 0], yx[:, 1]]
+    basis = 5 // 2
+
+    merged_segments = []
+    for yx_pt, max_indice, value in zip(yx, indices, topk_values):
+        y, x = yx_pt
+        if max_indice == -1 or value == 0:
+            continue
+        segment_list = []
+        for y_offset in range(-basis, basis + 1):
+            for x_offset in range(-basis, basis + 1):
+                indice = idx_map[y + y_offset, x + x_offset]
+                cnt = int(acc_map_np[y + y_offset, x + x_offset])
+                if indice != -1:
+                    segment_list.append(segments[indice])
+                if cnt > 1:
+                    check_cnt = 1
+                    current_hough = hough[indice]
+                    for new_indice, new_hough in enumerate(hough):
+                        if (current_hough == new_hough).all() and indice != new_indice:
+                            segment_list.append(segments[new_indice])
+                            check_cnt += 1
+                        if check_cnt == cnt:
+                            break
+        group_segments = np.array(segment_list).reshape([-1, 2])
+        sorted_group_segments = np.sort(group_segments, axis=0)
+        x_min, y_min = sorted_group_segments[0, :]
+        x_max, y_max = sorted_group_segments[-1, :]
+
+        deg = theta[max_indice]
+        if deg >= 90:
+            merged_segments.append([x_min, y_max, x_max, y_min])
+        else:
+            merged_segments.append([x_min, y_min, x_max, y_max])
+
+    # 2. get intersections
+    new_segments = np.array(merged_segments)  # (x1, y1, x2, y2)
+    start = new_segments[:, :2]  # (x1, y1)
+    end = new_segments[:, 2:]  # (x2, y2)
+    new_centers = (start + end) / 2.0
+    diff = start - end
+    dist_segments = np.sqrt(np.sum(diff ** 2, axis=-1))
+
+    # ax + by = c
+    a = diff[:, 1]
+    b = -diff[:, 0]
+    c = a * start[:, 0] + b * start[:, 1]
+    pre_det = a[:, None] * b[None, :]
+    det = pre_det - np.transpose(pre_det)
+
+    pre_inter_y = a[:, None] * c[None, :]
+    inter_y = (pre_inter_y - np.transpose(pre_inter_y)) / (det + 1e-10)
+    pre_inter_x = c[:, None] * b[None, :]
+    inter_x = (pre_inter_x - np.transpose(pre_inter_x)) / (det + 1e-10)
+    inter_pts = np.concatenate([inter_x[:, :, None], inter_y[:, :, None]], axis=-1).astype('int32')
+
+    # 3. get corner information
+    # 3.1 get distance
+    '''
+    dist_segments:
+        | dist(0), dist(1), dist(2), ...|
+    dist_inter_to_segment1:
+        | dist(inter,0), dist(inter,0), dist(inter,0), ... |
+        | dist(inter,1), dist(inter,1), dist(inter,1), ... |
+        ...
+    dist_inter_to_semgnet2:
+        | dist(inter,0), dist(inter,1), dist(inter,2), ... |
+        | dist(inter,0), dist(inter,1), dist(inter,2), ... |
+        ...
+    '''
+
+    dist_inter_to_segment1_start = np.sqrt(
+        np.sum(((inter_pts - start[:, None, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
+    dist_inter_to_segment1_end = np.sqrt(
+        np.sum(((inter_pts - end[:, None, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
+    dist_inter_to_segment2_start = np.sqrt(
+        np.sum(((inter_pts - start[None, :, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
+    dist_inter_to_segment2_end = np.sqrt(
+        np.sum(((inter_pts - end[None, :, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
+
+    # sort ascending
+    dist_inter_to_segment1 = np.sort(
+        np.concatenate([dist_inter_to_segment1_start, dist_inter_to_segment1_end], axis=-1),
+        axis=-1)  # [n_batch, n_batch, 2]
+    dist_inter_to_segment2 = np.sort(
+        np.concatenate([dist_inter_to_segment2_start, dist_inter_to_segment2_end], axis=-1),
+        axis=-1)  # [n_batch, n_batch, 2]
+
+    # 3.2 get degree
+    inter_to_start = new_centers[:, None, :] - inter_pts
+    deg_inter_to_start = np.arctan2(inter_to_start[:, :, 1], inter_to_start[:, :, 0]) * 180 / np.pi
+    deg_inter_to_start[deg_inter_to_start < 0.0] += 360
+    inter_to_end = new_centers[None, :, :] - inter_pts
+    deg_inter_to_end = np.arctan2(inter_to_end[:, :, 1], inter_to_end[:, :, 0]) * 180 / np.pi
+    deg_inter_to_end[deg_inter_to_end < 0.0] += 360
+
+    '''
+    B -- G
+    |    |
+    C -- R
+    B : blue / G: green / C: cyan / R: red
+
+    0 -- 1
+    |    |
+    3 -- 2
+    '''
+    # rename variables
+    deg1_map, deg2_map = deg_inter_to_start, deg_inter_to_end
+    # sort deg ascending
+    deg_sort = np.sort(np.concatenate([deg1_map[:, :, None], deg2_map[:, :, None]], axis=-1), axis=-1)
+
+    deg_diff_map = np.abs(deg1_map - deg2_map)
+    # we only consider the smallest degree of intersect
+    deg_diff_map[deg_diff_map > 180] = 360 - deg_diff_map[deg_diff_map > 180]
+
+    # define available degree range
+    deg_range = [60, 120]
+
+    corner_dict = {corner_info: [] for corner_info in range(4)}
+    inter_points = []
+    for i in range(inter_pts.shape[0]):
+        for j in range(i + 1, inter_pts.shape[1]):
+            # i, j > line index, always i < j
+            x, y = inter_pts[i, j, :]
+            deg1, deg2 = deg_sort[i, j, :]
+            deg_diff = deg_diff_map[i, j]
+
+            check_degree = deg_diff > deg_range[0] and deg_diff < deg_range[1]
+
+            outside_ratio = params['outside_ratio']  # over ratio >>> drop it!
+            inside_ratio = params['inside_ratio']  # over ratio >>> drop it!
+            check_distance = ((dist_inter_to_segment1[i, j, 1] >= dist_segments[i] and \
+                               dist_inter_to_segment1[i, j, 0] <= dist_segments[i] * outside_ratio) or \
+                              (dist_inter_to_segment1[i, j, 1] <= dist_segments[i] and \
+                               dist_inter_to_segment1[i, j, 0] <= dist_segments[i] * inside_ratio)) and \
+                             ((dist_inter_to_segment2[i, j, 1] >= dist_segments[j] and \
+                               dist_inter_to_segment2[i, j, 0] <= dist_segments[j] * outside_ratio) or \
+                              (dist_inter_to_segment2[i, j, 1] <= dist_segments[j] and \
+                               dist_inter_to_segment2[i, j, 0] <= dist_segments[j] * inside_ratio))
+
+            if check_degree and check_distance:
+                corner_info = None
+
+                if (deg1 >= 0 and deg1 <= 45 and deg2 >= 45 and deg2 <= 120) or \
+                        (deg2 >= 315 and deg1 >= 45 and deg1 <= 120):
+                    corner_info, color_info = 0, 'blue'
+                elif (deg1 >= 45 and deg1 <= 125 and deg2 >= 125 and deg2 <= 225):
+                    corner_info, color_info = 1, 'green'
+                elif (deg1 >= 125 and deg1 <= 225 and deg2 >= 225 and deg2 <= 315):
+                    corner_info, color_info = 2, 'black'
+                elif (deg1 >= 0 and deg1 <= 45 and deg2 >= 225 and deg2 <= 315) or \
+                        (deg2 >= 315 and deg1 >= 225 and deg1 <= 315):
+                    corner_info, color_info = 3, 'cyan'
+                else:
+                    corner_info, color_info = 4, 'red'  # we don't use it
+                    continue
+
+                corner_dict[corner_info].append([x, y, i, j])
+                inter_points.append([x, y])
+
+    square_list = []
+    connect_list = []
+    segments_list = []
+    for corner0 in corner_dict[0]:
+        for corner1 in corner_dict[1]:
+            connect01 = False
+            for corner0_line in corner0[2:]:
+                if corner0_line in corner1[2:]:
+                    connect01 = True
+                    break
+            if connect01:
+                for corner2 in corner_dict[2]:
+                    connect12 = False
+                    for corner1_line in corner1[2:]:
+                        if corner1_line in corner2[2:]:
+                            connect12 = True
+                            break
+                    if connect12:
+                        for corner3 in corner_dict[3]:
+                            connect23 = False
+                            for corner2_line in corner2[2:]:
+                                if corner2_line in corner3[2:]:
+                                    connect23 = True
+                                    break
+                            if connect23:
+                                for corner3_line in corner3[2:]:
+                                    if corner3_line in corner0[2:]:
+                                        # SQUARE!!!
+                                        '''
+                                        0 -- 1
+                                        |    |
+                                        3 -- 2
+                                        square_list:
+                                            order: 0 > 1 > 2 > 3
+                                            | x0, y0, x1, y1, x2, y2, x3, y3 |
+                                            | x0, y0, x1, y1, x2, y2, x3, y3 |
+                                            ...
+                                        connect_list:
+                                            order: 01 > 12 > 23 > 30
+                                            | line_idx01, line_idx12, line_idx23, line_idx30 |
+                                            | line_idx01, line_idx12, line_idx23, line_idx30 |
+                                            ...
+                                        segments_list:
+                                            order: 0 > 1 > 2 > 3
+                                            | line_idx0_i, line_idx0_j, line_idx1_i, line_idx1_j, line_idx2_i, line_idx2_j, line_idx3_i, line_idx3_j |
+                                            | line_idx0_i, line_idx0_j, line_idx1_i, line_idx1_j, line_idx2_i, line_idx2_j, line_idx3_i, line_idx3_j |
+                                            ...
+                                        '''
+                                        square_list.append(corner0[:2] + corner1[:2] + corner2[:2] + corner3[:2])
+                                        connect_list.append([corner0_line, corner1_line, corner2_line, corner3_line])
+                                        segments_list.append(corner0[2:] + corner1[2:] + corner2[2:] + corner3[2:])
+
+    def check_outside_inside(segments_info, connect_idx):
+        # return 'outside or inside', min distance, cover_param, peri_param
+        if connect_idx == segments_info[0]:
+            check_dist_mat = dist_inter_to_segment1
+        else:
+            check_dist_mat = dist_inter_to_segment2
+
+        i, j = segments_info
+        min_dist, max_dist = check_dist_mat[i, j, :]
+        connect_dist = dist_segments[connect_idx]
+        if max_dist > connect_dist:
+            return 'outside', min_dist, 0, 1
+        else:
+            return 'inside', min_dist, -1, -1
+
+    top_square = None
+
+    try:
+        map_size = input_shape[0] / 2
+        squares = np.array(square_list).reshape([-1, 4, 2])
+        score_array = []
+        connect_array = np.array(connect_list)
+        segments_array = np.array(segments_list).reshape([-1, 4, 2])
+
+        # get degree of corners:
+        squares_rollup = np.roll(squares, 1, axis=1)
+        squares_rolldown = np.roll(squares, -1, axis=1)
+        vec1 = squares_rollup - squares
+        normalized_vec1 = vec1 / (np.linalg.norm(vec1, axis=-1, keepdims=True) + 1e-10)
+        vec2 = squares_rolldown - squares
+        normalized_vec2 = vec2 / (np.linalg.norm(vec2, axis=-1, keepdims=True) + 1e-10)
+        inner_products = np.sum(normalized_vec1 * normalized_vec2, axis=-1)  # [n_squares, 4]
+        squares_degree = np.arccos(inner_products) * 180 / np.pi  # [n_squares, 4]
+
+        # get square score
+        overlap_scores = []
+        degree_scores = []
+        length_scores = []
+
+        for connects, segments, square, degree in zip(connect_array, segments_array, squares, squares_degree):
+            '''
+            0 -- 1
+            |    |
+            3 -- 2
+
+            # segments: [4, 2]
+            # connects: [4]
+            '''
+
+            ###################################### OVERLAP SCORES
+            cover = 0
+            perimeter = 0
+            # check 0 > 1 > 2 > 3
+            square_length = []
+
+            for start_idx in range(4):
+                end_idx = (start_idx + 1) % 4
+
+                connect_idx = connects[start_idx]  # segment idx of segment01
+                start_segments = segments[start_idx]
+                end_segments = segments[end_idx]
+
+                start_point = square[start_idx]
+                end_point = square[end_idx]
+
+                # check whether outside or inside
+                start_position, start_min, start_cover_param, start_peri_param = check_outside_inside(start_segments,
+                                                                                                      connect_idx)
+                end_position, end_min, end_cover_param, end_peri_param = check_outside_inside(end_segments, connect_idx)
+
+                cover += dist_segments[connect_idx] + start_cover_param * start_min + end_cover_param * end_min
+                perimeter += dist_segments[connect_idx] + start_peri_param * start_min + end_peri_param * end_min
+
+                square_length.append(
+                    dist_segments[connect_idx] + start_peri_param * start_min + end_peri_param * end_min)
+
+            overlap_scores.append(cover / perimeter)
+            ######################################
+            ###################################### DEGREE SCORES
+            '''
+            deg0 vs deg2
+            deg1 vs deg3
+            '''
+            deg0, deg1, deg2, deg3 = degree
+            deg_ratio1 = deg0 / deg2
+            if deg_ratio1 > 1.0:
+                deg_ratio1 = 1 / deg_ratio1
+            deg_ratio2 = deg1 / deg3
+            if deg_ratio2 > 1.0:
+                deg_ratio2 = 1 / deg_ratio2
+            degree_scores.append((deg_ratio1 + deg_ratio2) / 2)
+            ######################################
+            ###################################### LENGTH SCORES
+            '''
+            len0 vs len2
+            len1 vs len3
+            '''
+            len0, len1, len2, len3 = square_length
+            len_ratio1 = len0 / len2 if len2 > len0 else len2 / len0
+            len_ratio2 = len1 / len3 if len3 > len1 else len3 / len1
+            length_scores.append((len_ratio1 + len_ratio2) / 2)
+
+            ######################################
+
+        overlap_scores = np.array(overlap_scores)
+        overlap_scores /= np.max(overlap_scores)
+
+        degree_scores = np.array(degree_scores)
+        # degree_scores /= np.max(degree_scores)
+
+        length_scores = np.array(length_scores)
+
+        ###################################### AREA SCORES
+        area_scores = np.reshape(squares, [-1, 4, 2])
+        area_x = area_scores[:, :, 0]
+        area_y = area_scores[:, :, 1]
+        correction = area_x[:, -1] * area_y[:, 0] - area_y[:, -1] * area_x[:, 0]
+        area_scores = np.sum(area_x[:, :-1] * area_y[:, 1:], axis=-1) - np.sum(area_y[:, :-1] * area_x[:, 1:], axis=-1)
+        area_scores = 0.5 * np.abs(area_scores + correction)
+        area_scores /= (map_size * map_size)  # np.max(area_scores)
+        ######################################
+
+        ###################################### CENTER SCORES
+        centers = np.array([[256 // 2, 256 // 2]], dtype='float32')  # [1, 2]
+        # squares: [n, 4, 2]
+        square_centers = np.mean(squares, axis=1)  # [n, 2]
+        center2center = np.sqrt(np.sum((centers - square_centers) ** 2))
+        center_scores = center2center / (map_size / np.sqrt(2.0))
+
+        '''
+        score_w = [overlap, degree, area, center, length]
+        '''
+        score_w = [0.0, 1.0, 10.0, 0.5, 1.0]
+        score_array = params['w_overlap'] * overlap_scores \
+                      + params['w_degree'] * degree_scores \
+                      + params['w_area'] * area_scores \
+                      - params['w_center'] * center_scores \
+                      + params['w_length'] * length_scores
+
+        best_square = []
+
+        sorted_idx = np.argsort(score_array)[::-1]
+        score_array = score_array[sorted_idx]
+        squares = squares[sorted_idx]
+
+    except Exception as e:
+        pass
+
+    '''return list
+    merged_lines, squares, scores
+    '''
+
+    try:
+        new_segments[:, 0] = new_segments[:, 0] * 2 / input_shape[1] * original_shape[1]
+        new_segments[:, 1] = new_segments[:, 1] * 2 / input_shape[0] * original_shape[0]
+        new_segments[:, 2] = new_segments[:, 2] * 2 / input_shape[1] * original_shape[1]
+        new_segments[:, 3] = new_segments[:, 3] * 2 / input_shape[0] * original_shape[0]
+    except:
+        new_segments = []
+
+    try:
+        squares[:, :, 0] = squares[:, :, 0] * 2 / input_shape[1] * original_shape[1]
+        squares[:, :, 1] = squares[:, :, 1] * 2 / input_shape[0] * original_shape[0]
+    except:
+        squares = []
+        score_array = []
+
+    try:
+        inter_points = np.array(inter_points)
+        inter_points[:, 0] = inter_points[:, 0] * 2 / input_shape[1] * original_shape[1]
+        inter_points[:, 1] = inter_points[:, 1] * 2 / input_shape[0] * original_shape[0]
+    except:
+        inter_points = []
+
+    return new_segments, squares, score_array, inter_points

src/ControlNet/annotator/openpose/LICENSE

@@ -0,0 +1,108 @@
+OPENPOSE: MULTIPERSON KEYPOINT DETECTION
+SOFTWARE LICENSE AGREEMENT
+ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY
+
+BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS LICENSE AGREEMENT.  IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR DOWNLOAD THE SOFTWARE.
+
+This is a license agreement ("Agreement") between your academic institution or non-profit organization or self (called "Licensee" or "You" in this Agreement) and Carnegie Mellon University (called "Licensor" in this Agreement).  All rights not specifically granted to you in this Agreement are reserved for Licensor. 
+
+RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: 
+Licensor retains exclusive ownership of any copy of the Software (as defined below) licensed under this Agreement and hereby grants to Licensee a personal, non-exclusive, 
+non-transferable license to use the Software for noncommercial research purposes, without the right to sublicense, pursuant to the terms and conditions of this Agreement.  As used in this Agreement, the term "Software" means (i) the actual copy of all or any portion of code for program routines made accessible to Licensee by Licensor pursuant to this Agreement, inclusive of backups, updates, and/or merged copies permitted hereunder or subsequently supplied by Licensor,  including all or any file structures, programming instructions, user interfaces and screen formats and sequences as well as any and all documentation and instructions related to it, and (ii) all or any derivatives and/or modifications created or made by You to any of the items specified in (i).
+
+CONFIDENTIALITY: Licensee acknowledges that the Software is proprietary to Licensor, and as such, Licensee agrees to receive all such materials in confidence and use the Software only in accordance with the terms of this Agreement.  Licensee agrees to use reasonable effort to protect the Software from unauthorized use, reproduction, distribution, or publication.
+
+COPYRIGHT: The Software is owned by Licensor and is protected by United 
+States copyright laws and applicable international treaties and/or conventions.
+
+PERMITTED USES:  The Software may be used for your own noncommercial internal research purposes. You understand and agree that Licensor is not obligated to implement any suggestions and/or feedback you might provide regarding the Software, but to the extent Licensor does so, you are not entitled to any compensation related thereto.
+
+DERIVATIVES: You may create derivatives of or make modifications to the Software, however, You agree that all and any such derivatives and modifications will be owned by Licensor and become a part of the Software licensed to You under this Agreement.  You may only use such derivatives and modifications for your own noncommercial internal research purposes, and you may not otherwise use, distribute or copy such derivatives and modifications in violation of this Agreement.
+
+BACKUPS:  If Licensee is an organization, it may make that number of copies of the Software necessary for internal noncommercial use at a single site within its organization provided that all information appearing in or on the original labels, including the copyright and trademark notices are copied onto the labels of the copies.
+
+USES NOT PERMITTED:  You may not distribute, copy or use the Software except as explicitly permitted herein. Licensee has not been granted any trademark license as part of this Agreement and may not use the name or mark “OpenPose", "Carnegie Mellon" or any renditions thereof without the prior written permission of Licensor.
+
+You may not sell, rent, lease, sublicense, lend, time-share or transfer, in whole or in part, or provide third parties access to prior or present versions (or any parts thereof) of the Software.
+
+ASSIGNMENT: You may not assign this Agreement or your rights hereunder without the prior written consent of Licensor. Any attempted assignment without such consent shall be null and void.
+
+TERM: The term of the license granted by this Agreement is from Licensee's acceptance of this Agreement by downloading the Software or by using the Software until terminated as provided below.
+
+The Agreement automatically terminates without notice if you fail to comply with any provision of this Agreement.  Licensee may terminate this Agreement by ceasing using the Software.  Upon any termination of this Agreement, Licensee will delete any and all copies of the Software. You agree that all provisions which operate to protect the proprietary rights of Licensor shall remain in force should breach occur and that the obligation of confidentiality described in this Agreement is binding in perpetuity and, as such, survives the term of the Agreement.
+
+FEE: Provided Licensee abides completely by the terms and conditions of this Agreement, there is no fee due to Licensor for Licensee's use of the Software in accordance with this Agreement.
+
+DISCLAIMER OF WARRANTIES:  THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON-INFRINGEMENT.  LICENSEE BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND RELATED MATERIALS.
+
+SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is provided as part of this Agreement.  
+
+EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent permitted under applicable law, Licensor shall not be liable for direct, indirect, special, incidental, or consequential damages or lost profits related to Licensee's use of and/or inability to use the Software, even if Licensor is advised of the possibility of such damage.
+
+EXPORT REGULATION: Licensee agrees to comply with any and all applicable 
+U.S. export control laws, regulations, and/or other laws related to embargoes and sanction programs administered by the Office of Foreign Assets Control.
+
+SEVERABILITY: If any provision(s) of this Agreement shall be held to be invalid, illegal, or unenforceable by a court or other tribunal of competent jurisdiction, the validity, legality and enforceability of the remaining provisions shall not in any way be affected or impaired thereby.
+
+NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right or remedy under this Agreement shall be construed as a waiver of any future or other exercise of such right or remedy by Licensor.
+
+GOVERNING LAW: This Agreement shall be construed and enforced in accordance with the laws of the Commonwealth of Pennsylvania without reference to conflict of laws principles.  You consent to the personal jurisdiction of the courts of this County and waive their rights to venue outside of Allegheny County, Pennsylvania.
+
+ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and entire agreement between Licensee and Licensor as to the matter set forth herein and supersedes any previous agreements, understandings, and arrangements between the parties relating hereto.
+
+
+
+************************************************************************
+
+THIRD-PARTY SOFTWARE NOTICES AND INFORMATION
+
+This project incorporates material from the project(s) listed below (collectively, "Third Party Code").  This Third Party Code is licensed to you under their original license terms set forth below.  We reserves all other rights not expressly granted, whether by implication, estoppel or otherwise.
+ 
+1.	Caffe, version 1.0.0, (https://github.com/BVLC/caffe/)
+
+COPYRIGHT
+
+All contributions by the University of California:
+Copyright (c) 2014-2017 The Regents of the University of California (Regents)
+All rights reserved.
+
+All other contributions:
+Copyright (c) 2014-2017, the respective contributors
+All rights reserved.
+
+Caffe uses a shared copyright model: each contributor holds copyright over
+their contributions to Caffe. The project versioning records all such
+contribution and copyright details. If a contributor wants to further mark
+their specific copyright on a particular contribution, they should indicate
+their copyright solely in the commit message of the change when it is
+committed.
+
+LICENSE
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met: 
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer. 
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution. 
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+CONTRIBUTION AGREEMENT
+
+By contributing to the BVLC/caffe repository through pull-request, comment,
+or otherwise, the contributor releases their content to the
+license and copyright terms herein.
+
+************END OF THIRD-PARTY SOFTWARE NOTICES AND INFORMATION**********

src/ControlNet/annotator/openpose/__init__.py

@@ -0,0 +1,49 @@
+# Openpose
+# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose
+# 2nd Edited by https://github.com/Hzzone/pytorch-openpose
+# 3rd Edited by ControlNet
+
+import os
+os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
+
+import torch
+import numpy as np
+from . import util
+from .body import Body
+from .hand import Hand
+from annotator.util import annotator_ckpts_path
+
+
+body_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/body_pose_model.pth"
+hand_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/hand_pose_model.pth"
+
+
+class OpenposeDetector:
+    def __init__(self):
+        body_modelpath = os.path.join(annotator_ckpts_path, "body_pose_model.pth")
+        hand_modelpath = os.path.join(annotator_ckpts_path, "hand_pose_model.pth")
+
+        if not os.path.exists(hand_modelpath):
+            from basicsr.utils.download_util import load_file_from_url
+            load_file_from_url(body_model_path, model_dir=annotator_ckpts_path)
+            load_file_from_url(hand_model_path, model_dir=annotator_ckpts_path)
+
+        self.body_estimation = Body(body_modelpath)
+        self.hand_estimation = Hand(hand_modelpath)
+
+    def __call__(self, oriImg, hand=False):
+        oriImg = oriImg[:, :, ::-1].copy()
+        with torch.no_grad():
+            candidate, subset = self.body_estimation(oriImg)
+            canvas = np.zeros_like(oriImg)
+            canvas = util.draw_bodypose(canvas, candidate, subset)
+            if hand:
+                hands_list = util.handDetect(candidate, subset, oriImg)
+                all_hand_peaks = []
+                for x, y, w, is_left in hands_list:
+                    peaks = self.hand_estimation(oriImg[y:y+w, x:x+w, :])
+                    peaks[:, 0] = np.where(peaks[:, 0] == 0, peaks[:, 0], peaks[:, 0] + x)
+                    peaks[:, 1] = np.where(peaks[:, 1] == 0, peaks[:, 1], peaks[:, 1] + y)
+                    all_hand_peaks.append(peaks)
+                canvas = util.draw_handpose(canvas, all_hand_peaks)
+            return canvas, dict(candidate=candidate.tolist(), subset=subset.tolist())

src/ControlNet/annotator/openpose/body.py

@@ -0,0 +1,219 @@
+import cv2
+import numpy as np
+import math
+import time
+from scipy.ndimage.filters import gaussian_filter
+import matplotlib.pyplot as plt
+import matplotlib
+import torch
+from torchvision import transforms
+
+from . import util
+from .model import bodypose_model
+
+class Body(object):
+    def __init__(self, model_path):
+        self.model = bodypose_model()
+        if torch.cuda.is_available():
+            self.model = self.model.cuda()
+            print('cuda')
+        model_dict = util.transfer(self.model, torch.load(model_path))
+        self.model.load_state_dict(model_dict)
+        self.model.eval()
+
+    def __call__(self, oriImg):
+        # scale_search = [0.5, 1.0, 1.5, 2.0]
+        scale_search = [0.5]
+        boxsize = 368
+        stride = 8
+        padValue = 128
+        thre1 = 0.1
+        thre2 = 0.05
+        multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
+        heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
+        paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
+
+        for m in range(len(multiplier)):
+            scale = multiplier[m]
+            imageToTest = cv2.resize(oriImg, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
+            imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
+            im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
+            im = np.ascontiguousarray(im)
+
+            data = torch.from_numpy(im).float()
+            if torch.cuda.is_available():
+                data = data.cuda()
+            # data = data.permute([2, 0, 1]).unsqueeze(0).float()
+            with torch.no_grad():
+                Mconv7_stage6_L1, Mconv7_stage6_L2 = self.model(data)
+            Mconv7_stage6_L1 = Mconv7_stage6_L1.cpu().numpy()
+            Mconv7_stage6_L2 = Mconv7_stage6_L2.cpu().numpy()
+
+            # extract outputs, resize, and remove padding
+            # heatmap = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[1]].data), (1, 2, 0))  # output 1 is heatmaps
+            heatmap = np.transpose(np.squeeze(Mconv7_stage6_L2), (1, 2, 0))  # output 1 is heatmaps
+            heatmap = cv2.resize(heatmap, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
+            heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
+            heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)
+
+            # paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[0]].data), (1, 2, 0))  # output 0 is PAFs
+            paf = np.transpose(np.squeeze(Mconv7_stage6_L1), (1, 2, 0))  # output 0 is PAFs
+            paf = cv2.resize(paf, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
+            paf = paf[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
+            paf = cv2.resize(paf, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)
+
+            heatmap_avg += heatmap_avg + heatmap / len(multiplier)
+            paf_avg += + paf / len(multiplier)
+
+        all_peaks = []
+        peak_counter = 0
+
+        for part in range(18):
+            map_ori = heatmap_avg[:, :, part]
+            one_heatmap = gaussian_filter(map_ori, sigma=3)
+
+            map_left = np.zeros(one_heatmap.shape)
+            map_left[1:, :] = one_heatmap[:-1, :]
+            map_right = np.zeros(one_heatmap.shape)
+            map_right[:-1, :] = one_heatmap[1:, :]
+            map_up = np.zeros(one_heatmap.shape)
+            map_up[:, 1:] = one_heatmap[:, :-1]
+            map_down = np.zeros(one_heatmap.shape)
+            map_down[:, :-1] = one_heatmap[:, 1:]
+
+            peaks_binary = np.logical_and.reduce(
+                (one_heatmap >= map_left, one_heatmap >= map_right, one_heatmap >= map_up, one_heatmap >= map_down, one_heatmap > thre1))
+            peaks = list(zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]))  # note reverse
+            peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
+            peak_id = range(peak_counter, peak_counter + len(peaks))
+            peaks_with_score_and_id = [peaks_with_score[i] + (peak_id[i],) for i in range(len(peak_id))]
+
+            all_peaks.append(peaks_with_score_and_id)
+            peak_counter += len(peaks)
+
+        # find connection in the specified sequence, center 29 is in the position 15
+        limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
+                   [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
+                   [1, 16], [16, 18], [3, 17], [6, 18]]
+        # the middle joints heatmap correpondence
+        mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22], \
+                  [23, 24], [25, 26], [27, 28], [29, 30], [47, 48], [49, 50], [53, 54], [51, 52], \
+                  [55, 56], [37, 38], [45, 46]]
+
+        connection_all = []
+        special_k = []
+        mid_num = 10
+
+        for k in range(len(mapIdx)):
+            score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
+            candA = all_peaks[limbSeq[k][0] - 1]
+            candB = all_peaks[limbSeq[k][1] - 1]
+            nA = len(candA)
+            nB = len(candB)
+            indexA, indexB = limbSeq[k]
+            if (nA != 0 and nB != 0):
+                connection_candidate = []
+                for i in range(nA):
+                    for j in range(nB):
+                        vec = np.subtract(candB[j][:2], candA[i][:2])
+                        norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
+                        norm = max(0.001, norm)
+                        vec = np.divide(vec, norm)
+
+                        startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
+                                            np.linspace(candA[i][1], candB[j][1], num=mid_num)))
+
+                        vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
+                                          for I in range(len(startend))])
+                        vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
+                                          for I in range(len(startend))])
+
+                        score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
+                        score_with_dist_prior = sum(score_midpts) / len(score_midpts) + min(
+                            0.5 * oriImg.shape[0] / norm - 1, 0)
+                        criterion1 = len(np.nonzero(score_midpts > thre2)[0]) > 0.8 * len(score_midpts)
+                        criterion2 = score_with_dist_prior > 0
+                        if criterion1 and criterion2:
+                            connection_candidate.append(
+                                [i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])
+
+                connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
+                connection = np.zeros((0, 5))
+                for c in range(len(connection_candidate)):
+                    i, j, s = connection_candidate[c][0:3]
+                    if (i not in connection[:, 3] and j not in connection[:, 4]):
+                        connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
+                        if (len(connection) >= min(nA, nB)):
+                            break
+
+                connection_all.append(connection)
+            else:
+                special_k.append(k)
+                connection_all.append([])
+
+        # last number in each row is the total parts number of that person
+        # the second last number in each row is the score of the overall configuration
+        subset = -1 * np.ones((0, 20))
+        candidate = np.array([item for sublist in all_peaks for item in sublist])
+
+        for k in range(len(mapIdx)):
+            if k not in special_k:
+                partAs = connection_all[k][:, 0]
+                partBs = connection_all[k][:, 1]
+                indexA, indexB = np.array(limbSeq[k]) - 1
+
+                for i in range(len(connection_all[k])):  # = 1:size(temp,1)
+                    found = 0
+                    subset_idx = [-1, -1]
+                    for j in range(len(subset)):  # 1:size(subset,1):
+                        if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
+                            subset_idx[found] = j
+                            found += 1
+
+                    if found == 1:
+                        j = subset_idx[0]
+                        if subset[j][indexB] != partBs[i]:
+                            subset[j][indexB] = partBs[i]
+                            subset[j][-1] += 1
+                            subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
+                    elif found == 2:  # if found 2 and disjoint, merge them
+                        j1, j2 = subset_idx
+                        membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2]
+                        if len(np.nonzero(membership == 2)[0]) == 0:  # merge
+                            subset[j1][:-2] += (subset[j2][:-2] + 1)
+                            subset[j1][-2:] += subset[j2][-2:]
+                            subset[j1][-2] += connection_all[k][i][2]
+                            subset = np.delete(subset, j2, 0)
+                        else:  # as like found == 1
+                            subset[j1][indexB] = partBs[i]
+                            subset[j1][-1] += 1
+                            subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
+
+                    # if find no partA in the subset, create a new subset
+                    elif not found and k < 17:
+                        row = -1 * np.ones(20)
+                        row[indexA] = partAs[i]
+                        row[indexB] = partBs[i]
+                        row[-1] = 2
+                        row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2]
+                        subset = np.vstack([subset, row])
+        # delete some rows of subset which has few parts occur
+        deleteIdx = []
+        for i in range(len(subset)):
+            if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
+                deleteIdx.append(i)
+        subset = np.delete(subset, deleteIdx, axis=0)
+
+        # subset: n*20 array, 0-17 is the index in candidate, 18 is the total score, 19 is the total parts
+        # candidate: x, y, score, id
+        return candidate, subset
+
+if __name__ == "__main__":
+    body_estimation = Body('../model/body_pose_model.pth')
+
+    test_image = '../images/ski.jpg'
+    oriImg = cv2.imread(test_image)  # B,G,R order
+    candidate, subset = body_estimation(oriImg)
+    canvas = util.draw_bodypose(oriImg, candidate, subset)
+    plt.imshow(canvas[:, :, [2, 1, 0]])
+    plt.show()

src/ControlNet/annotator/openpose/hand.py

@@ -0,0 +1,86 @@
+import cv2
+import json
+import numpy as np
+import math
+import time
+from scipy.ndimage.filters import gaussian_filter
+import matplotlib.pyplot as plt
+import matplotlib
+import torch
+from skimage.measure import label
+
+from .model import handpose_model
+from . import util
+
+class Hand(object):
+    def __init__(self, model_path):
+        self.model = handpose_model()
+        if torch.cuda.is_available():
+            self.model = self.model.cuda()
+            print('cuda')
+        model_dict = util.transfer(self.model, torch.load(model_path))
+        self.model.load_state_dict(model_dict)
+        self.model.eval()
+
+    def __call__(self, oriImg):
+        scale_search = [0.5, 1.0, 1.5, 2.0]
+        # scale_search = [0.5]
+        boxsize = 368
+        stride = 8
+        padValue = 128
+        thre = 0.05
+        multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
+        heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 22))
+        # paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
+
+        for m in range(len(multiplier)):
+            scale = multiplier[m]
+            imageToTest = cv2.resize(oriImg, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
+            imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
+            im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
+            im = np.ascontiguousarray(im)
+
+            data = torch.from_numpy(im).float()
+            if torch.cuda.is_available():
+                data = data.cuda()
+            # data = data.permute([2, 0, 1]).unsqueeze(0).float()
+            with torch.no_grad():
+                output = self.model(data).cpu().numpy()
+                # output = self.model(data).numpy()q
+
+            # extract outputs, resize, and remove padding
+            heatmap = np.transpose(np.squeeze(output), (1, 2, 0))  # output 1 is heatmaps
+            heatmap = cv2.resize(heatmap, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
+            heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
+            heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)
+
+            heatmap_avg += heatmap / len(multiplier)
+
+        all_peaks = []
+        for part in range(21):
+            map_ori = heatmap_avg[:, :, part]
+            one_heatmap = gaussian_filter(map_ori, sigma=3)
+            binary = np.ascontiguousarray(one_heatmap > thre, dtype=np.uint8)
+            # 全部小于阈值
+            if np.sum(binary) == 0:
+                all_peaks.append([0, 0])
+                continue
+            label_img, label_numbers = label(binary, return_num=True, connectivity=binary.ndim)
+            max_index = np.argmax([np.sum(map_ori[label_img == i]) for i in range(1, label_numbers + 1)]) + 1
+            label_img[label_img != max_index] = 0
+            map_ori[label_img == 0] = 0
+
+            y, x = util.npmax(map_ori)
+            all_peaks.append([x, y])
+        return np.array(all_peaks)
+
+if __name__ == "__main__":
+    hand_estimation = Hand('../model/hand_pose_model.pth')
+
+    # test_image = '../images/hand.jpg'
+    test_image = '../images/hand.jpg'
+    oriImg = cv2.imread(test_image)  # B,G,R order
+    peaks = hand_estimation(oriImg)
+    canvas = util.draw_handpose(oriImg, peaks, True)
+    cv2.imshow('', canvas)
+    cv2.waitKey(0)

src/ControlNet/annotator/openpose/model.py

@@ -0,0 +1,219 @@
+import torch
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+
+def make_layers(block, no_relu_layers):
+    layers = []
+    for layer_name, v in block.items():
+        if 'pool' in layer_name:
+            layer = nn.MaxPool2d(kernel_size=v[0], stride=v[1],
+                                    padding=v[2])
+            layers.append((layer_name, layer))
+        else:
+            conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1],
+                               kernel_size=v[2], stride=v[3],
+                               padding=v[4])
+            layers.append((layer_name, conv2d))
+            if layer_name not in no_relu_layers:
+                layers.append(('relu_'+layer_name, nn.ReLU(inplace=True)))
+
+    return nn.Sequential(OrderedDict(layers))
+
+class bodypose_model(nn.Module):
+    def __init__(self):
+        super(bodypose_model, self).__init__()
+
+        # these layers have no relu layer
+        no_relu_layers = ['conv5_5_CPM_L1', 'conv5_5_CPM_L2', 'Mconv7_stage2_L1',\
+                          'Mconv7_stage2_L2', 'Mconv7_stage3_L1', 'Mconv7_stage3_L2',\
+                          'Mconv7_stage4_L1', 'Mconv7_stage4_L2', 'Mconv7_stage5_L1',\
+                          'Mconv7_stage5_L2', 'Mconv7_stage6_L1', 'Mconv7_stage6_L1']
+        blocks = {}
+        block0 = OrderedDict([
+                      ('conv1_1', [3, 64, 3, 1, 1]),
+                      ('conv1_2', [64, 64, 3, 1, 1]),
+                      ('pool1_stage1', [2, 2, 0]),
+                      ('conv2_1', [64, 128, 3, 1, 1]),
+                      ('conv2_2', [128, 128, 3, 1, 1]),
+                      ('pool2_stage1', [2, 2, 0]),
+                      ('conv3_1', [128, 256, 3, 1, 1]),
+                      ('conv3_2', [256, 256, 3, 1, 1]),
+                      ('conv3_3', [256, 256, 3, 1, 1]),
+                      ('conv3_4', [256, 256, 3, 1, 1]),
+                      ('pool3_stage1', [2, 2, 0]),
+                      ('conv4_1', [256, 512, 3, 1, 1]),
+                      ('conv4_2', [512, 512, 3, 1, 1]),
+                      ('conv4_3_CPM', [512, 256, 3, 1, 1]),
+                      ('conv4_4_CPM', [256, 128, 3, 1, 1])
+                  ])
+
+
+        # Stage 1
+        block1_1 = OrderedDict([
+                        ('conv5_1_CPM_L1', [128, 128, 3, 1, 1]),
+                        ('conv5_2_CPM_L1', [128, 128, 3, 1, 1]),
+                        ('conv5_3_CPM_L1', [128, 128, 3, 1, 1]),
+                        ('conv5_4_CPM_L1', [128, 512, 1, 1, 0]),
+                        ('conv5_5_CPM_L1', [512, 38, 1, 1, 0])
+                    ])
+
+        block1_2 = OrderedDict([
+                        ('conv5_1_CPM_L2', [128, 128, 3, 1, 1]),
+                        ('conv5_2_CPM_L2', [128, 128, 3, 1, 1]),
+                        ('conv5_3_CPM_L2', [128, 128, 3, 1, 1]),
+                        ('conv5_4_CPM_L2', [128, 512, 1, 1, 0]),
+                        ('conv5_5_CPM_L2', [512, 19, 1, 1, 0])
+                    ])
+        blocks['block1_1'] = block1_1
+        blocks['block1_2'] = block1_2
+
+        self.model0 = make_layers(block0, no_relu_layers)
+
+        # Stages 2 - 6
+        for i in range(2, 7):
+            blocks['block%d_1' % i] = OrderedDict([
+                    ('Mconv1_stage%d_L1' % i, [185, 128, 7, 1, 3]),
+                    ('Mconv2_stage%d_L1' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv3_stage%d_L1' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv4_stage%d_L1' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv5_stage%d_L1' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv6_stage%d_L1' % i, [128, 128, 1, 1, 0]),
+                    ('Mconv7_stage%d_L1' % i, [128, 38, 1, 1, 0])
+                ])
+
+            blocks['block%d_2' % i] = OrderedDict([
+                    ('Mconv1_stage%d_L2' % i, [185, 128, 7, 1, 3]),
+                    ('Mconv2_stage%d_L2' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv3_stage%d_L2' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv4_stage%d_L2' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv5_stage%d_L2' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv6_stage%d_L2' % i, [128, 128, 1, 1, 0]),
+                    ('Mconv7_stage%d_L2' % i, [128, 19, 1, 1, 0])
+                ])
+
+        for k in blocks.keys():
+            blocks[k] = make_layers(blocks[k], no_relu_layers)
+
+        self.model1_1 = blocks['block1_1']
+        self.model2_1 = blocks['block2_1']
+        self.model3_1 = blocks['block3_1']
+        self.model4_1 = blocks['block4_1']
+        self.model5_1 = blocks['block5_1']
+        self.model6_1 = blocks['block6_1']
+
+        self.model1_2 = blocks['block1_2']
+        self.model2_2 = blocks['block2_2']
+        self.model3_2 = blocks['block3_2']
+        self.model4_2 = blocks['block4_2']
+        self.model5_2 = blocks['block5_2']
+        self.model6_2 = blocks['block6_2']
+
+
+    def forward(self, x):
+
+        out1 = self.model0(x)
+
+        out1_1 = self.model1_1(out1)
+        out1_2 = self.model1_2(out1)
+        out2 = torch.cat([out1_1, out1_2, out1], 1)
+
+        out2_1 = self.model2_1(out2)
+        out2_2 = self.model2_2(out2)
+        out3 = torch.cat([out2_1, out2_2, out1], 1)
+
+        out3_1 = self.model3_1(out3)
+        out3_2 = self.model3_2(out3)
+        out4 = torch.cat([out3_1, out3_2, out1], 1)
+
+        out4_1 = self.model4_1(out4)
+        out4_2 = self.model4_2(out4)
+        out5 = torch.cat([out4_1, out4_2, out1], 1)
+
+        out5_1 = self.model5_1(out5)
+        out5_2 = self.model5_2(out5)
+        out6 = torch.cat([out5_1, out5_2, out1], 1)
+
+        out6_1 = self.model6_1(out6)
+        out6_2 = self.model6_2(out6)
+
+        return out6_1, out6_2
+
+class handpose_model(nn.Module):
+    def __init__(self):
+        super(handpose_model, self).__init__()
+
+        # these layers have no relu layer
+        no_relu_layers = ['conv6_2_CPM', 'Mconv7_stage2', 'Mconv7_stage3',\
+                          'Mconv7_stage4', 'Mconv7_stage5', 'Mconv7_stage6']
+        # stage 1
+        block1_0 = OrderedDict([
+                ('conv1_1', [3, 64, 3, 1, 1]),
+                ('conv1_2', [64, 64, 3, 1, 1]),
+                ('pool1_stage1', [2, 2, 0]),
+                ('conv2_1', [64, 128, 3, 1, 1]),
+                ('conv2_2', [128, 128, 3, 1, 1]),
+                ('pool2_stage1', [2, 2, 0]),
+                ('conv3_1', [128, 256, 3, 1, 1]),
+                ('conv3_2', [256, 256, 3, 1, 1]),
+                ('conv3_3', [256, 256, 3, 1, 1]),
+                ('conv3_4', [256, 256, 3, 1, 1]),
+                ('pool3_stage1', [2, 2, 0]),
+                ('conv4_1', [256, 512, 3, 1, 1]),
+                ('conv4_2', [512, 512, 3, 1, 1]),
+                ('conv4_3', [512, 512, 3, 1, 1]),
+                ('conv4_4', [512, 512, 3, 1, 1]),
+                ('conv5_1', [512, 512, 3, 1, 1]),
+                ('conv5_2', [512, 512, 3, 1, 1]),
+                ('conv5_3_CPM', [512, 128, 3, 1, 1])
+            ])
+
+        block1_1 = OrderedDict([
+            ('conv6_1_CPM', [128, 512, 1, 1, 0]),
+            ('conv6_2_CPM', [512, 22, 1, 1, 0])
+        ])
+
+        blocks = {}
+        blocks['block1_0'] = block1_0
+        blocks['block1_1'] = block1_1
+
+        # stage 2-6
+        for i in range(2, 7):
+            blocks['block%d' % i] = OrderedDict([
+                    ('Mconv1_stage%d' % i, [150, 128, 7, 1, 3]),
+                    ('Mconv2_stage%d' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv3_stage%d' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv4_stage%d' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv5_stage%d' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv6_stage%d' % i, [128, 128, 1, 1, 0]),
+                    ('Mconv7_stage%d' % i, [128, 22, 1, 1, 0])
+                ])
+
+        for k in blocks.keys():
+            blocks[k] = make_layers(blocks[k], no_relu_layers)
+
+        self.model1_0 = blocks['block1_0']
+        self.model1_1 = blocks['block1_1']
+        self.model2 = blocks['block2']
+        self.model3 = blocks['block3']
+        self.model4 = blocks['block4']
+        self.model5 = blocks['block5']
+        self.model6 = blocks['block6']
+
+    def forward(self, x):
+        out1_0 = self.model1_0(x)
+        out1_1 = self.model1_1(out1_0)
+        concat_stage2 = torch.cat([out1_1, out1_0], 1)
+        out_stage2 = self.model2(concat_stage2)
+        concat_stage3 = torch.cat([out_stage2, out1_0], 1)
+        out_stage3 = self.model3(concat_stage3)
+        concat_stage4 = torch.cat([out_stage3, out1_0], 1)
+        out_stage4 = self.model4(concat_stage4)
+        concat_stage5 = torch.cat([out_stage4, out1_0], 1)
+        out_stage5 = self.model5(concat_stage5)
+        concat_stage6 = torch.cat([out_stage5, out1_0], 1)
+        out_stage6 = self.model6(concat_stage6)
+        return out_stage6
+
+

src/ControlNet/annotator/openpose/util.py

@@ -0,0 +1,164 @@
+import math
+import numpy as np
+import matplotlib
+import cv2
+
+
+def padRightDownCorner(img, stride, padValue):
+    h = img.shape[0]
+    w = img.shape[1]
+
+    pad = 4 * [None]
+    pad[0] = 0 # up
+    pad[1] = 0 # left
+    pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down
+    pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right
+
+    img_padded = img
+    pad_up = np.tile(img_padded[0:1, :, :]*0 + padValue, (pad[0], 1, 1))
+    img_padded = np.concatenate((pad_up, img_padded), axis=0)
+    pad_left = np.tile(img_padded[:, 0:1, :]*0 + padValue, (1, pad[1], 1))
+    img_padded = np.concatenate((pad_left, img_padded), axis=1)
+    pad_down = np.tile(img_padded[-2:-1, :, :]*0 + padValue, (pad[2], 1, 1))
+    img_padded = np.concatenate((img_padded, pad_down), axis=0)
+    pad_right = np.tile(img_padded[:, -2:-1, :]*0 + padValue, (1, pad[3], 1))
+    img_padded = np.concatenate((img_padded, pad_right), axis=1)
+
+    return img_padded, pad
+
+# transfer caffe model to pytorch which will match the layer name
+def transfer(model, model_weights):
+    transfered_model_weights = {}
+    for weights_name in model.state_dict().keys():
+        transfered_model_weights[weights_name] = model_weights['.'.join(weights_name.split('.')[1:])]
+    return transfered_model_weights
+
+# draw the body keypoint and lims
+def draw_bodypose(canvas, candidate, subset):
+    stickwidth = 4
+    limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
+               [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
+               [1, 16], [16, 18], [3, 17], [6, 18]]
+
+    colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
+              [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
+              [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
+    for i in range(18):
+        for n in range(len(subset)):
+            index = int(subset[n][i])
+            if index == -1:
+                continue
+            x, y = candidate[index][0:2]
+            cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
+    for i in range(17):
+        for n in range(len(subset)):
+            index = subset[n][np.array(limbSeq[i]) - 1]
+            if -1 in index:
+                continue
+            cur_canvas = canvas.copy()
+            Y = candidate[index.astype(int), 0]
+            X = candidate[index.astype(int), 1]
+            mX = np.mean(X)
+            mY = np.mean(Y)
+            length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
+            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
+            polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
+            cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
+            canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
+    # plt.imsave("preview.jpg", canvas[:, :, [2, 1, 0]])
+    # plt.imshow(canvas[:, :, [2, 1, 0]])
+    return canvas
+
+
+# image drawed by opencv is not good.
+def draw_handpose(canvas, all_hand_peaks, show_number=False):
+    edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
+             [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
+
+    for peaks in all_hand_peaks:
+        for ie, e in enumerate(edges):
+            if np.sum(np.all(peaks[e], axis=1)==0)==0:
+                x1, y1 = peaks[e[0]]
+                x2, y2 = peaks[e[1]]
+                cv2.line(canvas, (x1, y1), (x2, y2), matplotlib.colors.hsv_to_rgb([ie/float(len(edges)), 1.0, 1.0])*255, thickness=2)
+
+        for i, keyponit in enumerate(peaks):
+            x, y = keyponit
+            cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
+            if show_number:
+                cv2.putText(canvas, str(i), (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 0), lineType=cv2.LINE_AA)
+    return canvas
+
+# detect hand according to body pose keypoints
+# please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp
+def handDetect(candidate, subset, oriImg):
+    # right hand: wrist 4, elbow 3, shoulder 2
+    # left hand: wrist 7, elbow 6, shoulder 5
+    ratioWristElbow = 0.33
+    detect_result = []
+    image_height, image_width = oriImg.shape[0:2]
+    for person in subset.astype(int):
+        # if any of three not detected
+        has_left = np.sum(person[[5, 6, 7]] == -1) == 0
+        has_right = np.sum(person[[2, 3, 4]] == -1) == 0
+        if not (has_left or has_right):
+            continue
+        hands = []
+        #left hand
+        if has_left:
+            left_shoulder_index, left_elbow_index, left_wrist_index = person[[5, 6, 7]]
+            x1, y1 = candidate[left_shoulder_index][:2]
+            x2, y2 = candidate[left_elbow_index][:2]
+            x3, y3 = candidate[left_wrist_index][:2]
+            hands.append([x1, y1, x2, y2, x3, y3, True])
+        # right hand
+        if has_right:
+            right_shoulder_index, right_elbow_index, right_wrist_index = person[[2, 3, 4]]
+            x1, y1 = candidate[right_shoulder_index][:2]
+            x2, y2 = candidate[right_elbow_index][:2]
+            x3, y3 = candidate[right_wrist_index][:2]
+            hands.append([x1, y1, x2, y2, x3, y3, False])
+
+        for x1, y1, x2, y2, x3, y3, is_left in hands:
+            # pos_hand = pos_wrist + ratio * (pos_wrist - pos_elbox) = (1 + ratio) * pos_wrist - ratio * pos_elbox
+            # handRectangle.x = posePtr[wrist*3] + ratioWristElbow * (posePtr[wrist*3] - posePtr[elbow*3]);
+            # handRectangle.y = posePtr[wrist*3+1] + ratioWristElbow * (posePtr[wrist*3+1] - posePtr[elbow*3+1]);
+            # const auto distanceWristElbow = getDistance(poseKeypoints, person, wrist, elbow);
+            # const auto distanceElbowShoulder = getDistance(poseKeypoints, person, elbow, shoulder);
+            # handRectangle.width = 1.5f * fastMax(distanceWristElbow, 0.9f * distanceElbowShoulder);
+            x = x3 + ratioWristElbow * (x3 - x2)
+            y = y3 + ratioWristElbow * (y3 - y2)
+            distanceWristElbow = math.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2)
+            distanceElbowShoulder = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
+            width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder)
+            # x-y refers to the center --> offset to topLeft point
+            # handRectangle.x -= handRectangle.width / 2.f;
+            # handRectangle.y -= handRectangle.height / 2.f;
+            x -= width / 2
+            y -= width / 2  # width = height
+            # overflow the image
+            if x < 0: x = 0
+            if y < 0: y = 0
+            width1 = width
+            width2 = width
+            if x + width > image_width: width1 = image_width - x
+            if y + width > image_height: width2 = image_height - y
+            width = min(width1, width2)
+            # the max hand box value is 20 pixels
+            if width >= 20:
+                detect_result.append([int(x), int(y), int(width), is_left])
+
+    '''
+    return value: [[x, y, w, True if left hand else False]].
+    width=height since the network require squared input.
+    x, y is the coordinate of top left 
+    '''
+    return detect_result
+
+# get max index of 2d array
+def npmax(array):
+    arrayindex = array.argmax(1)
+    arrayvalue = array.max(1)
+    i = arrayvalue.argmax()
+    j = arrayindex[i]
+    return i, j

src/ControlNet/annotator/util.py

@@ -0,0 +1,38 @@
+import numpy as np
+import cv2
+import os
+
+
+annotator_ckpts_path = os.path.join(os.path.dirname(__file__), 'ckpts')
+
+
+def HWC3(x):
+    assert x.dtype == np.uint8
+    if x.ndim == 2:
+        x = x[:, :, None]
+    assert x.ndim == 3
+    H, W, C = x.shape
+    assert C == 1 or C == 3 or C == 4
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    if C == 4:
+        color = x[:, :, 0:3].astype(np.float32)
+        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+        y = color * alpha + 255.0 * (1.0 - alpha)
+        y = y.clip(0, 255).astype(np.uint8)
+        return y
+
+
+def resize_image(input_image, resolution):
+    H, W, C = input_image.shape
+    H = float(H)
+    W = float(W)
+    k = float(resolution) / min(H, W)
+    H *= k
+    W *= k
+    H = int(np.round(H / 64.0)) * 64
+    W = int(np.round(W / 64.0)) * 64
+    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
+    return img

src/EGNet/README.md

@@ -0,0 +1,49 @@
+# EGNet
+EGNet:Edge Guidance Network for Salient Object Detection (ICCV 2019)
+
+We use the sal2edge.m to generate the edge label for training.
+### For training:
+1. Clone this code by `git clone https://github.com/JXingZhao/EGNet.git --recursive`, assume your source code directory is`$EGNet`;
+
+2. Download [training data](https://pan.baidu.com/s/1LaQoNRS8-11V7grAfFiHCg) (fsex) ([google drive](https://drive.google.com/open?id=1wduPbFMkxB_3W72LvJckD7N0hWbXsKsj));
+
+3. Download [initial model](https://pan.baidu.com/s/1dD2JOY_FBSLzjp5tUPBDBQ) (8ir7) ([google_drive](https://drive.google.com/open?id=1q7FtHWoarRzGNQQXTn9t7QSR8jJL8vk6)); 
+
+4. Change the image path and intial model path in run.py and dataset.py;
+
+5. Start to train with `python3 run.py --mode train`.
+
+### For testing:
+1. Download [pretrained model](https://pan.baidu.com/s/1s35ZyGDSNVzVIeVd7Aot0Q) (2cf5)  ([google drive](https://drive.google.com/open?id=17Ffc6V5EiujtcFKupsJXhtlQ3cLK5OGp)); 
+
+2. Change the test image path in dataset.py 
+
+3. Generate saliency maps for SOD dataset by `python3 run.py --mode test --sal_mode s`, PASCALS by `python3 run.py --mode test --sal_mode p` and so on;
+
+4. Testing code we use is the public open source code. (https://github.com/Andrew-Qibin/SalMetric)
+
+
+
+### Pretrained models, datasets and results:
+| [Page](https://mmcheng.net/jxzhao/) |
+| [Training Set](https://pan.baidu.com/s/1LaQoNRS8-11V7grAfFiHCg) (fsex) ([google drive](https://drive.google.com/open?id=1wduPbFMkxB_3W72LvJckD7N0hWbXsKsj)) |
+| [Pretrained models](https://pan.baidu.com/s/1s35ZyGDSNVzVIeVd7Aot0Q) (2cf5) |
+| [Saliency maps](https://pan.baidu.com/s/1M_dqPJ08oaYWge_zZnHSTQ) (54gi) ([google drive VGG](https://drive.google.com/open?id=1WEuEqNmqMePyxD8anGo0KA4rWK9Nyb9I)) ([google drive resnet](https://drive.google.com/open?id=1h5R8tT3Jq_2S3pLfXREpuWaKvFphQ4K9)) |
+
+
+### If you think this work is helpful, please cite
+```latex
+@inproceedings{zhao2019EGNet,
+ title={EGNet:Edge Guidance Network for Salient Object Detection},
+ author={Zhao, Jia-Xing and Liu, Jiang-Jiang and Fan, Deng-Ping and Cao, Yang and Yang, Jufeng and Cheng, Ming-Ming},
+ booktitle={The IEEE International Conference on Computer Vision (ICCV)},
+ month={Oct},
+ year={2019},
+}
+```
+
+### Other related work
+Contrast Prior and Fluid Pyramid Integration for RGBD Salient Object Detection. (CVPR2019) [page](https://mmcheng.net/rgbdsalpyr/)
+
+
+

src/EGNet/dataset.py

@@ -0,0 +1,283 @@
+import os
+from PIL import Image
+import cv2
+import torch
+from torch.utils import data
+from torchvision import transforms
+from torchvision.transforms import functional as F
+import numbers
+import numpy as np
+import random
+
+#re_size = (256, 256)
+#cr_size = (224, 224)
+
+class ImageDataTrain(data.Dataset):
+    def __init__(self):
+
+        self.sal_root = '/home/liuj/dataset/DUTS/DUTS-TR'
+        self.sal_source = '/home/liuj/dataset/DUTS/DUTS-TR/train_pair_edge.lst'
+
+        with open(self.sal_source, 'r') as f:
+            self.sal_list = [x.strip() for x in f.readlines()]
+
+        self.sal_num = len(self.sal_list)
+
+
+    def __getitem__(self, item):
+
+
+        sal_image = load_image(os.path.join(self.sal_root, self.sal_list[item%self.sal_num].split()[0]))
+        sal_label = load_sal_label(os.path.join(self.sal_root, self.sal_list[item%self.sal_num].split()[1]))
+        sal_edge = load_edge_label(os.path.join(self.sal_root, self.sal_list[item%self.sal_num].split()[2]))
+        sal_image, sal_label, sal_edge = cv_random_flip(sal_image, sal_label, sal_edge)
+        sal_image = torch.Tensor(sal_image)
+        sal_label = torch.Tensor(sal_label)
+        sal_edge = torch.Tensor(sal_edge)
+
+        sample = {'sal_image': sal_image, 'sal_label': sal_label, 'sal_edge': sal_edge}
+        return sample
+
+    def __len__(self):
+        # return max(max(self.edge_num, self.sal_num), self.skel_num)
+        return self.sal_num
+
+class ImageDataTest(data.Dataset):
+    def __init__(self, test_mode=1, sal_mode='e'):
+        if test_mode == 0:
+            # self.image_root = '/home/liuj/dataset/saliency_test/ECSSD/Imgs/'
+            # self.image_source = '/home/liuj/dataset/saliency_test/ECSSD/test.lst'
+            self.image_root = '/home/liuj/dataset/HED-BSDS_PASCAL/HED-BSDS/test/'
+            self.image_source = '/home/liuj/dataset/HED-BSDS_PASCAL/HED-BSDS/test.lst'
+            
+            
+        elif test_mode == 1:
+            if sal_mode == 'e':
+                self.image_root = '/home/liuj/dataset/saliency_test/ECSSD/Imgs/'
+                self.image_source = '/home/liuj/dataset/saliency_test/ECSSD/test.lst'
+                self.test_fold = '/media/ubuntu/disk/Result/saliency/ECSSD/'
+            elif sal_mode == 'p':
+                self.image_root = '/home/liuj/dataset/saliency_test/PASCALS/Imgs/'
+                self.image_source = '/home/liuj/dataset/saliency_test/PASCALS/test.lst'
+                self.test_fold = '/media/ubuntu/disk/Result/saliency/PASCALS/'
+            elif sal_mode == 'd':
+                self.image_root = '/home/liuj/dataset/saliency_test/DUTOMRON/Imgs/'
+                self.image_source = '/home/liuj/dataset/saliency_test/DUTOMRON/test.lst'
+                self.test_fold = '/media/ubuntu/disk/Result/saliency/DUTOMRON/'
+            elif sal_mode == 'h':
+                self.image_root = '/home/liuj/dataset/saliency_test/HKU-IS/Imgs/'
+                self.image_source = '/home/liuj/dataset/saliency_test/HKU-IS/test.lst'
+                self.test_fold = '/media/ubuntu/disk/Result/saliency/HKU-IS/'
+            elif sal_mode == 's':
+                self.image_root = '/home/liuj/dataset/saliency_test/SOD/Imgs/'
+                self.image_source = '/home/liuj/dataset/saliency_test/SOD/test.lst'
+                self.test_fold = '/media/ubuntu/disk/Result/saliency/SOD/'
+            elif sal_mode == 'm':
+                self.image_root = '/home/liuj/dataset/saliency_test/MSRA/Imgs/'
+                self.image_source = '/home/liuj/dataset/saliency_test/MSRA/test.lst'
+            elif sal_mode == 'o':
+                self.image_root = '/home/liuj/dataset/saliency_test/SOC/TestSet/Imgs/'
+                self.image_source = '/home/liuj/dataset/saliency_test/SOC/TestSet/test.lst'
+                self.test_fold = '/media/ubuntu/disk/Result/saliency/SOC/'
+            elif sal_mode == 't':
+                self.image_root = '/home/liuj/dataset/DUTS/DUTS-TE/DUTS-TE-Image/'
+                self.image_source = '/home/liuj/dataset/DUTS/DUTS-TE/test.lst'
+                self.test_fold = '/media/ubuntu/disk/Result/saliency/DUTS/'
+        elif test_mode == 2:
+
+            self.image_root = '/home/liuj/dataset/SK-LARGE/images/test/'
+            self.image_source = '/home/liuj/dataset/SK-LARGE/test.lst'
+
+        with open(self.image_source, 'r') as f:
+            self.image_list = [x.strip() for x in f.readlines()]
+
+        self.image_num = len(self.image_list)
+
+    def __getitem__(self, item):
+        image, im_size = load_image_test(os.path.join(self.image_root, self.image_list[item]))
+        image = torch.Tensor(image)
+
+        return {'image': image, 'name': self.image_list[item%self.image_num], 'size': im_size}
+    def save_folder(self):
+        return self.test_fold
+
+    def __len__(self):
+        # return max(max(self.edge_num, self.skel_num), self.sal_num)
+        return self.image_num
+
+
+# get the dataloader (Note: without data augmentation, except saliency with random flip)
+def get_loader(batch_size, mode='train', num_thread=1, test_mode=0, sal_mode='e'):
+    shuffle = False
+    if mode == 'train':
+        shuffle = True
+        dataset = ImageDataTrain()
+    else:
+        dataset = ImageDataTest(test_mode=test_mode, sal_mode=sal_mode)
+
+    data_loader = data.DataLoader(dataset=dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_thread)
+    return data_loader, dataset
+
+def load_image(pah):
+    if not os.path.exists(pah):
+        print('File Not Exists')
+    im = cv2.imread(pah)
+    in_ = np.array(im, dtype=np.float32)
+    # in_ = cv2.resize(in_, im_sz, interpolation=cv2.INTER_CUBIC)
+    # in_ = in_[:,:,::-1] # only if use PIL to load image
+    in_ -= np.array((104.00699, 116.66877, 122.67892))
+    in_ = in_.transpose((2,0,1))
+    return in_
+
+def load_image_test(pah):
+    if not os.path.exists(pah):
+        print('File Not Exists')
+    im = cv2.imread(pah)
+    in_ = np.array(im, dtype=np.float32)
+    im_size = tuple(in_.shape[:2])
+    # in_ = cv2.resize(in_, (cr_size[1], cr_size[0]), interpolation=cv2.INTER_LINEAR)
+    # in_ = in_[:,:,::-1] # only if use PIL to load image
+    in_ -= np.array((104.00699, 116.66877, 122.67892))
+    in_ = in_.transpose((2,0,1))
+    return in_, im_size
+
+def load_edge_label(pah):
+    """
+    pixels > 0.5 -> 1
+    Load label image as 1 x height x width integer array of label indices.
+    The leading singleton dimension is required by the loss.
+    """
+    if not os.path.exists(pah):
+        print('File Not Exists')
+    im = Image.open(pah)
+    label = np.array(im, dtype=np.float32)
+    if len(label.shape) == 3:
+        label = label[:,:,0]
+    # label = cv2.resize(label, im_sz, interpolation=cv2.INTER_NEAREST)
+    label = label / 255.
+    label[np.where(label > 0.5)] = 1.
+    label = label[np.newaxis, ...]
+    return label
+
+def load_skel_label(pah):
+    """
+    pixels > 0 -> 1
+    Load label image as 1 x height x width integer array of label indices.
+    The leading singleton dimension is required by the loss.
+    """
+    if not os.path.exists(pah):
+        print('File Not Exists')
+    im = Image.open(pah)
+    label = np.array(im, dtype=np.float32)
+    if len(label.shape) == 3:
+        label = label[:,:,0]
+    # label = cv2.resize(label, im_sz, interpolation=cv2.INTER_NEAREST)
+    label = label / 255.
+    label[np.where(label > 0.)] = 1.
+    label = label[np.newaxis, ...]
+    return label
+
+def load_sal_label(pah):
+    """
+    Load label image as 1 x height x width integer array of label indices.
+    The leading singleton dimension is required by the loss.
+    """
+    if not os.path.exists(pah):
+        print('File Not Exists')
+    im = Image.open(pah)
+    label = np.array(im, dtype=np.float32)
+    if len(label.shape) == 3:
+        label = label[:,:,0]
+    # label = cv2.resize(label, im_sz, interpolation=cv2.INTER_NEAREST)
+    label = label / 255.
+    label = label[np.newaxis, ...]
+    return label
+
+def load_sem_label(pah):
+    """
+    Load label image as 1 x height x width integer array of label indices.
+    The leading singleton dimension is required by the loss.
+    """
+    if not os.path.exists(pah):
+        print('File Not Exists')
+    im = Image.open(pah)
+    label = np.array(im, dtype=np.float32)
+    if len(label.shape) == 3:
+        label = label[:,:,0]
+    # label = cv2.resize(label, im_sz, interpolation=cv2.INTER_NEAREST)
+    # label = label / 255.
+    label = label[np.newaxis, ...]
+    return label
+
+def edge_thres_transform(x, thres):
+    # y0 = torch.zeros(x.size())
+    y1 = torch.ones(x.size())
+    x = torch.where(x >= thres, y1, x)
+    return x
+
+def skel_thres_transform(x, thres):
+    y0 = torch.zeros(x.size())
+    y1 = torch.ones(x.size())
+    x = torch.where(x > thres, y1, y0)
+    return x
+
+def cv_random_flip(img, label, edge):
+    flip_flag = random.randint(0, 1)
+    if flip_flag == 1:
+        img = img[:,:,::-1].copy()
+        label = label[:,:,::-1].copy()
+        edge = edge[:,:,::-1].copy()
+    return img, label, edge
+
+def cv_random_crop_flip(img, label, resize_size, crop_size, random_flip=True):
+    def get_params(img_size, output_size):
+        h, w = img_size
+        th, tw = output_size
+        if w == tw and h == th:
+            return 0, 0, h, w
+        i = random.randint(0, h - th)
+        j = random.randint(0, w - tw)
+        return i, j, th, tw
+    if random_flip:
+        flip_flag = random.randint(0, 1)
+    img = img.transpose((1,2,0)) # H, W, C
+    label = label[0,:,:] # H, W
+    img = cv2.resize(img, (resize_size[1], resize_size[0]), interpolation=cv2.INTER_LINEAR)
+    label = cv2.resize(label, (resize_size[1], resize_size[0]), interpolation=cv2.INTER_NEAREST)
+    i, j, h, w = get_params(resize_size, crop_size)
+    img = img[i:i+h, j:j+w, :].transpose((2,0,1)) # C, H, W
+    label = label[i:i+h, j:j+w][np.newaxis, ...] # 1, H, W
+    if flip_flag == 1:
+        img = img[:,:,::-1].copy()
+        label = label[:,:,::-1].copy()
+    return img, label
+
+def random_crop(img, label, size, padding=None, pad_if_needed=True, fill_img=(123, 116, 103), fill_label=0, padding_mode='constant'):
+
+    def get_params(img, output_size):
+        w, h = img.size
+        th, tw = output_size
+        if w == tw and h == th:
+            return 0, 0, h, w
+
+        i = random.randint(0, h - th)
+        j = random.randint(0, w - tw)
+        return i, j, th, tw
+
+    if isinstance(size, numbers.Number):
+        size = (int(size), int(size))
+    if padding is not None:
+        img = F.pad(img, padding, fill_img, padding_mode)
+        label = F.pad(label, padding, fill_label, padding_mode)
+
+    # pad the width if needed
+    if pad_if_needed and img.size[0] < size[1]:
+        img = F.pad(img, (int((1 + size[1] - img.size[0]) / 2), 0), fill_img, padding_mode)
+        label = F.pad(label, (int((1 + size[1] - label.size[0]) / 2), 0), fill_label, padding_mode)
+    # pad the height if needed
+    if pad_if_needed and img.size[1] < size[0]:
+        img = F.pad(img, (0, int((1 + size[0] - img.size[1]) / 2)), fill_img, padding_mode)
+        label = F.pad(label, (0, int((1 + size[0] - label.size[1]) / 2)), fill_label, padding_mode)
+
+    i, j, h, w = get_params(img, size)
+    return [F.crop(img, i, j, h, w), F.crop(label, i, j, h, w)]

src/EGNet/model.py

@@ -0,0 +1,208 @@
+import torch
+from torch import nn
+from torch.nn import init
+import torch.nn.functional as F
+import math
+from torch.autograd import Variable
+import numpy as np
+
+from resnet import resnet50
+from vgg import vgg16
+
+
+config_vgg = {'convert': [[128,256,512,512,512],[64,128,256,512,512]], 'merge1': [[128, 256, 128, 3,1], [256, 512, 256, 3, 1], [512, 0, 512, 5, 2], [512, 0, 512, 5, 2],[512, 0, 512, 7, 3]], 'merge2': [[128], [256, 512, 512, 512]]}  # no convert layer, no conv6
+
+config_resnet = {'convert': [[64,256,512,1024,2048],[128,256,512,512,512]], 'deep_pool': [[512, 512, 256, 256, 128], [512, 256, 256, 128, 128], [False, True, True, True, False], [True, True, True, True, False]], 'score': 256, 'edgeinfo':[[16, 16, 16, 16], 128, [16,8,4,2]],'edgeinfoc':[64,128], 'block': [[512, [16]], [256, [16]], [256, [16]], [128, [16]]], 'fuse': [[16, 16, 16, 16], True], 'fuse_ratio': [[16,1], [8,1], [4,1], [2,1]],  'merge1': [[128, 256, 128, 3,1], [256, 512, 256, 3, 1], [512, 0, 512, 5, 2], [512, 0, 512, 5, 2],[512, 0, 512, 7, 3]], 'merge2': [[128], [256, 512, 512, 512]]}
+
+
+class ConvertLayer(nn.Module):
+    def __init__(self, list_k):
+        super(ConvertLayer, self).__init__()
+        up0, up1, up2 = [], [], []
+        for i in range(len(list_k[0])):
+          
+            up0.append(nn.Sequential(nn.Conv2d(list_k[0][i], list_k[1][i], 1, 1, bias=False), nn.ReLU(inplace=True)))
+
+
+        self.convert0 = nn.ModuleList(up0)
+
+
+    def forward(self, list_x):
+        resl = []
+        for i in range(len(list_x)):
+            resl.append(self.convert0[i](list_x[i]))
+        return resl
+
+
+        
+
+class MergeLayer1(nn.Module): # list_k: [[64, 512, 64], [128, 512, 128], [256, 0, 256] ... ]
+    def __init__(self, list_k):
+        super(MergeLayer1, self).__init__()
+        self.list_k = list_k
+        trans, up, score = [], [], []
+        for ik in list_k:
+            if ik[1] > 0:
+                trans.append(nn.Sequential(nn.Conv2d(ik[1], ik[0], 1, 1, bias=False), nn.ReLU(inplace=True)))
+
+           
+            up.append(nn.Sequential(nn.Conv2d(ik[0], ik[2], ik[3], 1, ik[4]), nn.ReLU(inplace=True), nn.Conv2d(ik[2], ik[2], ik[3], 1, ik[4]), nn.ReLU(inplace=True), nn.Conv2d(ik[2], ik[2], ik[3], 1, ik[4]), nn.ReLU(inplace=True)))
+            score.append(nn.Conv2d(ik[2], 1, 3, 1, 1))
+        trans.append(nn.Sequential(nn.Conv2d(512, 128, 1, 1, bias=False), nn.ReLU(inplace=True)))
+        self.trans, self.up, self.score = nn.ModuleList(trans), nn.ModuleList(up), nn.ModuleList(score)
+        self.relu =nn.ReLU()
+
+    def forward(self, list_x, x_size):
+        up_edge, up_sal, edge_feature, sal_feature = [], [], [], []
+        
+        
+        num_f = len(list_x)
+        tmp = self.up[num_f - 1](list_x[num_f-1])
+        sal_feature.append(tmp)
+        U_tmp = tmp
+        up_sal.append(F.interpolate(self.score[num_f - 1](tmp), x_size, mode='bilinear', align_corners=True))
+        
+        for j in range(2, num_f ):
+            i = num_f - j
+             
+            if list_x[i].size()[1] < U_tmp.size()[1]:
+                U_tmp = list_x[i] + F.interpolate((self.trans[i](U_tmp)), list_x[i].size()[2:], mode='bilinear', align_corners=True)
+            else:
+                U_tmp = list_x[i] + F.interpolate((U_tmp), list_x[i].size()[2:], mode='bilinear', align_corners=True)
+            
+            
+            
+                
+            
+            tmp = self.up[i](U_tmp)
+            U_tmp = tmp
+            sal_feature.append(tmp)
+            up_sal.append(F.interpolate(self.score[i](tmp), x_size, mode='bilinear', align_corners=True))
+
+        U_tmp = list_x[0] + F.interpolate((self.trans[-1](sal_feature[0])), list_x[0].size()[2:], mode='bilinear', align_corners=True)
+        tmp = self.up[0](U_tmp)
+        edge_feature.append(tmp)
+       
+        up_edge.append(F.interpolate(self.score[0](tmp), x_size, mode='bilinear', align_corners=True)) 
+        return up_edge, edge_feature, up_sal, sal_feature        
+        
+class MergeLayer2(nn.Module): 
+    def __init__(self, list_k):
+        super(MergeLayer2, self).__init__()
+        self.list_k = list_k
+        trans, up, score = [], [], []
+        for i in list_k[0]:
+            tmp = []
+            tmp_up = []
+            tmp_score = []
+            feature_k = [[3,1],[5,2], [5,2], [7,3]]
+            for idx, j in enumerate(list_k[1]):
+                tmp.append(nn.Sequential(nn.Conv2d(j, i, 1, 1, bias=False), nn.ReLU(inplace=True)))
+
+                tmp_up.append(nn.Sequential(nn.Conv2d(i , i, feature_k[idx][0], 1, feature_k[idx][1]), nn.ReLU(inplace=True), nn.Conv2d(i, i,  feature_k[idx][0],1 , feature_k[idx][1]), nn.ReLU(inplace=True), nn.Conv2d(i, i, feature_k[idx][0], 1, feature_k[idx][1]), nn.ReLU(inplace=True)))
+                tmp_score.append(nn.Conv2d(i, 1, 3, 1, 1))
+            trans.append(nn.ModuleList(tmp))
+
+            up.append(nn.ModuleList(tmp_up))
+            score.append(nn.ModuleList(tmp_score))
+            
+
+        self.trans, self.up, self.score = nn.ModuleList(trans), nn.ModuleList(up), nn.ModuleList(score)       
+        self.final_score = nn.Sequential(nn.Conv2d(list_k[0][0], list_k[0][0], 5, 1, 2), nn.ReLU(inplace=True), nn.Conv2d(list_k[0][0], 1, 3, 1, 1))
+        self.relu =nn.ReLU()
+
+    def forward(self, list_x, list_y, x_size):
+        up_score, tmp_feature = [], []
+        list_y = list_y[::-1]
+
+        
+        for i, i_x in enumerate(list_x):
+            for j, j_x in enumerate(list_y):                              
+                tmp = F.interpolate(self.trans[i][j](j_x), i_x.size()[2:], mode='bilinear', align_corners=True) + i_x                
+                tmp_f = self.up[i][j](tmp)             
+                up_score.append(F.interpolate(self.score[i][j](tmp_f), x_size, mode='bilinear', align_corners=True))                  
+                tmp_feature.append(tmp_f)
+       
+        tmp_fea = tmp_feature[0]
+        for i_fea in range(len(tmp_feature) - 1):
+            tmp_fea = self.relu(torch.add(tmp_fea, F.interpolate((tmp_feature[i_fea+1]), tmp_feature[0].size()[2:], mode='bilinear', align_corners=True)))
+        up_score.append(F.interpolate(self.final_score(tmp_fea), x_size, mode='bilinear', align_corners=True))
+      
+
+
+        return up_score
+       
+
+
+# extra part
+def extra_layer(base_model_cfg, vgg):
+    if base_model_cfg == 'vgg':
+        config = config_vgg
+    elif base_model_cfg == 'resnet':
+        config = config_resnet
+    merge1_layers = MergeLayer1(config['merge1'])
+    merge2_layers = MergeLayer2(config['merge2'])
+
+    return vgg, merge1_layers, merge2_layers
+
+
+# TUN network
+class TUN_bone(nn.Module):
+    def __init__(self, base_model_cfg, base, merge1_layers, merge2_layers):
+        super(TUN_bone, self).__init__()
+        self.base_model_cfg = base_model_cfg
+        if self.base_model_cfg == 'vgg':
+
+            self.base = base
+            # self.base_ex = nn.ModuleList(base_ex)
+            self.merge1 = merge1_layers
+            self.merge2 = merge2_layers
+
+        elif self.base_model_cfg == 'resnet':
+            self.convert = ConvertLayer(config_resnet['convert'])
+            self.base = base
+            self.merge1 = merge1_layers
+            self.merge2 = merge2_layers
+
+    def forward(self, x):
+        x_size = x.size()[2:]
+        conv2merge = self.base(x)        
+        if self.base_model_cfg == 'resnet':            
+            conv2merge = self.convert(conv2merge)
+        up_edge, edge_feature, up_sal, sal_feature = self.merge1(conv2merge, x_size)
+        up_sal_final = self.merge2(edge_feature, sal_feature, x_size)
+        return up_edge, up_sal, up_sal_final
+
+
+# build the whole network
+def build_model(base_model_cfg='vgg'):
+    if base_model_cfg == 'vgg':
+        return TUN_bone(base_model_cfg, *extra_layer(base_model_cfg, vgg16()))
+    elif base_model_cfg == 'resnet':
+        return TUN_bone(base_model_cfg, *extra_layer(base_model_cfg, resnet50()))
+
+
+# weight init
+def xavier(param):
+    # init.xavier_uniform(param)
+    init.xavier_uniform_(param)
+
+
+def weights_init(m):
+    if isinstance(m, nn.Conv2d):
+        # xavier(m.weight.data)
+        m.weight.data.normal_(0, 0.01)
+        if m.bias is not None:
+            m.bias.data.zero_()
+
+if __name__ == '__main__':
+    from torch.autograd import Variable
+    net = TUN(*extra_layer(vgg(base['tun'], 3), vgg(base['tun_ex'], 512), config['merge_block'], config['fuse'])).cuda()
+    img = Variable(torch.randn((1, 3, 256, 256))).cuda()
+    out = net(img, mode = 2)
+    print(len(out))
+    print(len(out[0]))
+    print(out[0].shape)
+    print(len(out[1]))
+    # print(net)
+    input('Press Any to Continue...')

src/EGNet/resnet.py

@@ -0,0 +1,301 @@
+import torch.nn as nn
+import math
+# import torch.utils.model_zoo as model_zoo
+import torch
+import numpy as np
+import torch.nn.functional as F
+affine_par = True
+
+
+# def outS(i):
+#     i = int(i)
+#     i = (i+1)/2
+#     i = int(np.ceil((i+1)/2.0))
+#     i = (i+1)/2
+#     return i
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes, affine = affine_par)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes, affine = affine_par)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1,  dilation_ = 1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, bias=False) # change
+        self.bn1 = nn.BatchNorm2d(planes,affine = affine_par)
+        for i in self.bn1.parameters():
+            i.requires_grad = False
+        padding = 1
+        if dilation_ == 2:
+            padding = 2
+        elif dilation_ == 4:
+            padding = 4
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, # change
+                               padding=padding, bias=False, dilation = dilation_)
+        self.bn2 = nn.BatchNorm2d(planes,affine = affine_par)
+        for i in self.bn2.parameters():
+            i.requires_grad = False
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4, affine = affine_par)
+        for i in self.bn3.parameters():
+            i.requires_grad = False
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, layers):
+        self.inplanes = 64
+        super(ResNet, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64,affine = affine_par)
+        for i in self.bn1.parameters():
+            i.requires_grad = False
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=True) # change
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        # self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation__ = 2)
+        # self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation__ = 4)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation__ = 2)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, 0.01)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+        #        for i in m.parameters():
+        #            i.requires_grad = False
+
+    def _make_layer(self, block, planes, blocks, stride=1,dilation__ = 1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion or dilation__ == 2 or dilation__ == 4:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion,affine = affine_par),
+            )
+        for i in downsample._modules['1'].parameters():
+            i.requires_grad = False
+        layers = []
+        layers.append(block(self.inplanes, planes, stride,dilation_=dilation__, downsample = downsample ))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes,dilation_=dilation__))
+
+        return nn.Sequential(*layers)
+    # def _make_pred_layer(self,block, dilation_series, padding_series,NoLabels):
+    # return block(dilation_series,padding_series,NoLabels)
+
+    def forward(self, x):
+        tmp_x = []
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        tmp_x.append(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        tmp_x.append(x)
+        x = self.layer2(x)
+        tmp_x.append(x)
+        x = self.layer3(x)
+        tmp_x.append(x)
+        x = self.layer4(x)
+        tmp_x.append(x)
+
+        return tmp_x
+
+
+
+class ResNet_locate(nn.Module):
+    def __init__(self, block, layers):
+        super(ResNet_locate,self).__init__()
+        self.resnet = ResNet(block, layers)
+        self.in_planes = 512
+        self.out_planes = [512, 256, 256, 128]
+
+        self.ppms_pre = nn.Conv2d(2048, self.in_planes, 1, 1, bias=False)
+        ppms, infos = [], []
+        for ii in [1, 3, 5]:
+            ppms.append(nn.Sequential(nn.AdaptiveAvgPool2d(ii), nn.Conv2d(self.in_planes, self.in_planes, 1, 1, bias=False), nn.ReLU(inplace=True)))
+        self.ppms = nn.ModuleList(ppms)
+
+        self.ppm_cat = nn.Sequential(nn.Conv2d(self.in_planes * 4, self.in_planes, 3, 1, 1, bias=False), nn.ReLU(inplace=True))
+        # self.ppm_score = nn.Conv2d(self.in_planes, 1, 1, 1)
+        for ii in self.out_planes:
+            infos.append(nn.Sequential(nn.Conv2d(self.in_planes, ii, 3, 1, 1, bias=False), nn.ReLU(inplace=True)))
+        self.infos = nn.ModuleList(infos)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, 0.01)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def load_pretrained_model(self, model):
+        self.resnet.load_state_dict(model)
+
+    def forward(self, x):
+        x_size = x.size()[2:]
+        xs = self.resnet(x)
+
+        xs_1 = self.ppms_pre(xs[-1])
+        xls = [xs_1]
+        for k in range(len(self.ppms)):
+            xls.append(F.interpolate(self.ppms[k](xs_1), xs_1.size()[2:], mode='bilinear', align_corners=True))
+        xls = self.ppm_cat(torch.cat(xls, dim=1))
+        top_score = None
+        # top_score = F.interpolate(self.ppm_score(xls), x_size, mode='bilinear', align_corners=True)
+
+        infos = []
+        for k in range(len(self.infos)):
+            infos.append(self.infos[k](F.interpolate(xls, xs[len(self.infos) - 1 - k].size()[2:], mode='bilinear', align_corners=True)))
+
+        return xs, top_score, infos
+
+class BottleneckEZ(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1,  dilation_ = 1, downsample=None):
+        super(BottleneckEZ, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, bias=False) # change
+        # self.bn1 = nn.BatchNorm2d(planes,affine = affine_par)
+        # for i in self.bn1.parameters():
+        #     i.requires_grad = False
+        padding = 1
+        if dilation_ == 2:
+            padding = 2
+        elif dilation_ == 4:
+            padding = 4
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, # change
+                               padding=padding, bias=False, dilation = dilation_)
+        # self.bn2 = nn.BatchNorm2d(planes,affine = affine_par)
+        # for i in self.bn2.parameters():
+        #     i.requires_grad = False
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        # self.bn3 = nn.BatchNorm2d(planes * 4, affine = affine_par)
+        # for i in self.bn3.parameters():
+        #     i.requires_grad = False
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        # out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        # out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        # out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+
+def resnet50(pretrained=False):
+    """Constructs a ResNet-50 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on Places
+    """
+    # model = ResNet(Bottleneck, [3, 4, 6, 3])
+    model = ResNet(Bottleneck, [3, 4, 6, 3])
+    if pretrained:
+        model.load_state_dict(load_url(model_urls['resnet50']), strict=False)
+    return model
+
+def resnet101(pretrained=False):
+    """Constructs a ResNet-101 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    # model = ResNet(Bottleneck, [3, 4, 23, 3])
+    model = ResNet_locate(Bottleneck, [3, 4, 23, 3])
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
+    return model

src/EGNet/run.py

@@ -0,0 +1,68 @@
+import argparse
+import os
+from dataset import get_loader
+from solver import Solver
+
+
+def main(config):
+    if config.mode == 'train':
+        train_loader, dataset = get_loader(config.batch_size, num_thread=config.num_thread)
+        run = "nnet"
+        if not os.path.exists("%s/run-%s" % (config.save_fold, run)): 
+            os.mkdir("%s/run-%s" % (config.save_fold, run))
+            os.mkdir("%s/run-%s/logs" % (config.save_fold, run))
+            os.mkdir("%s/run-%s/models" % (config.save_fold, run))
+        config.save_fold = "%s/run-%s" % (config.save_fold, run)
+        train = Solver(train_loader, None, config)
+        train.train()
+    elif config.mode == 'test':
+        test_loader, dataset = get_loader(config.test_batch_size, mode='test',num_thread=config.num_thread, test_mode=config.test_mode, sal_mode=config.sal_mode)
+
+        test = Solver(None, test_loader, config, dataset.save_folder())
+        test.test(test_mode=config.test_mode)
+    else:
+        raise IOError("illegal input!!!")
+
+
+
+if __name__ == '__main__':
+
+    vgg_path = '/home/liuj/code/Messal/weights/vgg16_20M.pth'
+    resnet_path = '/home/liuj/code/Messal/weights/resnet50_caffe.pth'
+
+    parser = argparse.ArgumentParser()
+
+    # Hyper-parameters
+    parser.add_argument('--n_color', type=int, default=3)
+
+    parser.add_argument('--cuda', type=bool, default=True)
+
+    # Training settings
+    parser.add_argument('--vgg', type=str, default=vgg_path)
+    parser.add_argument('--resnet', type=str, default=resnet_path)
+    parser.add_argument('--epoch', type=int, default=30) # 12, now x3
+    parser.add_argument('--batch_size', type=int, default=1)
+    parser.add_argument('--test_batch_size', type=int, default=1)
+    parser.add_argument('--num_thread', type=int, default=4)
+    parser.add_argument('--load_bone', type=str, default='')
+    # parser.add_argument('--load_branch', type=str, default='')
+    parser.add_argument('--save_fold', type=str, default='./EGNet')
+    # parser.add_argument('--epoch_val', type=int, default=20)
+    parser.add_argument('--epoch_save', type=int, default=1) # 2, now x3
+    parser.add_argument('--epoch_show', type=int, default=1)
+    parser.add_argument('--pre_trained', type=str, default=None)
+
+    # Testing settings
+    parser.add_argument('--model', type=str, default='./epoch_resnet.pth')
+    parser.add_argument('--test_fold', type=str, default='./results/test')
+    parser.add_argument('--test_mode', type=int, default=1)
+    parser.add_argument('--sal_mode', type=str, default='t')
+
+    # Misc
+    parser.add_argument('--mode', type=str, default='train', choices=['train', 'test'])
+    parser.add_argument('--visdom', type=bool, default=False)
+    
+    config = parser.parse_args()
+  
+    if not os.path.exists(config.save_fold): os.mkdir(config.save_fold)
+    main(config)

src/EGNet/sal2edge.m

@@ -0,0 +1,34 @@
+data_root = '/home/liuj/dataset/DUTS/DUTS-TR/DUTS-TR-Mask';
+out_root = '/home/liuj/dataset/DUTS/DUTS-TR/DUTS-TR-Mask';
+lst_set = '/home/liuj/dataset/DUTS/DUTS-TR/train'
+index_file = fullfile([lst_set '.lst']);
+
+fileID = fopen(index_file);
+im_ids = textscan(fileID, '%s');
+im_ids = im_ids{1};
+fclose(fileID);
+
+
+num_images = length(im_ids);
+
+for im_id = 1:10
+
+    id = im_ids{im_id};
+    id = id(1:end-4);
+    
+%     img_path = fullfile(data_root, [id '.jpg']);
+%     image = imread(img_path);
+   
+    gt = imread(fullfile(data_root, [id '.png']));
+    gt = (gt > 128);
+    gt = double(gt);
+
+    [gy, gx] = gradient(gt);
+    temp_edge = gy.*gy + gx.*gx;
+    temp_edge(temp_edge~=0)=1;
+    bound = uint8(temp_edge*255);
+
+    save_path = fullfile(out_root, [id '_edge.png']);
+    imwrite(bound, save_path);
+
+end

src/EGNet/solver.py

@@ -0,0 +1,230 @@
+import torch
+from collections import OrderedDict
+from torch.nn import utils, functional as F
+from torch.optim import Adam, SGD
+from torch.autograd import Variable
+from torch.backends import cudnn
+from model import build_model, weights_init
+import scipy.misc as sm
+import numpy as np
+import os
+import torchvision.utils as vutils
+import cv2
+import torch.nn.functional as F
+import math
+import time
+import sys
+import PIL.Image
+import scipy.io
+import os
+import logging
+EPSILON = 1e-8
+p = OrderedDict()
+
+from dataset import get_loader
+base_model_cfg = 'resnet'
+p['lr_bone'] = 5e-5  # Learning rate resnet:5e-5, vgg:2e-5
+p['lr_branch'] = 0.025  # Learning rate
+p['wd'] = 0.0005  # Weight decay
+p['momentum'] = 0.90  # Momentum
+lr_decay_epoch = [15, 24] # [6, 9], now x3 #15
+nAveGrad = 10  # Update the weights once in 'nAveGrad' forward passes
+showEvery = 50
+tmp_path = 'tmp_see'
+
+
+class Solver(object):
+    def __init__(self, train_loader, test_loader, config, save_fold=None):
+        self.train_loader = train_loader
+        self.test_loader = test_loader
+        self.config = config
+        self.save_fold = save_fold
+        self.mean = torch.Tensor([123.68, 116.779, 103.939]).view(3, 1, 1) / 255.
+        # inference: choose the side map (see paper)
+        if config.visdom:
+            self.visual = Viz_visdom("trueUnify", 1)
+        self.build_model()
+        if self.config.pre_trained: self.net.load_state_dict(torch.load(self.config.pre_trained))
+        if config.mode == 'train':
+            self.log_output = open("%s/logs/log.txt" % config.save_fold, 'w')
+        else:
+            print('Loading pre-trained model from %s...' % self.config.model)
+            self.net_bone.load_state_dict(torch.load(self.config.model))
+            self.net_bone.eval()
+
+    def print_network(self, model, name):
+        num_params = 0
+        for p in model.parameters():
+            num_params += p.numel()
+        print(name)
+        print(model)
+        print("The number of parameters: {}".format(num_params))
+
+    def get_params(self, base_lr):
+        ml = []
+        for name, module in self.net_bone.named_children():
+            print(name)
+            if name == 'loss_weight':
+                ml.append({'params': module.parameters(), 'lr': p['lr_branch']})          
+            else:
+                ml.append({'params': module.parameters()})
+        return ml
+
+    # build the network
+    def build_model(self):
+        self.net_bone = build_model(base_model_cfg)
+        if self.config.cuda:
+            self.net_bone = self.net_bone.cuda()
+            
+        self.net_bone.eval()  # use_global_stats = True
+        self.net_bone.apply(weights_init)
+        if self.config.mode == 'train':
+            if self.config.load_bone == '':
+                if base_model_cfg == 'vgg':
+                    self.net_bone.base.load_pretrained_model(torch.load(self.config.vgg))
+                elif base_model_cfg == 'resnet':
+                    self.net_bone.base.load_state_dict(torch.load(self.config.resnet))
+            if self.config.load_bone != '': self.net_bone.load_state_dict(torch.load(self.config.load_bone))
+
+        self.lr_bone = p['lr_bone']
+        self.lr_branch = p['lr_branch']
+        self.optimizer_bone = Adam(filter(lambda p: p.requires_grad, self.net_bone.parameters()), lr=self.lr_bone, weight_decay=p['wd'])
+
+        self.print_network(self.net_bone, 'trueUnify bone part')
+
+    # update the learning rate
+    def update_lr(self, rate):
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = param_group['lr'] * rate
+
+
+    def test(self, test_mode=0):
+        EPSILON = 1e-8
+        img_num = len(self.test_loader)
+        time_t = 0.0
+        name_t = 'EGNet_ResNet50/'
+
+        if not os.path.exists(os.path.join(self.save_fold, name_t)):             
+            os.mkdir(os.path.join(self.save_fold, name_t))
+        for i, data_batch in enumerate(self.test_loader):
+            self.config.test_fold = self.save_fold
+            print(self.config.test_fold)
+            images_, name, im_size = data_batch['image'], data_batch['name'][0], np.asarray(data_batch['size'])
+            
+            with torch.no_grad():
+                
+                images = Variable(images_)
+                if self.config.cuda:
+                    images = images.cuda()
+                print(images.size())
+                time_start = time.time()
+                up_edge, up_sal, up_sal_f = self.net_bone(images)
+                torch.cuda.synchronize()
+                time_end = time.time()
+                print(time_end - time_start)
+                time_t = time_t + time_end - time_start                              
+                pred = np.squeeze(torch.sigmoid(up_sal_f[-1]).cpu().data.numpy())             
+                multi_fuse = 255 * pred
+                
+
+                
+                cv2.imwrite(os.path.join(self.config.test_fold,name_t, name[:-4] + '.png'), multi_fuse)
+          
+        print("--- %s seconds ---" % (time_t))
+        print('Test Done!')
+
+   
+    # training phase
+    def train(self):
+        iter_num = len(self.train_loader.dataset) // self.config.batch_size
+        aveGrad = 0
+        F_v = 0
+        if not os.path.exists(tmp_path): 
+            os.mkdir(tmp_path)
+        for epoch in range(self.config.epoch):                          
+            r_edge_loss, r_sal_loss, r_sum_loss= 0,0,0
+            self.net_bone.zero_grad()
+            for i, data_batch in enumerate(self.train_loader):
+                sal_image, sal_label, sal_edge = data_batch['sal_image'], data_batch['sal_label'], data_batch['sal_edge']
+                if sal_image.size()[2:] != sal_label.size()[2:]:
+                    print("Skip this batch")
+                    continue
+                sal_image, sal_label, sal_edge = Variable(sal_image), Variable(sal_label), Variable(sal_edge)
+                if self.config.cuda: 
+                    sal_image, sal_label, sal_edge = sal_image.cuda(), sal_label.cuda(), sal_edge.cuda()
+
+                up_edge, up_sal, up_sal_f = self.net_bone(sal_image)
+                # edge part
+                edge_loss = []
+                for ix in up_edge:
+                    edge_loss.append(bce2d_new(ix, sal_edge, reduction='sum'))
+                edge_loss = sum(edge_loss) / (nAveGrad * self.config.batch_size)
+                r_edge_loss += edge_loss.data
+                # sal part
+                sal_loss1= []
+                sal_loss2 = []
+                for ix in up_sal:
+                    sal_loss1.append(F.binary_cross_entropy_with_logits(ix, sal_label, reduction='sum'))
+
+                for ix in up_sal_f:
+                    sal_loss2.append(F.binary_cross_entropy_with_logits(ix, sal_label, reduction='sum'))
+                sal_loss = (sum(sal_loss1) + sum(sal_loss2)) / (nAveGrad * self.config.batch_size)
+              
+                r_sal_loss += sal_loss.data
+                loss = sal_loss + edge_loss
+                r_sum_loss += loss.data
+                loss.backward()
+                aveGrad += 1
+
+                if aveGrad % nAveGrad == 0:
+       
+                    self.optimizer_bone.step()
+                    self.optimizer_bone.zero_grad()           
+                    aveGrad = 0
+
+
+                if i % showEvery == 0:
+
+                    print('epoch: [%2d/%2d], iter: [%5d/%5d]  ||  Edge : %10.4f  ||  Sal : %10.4f  ||  Sum : %10.4f' % (
+                        epoch, self.config.epoch, i, iter_num,  r_edge_loss*(nAveGrad * self.config.batch_size)/showEvery,
+                                                                r_sal_loss*(nAveGrad * self.config.batch_size)/showEvery,
+                                                                r_sum_loss*(nAveGrad * self.config.batch_size)/showEvery))
+
+                    print('Learning rate: ' + str(self.lr_bone))
+                    r_edge_loss, r_sal_loss, r_sum_loss= 0,0,0
+
+                if i % 200 == 0:
+
+                    vutils.save_image(torch.sigmoid(up_sal_f[-1].data), tmp_path+'/iter%d-sal-0.jpg' % i, normalize=True, padding = 0)
+
+                    vutils.save_image(sal_image.data, tmp_path+'/iter%d-sal-data.jpg' % i, padding = 0)
+                    vutils.save_image(sal_label.data, tmp_path+'/iter%d-sal-target.jpg' % i, padding = 0)
+            
+            if (epoch + 1) % self.config.epoch_save == 0:
+                torch.save(self.net_bone.state_dict(), '%s/models/epoch_%d_bone.pth' % (self.config.save_fold, epoch + 1))
+                
+            if epoch in lr_decay_epoch:
+                self.lr_bone = self.lr_bone * 0.1  
+                self.optimizer_bone = Adam(filter(lambda p: p.requires_grad, self.net_bone.parameters()), lr=self.lr_bone, weight_decay=p['wd'])
+
+
+        torch.save(self.net_bone.state_dict(), '%s/models/final_bone.pth' % self.config.save_fold)
+        
+def bce2d_new(input, target, reduction=None):
+    assert(input.size() == target.size())
+    pos = torch.eq(target, 1).float()
+    neg = torch.eq(target, 0).float()
+    # ing = ((torch.gt(target, 0) & torch.lt(target, 1))).float()
+
+    num_pos = torch.sum(pos)
+    num_neg = torch.sum(neg)
+    num_total = num_pos + num_neg
+
+    alpha = num_neg  / num_total
+    beta = 1.1 * num_pos  / num_total
+    # target pixel = 1 -> weight beta
+    # target pixel = 0 -> weight 1-beta
+    weights = alpha * pos + beta * neg
+
+    return F.binary_cross_entropy_with_logits(input, target, weights, reduction=reduction)
+

src/EGNet/vgg.py

@@ -0,0 +1,273 @@
+import torch.nn as nn
+import math
+# import torch.utils.model_zoo as model_zoo
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+# vgg16
+def vgg(cfg, i, batch_norm=False):
+    layers = []
+    in_channels = i
+    stage = 1
+    for v in cfg:
+        if v == 'M':
+            stage += 1
+            if stage == 6:
+                layers += [nn.MaxPool2d(kernel_size=3, stride=2, padding=1)]
+            else:
+                layers += [nn.MaxPool2d(kernel_size=3, stride=2, padding=1)]
+        else:
+            if stage == 6:
+                # conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=4, dilation=4, bias=False)
+                conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
+            else:
+                conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
+            if batch_norm:
+                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
+            else:
+                layers += [conv2d, nn.ReLU(inplace=True)]
+            in_channels = v
+    return layers
+
+class vgg16(nn.Module):
+    def __init__(self):
+        super(vgg16, self).__init__()
+        self.cfg = {'tun': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 'tun_ex': [512, 512, 512]}
+        self.extract = [8, 15, 22, 29] # [3, 8, 15, 22, 29]
+        self.extract_ex = [5]
+        self.base = nn.ModuleList(vgg(self.cfg['tun'], 3))
+        self.base_ex = vgg_ex(self.cfg['tun_ex'], 512)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, 0.01)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def load_pretrained_model(self, model):
+        self.base.load_state_dict(model)
+
+    def forward(self, x, multi=0):
+        tmp_x = []
+        for k in range(len(self.base)):
+            x = self.base[k](x)
+            if k in self.extract:
+                tmp_x.append(x)
+        x = self.base_ex(x)
+        tmp_x.append(x)
+        if multi == 1:
+            tmp_y = []
+            tmp_y.append(tmp_x[0])
+            return tmp_y
+        else:
+            return tmp_x
+
+class vgg_ex(nn.Module):
+    def __init__(self, cfg, incs=512, padding=1, dilation=1):
+        super(vgg_ex, self).__init__()
+        self.cfg = cfg
+        layers = []
+        for v in self.cfg:
+            # conv2d = nn.Conv2d(incs, v, kernel_size=3, padding=4, dilation=4, bias=False)
+            conv2d = nn.Conv2d(incs, v, kernel_size=3, padding=padding, dilation=dilation, bias=False)
+            layers += [conv2d, nn.ReLU(inplace=True)]
+            incs = v
+        self.ex = nn.Sequential(*layers)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, 0.01)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def forward(self, x):
+        x = self.ex(x)
+        return x
+
+# class vgg16_locate(nn.Module):
+#     def __init__(self):
+#         super(vgg16_locate,self).__init__()
+#         self.cfg = [512, 512, 512]
+#         self.vgg16 = vgg16()
+#         # self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+#         # self.maxpool3 = nn.MaxPool2d(kernel_size=3, stride=3, padding=1)
+#         self.layer61 = vgg_ex(self.cfg, 512, 3, 3)
+#         self.layer62 = vgg_ex(self.cfg, 512, 6, 6)
+#         self.layer63 = vgg_ex(self.cfg, 512, 9, 9)
+#         self.layer64 = vgg_ex(self.cfg, 512, 12, 12)
+#
+#
+#         # self.layer6_convert, self.layer6_trans, self.layer6_score = [],[],[]
+#         # for ii in range(3):
+#         #     self.layer6_convert.append(nn.Conv2d(1024, 512, 3, 1, 1, bias=False))
+#         #     self.layer6_trans.append(nn.Conv2d(512, 512, 1, 1, bias=False))
+#         #     self.layer6_score.append(nn.Conv2d(512, 1, 1, 1))
+#         # self.layer6_convert, self.layer6_trans, self.layer6_score = nn.ModuleList(self.layer6_convert), nn.ModuleList(self.layer6_trans), nn.ModuleList(self.layer6_score)
+#         self.trans = nn.Conv2d(512*5, 512, 3, 1, 1, bias=False)
+#         # self.score = nn.Conv2d(3, 1, 1, 1)
+#         # self.score = nn.Conv2d(1, 1, 1, 1)
+#         self.relu = nn.ReLU(inplace=True)
+#
+#         for m in self.modules():
+#             if isinstance(m, nn.Conv2d):
+#                 n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+#                 m.weight.data.normal_(0, 0.01)
+#             elif isinstance(m, nn.BatchNorm2d):
+#                 m.weight.data.fill_(1)
+#                 m.bias.data.zero_()
+#
+#     def load_pretrained_model(self, model):
+#         self.vgg16.load_pretrained_model(model)
+#
+#     def forward(self, x):
+#         x_size = x.size()[2:]
+#         xs = self.vgg16(x)
+#
+#         xls = [xs[-1]]
+#         xls.append(self.layer61(xs[-2]))
+#         xls.append(self.layer62(xs[-2]))
+#         xls.append(self.layer63(xs[-2]))
+#         xls.append(self.layer64(xs[-2]))
+#
+#         # xls_tmp = [self.layer6_convert[0](xls[0])]
+#         # for ii in range(1, 3):
+#         #     xls_tmp.append(F.interpolate(self.layer6_convert[ii](xls[ii]), xls_tmp[0].size()[2:], mode='bilinear', align_corners=True))
+#         #
+#         # xls_trans = self.layer6_trans[0](xls_tmp[0])
+#         # for ii in range(1, 3):
+#         #     xls_trans = torch.add(xls_trans, self.layer6_trans[ii](xls_tmp[ii]))
+#         score, score_fuse = [], None
+#         # for ii in range(3):
+#         #     score.append(self.layer6_score[ii](xls_tmp[ii]))
+#
+#         xls_trans = self.trans(self.relu(torch.cat(xls, dim=1)))
+#         xs[-1] = xls_trans
+#         # score_fuse = F.interpolate(self.score(torch.cat(score, dim=1)), x_size, mode='bilinear', align_corners=True)
+#         # score_fuse = F.interpolate(self.score(torch.add(torch.add(score[0], score[1]), score[2])), x_size, mode='bilinear', align_corners=True)
+#
+#         # score = [F.interpolate(ss, x_size, mode='bilinear', align_corners=True) for ss in score]
+#
+#         return xs, score_fuse, score
+
+class vgg16_locate(nn.Module):
+    def __init__(self):
+        super(vgg16_locate,self).__init__()
+        self.vgg16 = vgg16()
+        self.in_planes = 512
+        # self.out_planes = [512, 256, 128, 64] #  with convert layer, with conv6
+        # self.out_planes = [512, 512, 256, 128] #  no convert layer, with conv6
+        self.out_planes = [512, 256, 128] #  no convert layer, no conv6
+
+        ppms, infos = [], []
+        # for ii in [3, 6, 12]:
+        #     if ii <= 8:
+        #         ppms.append(nn.Sequential(nn.AvgPool2d(kernel_size=ii, stride=ii), nn.Conv2d(self.in_planes, self.in_planes, 1, 1, bias=False), nn.ReLU(inplace=True)))
+        #     else:
+        #         ppms.append(nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Conv2d(self.in_planes, self.in_planes, 1, 1, bias=False), nn.ReLU(inplace=True)))
+        for ii in [1, 3, 5]:
+            ppms.append(nn.Sequential(nn.AdaptiveAvgPool2d(ii), nn.Conv2d(self.in_planes, self.in_planes, 1, 1, bias=False), nn.ReLU(inplace=True)))
+        self.ppms = nn.ModuleList(ppms)
+
+        self.ppm_cat = nn.Sequential(nn.Conv2d(self.in_planes * 4, self.in_planes, 3, 1, 1, bias=False), nn.ReLU(inplace=True))
+        #self.ppm_cat = nn.Sequential(nn.Conv2d(self.in_planes, self.in_planes, 3, 1, 1, bias=False), nn.ReLU(inplace=True))
+        # self.ppm_score = nn.Conv2d(self.in_planes, 1, 1, 1)
+        for ii in self.out_planes:
+            infos.append(nn.Sequential(nn.Conv2d(self.in_planes, ii, 3, 1, 1, bias=False), nn.ReLU(inplace=True)))
+        self.infos = nn.ModuleList(infos)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, 0.01)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def load_pretrained_model(self, model):
+        self.vgg16.load_pretrained_model(model)
+
+    def forward(self, x):
+        x_size = x.size()[2:]
+        xs = self.vgg16(x)
+
+        xls = [xs[-1]]
+        #xls = xs[-1]
+        for k in range(len(self.ppms)):
+            xls.append(F.interpolate(self.ppms[k](xs[-1]), xs[-1].size()[2:], mode='bilinear', align_corners=True))
+            #xls = torch.add(xls, F.interpolate(self.ppms[k](xs[-1]), xs[-1].size()[2:], mode='bilinear', align_corners=True))
+        xls = self.ppm_cat(torch.cat(xls, dim=1))
+        #xls = self.ppm_cat(xls)
+        top_score = None
+        # top_score = F.interpolate(self.ppm_score(xls), x_size, mode='bilinear', align_corners=True)
+
+        infos = []
+        for k in range(len(self.infos)):
+            infos.append(self.infos[k](F.interpolate(xls, xs[len(self.infos) - 1 - k].size()[2:], mode='bilinear', align_corners=True)))
+
+        return xs, top_score, infos
+
+# class vgg16_locate(nn.Module):
+#     def __init__(self):
+#         super(vgg16_locate,self).__init__()
+#         self.cfg = [1024, 1024, 1024]
+#         self.vgg16 = vgg16()
+#         self.maxpool4 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+#         self.maxpool5 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+#         self.maxpool6 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+#
+#         self.layer5 = vgg_ex(self.cfg, 1024)
+#         self.layer6 = vgg_ex(self.cfg, 1024)
+#         self.layer7 = vgg_ex(self.cfg, 1024)
+#
+#         self.layer71 = nn.Conv2d(1024, 512, 1, 1, bias=False)
+#         self.layer61 = nn.Conv2d(1024, 512, 1, 1, bias=False)
+#         self.layer51 = nn.Conv2d(1024, 512, 1, 1, bias=False)
+#         self.layer41 = nn.Conv2d(1024, 512, 1, 1, bias=False)
+#
+#         self.layer76 = nn.Conv2d(512, 512, 3, 1, 1, bias=False)
+#         self.layer65 = nn.Conv2d(512, 512, 3, 1, 1, bias=False)
+#         self.layer54 = nn.Sequential(nn.Conv2d(512, 512, 3, 1, 1, bias=False), nn.ReLU(inplace=True), nn.Conv2d(512, 512, 1, 1, bias=False))
+#         # self.layer54 = nn.Conv2d(512, 512, 3, 1, 1, bias=False)
+#         # self.layer54_ = nn.Sequential(nn.ReLU(inplace=True), nn.Conv2d(512, 512, 1, 1, bias=False))
+#         # self.score = nn.Conv2d(512, 1, 1, 1)
+#
+#         self.relu = nn.ReLU(inplace=True)
+#
+#         for m in self.modules():
+#             if isinstance(m, nn.Conv2d):
+#                 n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+#                 m.weight.data.normal_(0, 0.01)
+#             elif isinstance(m, nn.BatchNorm2d):
+#                 m.weight.data.fill_(1)
+#                 m.bias.data.zero_()
+#
+#     def load_pretrained_model(self, model):
+#         self.vgg16.load_pretrained_model(model)
+#
+#     def forward(self, x):
+#         x_size = x.size()[2:]
+#         score_fuse, score = None, None
+#         xs = self.vgg16(x)
+#
+#         x5 = self.layer5(self.maxpool4(xs[-1]))
+#         x6 = self.layer6(self.maxpool5(x5))
+#         x7 = self.layer7(self.maxpool6(x6))
+#
+#         x8 = self.layer76(self.relu(torch.add(F.interpolate(self.layer71(x7) , x6.size()[2:], mode='bilinear', align_corners=True), self.layer61(x6))))
+#         x8 = self.layer65(self.relu(torch.add(F.interpolate(x8 , x5.size()[2:], mode='bilinear', align_corners=True), self.layer51(x5))))
+#         x8 = self.layer54(self.relu(torch.add(F.interpolate(x8 , xs[-1].size()[2:], mode='bilinear', align_corners=True), self.layer41(xs[-1]))))
+#         xs[-1] = x8
+#
+#         # x8 = self.layer76(self.relu(torch.add(F.interpolate(self.layer71(x7) , x6.size()[2:], mode='bilinear', align_corners=True), self.layer61(x6))))
+#         # x9 = self.layer65(self.relu(torch.add(F.interpolate(x8 , x5.size()[2:], mode='bilinear', align_corners=True), self.layer51(x5))))
+#         # x10 = self.layer54(self.relu(torch.add(F.interpolate(x9 , xs[-1].size()[2:], mode='bilinear', align_corners=True), self.layer41(xs[-1]))))
+#         # score_fuse = F.interpolate(self.score(self.relu(torch.add(torch.add(F.interpolate(x8 , x10.size()[2:], mode='bilinear', align_corners=True),
+#         #                        F.interpolate(x9 , x10.size()[2:], mode='bilinear', align_corners=True)), x10))), x_size, mode='bilinear', align_corners=True)
+#         # xs[-1] = self.layer54_(x10)
+#
+#         return xs, score_fuse, score

src/diffusion_hacked.py

@@ -0,0 +1,957 @@
+from einops import rearrange, reduce, repeat
+import torch.nn.functional as F
+import torch
+import gc
+from src.utils import *
+from src.flow_utils import get_mapping_ind, warp_tensor
+from diffusers.models.unet_2d_condition import UNet2DConditionOutput
+from diffusers.models.attention_processor import AttnProcessor2_0
+from typing import Any, Dict, List, Optional, Tuple, Union
+import sys
+sys.path.append("./src/ebsynth/deps/gmflow/")
+from gmflow.geometry import flow_warp, forward_backward_consistency_check
+
+"""
+==========================================================================
+PART I - FRESCO-based attention
+* Class AttentionControl: Control the function of FRESCO-based attention
+* Class FRESCOAttnProcessor2_0: FRESCO-based attention
+* apply_FRESCO_attn(): Apply FRESCO-based attention to a StableDiffusionPipeline
+==========================================================================
+"""
+
+class AttentionControl():
+    """
+    Control FRESCO-based attention
+    * enable/diable spatial-guided attention
+    * enable/diable temporal-guided attention
+    * enable/diable cross-frame attention
+    * collect intermediate attention feature (for spatial-guided attention)
+    """
+    def __init__(self):
+        self.stored_attn = self.get_empty_store()
+        self.store = False
+        self.index = 0
+        self.attn_mask = None
+        self.interattn_paras = None
+        self.use_interattn = False
+        self.use_cfattn = False
+        self.use_intraattn = False
+        self.intraattn_bias = 0
+        self.intraattn_scale_factor = 0.2
+        self.interattn_scale_factor = 0.2
+    
+    @staticmethod
+    def get_empty_store():
+        return {
+            'decoder_attn': [],
+        }
+    
+    def clear_store(self):
+        del self.stored_attn
+        torch.cuda.empty_cache()
+        gc.collect()
+        self.stored_attn = self.get_empty_store()
+        self.disable_intraattn()
+
+    # store attention feature of the input frame for spatial-guided attention
+    def enable_store(self):
+        self.store = True
+        
+    def disable_store(self):
+        self.store = False  
+
+    # spatial-guided attention
+    def enable_intraattn(self):
+        self.index = 0
+        self.use_intraattn = True
+        self.disable_store()
+        if len(self.stored_attn['decoder_attn']) == 0:
+            self.use_intraattn = False
+        
+    def disable_intraattn(self):
+        self.index = 0
+        self.use_intraattn = False
+        self.disable_store()
+
+    def disable_cfattn(self):
+        self.use_cfattn = False        
+
+    # cross frame attention
+    def enable_cfattn(self, attn_mask=None):
+        if attn_mask:
+            if self.attn_mask:
+                del self.attn_mask
+                torch.cuda.empty_cache()
+            self.attn_mask = attn_mask
+            self.use_cfattn = True  
+        else:
+            if self.attn_mask:
+                self.use_cfattn = True
+            else:
+                print('Warning: no valid cross-frame attention parameters available!')
+                self.disable_cfattn()       
+        
+    def disable_interattn(self):
+        self.use_interattn = False
+
+    # temporal-guided attention
+    def enable_interattn(self, interattn_paras=None):
+        if interattn_paras:
+            if self.interattn_paras:
+                del self.interattn_paras
+                torch.cuda.empty_cache()
+            self.interattn_paras = interattn_paras
+            self.use_interattn = True
+        else:
+            if self.interattn_paras:
+                self.use_interattn = True
+            else:
+                print('Warning: no valid temporal-guided attention parameters available!')
+                self.disable_interattn()
+    
+    def disable_controller(self):
+        self.disable_intraattn()
+        self.disable_interattn()
+        self.disable_cfattn()
+    
+    def enable_controller(self, interattn_paras=None, attn_mask=None):
+        self.enable_intraattn()
+        self.enable_interattn(interattn_paras)
+        self.enable_cfattn(attn_mask)    
+    
+    def forward(self, context):
+        if self.store:
+            self.stored_attn['decoder_attn'].append(context.detach())
+        if self.use_intraattn and len(self.stored_attn['decoder_attn']) > 0:
+            tmp = self.stored_attn['decoder_attn'][self.index]
+            self.index = self.index + 1
+            if self.index >= len(self.stored_attn['decoder_attn']):
+                self.index = 0
+                self.disable_store()
+            return tmp
+        return context
+    
+    def __call__(self, context):
+        context = self.forward(context)
+        return context
+
+
+#import xformers
+#import importlib
+class FRESCOAttnProcessor2_0:
+    """
+    Hack self attention to FRESCO-based attention
+    * adding spatial-guided attention
+    * adding temporal-guided attention
+    * adding cross-frame attention
+    
+    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+    Usage
+    frescoProc = FRESCOAttnProcessor2_0(2, attn_mask)
+    attnProc = AttnProcessor2_0()
+    
+    attn_processor_dict = {}
+    for k in pipe.unet.attn_processors.keys():
+        if k.startswith("up_blocks.2") or k.startswith("up_blocks.3"):
+            attn_processor_dict[k] = frescoProc
+        else:
+            attn_processor_dict[k] = attnProc
+    pipe.unet.set_attn_processor(attn_processor_dict)
+    """
+
+    def __init__(self, unet_chunk_size=2, controller=None):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+        self.unet_chunk_size = unet_chunk_size
+        self.controller = controller
+            
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        crossattn = False
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+            if self.controller and self.controller.store:
+                self.controller(hidden_states.detach().clone())
+        else:
+            crossattn = True
+            if attn.norm_cross:
+                encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+            
+        # BC * HW * 8D
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+        
+        query_raw, key_raw = None, None
+        if self.controller and self.controller.use_interattn and (not crossattn):
+            query_raw, key_raw = query.clone(), key.clone()
+
+        inner_dim = key.shape[-1] # 8D
+        head_dim = inner_dim // attn.heads # D
+        
+        '''for efficient cross-frame attention'''
+        if self.controller and self.controller.use_cfattn and (not crossattn):
+            video_length = key.size()[0] // self.unet_chunk_size
+            former_frame_index = [0] * video_length
+            attn_mask = None
+            if self.controller.attn_mask is not None:
+                for m in self.controller.attn_mask:
+                    if m.shape[1] == key.shape[1]:
+                        attn_mask = m
+            # BC * HW * 8D --> B * C * HW * 8D
+            key = rearrange(key, "(b f) d c -> b f d c", f=video_length)
+            # B * C * HW * 8D --> B * C * HW * 8D
+            if attn_mask is None:
+                key = key[:, former_frame_index]
+            else:
+                key = repeat(key[:, attn_mask], "b d c -> b f d c", f=video_length)
+            # B * C * HW * 8D --> BC * HW * 8D 
+            key = rearrange(key, "b f d c -> (b f) d c").detach()
+            value = rearrange(value, "(b f) d c -> b f d c", f=video_length)
+            if attn_mask is None:
+                value = value[:, former_frame_index]
+            else:
+                value = repeat(value[:, attn_mask], "b d c -> b f d c", f=video_length)              
+            value = rearrange(value, "b f d c -> (b f) d c").detach()
+        
+        # BC * HW * 8D --> BC * HW * 8 * D --> BC * 8 * HW * D
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        # BC * 8 * HW2 * D
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        # BC * 8 * HW2 * D2
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        
+        '''for spatial-guided intra-frame attention'''
+        if self.controller and self.controller.use_intraattn and (not crossattn): 
+            ref_hidden_states = self.controller(None)
+            assert ref_hidden_states.shape == encoder_hidden_states.shape
+            query_ = attn.to_q(ref_hidden_states)
+            key_ = attn.to_k(ref_hidden_states) 
+            
+            ''' 
+            # for xformers implementation 
+            if importlib.util.find_spec("xformers") is not None:
+                # BC * HW * 8D --> BC * HW * 8 * D
+                query_ = rearrange(query_, "b d (h c) -> b d h c", h=attn.heads)
+                key_ = rearrange(key_, "b d (h c) -> b d h c", h=attn.heads)
+                # BC * 8 * HW * D --> 8BC * HW * D
+                query = rearrange(query, "b h d c -> b d h c")
+                query = xformers.ops.memory_efficient_attention(
+                    query_, key_ * self.sattn_scale_factor, query, 
+                    attn_bias=torch.eye(query_.size(1), key_.size(1), 
+                    dtype=query.dtype, device=query.device) * self.bias_weight, op=None
+                )
+                query = rearrange(query, "b d h c -> b h d c").detach()
+            '''
+            # BC * 8 * HW * D
+            query_ = query_.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            key_ = key_.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            query = F.scaled_dot_product_attention(
+                query_, key_ * self.controller.intraattn_scale_factor, query, 
+                attn_mask = torch.eye(query_.size(-2), key_.size(-2), 
+                                      dtype=query.dtype, device=query.device) * self.controller.intraattn_bias,
+            ).detach()
+            #print('intra: ', GPU.getGPUs()[1].memoryUsed)
+            del query_, key_
+            torch.cuda.empty_cache()
+        
+        '''
+        # for xformers implementation
+        if importlib.util.find_spec("xformers") is not None:
+            hidden_states = xformers.ops.memory_efficient_attention(
+                    rearrange(query, "b h d c -> b d h c"), rearrange(key, "b h d c -> b d h c"), 
+                    rearrange(value, "b h d c -> b d h c"), 
+                    attn_bias=attention_mask, op=None
+                )
+            hidden_states = rearrange(hidden_states, "b d h c -> b h d c", h=attn.heads)
+        '''
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        # output: BC * 8 * HW * D2      
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+        #print('cross: ', GPU.getGPUs()[1].memoryUsed)
+        
+        '''for temporal-guided inter-frame attention (FLATTEN)'''
+        if self.controller and self.controller.use_interattn and (not crossattn):
+            del query, key, value
+            torch.cuda.empty_cache()
+            bwd_mapping = None
+            fwd_mapping = None
+            flattn_mask = None
+            for i, f in enumerate(self.controller.interattn_paras['fwd_mappings']):
+                if f.shape[2] == hidden_states.shape[2]:
+                    fwd_mapping = f
+                    bwd_mapping = self.controller.interattn_paras['bwd_mappings'][i]
+                    interattn_mask = self.controller.interattn_paras['interattn_masks'][i]
+            video_length = key_raw.size()[0] // self.unet_chunk_size
+            # BC * HW * 8D --> C * 8BD * HW
+            key = rearrange(key_raw, "(b f) d c -> f (b c) d", f=video_length)
+            query = rearrange(query_raw, "(b f) d c -> f (b c) d", f=video_length)
+            # BC * 8 * HW * D --> C * 8BD * HW
+            #key = rearrange(hidden_states, "(b f) h d c -> f (b h c) d", f=video_length) ########
+            #query = rearrange(hidden_states, "(b f) h d c -> f (b h c) d", f=video_length) #######
+            
+            value = rearrange(hidden_states, "(b f) h d c -> f (b h c) d", f=video_length)
+            key = torch.gather(key, 2, fwd_mapping.expand(-1,key.shape[1],-1))
+            query = torch.gather(query, 2, fwd_mapping.expand(-1,query.shape[1],-1))
+            value = torch.gather(value, 2, fwd_mapping.expand(-1,value.shape[1],-1))
+            # C * 8BD * HW --> BHW, C, 8D
+            key = rearrange(key, "f (b c) d -> (b d) f c", b=self.unet_chunk_size)
+            query = rearrange(query, "f (b c) d -> (b d) f c", b=self.unet_chunk_size)
+            value = rearrange(value, "f (b c) d -> (b d) f c", b=self.unet_chunk_size) 
+            '''
+            # for xformers implementation 
+            if importlib.util.find_spec("xformers") is not None:
+                # BHW * C * 8D --> BHW * C * 8 * D
+                query = rearrange(query, "b d (h c) -> b d h c", h=attn.heads)
+                key = rearrange(key, "b d (h c) -> b d h c", h=attn.heads)
+                value = rearrange(value, "b d (h c) -> b d h c", h=attn.heads)
+                B, D, C, _ = flattn_mask.shape
+                C1 = int(np.ceil(C / 4) * 4)
+                attn_bias = torch.zeros(B, D, C, C1, dtype=value.dtype, device=value.device) # HW * 1 * C * C
+                attn_bias[:,:,:,:C].masked_fill_(interattn_mask.logical_not(), float("-inf")) # BHW * C * C
+                hidden_states_ = xformers.ops.memory_efficient_attention(
+                    query, key * self.controller.interattn_scale_factor, value, 
+                    attn_bias=attn_bias.squeeze(1).repeat(self.unet_chunk_size*attn.heads,1,1)[:,:,:C], op=None
+                )
+                hidden_states_ = rearrange(hidden_states_, "b d h c -> b h d c", h=attn.heads).detach()
+            '''
+            # BHW * C * 8D --> BHW * C * 8 * D--> BHW * 8 * C * D
+            query = query.view(-1, video_length, attn.heads, head_dim).transpose(1, 2).detach()
+            key = key.view(-1, video_length, attn.heads, head_dim).transpose(1, 2).detach()
+            value = value.view(-1, video_length, attn.heads, head_dim).transpose(1, 2).detach()
+            hidden_states_ = F.scaled_dot_product_attention(
+                query, key * self.controller.interattn_scale_factor, value, 
+                attn_mask = (interattn_mask.repeat(self.unet_chunk_size,1,1,1))#.to(query.dtype)-1.0) * 1e6 -
+                #torch.eye(interattn_mask.shape[2]).to(query.device).to(query.dtype) * 1e4,
+            )
+                
+            # BHW * 8 * C * D --> C * 8BD * HW
+            hidden_states_ = rearrange(hidden_states_, "(b d) h f c -> f (b h c) d", b=self.unet_chunk_size)
+            hidden_states_ = torch.gather(hidden_states_, 2, bwd_mapping.expand(-1,hidden_states_.shape[1],-1)).detach()
+            # C * 8BD * HW --> BC * 8 * HW * D
+            hidden_states = rearrange(hidden_states_, "f (b h c) d -> (b f) h d c", b=self.unet_chunk_size, h=attn.heads)
+            #print('inter: ', GPU.getGPUs()[1].memoryUsed)
+            
+        # BC * 8 * HW * D --> BC * HW * 8D
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+def apply_FRESCO_attn(pipe):
+    """
+    Apply FRESCO-guided attention to a StableDiffusionPipeline
+    """    
+    frescoProc = FRESCOAttnProcessor2_0(2, AttentionControl())
+    attnProc = AttnProcessor2_0()
+    attn_processor_dict = {}
+    for k in pipe.unet.attn_processors.keys():
+        if k.startswith("up_blocks.2") or k.startswith("up_blocks.3"):
+            attn_processor_dict[k] = frescoProc
+        else:
+            attn_processor_dict[k] = attnProc
+    pipe.unet.set_attn_processor(attn_processor_dict)
+    return frescoProc
+    
+
+"""
+==========================================================================
+PART II - FRESCO-based optimization
+* optimize_feature(): function to optimze latent feature 
+* my_forward(): hacked pipe.unet.forward(), adding feature optimization
+* apply_FRESCO_opt(): function to apply FRESCO-based optimization to a StableDiffusionPipeline
+* disable_FRESCO_opt(): function to disable the FRESCO-based optimization
+==========================================================================
+"""
+
+def optimize_feature(sample, flows, occs, correlation_matrix=[], 
+                     intra_weight = 1e2, iters=20, unet_chunk_size=2, optimize_temporal = True):
+    """
+    FRESO-guided latent feature optimization
+    * optimize spatial correspondence (match correlation_matrix)
+    * optimize temporal correspondence (match warped_image)
+    """
+    if (flows is None or occs is None or (not optimize_temporal)) and (intra_weight == 0 or len(correlation_matrix) == 0):
+        return sample
+    # flows=[fwd_flows, bwd_flows]: (N-1)*2*H1*W1
+    # occs=[fwd_occs, bwd_occs]: (N-1)*H1*W1
+    # sample: 2N*C*H*W
+    torch.cuda.empty_cache()
+    video_length = sample.shape[0] // unet_chunk_size
+    latent = rearrange(sample.to(torch.float32), "(b f) c h w -> b f c h w", f=video_length)
+    
+    cs = torch.nn.Parameter((latent.detach().clone()))
+    optimizer = torch.optim.Adam([cs], lr=0.2)
+
+    # unify resolution
+    if flows is not None and occs is not None:
+        scale = sample.shape[2] * 1.0 / flows[0].shape[2]
+        kernel = int(1 / scale)
+        bwd_flow_ = F.interpolate(flows[1] * scale, scale_factor=scale, mode='bilinear').repeat(unet_chunk_size,1,1,1)
+        bwd_occ_ = F.max_pool2d(occs[1].unsqueeze(1), kernel_size=kernel).repeat(unet_chunk_size,1,1,1) # 2(N-1)*1*H1*W1
+        fwd_flow_ = F.interpolate(flows[0] * scale, scale_factor=scale, mode='bilinear').repeat(unet_chunk_size,1,1,1)
+        fwd_occ_ = F.max_pool2d(occs[0].unsqueeze(1), kernel_size=kernel).repeat(unet_chunk_size,1,1,1) # 2(N-1)*1*H1*W1
+        # match frame 0,1,2,3 and frame 1,2,3,0
+        reshuffle_list = list(range(1,video_length))+[0]
+        
+    # attention_probs is the GRAM matrix of the normalized feature 
+    attention_probs = None
+    for tmp in correlation_matrix:
+        if sample.shape[2] * sample.shape[3] == tmp.shape[1]:
+            attention_probs = tmp # 2N*HW*HW
+            break
+     
+    n_iter=[0]
+    while n_iter[0] < iters:
+        def closure():
+            optimizer.zero_grad()
+            
+            loss = 0
+
+            # temporal consistency loss 
+            if optimize_temporal and flows is not None and occs is not None:
+                c1 = rearrange(cs[:,:], "b f c h w -> (b f) c h w")
+                c2 = rearrange(cs[:,reshuffle_list], "b f c h w -> (b f) c h w")
+                warped_image1 = flow_warp(c1, bwd_flow_)
+                warped_image2 = flow_warp(c2, fwd_flow_)
+                loss = (abs((c2-warped_image1)*(1-bwd_occ_)) + abs((c1-warped_image2)*(1-fwd_occ_))).mean() * 2
+                
+            # spatial consistency loss
+            if attention_probs is not None and intra_weight > 0:
+                cs_vector = rearrange(cs, "b f c h w -> (b f) (h w) c")
+                #attention_scores = torch.bmm(cs_vector, cs_vector.transpose(-1, -2))
+                #cs_attention_probs = attention_scores.softmax(dim=-1)
+                cs_vector = cs_vector / ((cs_vector ** 2).sum(dim=2, keepdims=True) ** 0.5)
+                cs_attention_probs = torch.bmm(cs_vector, cs_vector.transpose(-1, -2))
+                tmp = F.l1_loss(cs_attention_probs, attention_probs) * intra_weight
+                loss = tmp + loss
+                
+            loss.backward()
+            n_iter[0]+=1
+            
+            
+            if False: # for debug
+                print('Iteration: %d, loss: %f'%(n_iter[0]+1, loss.data.mean()))
+            return loss
+        optimizer.step(closure)
+
+    torch.cuda.empty_cache()
+    return adaptive_instance_normalization(rearrange(cs.data.to(sample.dtype), "b f c h w -> (b f) c h w"), sample)
+
+
+def my_forward(self, steps = [], layers = [0,1,2,3], flows = None, occs = None, 
+               correlation_matrix=[], intra_weight = 1e2, iters=20, optimize_temporal = True, saliency = None):
+    """
+    Hacked pipe.unet.forward()
+    copied from https://github.com/huggingface/diffusers/blob/v0.19.3/src/diffusers/models/unet_2d_condition.py#L700
+    if you are using a new version of diffusers, please copy the source code and modify it accordingly (find [HACK] in the code)
+    * restore and return the decoder features
+    * optimize the decoder features
+    * perform background smoothing
+    """    
+    def forward(
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        class_labels: Optional[torch.Tensor] = None,
+        timestep_cond: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[UNet2DConditionOutput, Tuple]:
+        r"""
+        The [`UNet2DConditionModel`] forward method.
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The noisy input tensor with the following shape `(batch, channel, height, width)`.
+            timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
+            encoder_hidden_states (`torch.FloatTensor`):
+                The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.
+            encoder_attention_mask (`torch.Tensor`):
+                A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If
+                `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,
+                which adds large negative values to the attention scores corresponding to "discard" tokens.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
+                tuple.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the [`AttnProcessor`].
+            added_cond_kwargs: (`dict`, *optional*):
+                A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that
+                are passed along to the UNet blocks.
+
+        Returns:
+            [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
+                If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise
+                a `tuple` is returned where the first element is the sample tensor.
+        """
+        # By default samples have to be AT least a multiple of the overall upsampling factor.
+        # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
+        # However, the upsampling interpolation output size can be forced to fit any upsampling size
+        # on the fly if necessary.
+        default_overall_up_factor = 2**self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+
+        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
+            logger.info("Forward upsample size to force interpolation output size.")
+            forward_upsample_size = True
+
+        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension
+        # expects mask of shape:
+        #   [batch, key_tokens]
+        # adds singleton query_tokens dimension:
+        #   [batch,                    1, key_tokens]
+        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+        if attention_mask is not None:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #       (keep = +0,     discard = -10000.0)
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        # 0. center input if necessary
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+
+        # `Timesteps` does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=sample.dtype)
+
+        emb = self.time_embedding(t_emb, timestep_cond)
+        aug_emb = None
+
+        if self.class_embedding is not None:
+            if class_labels is None:
+                raise ValueError("class_labels should be provided when num_class_embeds > 0")
+
+            if self.config.class_embed_type == "timestep":
+                class_labels = self.time_proj(class_labels)
+
+                # `Timesteps` does not contain any weights and will always return f32 tensors
+                # there might be better ways to encapsulate this.
+                class_labels = class_labels.to(dtype=sample.dtype)
+
+            class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)
+
+            if self.config.class_embeddings_concat:
+                emb = torch.cat([emb, class_emb], dim=-1)
+            else:
+                emb = emb + class_emb
+
+        if self.config.addition_embed_type == "text":
+            aug_emb = self.add_embedding(encoder_hidden_states)
+        elif self.config.addition_embed_type == "text_image":
+            # Kandinsky 2.1 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
+                )
+
+            image_embs = added_cond_kwargs.get("image_embeds")
+            text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)
+            aug_emb = self.add_embedding(text_embs, image_embs)
+        elif self.config.addition_embed_type == "text_time":
+            # SDXL - style
+            if "text_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
+                )
+            text_embeds = added_cond_kwargs.get("text_embeds")
+            if "time_ids" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
+                )
+            time_ids = added_cond_kwargs.get("time_ids")
+            time_embeds = self.add_time_proj(time_ids.flatten())
+            time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
+
+            add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
+            add_embeds = add_embeds.to(emb.dtype)
+            aug_emb = self.add_embedding(add_embeds)
+        elif self.config.addition_embed_type == "image":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
+                )
+            image_embs = added_cond_kwargs.get("image_embeds")
+            aug_emb = self.add_embedding(image_embs)
+        elif self.config.addition_embed_type == "image_hint":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`"
+                )
+            image_embs = added_cond_kwargs.get("image_embeds")
+            hint = added_cond_kwargs.get("hint")
+            aug_emb, hint = self.add_embedding(image_embs, hint)
+            sample = torch.cat([sample, hint], dim=1)
+
+        emb = emb + aug_emb if aug_emb is not None else emb
+
+        if self.time_embed_act is not None:
+            emb = self.time_embed_act(emb)
+
+        if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":
+            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj":
+            # Kadinsky 2.1 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            encoder_hidden_states = self.encoder_hid_proj(image_embeds)
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        # 3. down
+
+        is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
+        is_adapter = mid_block_additional_residual is None and down_block_additional_residuals is not None
+
+        down_block_res_samples = (sample,)
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                # For t2i-adapter CrossAttnDownBlock2D
+                additional_residuals = {}
+                if is_adapter and len(down_block_additional_residuals) > 0:
+                    additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0)
+
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                    **additional_residuals,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+                if is_adapter and len(down_block_additional_residuals) > 0:
+                    sample += down_block_additional_residuals.pop(0)
+            down_block_res_samples += res_samples
+
+        if is_controlnet:
+            new_down_block_res_samples = ()
+
+            for down_block_res_sample, down_block_additional_residual in zip(
+                down_block_res_samples, down_block_additional_residuals
+            ):
+                down_block_res_sample = down_block_res_sample + down_block_additional_residual
+                new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
+
+            down_block_res_samples = new_down_block_res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            sample = self.mid_block(
+                sample,
+                emb,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=attention_mask,
+                cross_attention_kwargs=cross_attention_kwargs,
+                encoder_attention_mask=encoder_attention_mask,
+            )
+
+        if is_controlnet:
+            sample = sample + mid_block_additional_residual
+        
+        # 5. up
+        '''
+        [HACK] restore the decoder features in up_samples
+        '''
+        up_samples = ()
+        #down_samples = ()
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
+
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            '''
+            [HACK] restore the decoder features in up_samples
+            [HACK] optimize the decoder features
+            [HACK] perform background smoothing
+            '''
+            if i in layers:
+                up_samples += (sample, )
+            if timestep in steps and i in layers: 
+                sample = optimize_feature(sample, flows, occs, correlation_matrix, 
+                                          intra_weight, iters, optimize_temporal = optimize_temporal)
+                if saliency is not None:
+                    sample = warp_tensor(sample, flows, occs, saliency, 2)
+            
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
+
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    upsample_size=upsample_size,
+                    attention_mask=attention_mask,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
+                )
+
+        # 6. post-process
+        if self.conv_norm_out:
+            sample = self.conv_norm_out(sample)
+            sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        '''
+        [HACK] return the output feature as well as the decoder features
+        '''        
+        if not return_dict:
+            return (sample, ) + up_samples
+
+        return UNet2DConditionOutput(sample=sample)
+    
+    return forward    
+
+
+def apply_FRESCO_opt(pipe, steps = [], layers = [0,1,2,3], flows = None, occs = None, 
+               correlation_matrix=[], intra_weight = 1e2, iters=20, optimize_temporal = True, saliency = None):
+    """
+    Apply FRESCO-based optimization to a StableDiffusionPipeline
+    """  
+    pipe.unet.forward = my_forward(pipe.unet, steps, layers, flows, occs, 
+               correlation_matrix, intra_weight, iters, optimize_temporal, saliency)
+
+def disable_FRESCO_opt(pipe):
+    """
+    Disable the FRESCO-based optimization
+    """  
+    apply_FRESCO_opt(pipe)
+
+
+"""
+=====================================================================================
+PART III - Prepare parameters for FRESCO-guided attention/optimization 
+* get_intraframe_paras(): get parameters for spatial-guided attention/optimization 
+* get_flow_and_interframe_paras(): get parameters for temporal-guided attention/optimization 
+=====================================================================================
+"""
+
+@torch.no_grad()
+def get_intraframe_paras(pipe, imgs, frescoProc, 
+                         prompt_embeds, do_classifier_free_guidance=True, seed=0):
+    """
+    Get parameters for spatial-guided attention and optimization
+    * perform one step denoising 
+    * collect attention feature, stored in frescoProc.controller.stored_attn['decoder_attn']
+    * compute the gram matrix of the normalized feature for spatial consistency loss
+    """
+
+    noise_scheduler = pipe.scheduler 
+    timestep = noise_scheduler.timesteps[-1]
+    device = pipe._execution_device
+    generator = torch.Generator(device=device).manual_seed(seed)
+    B, C, H, W = imgs.shape
+    
+    frescoProc.controller.disable_controller()
+    disable_FRESCO_opt(pipe)
+    frescoProc.controller.clear_store()
+    frescoProc.controller.enable_store()
+    
+    latents = pipe.prepare_latents(
+        B,
+        pipe.unet.config.in_channels,
+        H,
+        W,
+        prompt_embeds.dtype,
+        device,
+        generator,
+        latents = None,
+    )
+
+    latent_x0 = pipe.vae.config.scaling_factor * pipe.vae.encode(imgs.to(pipe.unet.dtype)).latent_dist.sample()
+    latents = noise_scheduler.add_noise(latent_x0, latents, timestep).detach()
+
+    latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents 
+    model_output = pipe.unet(
+        latent_model_input,
+        timestep,
+        encoder_hidden_states=prompt_embeds,
+        cross_attention_kwargs=None,
+        return_dict=False,
+    )
+    
+    frescoProc.controller.disable_store()
+
+    # gram matrix of the normalized feature for spatial consistency loss
+    correlation_matrix = []
+    for tmp in model_output[1:]:
+        latent_vector = rearrange(tmp, "b c h w -> b (h w) c")
+        latent_vector = latent_vector / ((latent_vector ** 2).sum(dim=2, keepdims=True) ** 0.5)
+        attention_probs = torch.bmm(latent_vector, latent_vector.transpose(-1, -2))
+        correlation_matrix += [attention_probs.detach().clone().to(torch.float32)]
+        del attention_probs, latent_vector, tmp
+    del model_output    
+    
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    return correlation_matrix
+
+
+@torch.no_grad()
+def get_flow_and_interframe_paras(flow_model, imgs, visualize_pipeline=False):
+    """
+    Get parameters for temporal-guided attention and optimization
+    * predict optical flow and occlusion mask
+    * compute pixel index correspondence for FLATTEN
+    """
+    images = torch.stack([torch.from_numpy(img).permute(2, 0, 1).float() for img in imgs], dim=0).cuda()
+    imgs_torch = torch.cat([numpy2tensor(img) for img in imgs], dim=0)
+    
+    reshuffle_list = list(range(1,len(images)))+[0]
+    
+    results_dict = flow_model(images, images[reshuffle_list], attn_splits_list=[2], 
+                              corr_radius_list=[-1], prop_radius_list=[-1], pred_bidir_flow=True)
+    flow_pr = results_dict['flow_preds'][-1]  # [2*B, 2, H, W]
+    fwd_flows, bwd_flows = flow_pr.chunk(2)   # [B, 2, H, W]
+    fwd_occs, bwd_occs = forward_backward_consistency_check(fwd_flows, bwd_flows) # [B, H, W]
+    
+    warped_image1 = flow_warp(images, bwd_flows)
+    bwd_occs = torch.clamp(bwd_occs + (abs(images[reshuffle_list]-warped_image1).mean(dim=1)>255*0.25).float(), 0 ,1)
+    
+    warped_image2 = flow_warp(images[reshuffle_list], fwd_flows)
+    fwd_occs = torch.clamp(fwd_occs + (abs(images-warped_image2).mean(dim=1)>255*0.25).float(), 0 ,1)    
+    
+    if visualize_pipeline:
+        print('visualized occlusion masks based on optical flows')
+        viz = torchvision.utils.make_grid(imgs_torch * (1-fwd_occs.unsqueeze(1)), len(images), 1)
+        visualize(viz.cpu(), 90)
+        viz = torchvision.utils.make_grid(imgs_torch[reshuffle_list] * (1-bwd_occs.unsqueeze(1)), len(images), 1)
+        visualize(viz.cpu(), 90) 
+        
+    attn_mask = []
+    for scale in [8.0, 16.0, 32.0]:
+        bwd_occs_ = F.interpolate(bwd_occs[:-1].unsqueeze(1), scale_factor=1./scale, mode='bilinear')
+        attn_mask += [torch.cat((bwd_occs_[0:1].reshape(1,-1)>-1, bwd_occs_.reshape(bwd_occs_.shape[0],-1)>0.5), dim=0)]   
+        
+    fwd_mappings = []
+    bwd_mappings = []
+    interattn_masks = []
+    for scale in [8.0, 16.0]:
+        fwd_mapping, bwd_mapping, interattn_mask = get_mapping_ind(bwd_flows, bwd_occs, imgs_torch, scale=scale)
+        fwd_mappings += [fwd_mapping]
+        bwd_mappings += [bwd_mapping]
+        interattn_masks += [interattn_mask]  
+    
+    interattn_paras = {}
+    interattn_paras['fwd_mappings'] = fwd_mappings
+    interattn_paras['bwd_mappings'] = bwd_mappings
+    interattn_paras['interattn_masks'] = interattn_masks    
+
+    gc.collect()
+    torch.cuda.empty_cache()
+    
+    return [fwd_flows, bwd_flows], [fwd_occs, bwd_occs], attn_mask, interattn_paras

src/ebsynth/blender/guide.py

@@ -0,0 +1,104 @@
+import os
+
+import cv2
+import numpy as np
+
+from flow.flow_utils import flow_calc, read_flow, read_mask
+
+
+class BaseGuide:
+
+    def __init__(self):
+        ...
+
+    def get_cmd(self, i, weight) -> str:
+        return (f'-guide {os.path.abspath(self.imgs[0])} '
+                f'{os.path.abspath(self.imgs[i])} -weight {weight}')
+
+
+class ColorGuide(BaseGuide):
+
+    def __init__(self, imgs):
+        super().__init__()
+        self.imgs = imgs
+
+
+class PositionalGuide(BaseGuide):
+
+    def __init__(self, flow_paths, save_paths):
+        super().__init__()
+        flows = [read_flow(f) for f in flow_paths]
+        masks = [read_mask(f) for f in flow_paths]
+        # TODO: modify the format of flow to numpy
+        H, W = flows[0].shape[2:]
+        first_img = PositionalGuide.__generate_first_img(H, W)
+        prev_img = first_img
+        imgs = [first_img]
+        cid = 0
+        for flow, mask in zip(flows, masks):
+            cur_img = flow_calc.warp(prev_img, flow,
+                                     'nearest').astype(np.uint8)
+            cur_img = cv2.inpaint(cur_img, mask, 30, cv2.INPAINT_TELEA)
+            prev_img = cur_img
+            imgs.append(cur_img)
+            cid += 1
+            cv2.imwrite(f'guide/{cid}.jpg', mask)
+
+        for path, img in zip(save_paths, imgs):
+            cv2.imwrite(path, img)
+        self.imgs = save_paths
+
+    @staticmethod
+    def __generate_first_img(H, W):
+        Hs = np.linspace(0, 1, H)
+        Ws = np.linspace(0, 1, W)
+        i, j = np.meshgrid(Hs, Ws, indexing='ij')
+        r = (i * 255).astype(np.uint8)
+        g = (j * 255).astype(np.uint8)
+        b = np.zeros(r.shape)
+        res = np.stack((b, g, r), 2)
+        return res
+
+
+class EdgeGuide(BaseGuide):
+
+    def __init__(self, imgs, save_paths):
+        super().__init__()
+        edges = [EdgeGuide.__generate_edge(cv2.imread(img)) for img in imgs]
+        for path, img in zip(save_paths, edges):
+            cv2.imwrite(path, img)
+        self.imgs = save_paths
+
+    @staticmethod
+    def __generate_edge(img):
+        filter = np.array([[0, -1, 0], [-1, 4, -1], [0, -1, 0]])
+        res = cv2.filter2D(img, -1, filter)
+        return res
+
+
+class TemporalGuide(BaseGuide):
+
+    def __init__(self, key_img, stylized_imgs, flow_paths, save_paths):
+        super().__init__()
+        self.flows = [read_flow(f) for f in flow_paths]
+        self.masks = [read_mask(f) for f in flow_paths]
+        self.stylized_imgs = stylized_imgs
+        self.imgs = save_paths
+
+        first_img = cv2.imread(key_img)
+        cv2.imwrite(self.imgs[0], first_img)
+
+    def get_cmd(self, i, weight) -> str:
+        if i == 0:
+            warped_img = self.stylized_imgs[0]
+        else:
+            prev_img = cv2.imread(self.stylized_imgs[i - 1])
+            warped_img = flow_calc.warp(prev_img, self.flows[i - 1],
+                                        'nearest').astype(np.uint8)
+
+            warped_img = cv2.inpaint(warped_img, self.masks[i - 1], 30,
+                                     cv2.INPAINT_TELEA)
+
+            cv2.imwrite(self.imgs[i], warped_img)
+
+        return super().get_cmd(i, weight)

src/ebsynth/blender/histogram_blend.py

@@ -0,0 +1,50 @@
+import cv2
+import numpy as np
+
+
+def histogram_transform(img: np.ndarray, means: np.ndarray, stds: np.ndarray,
+                        target_means: np.ndarray, target_stds: np.ndarray):
+    means = means.reshape((1, 1, 3))
+    stds = stds.reshape((1, 1, 3))
+    target_means = target_means.reshape((1, 1, 3))
+    target_stds = target_stds.reshape((1, 1, 3))
+    x = img.astype(np.float32)
+    x = (x - means) * target_stds / stds + target_means
+    # x = np.round(x)
+    # x = np.clip(x, 0, 255)
+    # x = x.astype(np.uint8)
+    return x
+
+
+def blend(a: np.ndarray,
+          b: np.ndarray,
+          min_error: np.ndarray,
+          weight1=0.5,
+          weight2=0.5):
+    a = cv2.cvtColor(a, cv2.COLOR_BGR2Lab)
+    b = cv2.cvtColor(b, cv2.COLOR_BGR2Lab)
+    min_error = cv2.cvtColor(min_error, cv2.COLOR_BGR2Lab)
+    a_mean = np.mean(a, axis=(0, 1))
+    a_std = np.std(a, axis=(0, 1))
+    b_mean = np.mean(b, axis=(0, 1))
+    b_std = np.std(b, axis=(0, 1))
+    min_error_mean = np.mean(min_error, axis=(0, 1))
+    min_error_std = np.std(min_error, axis=(0, 1))
+
+    t_mean_val = 0.5 * 256
+    t_std_val = (1 / 36) * 256
+    t_mean = np.ones([3], dtype=np.float32) * t_mean_val
+    t_std = np.ones([3], dtype=np.float32) * t_std_val
+    a = histogram_transform(a, a_mean, a_std, t_mean, t_std)
+
+    b = histogram_transform(b, b_mean, b_std, t_mean, t_std)
+    ab = (a * weight1 + b * weight2 - t_mean_val) / 0.5 + t_mean_val
+    ab_mean = np.mean(ab, axis=(0, 1))
+    ab_std = np.std(ab, axis=(0, 1))
+    ab = histogram_transform(ab, ab_mean, ab_std, min_error_mean,
+                             min_error_std)
+    ab = np.round(ab)
+    ab = np.clip(ab, 0, 255)
+    ab = ab.astype(np.uint8)
+    ab = cv2.cvtColor(ab, cv2.COLOR_Lab2BGR)
+    return ab

src/ebsynth/blender/poisson_fusion.py

@@ -0,0 +1,93 @@
+import cv2
+import numpy as np
+import scipy
+
+As = None
+prev_states = None
+
+
+def construct_A(h, w, grad_weight):
+    indgx_x = []
+    indgx_y = []
+    indgy_x = []
+    indgy_y = []
+    vdx = []
+    vdy = []
+    for i in range(h):
+        for j in range(w):
+            if i < h - 1:
+                indgx_x += [i * w + j]
+                indgx_y += [i * w + j]
+                vdx += [1]
+                indgx_x += [i * w + j]
+                indgx_y += [(i + 1) * w + j]
+                vdx += [-1]
+            if j < w - 1:
+                indgy_x += [i * w + j]
+                indgy_y += [i * w + j]
+                vdy += [1]
+                indgy_x += [i * w + j]
+                indgy_y += [i * w + j + 1]
+                vdy += [-1]
+    Ix = scipy.sparse.coo_array(
+        (np.ones(h * w), (np.arange(h * w), np.arange(h * w))),
+        shape=(h * w, h * w)).tocsc()
+    Gx = scipy.sparse.coo_array(
+        (np.array(vdx), (np.array(indgx_x), np.array(indgx_y))),
+        shape=(h * w, h * w)).tocsc()
+    Gy = scipy.sparse.coo_array(
+        (np.array(vdy), (np.array(indgy_x), np.array(indgy_y))),
+        shape=(h * w, h * w)).tocsc()
+    As = []
+    for i in range(3):
+        As += [
+            scipy.sparse.vstack([Gx * grad_weight[i], Gy * grad_weight[i], Ix])
+        ]
+    return As
+
+
+# blendI, I1, I2, mask should be RGB unit8 type
+# return poissson fusion result (RGB unit8 type)
+# I1 and I2: propagated results from previous and subsequent key frames
+# mask: pixel selection mask
+# blendI: contrastive-preserving blending results of I1 and I2
+def poisson_fusion(blendI, I1, I2, mask, grad_weight=[2.5, 0.5, 0.5]):
+    global As
+    global prev_states
+
+    Iab = cv2.cvtColor(blendI, cv2.COLOR_BGR2LAB).astype(float)
+    Ia = cv2.cvtColor(I1, cv2.COLOR_BGR2LAB).astype(float)
+    Ib = cv2.cvtColor(I2, cv2.COLOR_BGR2LAB).astype(float)
+    m = (mask > 0).astype(float)[:, :, np.newaxis]
+    h, w, c = Iab.shape
+
+    # fuse the gradient of I1 and I2 with mask
+    gx = np.zeros_like(Ia)
+    gy = np.zeros_like(Ia)
+    gx[:-1, :, :] = (Ia[:-1, :, :] - Ia[1:, :, :]) * (1 - m[:-1, :, :]) + (
+        Ib[:-1, :, :] - Ib[1:, :, :]) * m[:-1, :, :]
+    gy[:, :-1, :] = (Ia[:, :-1, :] - Ia[:, 1:, :]) * (1 - m[:, :-1, :]) + (
+        Ib[:, :-1, :] - Ib[:, 1:, :]) * m[:, :-1, :]
+
+    # construct A for solving Ax=b
+    crt_states = (h, w, grad_weight)
+    if As is None or crt_states != prev_states:
+        As = construct_A(*crt_states)
+        prev_states = crt_states
+
+    final = []
+    for i in range(3):
+        weight = grad_weight[i]
+        im_dx = np.clip(gx[:, :, i].reshape(h * w, 1), -100, 100)
+        im_dy = np.clip(gy[:, :, i].reshape(h * w, 1), -100, 100)
+        im = Iab[:, :, i].reshape(h * w, 1)
+        im_mean = im.mean()
+        im = im - im_mean
+        A = As[i]
+        b = np.vstack([im_dx * weight, im_dy * weight, im])
+        out = scipy.sparse.linalg.lsqr(A, b)
+        out_im = (out[0] + im_mean).reshape(h, w, 1)
+        final += [out_im]
+
+    final = np.clip(np.concatenate(final, axis=2), 0, 255)
+    return cv2.cvtColor(final.astype(np.uint8), cv2.COLOR_LAB2BGR)

src/ebsynth/blender/video_sequence.py

@@ -0,0 +1,187 @@
+import os
+import shutil
+
+
+class VideoSequence:
+
+    def __init__(self,
+                 base_dir,
+                 key_ind,
+                 input_subdir='videos',
+                 key_subdir='keys0',
+                 tmp_subdir='tmp',
+                 input_format='frame%04d.jpg',
+                 key_format='%04d.jpg',
+                 out_subdir_format='out_%d',
+                 blending_out_subdir='blend',
+                 output_format='%04d.jpg'):
+        #if (end_frame - beg_frame) % interval != 0:
+        #    end_frame -= (end_frame - beg_frame) % interval
+
+        self.__base_dir = base_dir
+        self.__input_dir = os.path.join(base_dir, input_subdir)
+        self.__key_dir = os.path.join(base_dir, key_subdir)
+        self.__tmp_dir = os.path.join(base_dir, tmp_subdir)
+        self.__input_format = input_format
+        self.__blending_out_dir = os.path.join(base_dir, blending_out_subdir)
+        self.__key_format = key_format
+        self.__out_subdir_format = out_subdir_format
+        self.__output_format = output_format
+        self.__key_ind = key_ind
+        #self.__beg_frame = beg_frame
+        #self.__end_frame = end_frame
+        #self.__interval = interval
+        self.__n_seq = len(key_ind)-1#(end_frame - beg_frame) // interval
+        self.__make_out_dirs()
+        os.makedirs(self.__tmp_dir, exist_ok=True)
+
+    @property
+    def beg_frame(self):
+        return self.__key_ind[0]#self.__beg_frame
+
+    @property
+    def end_frame(self):
+        return self.__key_ind[-1]#self.__end_frame
+
+    @property
+    def n_seq(self):
+        return self.__n_seq
+
+    @property
+    def blending_dir(self):
+        return os.path.abspath(self.__blending_out_dir)
+
+    def interval(self, i):
+        return self.get_sequence_beg_id(i + 1) - self.get_sequence_beg_id(i)    
+    
+    def remove_out_and_tmp(self):
+        for i in range(self.n_seq + 1):
+            out_dir = self.__get_out_subdir(i)
+            shutil.rmtree(out_dir)
+        shutil.rmtree(self.__tmp_dir)
+
+    def get_input_sequence(self, i, is_forward=True):
+        beg_id = self.get_sequence_beg_id(i)
+        end_id = self.get_sequence_beg_id(i + 1)
+        if is_forward:
+            id_list = list(range(beg_id, end_id))
+        else:
+            id_list = list(range(end_id, beg_id, -1))
+        path_dir = [
+            os.path.join(self.__input_dir, self.__input_format % id)
+            for id in id_list
+        ]
+        return path_dir
+
+    def get_output_sequence(self, i, is_forward=True):
+        beg_id = self.get_sequence_beg_id(i)
+        end_id = self.get_sequence_beg_id(i + 1)
+        if is_forward:
+            id_list = list(range(beg_id, end_id))
+        else:
+            i += 1
+            id_list = list(range(end_id, beg_id, -1))
+        out_subdir = self.__get_out_subdir(i)
+        path_dir = [
+            os.path.join(out_subdir, self.__output_format % id)
+            for id in id_list
+        ]
+        return path_dir
+
+    def get_temporal_sequence(self, i, is_forward=True):
+        beg_id = self.get_sequence_beg_id(i)
+        end_id = self.get_sequence_beg_id(i + 1)
+        if is_forward:
+            id_list = list(range(beg_id, end_id))
+        else:
+            i += 1
+            id_list = list(range(end_id, beg_id, -1))
+        tmp_dir = self.__get_tmp_out_subdir(i)
+        path_dir = [
+            os.path.join(tmp_dir, 'temporal_' + self.__output_format % id)
+            for id in id_list
+        ]
+        return path_dir
+
+    def get_edge_sequence(self, i, is_forward=True):
+        beg_id = self.get_sequence_beg_id(i)
+        end_id = self.get_sequence_beg_id(i + 1)
+        if is_forward:
+            id_list = list(range(beg_id, end_id))
+        else:
+            i += 1
+            id_list = list(range(end_id, beg_id, -1))
+        tmp_dir = self.__get_tmp_out_subdir(i)
+        path_dir = [
+            os.path.join(tmp_dir, 'edge_' + self.__output_format % id)
+            for id in id_list
+        ]
+        return path_dir
+
+    def get_pos_sequence(self, i, is_forward=True):
+        beg_id = self.get_sequence_beg_id(i)
+        end_id = self.get_sequence_beg_id(i + 1)
+        if is_forward:
+            id_list = list(range(beg_id, end_id))
+        else:
+            i += 1
+            id_list = list(range(end_id, beg_id, -1))
+        tmp_dir = self.__get_tmp_out_subdir(i)
+        path_dir = [
+            os.path.join(tmp_dir, 'pos_' + self.__output_format % id)
+            for id in id_list
+        ]
+        return path_dir
+
+    def get_flow_sequence(self, i, is_forward=True):
+        beg_id = self.get_sequence_beg_id(i)
+        end_id = self.get_sequence_beg_id(i + 1)
+        if is_forward:
+            id_list = list(range(beg_id, end_id - 1))
+            path_dir = [
+                os.path.join(self.__tmp_dir, 'flow_f_%04d.npy' % id)
+                for id in id_list
+            ]
+        else:
+            id_list = list(range(end_id, beg_id + 1, -1))
+            path_dir = [
+                os.path.join(self.__tmp_dir, 'flow_b_%04d.npy' % id)
+                for id in id_list
+            ]
+
+        return path_dir
+
+    def get_input_img(self, i):
+        return os.path.join(self.__input_dir, self.__input_format % i)
+
+    def get_key_img(self, i):
+        sequence_beg_id = self.get_sequence_beg_id(i)
+        return os.path.join(self.__key_dir,
+                            self.__key_format % sequence_beg_id)
+
+    def get_blending_img(self, i):
+        return os.path.join(self.__blending_out_dir, self.__output_format % i)
+
+    def get_sequence_beg_id(self, i):
+        return self.__key_ind[i]#i * self.__interval + self.__beg_frame
+
+    def __get_out_subdir(self, i):
+        dir_id = self.get_sequence_beg_id(i)
+        out_subdir = os.path.join(self.__base_dir,
+                                  self.__out_subdir_format % dir_id)
+        return out_subdir
+
+    def __get_tmp_out_subdir(self, i):
+        dir_id = self.get_sequence_beg_id(i)
+        tmp_out_subdir = os.path.join(self.__tmp_dir,
+                                      self.__out_subdir_format % dir_id)
+        return tmp_out_subdir
+
+    def __make_out_dirs(self):
+        os.makedirs(self.__base_dir, exist_ok=True)
+        os.makedirs(self.__blending_out_dir, exist_ok=True)
+        for i in range(self.__n_seq + 1):
+            out_subdir = self.__get_out_subdir(i)
+            tmp_subdir = self.__get_tmp_out_subdir(i)
+            os.makedirs(out_subdir, exist_ok=True)
+            os.makedirs(tmp_subdir, exist_ok=True)