# Documentation

## Python scripts

These files are for our monocular 3D Tracking pipeline:

`main.py` Execute 3D photo inpainting

`mesh.py` Functions about context-aware depth inpainting

`mesh_tools.py` Some common functions used in `mesh.py`

`utils.py` Some common functions used in image preprocessing, data loading

`networks.py` Network architectures of inpainting model


MiDaS/

`run.py` Execute depth estimation

`monodepth_net.py` Network architecture of depth estimation model

`MiDaS_utils.py` Some common functions in depth estimation


## Configuration

```bash
argument.yml
```

- `depth_edge_model_ckpt: checkpoints/EdgeModel.pth`
    - Pretrained model of depth-edge inpainting
- `depth_feat_model_ckpt: checkpoints/DepthModel.pth`
    - Pretrained model of depth inpainting
- `rgb_feat_model_ckpt: checkpoints/ColorModel.pth`
    - Pretrained model of color inpainting
- `MiDaS_model_ckpt: MiDaS/model.pt`
    - Pretrained model of depth estimation
- `use_boostmonodepth: True`
    - Use [BoostMonocularDepth](https://github.com/compphoto/BoostingMonocularDepth) to get sharper monocular depth estimation
- `fps: 40`
    - Frame per second of output rendered video
- `num_frames: 240`
    - Total number of frames in output rendered video
- `x_shift_range: [-0.03, -0.03, -0.03]`
    - The translations on x-axis of output rendered videos.
    - This parameter is a list. Each element corresponds to a specific camera motion.
- `y_shift_range: [-0.00, -0.00, -0.03]`
    - The translations on y-axis of output rendered videos.
    - This parameter is a list. Each element corresponds to a specific camera motion.
- `z_shift_range: [-0.07, -0.07, -0.07]`
    - The translations on z-axis of output rendered videos.
    - This parameter is a list. Each element corresponds to a specific camera motion.
- `traj_types: ['straight-line', 'circle', 'circle']`
    - The type of camera trajectory.
    - This parameter is a list.
    - Currently, we only privode `straight-line` and `circle`.
-  `video_postfix: ['zoom-in', 'swing', 'circle']`
    - The postfix of video.
    - This parameter is a list.
- Note that the number of elements in `x_shift_range`,  `y_shift_range`, `z_shift_range`, `traj_types` and `video_postfix` should be equal.
- `specific: '' `
    - The specific image name, use this to specify the image to be executed. By default, all the image in the folder will    be executed.
- `longer_side_len: 960`
    - The length of larger dimension in output resolution.
- `src_folder: image`
    - Input image directory. 
- `depth_folder: depth`
    - Estimated depth directory.
- `mesh_folder: mesh`
    - Output 3-D mesh directory.
- `video_folder: video`
    - Output rendered video directory
- `load_ply: False`
    - Action to load existed mesh (.ply) file
- `save_ply: True`
    - Action to store the output mesh (.ply) file
    - Disable this option `save_ply: False` to reduce the computational time.
- `inference_video: True`
    - Action to rendered the output video
- `gpu_ids: 0`
    - The ID of working GPU. Leave it blank or negative to use CPU.
- `offscreen_rendering: True`
    - If you're executing the process in a remote server (via ssh), please switch on this flag. 
    - Sometimes, using off-screen rendering result in longer execution time.
- `img_format: '.jpg'`
    - Input image format.
- `depth_format: '.npy'`
    - Input depth (disparity) format. Use NumPy array file as default.
    - If the user wants to edit the depth (disparity) map manually, we provide `.png` format depth (disparity) map.
        - Remember to switch this parameter from `.npy` to `.png` when using depth (disparity) map with `.png` format.
- `require_midas: True`
    - Set it to `True` if the user wants to use depth map estimated by `MiDaS`.
    - Set it to `False` if the user wants to use manually edited depth map.
    - If the user wants to edit the depth (disparity) map manually, we provide `.png` format depth (disparity) map.
        - Remember to switch this parameter from `True` to `False` when using manually edited depth map.
- `depth_threshold: 0.04`
    - A threshold in disparity, adjacent two pixels are discontinuity pixels 
      if the difference between them excceed this number.
- `ext_edge_threshold: 0.002`
    - The threshold to define inpainted depth edge. A pixel in inpainted edge 
      map belongs to extended depth edge if the value of that pixel exceeds this number,
- `sparse_iter: 5`
    - Total iteration numbers of bilateral median filter
- `filter_size: [7, 7, 5, 5, 5]`
    - Window size of bilateral median filter in each iteration.
- `sigma_s: 4.0`
    - Intensity term of bilateral median filter
- `sigma_r: 0.5`
    - Spatial term of bilateral median filter
- `redundant_number: 12`
    - The number defines short segments. If a depth edge is shorter than this number, 
      it is a short segment and removed.
- `background_thickness: 70`
    - The thickness of synthesis area.
- `context_thickness: 140`
    - The thickness of context area.
- `background_thickness_2: 70`
    - The thickness of synthesis area when inpaint second time.
- `context_thickness_2: 70`
    - The thickness of context area when inpaint second time.
- `discount_factor: 1.00`
- `log_depth: True`
    - The scale of depth inpainting. If true, performing inpainting in log scale. 
      Otherwise, performing in linear scale.
- `largest_size: 512`
    - The largest size of inpainted image patch.
- `depth_edge_dilate: 10`
    - The thickness of dilated synthesis area.
- `depth_edge_dilate_2: 5`
    - The thickness of dilated synthesis area when inpaint second time.
- `extrapolate_border: True`
    - Action to extrapolate out-side the border.
- `extrapolation_thickness: 60`
    - The thickness of extrapolated area.
- `repeat_inpaint_edge: True`
    - Action to apply depth edge inpainting model repeatedly. Sometimes inpainting depth 
      edge once results in short inpinated edge, apply depth edge inpainting repeatedly 
      could help you prolong the inpainted depth edge. 
- `crop_border: [0.03, 0.03, 0.05, 0.03]`
    - The fraction of pixels to crop out around the borders `[top, left, bottom, right]`.
- `anti_flickering: True`
    - Action to avoid flickering effect in the output video. 
    - This may result in longer computational time in rendering phase.                                                                                                                                                                                                                                                                                                                                                                                                   # [CVPR 2020] 3D Photography using Context-aware Layered Depth Inpainting

[![Open 3DPhotoInpainting in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1706ToQrkIZshRSJSHvZ1RuCiM__YX3Bz)

### [[Paper](https://arxiv.org/abs/2004.04727)] [[Project Website](https://shihmengli.github.io/3D-Photo-Inpainting/)] [[Google Colab](https://colab.research.google.com/drive/1706ToQrkIZshRSJSHvZ1RuCiM__YX3Bz)]

<p align='center'>
<img src='https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/3DPhoto_teaser.jpg' width='900'/>
</p>

We propose a method for converting a single RGB-D input image into a 3D photo, i.e., a multi-layer representation for novel view synthesis that contains hallucinated color and depth structures in regions occluded in the original view. We use a Layered Depth Image with explicit pixel connectivity as underlying representation, and present a learning-based inpainting model that iteratively synthesizes new local color-and-depth content into the occluded region in a spatial context-aware manner. The resulting 3D photos can be efficiently rendered with motion parallax using standard graphics engines. We validate the effectiveness of our method on a wide range of challenging everyday scenes and show fewer artifacts when compared with the state-of-the-arts.
<br/>

**3D Photography using Context-aware Layered Depth Inpainting**
<br/>
[Meng-Li Shih](https://shihmengli.github.io/), 
[Shih-Yang Su](https://lemonatsu.github.io/), 
[Johannes Kopf](https://johanneskopf.de/), and
[Jia-Bin Huang](https://filebox.ece.vt.edu/~jbhuang/)
<br/>
In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.


## Prerequisites

- Linux (tested on Ubuntu 18.04.4 LTS)
- Anaconda
- Python 3.7 (tested on 3.7.4)
- PyTorch 1.4.0 (tested on 1.4.0 for execution)

and the Python dependencies listed in [requirements.txt](requirements.txt)
- To get started, please run the following commands:
    ```bash
    conda create -n 3DP python=3.7 anaconda
    conda activate 3DP
    pip install -r requirements.txt
    conda install pytorch==1.4.0 torchvision==0.5.0 cudatoolkit==10.1.243 -c pytorch
    ```
- Next, please download the model weight using the following command:
    ```bash
    chmod +x download.sh
    ./download.sh
    ```    

## Quick start
Please follow the instructions in this section. 
This should allow to execute our results.
For more detailed instructions, please refer to [`DOCUMENTATION.md`](DOCUMENTATION.md).

## Execute
1. Put ```.jpg``` files (e.g., test.jpg) into the ```image``` folder. 
    - E.g., `image/moon.jpg`
2. Run the following command
    ```bash
    python main.py --config argument.yml
    ```
    - Note: The 3D photo generation process usually takes about 2-3 minutes depending on the available computing resources.
3. The results are stored in the following directories:
    - Corresponding depth map estimated by [MiDaS](https://github.com/intel-isl/MiDaS.git) 
        - E.g. ```depth/moon.npy```, ```depth/moon.png```
        - User could edit ```depth/moon.png``` manually. 
            - Remember to set the following two flags as listed below if user wants to use manually edited ```depth/moon.png``` as input for 3D Photo.
                - `depth_format: '.png'`
                - `require_midas: False`
    - Inpainted 3D mesh (Optional: User need to switch on the flag `save_ply`)
        - E.g. ```mesh/moon.ply```
    - Rendered videos with zoom-in motion
        - E.g. ```video/moon_zoom-in.mp4```
    - Rendered videos with swing motion
        - E.g. ```video/moon_swing.mp4```
    - Rendered videos with circle motion
        - E.g. ```video/moon_circle.mp4```         
    - Rendered videos with dolly zoom-in effect
        - E.g. ```video/moon_dolly-zoom-in.mp4```
        - Note: We assume that the object of focus is located at the center of the image.
4. (Optional) If you want to change the default configuration. Please read [`DOCUMENTATION.md`](DOCUMENTATION.md) and modified ```argument.yml```.


## License
This work is licensed under MIT License. See [LICENSE](LICENSE) for details. 

If you find our code/models useful, please consider citing our paper:
```
@inproceedings{Shih3DP20,
  author = {Shih, Meng-Li and Su, Shih-Yang and Kopf, Johannes and Huang, Jia-Bin},
  title = {3D Photography using Context-aware Layered Depth Inpainting},
  booktitle = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  year = {2020}
}
```

## Acknowledgments
- We thank Pratul Srinivasan for providing clarification of the method [Srinivasan et al. CVPR 2019](https://people.eecs.berkeley.edu/~pratul/publication/mpi_extrapolation/).
- We thank the author of [Zhou et al. 2018](https://people.eecs.berkeley.edu/~tinghuiz/projects/mpi/), [Choi et al. 2019](https://github.com/NVlabs/extreme-view-synth/), [Mildenhall et al. 2019](https://github.com/Fyusion/LLFF), [Srinivasan et al. 2019](https://github.com/google-research/google-research/tree/ac9b04e1dbdac468fda53e798a326fe9124e49fe/mpi_extrapolation), [Wiles et al. 2020](http://www.robots.ox.ac.uk/~ow/synsin.html), [Niklaus et al. 2019](https://github.com/sniklaus/3d-ken-burns) for providing their implementations online.
- Our code builds upon [EdgeConnect](https://github.com/knazeri/edge-connect), [MiDaS](https://github.com/intel-isl/MiDaS.git) and [pytorch-inpainting-with-partial-conv](https://github.com/naoto0804/pytorch-inpainting-with-partial-conv)                                                       depth_edge_model_ckpt: checkpoints/edge-model.pth
depth_feat_model_ckpt: checkpoints/depth-model.pth
rgb_feat_model_ckpt: checkpoints/color-model.pth
MiDaS_model_ckpt: MiDaS/model.pt
use_boostmonodepth: True
fps: 40
num_frames: 240
x_shift_range: [0.00, 0.00, -0.015, -0.015]
y_shift_range: [0.00, 0.00, -0.015, -0.00]
z_shift_range: [-0.05, -0.05, -0.05, -0.05]
traj_types: ['double-straight-line', 'double-straight-line', 'circle', 'circle']
video_postfix: ['dolly-zoom-in', 'zoom-in', 'circle', 'swing']
specific: ''
longer_side_len: 960
src_folder: image
depth_folder: depth
mesh_folder: mesh
video_folder: video
load_ply: False
save_ply: True
inference_video: True
gpu_ids: 0
offscreen_rendering: False
img_format: '.jpg'
depth_format: '.npy'
require_midas: True
depth_threshold: 0.04
ext_edge_threshold: 0.002
sparse_iter: 5
filter_size: [7, 7, 5, 5, 5]
sigma_s: 4.0
sigma_r: 0.5
redundant_number: 12
background_thickness: 70
context_thickness: 140
background_thickness_2: 70
context_thickness_2: 70
discount_factor: 1.00
log_depth: True
largest_size: 512
depth_edge_dilate: 10
depth_edge_dilate_2: 5
extrapolate_border: True
extrapolation_thickness: 60
repeat_inpaint_edge: True
crop_border: [0.03, 0.03, 0.05, 0.03]
anti_flickering: True                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                 import numpy as np
from functools import reduce

def sparse_bilateral_filtering(
    depth, image, config, HR=False, mask=None, gsHR=True, edge_id=None, num_iter=None, num_gs_iter=None, spdb=False
):
    """
    config:
    - filter_size
    """
    import time

    save_images = []
    save_depths = []
    save_discontinuities = []
    vis_depth = depth.copy()
    backup_vis_depth = vis_depth.copy()

    depth_max = vis_depth.max()
    depth_min = vis_depth.min()
    vis_image = image.copy()
    for i in range(num_iter):
        if isinstance(config["filter_size"], list):
            window_size = config["filter_size"][i]
        else:
            window_size = config["filter_size"]
        vis_image = image.copy()
        save_images.append(vis_image)
        save_depths.append(vis_depth)
        u_over, b_over, l_over, r_over = vis_depth_discontinuity(vis_depth, config, mask=mask)
        vis_image[u_over > 0] = np.array([0, 0, 0])
        vis_image[b_over > 0] = np.array([0, 0, 0])
        vis_image[l_over > 0] = np.array([0, 0, 0])
        vis_image[r_over > 0] = np.array([0, 0, 0])

        discontinuity_map = (u_over + b_over + l_over + r_over).clip(0.0, 1.0)
        discontinuity_map[depth == 0] = 1
        save_discontinuities.append(discontinuity_map)
        if mask is not None:
            discontinuity_map[mask == 0] = 0
        vis_depth = bilateral_filter(
            vis_depth, config, discontinuity_map=discontinuity_map, HR=HR, mask=mask, window_size=window_size
        )

    return save_images, save_depths


def vis_depth_discontinuity(depth, config, vis_diff=False, label=False, mask=None):
    """
    config:
    - 
    """
    if label == False:
        disp = 1./depth
        u_diff = (disp[1:, :] - disp[:-1, :])[:-1, 1:-1]
        b_diff = (disp[:-1, :] - disp[1:, :])[1:, 1:-1]
        l_diff = (disp[:, 1:] - disp[:, :-1])[1:-1, :-1]
        r_diff = (disp[:, :-1] - disp[:, 1:])[1:-1, 1:]
        if mask is not None:
            u_mask = (mask[1:, :] * mask[:-1, :])[:-1, 1:-1]
            b_mask = (mask[:-1, :] * mask[1:, :])[1:, 1:-1]
            l_mask = (mask[:, 1:] * mask[:, :-1])[1:-1, :-1]
            r_mask = (mask[:, :-1] * mask[:, 1:])[1:-1, 1:]
            u_diff = u_diff * u_mask
            b_diff = b_diff * b_mask
            l_diff = l_diff * l_mask
            r_diff = r_diff * r_mask
        u_over = (np.abs(u_diff) > config['depth_threshold']).astype(np.float32)
        b_over = (np.abs(b_diff) > config['depth_threshold']).astype(np.float32)
        l_over = (np.abs(l_diff) > config['depth_threshold']).astype(np.float32)
        r_over = (np.abs(r_diff) > config['depth_threshold']).astype(np.float32)
    else:
        disp = depth
        u_diff = (disp[1:, :] * disp[:-1, :])[:-1, 1:-1]
        b_diff = (disp[:-1, :] * disp[1:, :])[1:, 1:-1]
        l_diff = (disp[:, 1:] * disp[:, :-1])[1:-1, :-1]
        r_diff = (disp[:, :-1] * disp[:, 1:])[1:-1, 1:]
        if mask is not None:
            u_mask = (mask[1:, :] * mask[:-1, :])[:-1, 1:-1]
            b_mask = (mask[:-1, :] * mask[1:, :])[1:, 1:-1]
            l_mask = (mask[:, 1:] * mask[:, :-1])[1:-1, :-1]
            r_mask = (mask[:, :-1] * mask[:, 1:])[1:-1, 1:]
            u_diff = u_diff * u_mask
            b_diff = b_diff * b_mask
            l_diff = l_diff * l_mask
            r_diff = r_diff * r_mask
        u_over = (np.abs(u_diff) > 0).astype(np.float32)
        b_over = (np.abs(b_diff) > 0).astype(np.float32)
        l_over = (np.abs(l_diff) > 0).astype(np.float32)
        r_over = (np.abs(r_diff) > 0).astype(np.float32)
    u_over = np.pad(u_over, 1, mode='constant')
    b_over = np.pad(b_over, 1, mode='constant')
    l_over = np.pad(l_over, 1, mode='constant')
    r_over = np.pad(r_over, 1, mode='constant')
    u_diff = np.pad(u_diff, 1, mode='constant')
    b_diff = np.pad(b_diff, 1, mode='constant')
    l_diff = np.pad(l_diff, 1, mode='constant')
    r_diff = np.pad(r_diff, 1, mode='constant')

    if vis_diff:
        return [u_over, b_over, l_over, r_over], [u_diff, b_diff, l_diff, r_diff]
    else:
        return [u_over, b_over, l_over, r_over]

def bilateral_filter(depth, config, discontinuity_map=None, HR=False, mask=None, window_size=False):
    sort_time = 0
    replace_time = 0
    filter_time = 0
    init_time = 0
    filtering_time = 0
    sigma_s = config['sigma_s']
    sigma_r = config['sigma_r']
    if window_size == False:
        window_size = config['filter_size']
    midpt = window_size//2
    ax = np.arange(-midpt, midpt+1.)
    xx, yy = np.meshgrid(ax, ax)
    if discontinuity_map is not None:
        spatial_term = np.exp(-(xx**2 + yy**2) / (2. * sigma_s**2))

    # padding
    depth = depth[1:-1, 1:-1]
    depth = np.pad(depth, ((1,1), (1,1)), 'edge')
    pad_depth = np.pad(depth, (midpt,midpt), 'edge')
    if discontinuity_map is not None:
        discontinuity_map = discontinuity_map[1:-1, 1:-1]
        discontinuity_map = np.pad(discontinuity_map, ((1,1), (1,1)), 'edge')
        pad_discontinuity_map = np.pad(discontinuity_map, (midpt,midpt), 'edge')
        pad_discontinuity_hole = 1 - pad_discontinuity_map
    # filtering
    output = depth.copy()
    pad_depth_patches = rolling_window(pad_depth, [window_size, window_size], [1,1])
    if discontinuity_map is not None:
        pad_discontinuity_patches = rolling_window(pad_discontinuity_map, [window_size, window_size], [1,1])
        pad_discontinuity_hole_patches = rolling_window(pad_discontinuity_hole, [window_size, window_size], [1,1])

    if mask is not None:
        pad_mask = np.pad(mask, (midpt,midpt), 'constant')
        pad_mask_patches = rolling_window(pad_mask, [window_size, window_size], [1,1])
    from itertools import product
    if discontinuity_map is not None:
        pH, pW = pad_depth_patches.shape[:2]
        for pi in range(pH):
            for pj in range(pW):
                if mask is not None and mask[pi, pj] == 0:
                    continue
                if discontinuity_map is not None:
                    if bool(pad_discontinuity_patches[pi, pj].any()) is False:
                        continue
                    discontinuity_patch = pad_discontinuity_patches[pi, pj]
                    discontinuity_holes = pad_discontinuity_hole_patches[pi, pj]
                depth_patch = pad_depth_patches[pi, pj]
                depth_order = depth_patch.ravel().argsort()
                patch_midpt = depth_patch[window_size//2, window_size//2]
                if discontinuity_map is not None:
                    coef = discontinuity_holes.astype(np.float32)
                    if mask is not None:
                        coef = coef * pad_mask_patches[pi, pj]
                else:
                    range_term = np.exp(-(depth_patch-patch_midpt)**2 / (2. * sigma_r**2))
                    coef = spatial_term * range_term
                if coef.max() == 0:
                    output[pi, pj] = patch_midpt
                    continue
                if discontinuity_map is not None and (coef.max() == 0):
                    output[pi, pj] = patch_midpt
                else:
                    coef = coef/(coef.sum())
                    coef_order = coef.ravel()[depth_order]
                    cum_coef = np.cumsum(coef_order)
                    ind = np.digitize(0.5, cum_coef)
                    output[pi, pj] = depth_patch.ravel()[depth_order][ind]
    else:
        pH, pW = pad_depth_patches.shape[:2]
        for pi in range(pH):
            for pj in range(pW):
                if discontinuity_map is not None:
                    if pad_discontinuity_patches[pi, pj][window_size//2, window_size//2] == 1:
                        continue
                    discontinuity_patch = pad_discontinuity_patches[pi, pj]
                    discontinuity_holes = (1. - discontinuity_patch)
                depth_patch = pad_depth_patches[pi, pj]
                depth_order = depth_patch.ravel().argsort()
                patch_midpt = depth_patch[window_size//2, window_size//2]
                range_term = np.exp(-(depth_patch-patch_midpt)**2 / (2. * sigma_r**2))
                if discontinuity_map is not None:
                    coef = spatial_term * range_term * discontinuity_holes
                else:
                    coef = spatial_term * range_term
                if coef.sum() == 0:
                    output[pi, pj] = patch_midpt
                    continue
                if discontinuity_map is not None and (coef.sum() == 0):
                    output[pi, pj] = patch_midpt
                else:
                    coef = coef/(coef.sum())
                    coef_order = coef.ravel()[depth_order]
                    cum_coef = np.cumsum(coef_order)
                    ind = np.digitize(0.5, cum_coef)
                    output[pi, pj] = depth_patch.ravel()[depth_order][ind]

    return output

def rolling_window(a, window, strides):
    assert len(a.shape)==len(window)==len(strides), "\'a\', \'window\', \'strides\' dimension mismatch"
    shape_fn = lambda i,w,s: (a.shape[i]-w)//s + 1
    shape = [shape_fn(i,w,s) for i,(w,s) in enumerate(zip(window, strides))] + list(window)
    def acc_shape(i):
        if i+1>=len(a.shape):
            return 1
        else:
            return reduce(lambda x,y:x*y, a.shape[i+1:])
    _strides = [acc_shape(i)*s*a.itemsize for i,s in enumerate(strides)] + list(a.strides)

    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=_strides)                                                                                                                                                                                                                                                                                                                                                                                import os
import cv2
import glob
import numpy as np
import imageio
from MiDaS.MiDaS_utils import write_depth

BOOST_BASE = 'BoostingMonocularDepth'

BOOST_INPUTS = 'inputs'
BOOST_OUTPUTS = 'outputs'

def run_boostmonodepth(img_names, src_folder, depth_folder):

    if not isinstance(img_names, list):
        img_names = [img_names]

    # remove irrelevant files first
    clean_folder(os.path.join(BOOST_BASE, BOOST_INPUTS))
    clean_folder(os.path.join(BOOST_BASE, BOOST_OUTPUTS))

    tgt_names = []
    for img_name in img_names:
        base_name = os.path.basename(img_name)
        tgt_name = os.path.join(BOOST_BASE, BOOST_INPUTS, base_name)
        os.system(f'cp {img_name} {tgt_name}')

        # keep only the file name here.
        # they save all depth as .png file
        tgt_names.append(os.path.basename(tgt_name).replace('.jpg', '.png'))

    os.system(f'cd {BOOST_BASE} && python run.py --Final --data_dir {BOOST_INPUTS}/  --output_dir {BOOST_OUTPUTS} --depthNet 0')

    for i, (img_name, tgt_name) in enumerate(zip(img_names, tgt_names)):
        img = imageio.imread(img_name)
        H, W = img.shape[:2]
        scale = 640. / max(H, W)

        # resize and save depth
        target_height, target_width = int(round(H * scale)), int(round(W * scale))
        depth = imageio.imread(os.path.join(BOOST_BASE, BOOST_OUTPUTS, tgt_name))
        depth = np.array(depth).astype(np.float32)
        depth = resize_depth(depth, target_width, target_height)
        np.save(os.path.join(depth_folder, tgt_name.replace('.png', '.npy')), depth / 32768. - 1.)
        write_depth(os.path.join(depth_folder, tgt_name.replace('.png', '')), depth)

def clean_folder(folder, img_exts=['.png', '.jpg', '.npy']):

    for img_ext in img_exts:
        paths_to_check = os.path.join(folder, f'*{img_ext}')
        if len(glob.glob(paths_to_check)) == 0:
            continue
        print(paths_to_check)
        os.system(f'rm {paths_to_check}')

def resize_depth(depth, width, height):
    """Resize numpy (or image read by imageio) depth map

    Args:
        depth (numpy): depth
        width (int): image width
        height (int): image height

    Returns:
        array: processed depth
    """
    depth = cv2.blur(depth, (3, 3))
    return cv2.resize(depth, (width, height), interpolation=cv2.INTER_AREA)                                                                                                                                                                                                                                                                                                                                                                              #!/bin/sh
fb_status=$(wget --spider -S https://filebox.ece.vt.edu/ 2>&1 | grep  "HTTP/1.1 200 OK")

mkdir checkpoints

echo "downloading from filebox ..."
wget https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/color-model.pth
wget https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/depth-model.pth
wget https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/edge-model.pth
wget https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/model.pt

mv color-model.pth checkpoints/.
mv depth-model.pth checkpoints/.
mv edge-model.pth checkpoints/.
mv model.pt MiDaS/.

echo "cloning from BoostingMonocularDepth ..."
git clone https://github.com/compphoto/BoostingMonocularDepth.git
mkdir -p BoostingMonocularDepth/pix2pix/checkpoints/mergemodel/

echo "downloading mergenet weights ..."
wget https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/latest_net_G.pth
mv latest_net_G.pth BoostingMonocularDepth/pix2pix/checkpoints/mergemodel/
wget https://github.com/intel-isl/MiDaS/releases/download/v2/model-f46da743.pt
mv model-f46da743.pt BoostingMonocularDepth/midas/model.pt                                                                                                                                                                                                                                                                                                                                                                                             import numpy as np
import argparse
import glob
import os
from functools import partial
import vispy
import scipy.misc as misc
from tqdm import tqdm
import yaml
import time
import sys
from mesh import write_ply, read_ply, output_3d_photo
from utils import get_MiDaS_samples, read_MiDaS_depth
import torch
import cv2
from skimage.transform import resize
import imageio
import copy
from networks import Inpaint_Color_Net, Inpaint_Depth_Net, Inpaint_Edge_Net
from MiDaS.run import run_depth
from boostmonodepth_utils import run_boostmonodepth
from MiDaS.monodepth_net import MonoDepthNet
import MiDaS.MiDaS_utils as MiDaS_utils
from bilateral_filtering import sparse_bilateral_filtering

parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, default='argument.yml',help='Configure of post processing')
args = parser.parse_args()
config = yaml.load(open(args.config, 'r'))
if config['offscreen_rendering'] is True:
    vispy.use(app='egl')
os.makedirs(config['mesh_folder'], exist_ok=True)
os.makedirs(config['video_folder'], exist_ok=True)
os.makedirs(config['depth_folder'], exist_ok=True)
sample_list = get_MiDaS_samples(config['src_folder'], config['depth_folder'], config, config['specific'])
normal_canvas, all_canvas = None, None

if isinstance(config["gpu_ids"], int) and (config["gpu_ids"] >= 0):
    device = config["gpu_ids"]
else:
    device = "cpu"

print(f"running on device {device}")

for idx in tqdm(range(len(sample_list))):
    depth = None
    sample = sample_list[idx]
    print("Current Source ==> ", sample['src_pair_name'])
    mesh_fi = os.path.join(config['mesh_folder'], sample['src_pair_name'] +'.ply')
    image = imageio.imread(sample['ref_img_fi'])

    print(f"Running depth extraction at {time.time()}")
    if config['use_boostmonodepth'] is True:
        run_boostmonodepth(sample['ref_img_fi'], config['src_folder'], config['depth_folder'])
    elif config['require_midas'] is True:
        run_depth([sample['ref_img_fi']], config['src_folder'], config['depth_folder'],
                  config['MiDaS_model_ckpt'], MonoDepthNet, MiDaS_utils, target_w=640)

    if 'npy' in config['depth_format']:
        config['output_h'], config['output_w'] = np.load(sample['depth_fi']).shape[:2]
    else:
        config['output_h'], config['output_w'] = imageio.imread(sample['depth_fi']).shape[:2]
    frac = config['longer_side_len'] / max(config['output_h'], config['output_w'])
    config['output_h'], config['output_w'] = int(config['output_h'] * frac), int(config['output_w'] * frac)
    config['original_h'], config['original_w'] = config['output_h'], config['output_w']
    if image.ndim == 2:
        image = image[..., None].repeat(3, -1)
    if np.sum(np.abs(image[..., 0] - image[..., 1])) == 0 and np.sum(np.abs(image[..., 1] - image[..., 2])) == 0:
        config['gray_image'] = True
    else:
        config['gray_image'] = False
    image = cv2.resize(image, (config['output_w'], config['output_h']), interpolation=cv2.INTER_AREA)
    depth = read_MiDaS_depth(sample['depth_fi'], 3.0, config['output_h'], config['output_w'])
    mean_loc_depth = depth[depth.shape[0]//2, depth.shape[1]//2]
    if not(config['load_ply'] is True and os.path.exists(mesh_fi)):
        vis_photos, vis_depths = sparse_bilateral_filtering(depth.copy(), image.copy(), config, num_iter=config['sparse_iter'], spdb=False)
        depth = vis_depths[-1]
        model = None
        torch.cuda.empty_cache()
        print("Start Running 3D_Photo ...")
        print(f"Loading edge model at {time.time()}")
        depth_edge_model = Inpaint_Edge_Net(init_weights=True)
        depth_edge_weight = torch.load(config['depth_edge_model_ckpt'],
                                       map_location=torch.device(device))
        depth_edge_model.load_state_dict(depth_edge_weight)
        depth_edge_model = depth_edge_model.to(device)
        depth_edge_model.eval()

        print(f"Loading depth model at {time.time()}")
        depth_feat_model = Inpaint_Depth_Net()
        depth_feat_weight = torch.load(config['depth_feat_model_ckpt'],
                                       map_location=torch.device(device))
        depth_feat_model.load_state_dict(depth_feat_weight, strict=True)
        depth_feat_model = depth_feat_model.to(device)
        depth_feat_model.eval()
        depth_feat_model = depth_feat_model.to(device)
        print(f"Loading rgb model at {time.time()}")
        rgb_model = Inpaint_Color_Net()
        rgb_feat_weight = torch.load(config['rgb_feat_model_ckpt'],
                                     map_location=torch.device(device))
        rgb_model.load_state_dict(rgb_feat_weight)
        rgb_model.eval()
        rgb_model = rgb_model.to(device)
        graph = None


        print(f"Writing depth ply (and basically doing everything) at {time.time()}")
        rt_info = write_ply(image,
                              depth,
                              sample['int_mtx'],
                              mesh_fi,
                              config,
                              rgb_model,
                              depth_edge_model,
                              depth_edge_model,
                              depth_feat_model)

        if rt_info is False:
            continue
        rgb_model = None
        color_feat_model = None
        depth_edge_model = None
        depth_feat_model = None
        torch.cuda.empty_cache()
    if config['save_ply'] is True or config['load_ply'] is True:
        verts, colors, faces, Height, Width, hFov, vFov = read_ply(mesh_fi)
    else:
        verts, colors, faces, Height, Width, hFov, vFov = rt_info


    print(f"Making video at {time.time()}")
    videos_poses, video_basename = copy.deepcopy(sample['tgts_poses']), sample['tgt_name']
    top = (config.get('original_h') // 2 - sample['int_mtx'][1, 2] * config['output_h'])
    left = (config.get('original_w') // 2 - sample['int_mtx'][0, 2] * config['output_w'])
    down, right = top + config['output_h'], left + config['output_w']
    border = [int(xx) for xx in [top, down, left, right]]
    normal_canvas, all_canvas = output_3d_photo(verts.copy(), colors.copy(), faces.copy(), copy.deepcopy(Height), copy.deepcopy(Width), copy.deepcopy(hFov), copy.deepcopy(vFov),
                        copy.deepcopy(sample['tgt_pose']), sample['video_postfix'], copy.deepcopy(sample['ref_pose']), copy.deepcopy(config['video_folder']),
                        image.copy(), copy.deepcopy(sample['int_mtx']), config, image,
                        videos_poses, video_basename, config.get('original_h'), config.get('original_w'), border=border, depth=depth, normal_canvas=normal_canvas, all_canvas=all_canvas,
                        mean_loc_depth=mean_loc_depth)                                                                                                                                                                                                                                                                                                                                                                                                                                      import os
import numpy as np
try:
    import cynetworkx as netx
except ImportError:
    import networkx as netx
import matplotlib.pyplot as plt
from functools import partial
from vispy import scene, io
from vispy.scene import visuals
from vispy.visuals.filters import Alpha
import cv2
from moviepy.editor import ImageSequenceClip
from skimage.transform import resize
import time
import copy
import torch
import os
from utils import path_planning, open_small_mask, clean_far_edge, refine_depth_around_edge
from utils import refine_color_around_edge, filter_irrelevant_edge_new, require_depth_edge, clean_far_edge_new
from utils import create_placeholder, refresh_node, find_largest_rect
from mesh_tools import get_depth_from_maps, get_map_from_ccs, get_edge_from_nodes, get_depth_from_nodes, get_rgb_from_nodes, crop_maps_by_size, convert2tensor, recursive_add_edge, update_info, filter_edge, relabel_node, depth_inpainting
from mesh_tools import refresh_bord_depth, enlarge_border, fill_dummy_bord, extrapolate, fill_missing_node, incomplete_node, get_valid_size, dilate_valid_size, size_operation
import transforms3d
import random
from functools import reduce

def create_mesh(depth, image, int_mtx, config):
    H, W, C = image.shape
    ext_H, ext_W = H + 2 * config['extrapolation_thickness'], W + 2 * config['extrapolation_thickness']
    LDI = netx.Graph(H=ext_H, W=ext_W, noext_H=H, noext_W=W, cam_param=int_mtx)
    xy2depth = {}
    int_mtx_pix = int_mtx * np.array([[W], [H], [1.]])
    LDI.graph['cam_param_pix'], LDI.graph['cam_param_pix_inv'] = int_mtx_pix, np.linalg.inv(int_mtx_pix)
    disp = 1. / (-depth)
    LDI.graph['hoffset'], LDI.graph['woffset'] = config['extrapolation_thickness'], config['extrapolation_thickness']
    LDI.graph['bord_up'], LDI.graph['bord_down'] = LDI.graph['hoffset'] + 0, LDI.graph['hoffset'] + H
    LDI.graph['bord_left'], LDI.graph['bord_right'] = LDI.graph['woffset'] + 0, LDI.graph['woffset'] + W
    for idx in range(H):
        for idy in range(W):
            x, y = idx + LDI.graph['hoffset'], idy + LDI.graph['woffset']
            LDI.add_node((x, y, -depth[idx, idy]),
                         color=image[idx, idy],
                         disp=disp[idx, idy],
                         synthesis=False,
                         cc_id=set())
            xy2depth[(x, y)] = [-depth[idx, idy]]
    for x, y, d in LDI.nodes:
        two_nes = [ne for ne in [(x+1, y), (x, y+1)] if ne[0] < LDI.graph['bord_down'] and ne[1] < LDI.graph['bord_right']]
        [LDI.add_edge((ne[0], ne[1], xy2depth[ne][0]), (x, y, d)) for ne in two_nes]
    LDI = calculate_fov(LDI)
    image = np.pad(image,
                    pad_width=((config['extrapolation_thickness'], config['extrapolation_thickness']),
                               (config['extrapolation_thickness'], config['extrapolation_thickness']),
                               (0, 0)),
                    mode='constant')
    depth = np.pad(depth,
                    pad_width=((config['extrapolation_thickness'], config['extrapolation_thickness']),
                               (config['extrapolation_thickness'], config['extrapolation_thickness'])),
                    mode='constant')

    return LDI, xy2depth, image, depth


def tear_edges(mesh, threshold = 0.00025, xy2depth=None):
    remove_edge_list = []
    remove_horizon, remove_vertical = np.zeros((2, mesh.graph['H'], mesh.graph['W']))
    mesh_nodes = mesh.nodes
    for edge in mesh.edges:
        if abs(mesh_nodes[edge[0]]['disp'] - mesh_nodes[edge[1]]['disp']) > threshold:
            remove_edge_list.append((edge[0], edge[1]))

            near, far = edge if abs(edge[0][2]) < abs(edge[1][2]) else edge[::-1]

            mesh_nodes[far]['near'] = [] if mesh_nodes[far].get('near') is None else mesh_nodes[far]['near'].append(near)
            mesh_nodes[near]['far'] = [] if mesh_nodes[near].get('far') is None else mesh_nodes[near]['far'].append(far)

            if near[0] == far[0]:
                remove_horizon[near[0], np.minimum(near[1], far[1])] = 1
            elif near[1] == far[1]:
                remove_vertical[np.minimum(near[0], far[0]), near[1]] = 1
    mesh.remove_edges_from(remove_edge_list)

    remove_edge_list = []

    dang_horizon = np.where(np.roll(remove_horizon, 1, 0) + np.roll(remove_horizon, -1, 0) - remove_horizon == 2)
    dang_vertical = np.where(np.roll(remove_vertical, 1, 1) + np.roll(remove_vertical, -1, 1) - remove_vertical == 2)

    horizon_condition = lambda x, y: mesh.graph['bord_up'] + 1 <= x < mesh.graph['bord_down'] - 1
    vertical_condition = lambda x, y: mesh.graph['bord_left'] + 1 <= y < mesh.graph['bord_right'] - 1

    prjto3d = lambda x, y: (x, y, xy2depth[(x, y)][0])

    node_existence = lambda x, y: mesh.has_node(prjto3d(x, y))

    for x, y in zip(dang_horizon[0], dang_horizon[1]):
        if horizon_condition(x, y) and node_existence(x, y) and node_existence(x, y+1):
            remove_edge_list.append((prjto3d(x, y), prjto3d(x, y+1)))
    for x, y in zip(dang_vertical[0], dang_vertical[1]):
        if vertical_condition(x, y) and node_existence(x, y) and node_existence(x+1, y):
            remove_edge_list.append((prjto3d(x, y), prjto3d(x+1, y)))
    mesh.remove_edges_from(remove_edge_list)

    return mesh

def calculate_fov(mesh):
    k = mesh.graph['cam_param']
    mesh.graph['hFov'] = 2 * np.arctan(1. / (2*k[0, 0]))
    mesh.graph['vFov'] = 2 * np.arctan(1. / (2*k[1, 1]))
    mesh.graph['aspect'] = mesh.graph['noext_H'] / mesh.graph['noext_W']

    return mesh

def calculate_fov_FB(mesh):
    mesh.graph['aspect'] = mesh.graph['H'] / mesh.graph['W']
    if mesh.graph['H'] > mesh.graph['W']:
        mesh.graph['hFov'] = 0.508015513
        half_short = np.tan(mesh.graph['hFov']/2.0)
        half_long = half_short * mesh.graph['aspect']
        mesh.graph['vFov'] = 2.0 * np.arctan(half_long)
    else:
        mesh.graph['vFov'] = 0.508015513
        half_short = np.tan(mesh.graph['vFov']/2.0)
        half_long = half_short / mesh.graph['aspect']
        mesh.graph['hFov'] = 2.0 * np.arctan(half_long)

    return mesh

def reproject_3d_int_detail(sx, sy, z, k_00, k_02, k_11, k_12, w_offset, h_offset):
    abs_z = abs(z)
    return [abs_z * ((sy+0.5-w_offset) * k_00 + k_02), abs_z * ((sx+0.5-h_offset) * k_11 + k_12), abs_z]

def reproject_3d_int_detail_FB(sx, sy, z, w_offset, h_offset, mesh):
    if mesh.graph.get('tan_hFov') is None:
        mesh.graph['tan_hFov'] = np.tan(mesh.graph['hFov'] / 2.)
    if mesh.graph.get('tan_vFov') is None:
        mesh.graph['tan_vFov'] = np.tan(mesh.graph['vFov'] / 2.)

    ray = np.array([(-1. + 2. * ((sy+0.5-w_offset)/(mesh.graph['W'] - 1))) * mesh.graph['tan_hFov'],
                    (1. - 2. * (sx+0.5-h_offset)/(mesh.graph['H'] - 1)) * mesh.graph['tan_vFov'],
                    -1])
    point_3d = ray * np.abs(z)

    return point_3d


def reproject_3d_int(sx, sy, z, mesh):
    k = mesh.graph['cam_param_pix_inv'].copy()
    if k[0, 2] > 0:
        k = np.linalg.inv(k)
    ray = np.dot(k, np.array([sy-mesh.graph['woffset'], sx-mesh.graph['hoffset'], 1]).reshape(3, 1))

    point_3d = ray * np.abs(z)
    point_3d = point_3d.flatten()

    return point_3d

def generate_init_node(mesh, config, min_node_in_cc):
    mesh_nodes = mesh.nodes

    info_on_pix = {}

    ccs = sorted(netx.connected_components(mesh), key = len, reverse=True)
    remove_nodes = []

    for cc in ccs:

        remove_flag = True if len(cc) < min_node_in_cc else False
        if remove_flag is False:
            for (nx, ny, nd) in cc:
                info_on_pix[(nx, ny)] = [{'depth':nd,
                                          'color':mesh_nodes[(nx, ny, nd)]['color'],
                                          'synthesis':False,
                                          'disp':mesh_nodes[(nx, ny, nd)]['disp']}]
        else:
            [remove_nodes.append((nx, ny, nd)) for (nx, ny, nd) in cc]

    for node in remove_nodes:
        far_nodes = [] if mesh_nodes[node].get('far') is None else mesh_nodes[node]['far']
        for far_node in far_nodes:
            if mesh.has_node(far_node) and mesh_nodes[far_node].get('near') is not None and node in mesh_nodes[far_node]['near']:
                mesh_nodes[far_node]['near'].remove(node)
        near_nodes = [] if mesh_nodes[node].get('near') is None else mesh_nodes[node]['near']
        for near_node in near_nodes:
            if mesh.has_node(near_node) and mesh_nodes[near_node].get('far') is not None and node in mesh_nodes[near_node]['far']:
                mesh_nodes[near_node]['far'].remove(node)

    [mesh.remove_node(node) for node in remove_nodes]

    return mesh, info_on_pix

def get_neighbors(mesh, node):
    return [*mesh.neighbors(node)]

def generate_face(mesh, info_on_pix, config):
    H, W = mesh.graph['H'], mesh.graph['W']
    str_faces = []
    num_node = len(mesh.nodes)
    ply_flag = config.get('save_ply')
    def out_fmt(input, cur_id_b, cur_id_self, cur_id_a, ply_flag):
        if ply_flag is True:
            input.append(' '.join(['3', cur_id_b, cur_id_self, cur_id_a]) + '\n')
        else:
            input.append([cur_id_b, cur_id_self, cur_id_a])
    mesh_nodes = mesh.nodes
    for node in mesh_nodes:
        cur_id_self = mesh_nodes[node]['cur_id']
        ne_nodes = get_neighbors(mesh, node)
        four_dir_nes = {'up': [], 'left': [],
                        'down': [], 'right': []}
        for ne_node in ne_nodes:
            store_tuple = [ne_node, mesh_nodes[ne_node]['cur_id']]
            if ne_node[0] == node[0]:
                if ne_node[1] == ne_node[1] - 1:
                    four_dir_nes['left'].append(store_tuple)
                else:
                    four_dir_nes['right'].append(store_tuple)
            else:
                if ne_node[0] == ne_node[0] - 1:
                    four_dir_nes['up'].append(store_tuple)
                else:
                    four_dir_nes['down'].append(store_tuple)
        for node_a, cur_id_a in four_dir_nes['up']:
            for node_b, cur_id_b in four_dir_nes['right']:
                out_fmt(str_faces, cur_id_b, cur_id_self, cur_id_a, ply_flag)
        for node_a, cur_id_a in four_dir_nes['right']:
            for node_b, cur_id_b in four_dir_nes['down']:
                out_fmt(str_faces, cur_id_b, cur_id_self, cur_id_a, ply_flag)
        for node_a, cur_id_a in four_dir_nes['down']:
            for node_b, cur_id_b in four_dir_nes['left']:
                out_fmt(str_faces, cur_id_b, cur_id_self, cur_id_a, ply_flag)
        for node_a, cur_id_a in four_dir_nes['left']:
            for node_b, cur_id_b in four_dir_nes['up']:
                out_fmt(str_faces, cur_id_b, cur_id_self, cur_id_a, ply_flag)

    return str_faces

def reassign_floating_island(mesh, info_on_pix, image, depth):
    H, W = mesh.graph['H'], mesh.graph['W'],
    mesh_nodes = mesh.nodes
    bord_up, bord_down = mesh.graph['bord_up'], mesh.graph['bord_down']
    bord_left, bord_right = mesh.graph['bord_left'], mesh.graph['bord_right']
    W = mesh.graph['W']
    lost_map = np.zeros((H, W))

    '''
    (5) is_inside(x, y, xmin, xmax, ymin, ymax) : Check if a pixel(x, y) is inside the border.
    (6) get_cross_nes(x, y) : Get the four cross neighbors of pixel(x, y).
    '''
    key_exist = lambda d, k: k in d
    is_inside = lambda x, y, xmin, xmax, ymin, ymax: xmin <= x < xmax and ymin <= y < ymax
    get_cross_nes = lambda x, y: [(x + 1, y), (x - 1, y), (x, y - 1), (x, y + 1)]
    '''
    (A) Highlight the pixels on isolated floating island.
    (B) Number those isolated floating islands with connected component analysis.
    (C) For each isolated island:
        (1) Find its longest surrounded depth edge.
        (2) Propogate depth from that depth edge to the pixels on the isolated island.
        (3) Build the connection between the depth edge and that isolated island.
    '''
    for x in range(H):
        for y in range(W):
            if is_inside(x, y, bord_up, bord_down, bord_left, bord_right) and not(key_exist(info_on_pix, (x, y))):
                lost_map[x, y] = 1
    _, label_lost_map = cv2.connectedComponents(lost_map.astype(np.uint8), connectivity=4)
    mask = np.zeros((H, W))
    mask[bord_up:bord_down, bord_left:bord_right] = 1
    label_lost_map = (label_lost_map * mask).astype(np.int)

    for i in range(1, label_lost_map.max()+1):
        lost_xs, lost_ys = np.where(label_lost_map == i)
        surr_edge_ids = {}
        for lost_x, lost_y in zip(lost_xs, lost_ys):
            if (lost_x, lost_y) == (295, 389) or (lost_x, lost_y) == (296, 389):
                import pdb; pdb.set_trace()
            for ne in get_cross_nes(lost_x, lost_y):
                if key_exist(info_on_pix, ne):
                    for info in info_on_pix[ne]:
                        ne_node = (ne[0], ne[1], info['depth'])
                        if key_exist(mesh_nodes[ne_node], 'edge_id'):
                            edge_id = mesh_nodes[ne_node]['edge_id']
                            surr_edge_ids[edge_id] = surr_edge_ids[edge_id] + [ne_node] if \
                                                key_exist(surr_edge_ids, edge_id) else [ne_node]
        if len(surr_edge_ids) == 0:
            continue
        edge_id, edge_nodes = sorted([*surr_edge_ids.items()], key=lambda x: len(x[1]), reverse=True)[0]
        edge_depth_map = np.zeros((H, W))
        for node in edge_nodes:
            edge_depth_map[node[0], node[1]] = node[2]
        lost_xs, lost_ys = np.where(label_lost_map == i)
        while lost_xs.shape[0] > 0:
            lost_xs, lost_ys = np.where(label_lost_map == i)
            for lost_x, lost_y in zip(lost_xs, lost_ys):
                propagated_depth = []
                real_nes = []
                for ne in get_cross_nes(lost_x, lost_y):
                    if not(is_inside(ne[0], ne[1], bord_up, bord_down, bord_left, bord_right)) or \
                       edge_depth_map[ne[0], ne[1]] == 0:
                        continue
                    propagated_depth.append(edge_depth_map[ne[0], ne[1]])
                    real_nes.append(ne)
                if len(real_nes) == 0:
                    continue
                reassign_depth = np.mean(propagated_depth)
                label_lost_map[lost_x, lost_y] = 0
                edge_depth_map[lost_x, lost_y] = reassign_depth
                depth[lost_x, lost_y] = -reassign_depth
                mesh.add_node((lost_x, lost_y, reassign_depth), color=image[lost_x, lost_y],
                                                            synthesis=False,
                                                            disp=1./reassign_depth,
                                                            cc_id=set())
                info_on_pix[(lost_x, lost_y)] = [{'depth':reassign_depth,
                                                  'color':image[lost_x, lost_y],
                                                  'synthesis':False,
                                                  'disp':1./reassign_depth}]
                new_connections = [((lost_x, lost_y, reassign_depth),
                                    (ne[0], ne[1], edge_depth_map[ne[0], ne[1]])) for ne in real_nes]
                mesh.add_edges_from(new_connections)

    return mesh, info_on_pix, depth

def remove_node_feat(mesh, *feats):
    mesh_nodes = mesh.nodes
    for node in mesh_nodes:
        for feat in feats:
            mesh_nodes[node][feat] = None

    return mesh

def update_status(mesh, info_on_pix, depth=None):
    '''
    (2) clear_node_feat(G, *fts) : Clear all the node feature on graph G.
    (6) get_cross_nes(x, y) : Get the four cross neighbors of pixel(x, y).
    '''
    key_exist = lambda d, k: d.get(k) is not None
    is_inside = lambda x, y, xmin, xmax, ymin, ymax: xmin <= x < xmax and ymin <= y < ymax
    get_cross_nes = lambda x, y: [(x + 1, y), (x - 1, y), (x, y - 1), (x, y + 1)]
    append_element = lambda d, k, x: d[k] + [x] if key_exist(d, k) else [x]

    def clear_node_feat(G, fts):
        le_nodes = G.nodes
        for k in le_nodes:
            v = le_nodes[k]
            for ft in fts:
                if ft in v:
                    v[ft] = None

    clear_node_feat(mesh, ['edge_id', 'far', 'near'])
    bord_up, bord_down = mesh.graph['bord_up'], mesh.graph['bord_down']
    bord_left, bord_right = mesh.graph['bord_left'], mesh.graph['bord_right']

    le_nodes = mesh.nodes

    for node_key in le_nodes:
        if mesh.neighbors(node_key).__length_hint__() == 4:
            continue
        four_nes = [xx for xx in get_cross_nes(node_key[0], node_key[1]) if
                    is_inside(xx[0], xx[1], bord_up, bord_down, bord_left, bord_right) and
                    xx in info_on_pix]
        [four_nes.remove((ne_node[0], ne_node[1])) for ne_node in mesh.neighbors(node_key)]
        for ne in four_nes:
            for info in info_on_pix[ne]:
                assert mesh.has_node((ne[0], ne[1], info['depth'])), "No node_key"
                ind_node = le_nodes[node_key]
                if abs(node_key[2]) > abs(info['depth']):
                    ind_node['near'] = append_element(ind_node, 'near', (ne[0], ne[1], info['depth']))
                else:
                    ind_node['far'] = append_element(ind_node, 'far', (ne[0], ne[1], info['depth']))
    if depth is not None:
        for key, value in info_on_pix.items():
            if depth[key[0], key[1]] != abs(value[0]['depth']):
                value[0]['disp'] = 1. / value[0]['depth']
                depth[key[0], key[1]] = abs(value[0]['depth'])

        return mesh, depth, info_on_pix
    else:
        return mesh

def group_edges(LDI, config, image, remove_conflict_ordinal, spdb=False):

    '''
    (1) add_new_node(G, node) : add "node" to graph "G"
    (2) add_new_edge(G, node_a, node_b) : add edge "node_a--node_b" to graph "G"
    (3) exceed_thre(x, y, thre) : Check if difference between "x" and "y" exceed threshold "thre"
    (4) key_exist(d, k) : Check if key "k' exists in dictionary "d"
    (5) comm_opp_bg(G, x, y) : Check if node "x" and "y" in graph "G" treat the same opposite node as background
    (6) comm_opp_fg(G, x, y) : Check if node "x" and "y" in graph "G" treat the same opposite node as foreground
    '''
    add_new_node = lambda G, node: None if G.has_node(node) else G.add_node(node)
    add_new_edge = lambda G, node_a, node_b: None if G.has_edge(node_a, node_b) else G.add_edge(node_a, node_b)
    exceed_thre = lambda x, y, thre: (abs(x) - abs(y)) > thre
    key_exist = lambda d, k: d.get(k) is not None
    comm_opp_bg = lambda G, x, y: key_exist(G.nodes[x], 'far') and key_exist(G.nodes[y], 'far') and \
                                    not(set(G.nodes[x]['far']).isdisjoint(set(G.nodes[y]['far'])))
    comm_opp_fg = lambda G, x, y: key_exist(G.nodes[x], 'near') and key_exist(G.nodes[y], 'near') and \
                                    not(set(G.nodes[x]['near']).isdisjoint(set(G.nodes[y]['near'])))
    discont_graph = netx.Graph()
    '''
    (A) Skip the pixel at image boundary, we don't want to deal with them.
    (B) Identify discontinuity by the number of its neighbor(degree).
        If the degree < 4(up/right/buttom/left). We will go through following steps:
        (1) Add the discontinuity pixel "node" to graph "discont_graph".
        (2) Find "node"'s cross neighbor(up/right/buttom/left) "ne_node".
            - If the cross neighbor "ne_node" is a discontinuity pixel(degree("ne_node") < 4),
                (a) add it to graph "discont_graph" and build the connection between "ne_node" and "node".
                (b) label its cross neighbor as invalid pixels "inval_diag_candi" to avoid building
                    connection between original discontinuity pixel "node" and "inval_diag_candi".
            - Otherwise, find "ne_node"'s cross neighbors, called diagonal candidate "diag_candi".
                - The "diag_candi" is diagonal to the original discontinuity pixel "node".
                - If "diag_candi" exists, go to step(3).
        (3) A diagonal candidate "diag_candi" will be :
            - added to the "discont_graph" if its degree < 4.
            - connected to the original discontinuity pixel "node" if it satisfied either
                one of following criterion:
                (a) the difference of disparity between "diag_candi" and "node" is smaller than default threshold.
                (b) the "diag_candi" and "node" face the same opposite pixel. (See. function "tear_edges")
                (c) Both of "diag_candi" and "node" must_connect to each other. (See. function "combine_end_node")
    (C) Aggregate each connected part in "discont_graph" into "discont_ccs" (A.K.A. depth edge).
    '''
    for node in LDI.nodes:
        if not(LDI.graph['bord_up'] + 1 <= node[0] <= LDI.graph['bord_down'] - 2 and \
               LDI.graph['bord_left'] + 1 <= node[1] <= LDI.graph['bord_right'] - 2):
            continue
        neighbors = [*LDI.neighbors(node)]
        if len(neighbors) < 4:
            add_new_node(discont_graph, node)
            diag_candi_anc, inval_diag_candi, discont_nes = set(), set(), set()
            for ne_node in neighbors:
                if len([*LDI.neighbors(ne_node)]) < 4:
                    add_new_node(discont_graph, ne_node)
                    add_new_edge(discont_graph, ne_node, node)
                    discont_nes.add(ne_node)
                else:
                    diag_candi_anc.add(ne_node)
            inval_diag_candi = set([inval_diagonal for ne_node in discont_nes for inval_diagonal in LDI.neighbors(ne_node) if \
                                     abs(inval_diagonal[0] - node[0]) < 2 and abs(inval_diagonal[1] - node[1]) < 2])
            for ne_node in diag_candi_anc:
                if ne_node[0] == node[0]:
                    diagonal_xys = [[ne_node[0] + 1, ne_node[1]], [ne_node[0] - 1, ne_node[1]]]
                elif ne_node[1] == node[1]:
                    diagonal_xys = [[ne_node[0], ne_node[1] + 1], [ne_node[0], ne_node[1] - 1]]
                for diag_candi in LDI.neighbors(ne_node):
                    if [diag_candi[0], diag_candi[1]] in diagonal_xys and LDI.degree(diag_candi) < 4:
                        if diag_candi not in inval_diag_candi:
                            if not exceed_thre(1./node[2], 1./diag_candi[2], config['depth_threshold']) or \
                               (comm_opp_bg(LDI, diag_candi, node) and comm_opp_fg(LDI, diag_candi, node)):
                                add_new_node(discont_graph, diag_candi)
                                add_new_edge(discont_graph, diag_candi, node)
                        if key_exist(LDI.nodes[diag_candi], 'must_connect') and node in LDI.nodes[diag_candi]['must_connect'] and \
                            key_exist(LDI.nodes[node], 'must_connect') and diag_candi in LDI.nodes[node]['must_connect']:
                            add_new_node(discont_graph, diag_candi)
                            add_new_edge(discont_graph, diag_candi, node)
    if spdb == True:
        import pdb; pdb.set_trace()
    discont_ccs = [*netx.connected_components(discont_graph)]
    '''
    In some corner case, a depth edge "discont_cc" will contain both
    foreground(FG) and background(BG) pixels. This violate the assumption that
    a depth edge can only composite by one type of pixel(FG or BG).
    We need to further divide this depth edge into several sub-part so that the
    assumption is satisfied.
    (A) A depth edge is invalid if both of its "far_flag"(BG) and
        "near_flag"(FG) are True.
    (B) If the depth edge is invalid, we need to do:
        (1) Find the role("oridinal") of each pixel on the depth edge.
            "-1" --> Its opposite pixels has smaller depth(near) than it.
                     It is a backgorund pixel.
            "+1" --> Its opposite pixels has larger depth(far) than it.
                     It is a foregorund pixel.
            "0"  --> Some of opposite pixels has larger depth(far) than it,
                     and some has smaller pixel than it.
                     It is an ambiguous pixel.
        (2) For each pixel "discont_node", check if its neigbhors' roles are consistent.
            - If not, break the connection between the neighbor "ne_node" that has a role
              different from "discont_node".
            - If yes, remove all the role that are inconsistent to its neighbors "ne_node".
        (3) Connected component analysis to re-identified those divided depth edge.
    (C) Aggregate each connected part in "discont_graph" into "discont_ccs" (A.K.A. depth edge).
    '''
    if remove_conflict_ordinal:
        new_discont_ccs = []
        num_new_cc = 0
        for edge_id, discont_cc in enumerate(discont_ccs):
            near_flag = False
            far_flag = False
            for discont_node in discont_cc:
                near_flag = True if key_exist(LDI.nodes[discont_node], 'far') else near_flag
                far_flag = True if key_exist(LDI.nodes[discont_node], 'near') else far_flag
                if far_flag and near_flag:
                    break
            if far_flag and near_flag:
                for discont_node in discont_cc:
                    discont_graph.nodes[discont_node]['ordinal'] = \
                        np.array([key_exist(LDI.nodes[discont_node], 'far'),
                                  key_exist(LDI.nodes[discont_node], 'near')]) * \
                        np.array([-1, 1])
                    discont_graph.nodes[discont_node]['ordinal'] = \
                        np.sum(discont_graph.nodes[discont_node]['ordinal'])
                remove_nodes, remove_edges = [], []
                for discont_node in discont_cc:
                    ordinal_relation = np.sum([discont_graph.nodes[xx]['ordinal'] \
                                               for xx in discont_graph.neighbors(discont_node)])
                    near_side = discont_graph.nodes[discont_node]['ordinal'] <= 0
                    if abs(ordinal_relation) < len([*discont_graph.neighbors(discont_node)]):
                        remove_nodes.append(discont_node)
                        for ne_node in discont_graph.neighbors(discont_node):
                            remove_flag = (near_side and not(key_exist(LDI.nodes[ne_node], 'far'))) or \
                                          (not near_side and not(key_exist(LDI.nodes[ne_node], 'near')))
                            remove_edges += [(discont_node, ne_node)] if remove_flag else []
                    else:
                        if near_side and key_exist(LDI.nodes[discont_node], 'near'):
                            LDI.nodes[discont_node].pop('near')
                        elif not(near_side) and key_exist(LDI.nodes[discont_node], 'far'):
                            LDI.nodes[discont_node].pop('far')
                discont_graph.remove_edges_from(remove_edges)
                sub_mesh = discont_graph.subgraph(list(discont_cc)).copy()
                sub_discont_ccs = [*netx.connected_components(sub_mesh)]
                is_redun_near = lambda xx: len(xx) == 1 and xx[0] in remove_nodes and key_exist(LDI.nodes[xx[0]], 'far')
                for sub_discont_cc in sub_discont_ccs:
                    if is_redun_near(list(sub_discont_cc)):
                        LDI.nodes[list(sub_discont_cc)[0]].pop('far')
                    new_discont_ccs.append(sub_discont_cc)
            else:
                new_discont_ccs.append(discont_cc)
        discont_ccs = new_discont_ccs
        new_discont_ccs = None
    if spdb == True:
        import pdb; pdb.set_trace()

    for edge_id, edge_cc in enumerate(discont_ccs):
        for node in edge_cc:
            LDI.nodes[node]['edge_id'] = edge_id

    return discont_ccs, LDI, discont_graph

def combine_end_node(mesh, edge_mesh, edge_ccs, depth):
    import collections
    mesh_nodes = mesh.nodes
    connect_dict = dict()
    for valid_edge_id, valid_edge_cc in enumerate(edge_ccs):
        connect_info = []
        for valid_edge_node in valid_edge_cc:
            single_connect = set()
            for ne_node in mesh.neighbors(valid_edge_node):
                if mesh_nodes[ne_node].get('far') is not None:
                    for fn in mesh_nodes[ne_node].get('far'):
                        if mesh.has_node(fn) and mesh_nodes[fn].get('edge_id') is not None:
                            single_connect.add(mesh_nodes[fn]['edge_id'])
                if mesh_nodes[ne_node].get('near') is not None:
                    for fn in mesh_nodes[ne_node].get('near'):
                        if mesh.has_node(fn) and mesh_nodes[fn].get('edge_id') is not None:
                            single_connect.add(mesh_nodes[fn]['edge_id'])
            connect_info.extend([*single_connect])
        connect_dict[valid_edge_id] = collections.Counter(connect_info)

    end_maps = np.zeros((mesh.graph['H'], mesh.graph['W']))
    edge_maps = np.zeros((mesh.graph['H'], mesh.graph['W'])) - 1
    for valid_edge_id, valid_edge_cc in enumerate(edge_ccs):
        for valid_edge_node in valid_edge_cc:
            edge_maps[valid_edge_node[0], valid_edge_node[1]] = valid_edge_id
            if len([*edge_mesh.neighbors(valid_edge_node)]) == 1:
                num_ne = 1
                if num_ne == 1:
                    end_maps[valid_edge_node[0], valid_edge_node[1]] = valid_edge_node[2]
    nxs, nys = np.where(end_maps != 0)
    invalid_nodes = set()
    for nx, ny in zip(nxs, nys):
        if mesh.has_node((nx, ny, end_maps[nx, ny])) is False:
            invalid_nodes.add((nx, ny))
            continue
        four_nes = [xx for xx in [(nx - 1, ny), (nx + 1, ny), (nx, ny - 1), (nx, ny + 1)] \
                        if 0 <= xx[0] < mesh.graph['H'] and 0 <= xx[1] < mesh.graph['W'] and \
                        end_maps[xx[0], xx[1]] != 0]
        mesh_nes = [*mesh.neighbors((nx, ny, end_maps[nx, ny]))]
        remove_num = 0
        for fne in four_nes:
            if (fne[0], fne[1], end_maps[fne[0], fne[1]]) in mesh_nes:
                remove_num += 1
        if remove_num == len(four_nes):
            invalid_nodes.add((nx, ny))
    for invalid_node in invalid_nodes:
        end_maps[invalid_node[0], invalid_node[1]] = 0

    nxs, nys = np.where(end_maps != 0)
    invalid_nodes = set()
    for nx, ny in zip(nxs, nys):
        if mesh_nodes[(nx, ny, end_maps[nx, ny])].get('edge_id') is None:
            continue
        else:
            self_id = mesh_nodes[(nx, ny, end_maps[nx, ny])].get('edge_id')
            self_connect = connect_dict[self_id] if connect_dict.get(self_id) is not None else dict()
        four_nes = [xx for xx in [(nx - 1, ny), (nx + 1, ny), (nx, ny - 1), (nx, ny + 1)] \
                        if 0 <= xx[0] < mesh.graph['H'] and 0 <= xx[1] < mesh.graph['W'] and \
                        end_maps[xx[0], xx[1]] != 0]
        for fne in four_nes:
            if mesh_nodes[(fne[0], fne[1], end_maps[fne[0], fne[1]])].get('edge_id') is None:
                continue
            else:
                ne_id = mesh_nodes[(fne[0], fne[1], end_maps[fne[0], fne[1]])]['edge_id']
                if self_connect.get(ne_id) is None or self_connect.get(ne_id) == 1:
                    continue
                else:
                    invalid_nodes.add((nx, ny))
    for invalid_node in invalid_nodes:
        end_maps[invalid_node[0], invalid_node[1]] = 0
    nxs, nys = np.where(end_maps != 0)
    invalid_nodes = set()
    for nx, ny in zip(nxs, nys):
        four_nes = [xx for xx in [(nx - 1, ny), (nx + 1, ny), (nx, ny - 1), (nx, ny + 1)] \
                        if 0 <= xx[0] < mesh.graph['H'] and 0 <= xx[1] < mesh.graph['W'] and \
                        end_maps[xx[0], xx[1]] != 0]
        for fne in four_nes:
            if mesh.has_node((fne[0], fne[1], end_maps[fne[0], fne[1]])):
                node_a, node_b = (fne[0], fne[1], end_maps[fne[0], fne[1]]), (nx, ny, end_maps[nx, ny])
                mesh.add_edge(node_a, node_b)
                mesh_nodes[node_b]['must_connect'] = set() if mesh_nodes[node_b].get('must_connect') is None else mesh_nodes[node_b]['must_connect']
                mesh_nodes[node_b]['must_connect'].add(node_a)
                mesh_nodes[node_b]['must_connect'] |= set([xx for xx in [*edge_mesh.neighbors(node_a)] if \
                                                            (xx[0] - node_b[0]) < 2 and (xx[1] - node_b[1]) < 2])
                mesh_nodes[node_a]['must_connect'] = set() if mesh_nodes[node_a].get('must_connect') is None else mesh_nodes[node_a]['must_connect']
                mesh_nodes[node_a]['must_connect'].add(node_b)
                mesh_nodes[node_a]['must_connect'] |= set([xx for xx in [*edge_mesh.neighbors(node_b)] if \
                                                            (xx[0] - node_a[0]) < 2 and (xx[1] - node_a[1]) < 2])
                invalid_nodes.add((nx, ny))
    for invalid_node in invalid_nodes:
        end_maps[invalid_node[0], invalid_node[1]] = 0

    return mesh

def remove_redundant_edge(mesh, edge_mesh, edge_ccs, info_on_pix, config, redundant_number=1000, invalid=False, spdb=False):
    point_to_amount = {}
    point_to_id = {}
    end_maps = np.zeros((mesh.graph['H'], mesh.graph['W'])) - 1
    for valid_edge_id, valid_edge_cc in enumerate(edge_ccs):
        for valid_edge_node in valid_edge_cc:
            point_to_amount[valid_edge_node] = len(valid_edge_cc)
            point_to_id[valid_edge_node] = valid_edge_id
            if edge_mesh.has_node(valid_edge_node) is True:
                if len([*edge_mesh.neighbors(valid_edge_node)]) == 1:
                    end_maps[valid_edge_node[0], valid_edge_node[1]] = valid_edge_id
    nxs, nys = np.where(end_maps > -1)
    point_to_adjoint = {}
    for nx, ny in zip(nxs, nys):
        adjoint_edges = set([end_maps[x, y] for x, y in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)] if end_maps[x, y] != -1])
        point_to_adjoint[end_maps[nx, ny]] = (point_to_adjoint[end_maps[nx, ny]] | adjoint_edges) if point_to_adjoint.get(end_maps[nx, ny]) is not None else adjoint_edges
    valid_edge_ccs = filter_edge(mesh, edge_ccs, config, invalid=invalid)
    edge_canvas = np.zeros((mesh.graph['H'], mesh.graph['W'])) - 1
    for valid_edge_id, valid_edge_cc in enumerate(valid_edge_ccs):
        for valid_edge_node in valid_edge_cc:
            edge_canvas[valid_edge_node[0], valid_edge_node[1]] = valid_edge_id
    if spdb is True:
        plt.imshow(edge_canvas); plt.show()
        import pdb; pdb.set_trace()
    for valid_edge_id, valid_edge_cc in enumerate(valid_edge_ccs):
        end_number = 0
        four_end_number = 0
        eight_end_number = 0
        db_eight_end_number = 0
        if len(valid_edge_cc) > redundant_number:
            continue
        for valid_edge_node in valid_edge_cc:
            if len([*edge_mesh.neighbors(valid_edge_node)]) == 3:
                break
            elif len([*edge_mesh.neighbors(valid_edge_node)]) == 1:
                hx, hy, hz = valid_edge_node
                if invalid is False:
                    eight_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1),
                                                     (hx + 1, hy + 1), (hx - 1, hy - 1), (hx - 1, hy + 1), (hx + 1, hy - 1)] \
                                            if info_on_pix.get((x, y)) is not None and edge_canvas[x, y] != -1 and edge_canvas[x, y] != valid_edge_id]
                    if len(eight_nes) == 0:
                        end_number += 1
                if invalid is True:
                    four_nes = []; eight_nes = []; db_eight_nes = []
                    four_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] \
                                            if info_on_pix.get((x, y)) is not None and edge_canvas[x, y] != -1 and edge_canvas[x, y] != valid_edge_id]
                    eight_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1), \
                                                    (hx + 1, hy + 1), (hx - 1, hy - 1), (hx - 1, hy + 1), (hx + 1, hy - 1)] \
                                            if info_on_pix.get((x, y)) is not None and edge_canvas[x, y] != -1 and edge_canvas[x, y] != valid_edge_id]
                    db_eight_nes = [(x, y) for x in range(hx - 2, hx + 3) for y in range(hy - 2, hy + 3) \
                                    if info_on_pix.get((x, y)) is not None and edge_canvas[x, y] != -1 and edge_canvas[x, y] != valid_edge_id and (x, y) != (hx, hy)]
                    if len(four_nes) == 0 or len(eight_nes) == 0:
                        end_number += 1
                        if len(four_nes) == 0:
                            four_end_number += 1
                        if len(eight_nes) == 0:
                            eight_end_number += 1
                        if len(db_eight_nes) == 0:
                            db_eight_end_number += 1
            elif len([*edge_mesh.neighbors(valid_edge_node)]) == 0:
                hx, hy, hz = valid_edge_node
                four_nes = [(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] \
                                if info_on_pix.get((x, y)) is not None and \
                                    mesh.has_edge(valid_edge_node, (x, y, info_on_pix[(x, y)][0]['depth'])) is False]
                for ne in four_nes:
                    try:
                        if invalid is True or (point_to_amount.get(ne) is None or point_to_amount[ne] < redundant_number) or \
                            point_to_id[ne] in point_to_adjoint.get(point_to_id[valid_edge_node], set()):
                            mesh.add_edge(valid_edge_node, ne)
                    except:
                        import pdb; pdb.set_trace()
        if (invalid is not True and end_number >= 1) or (invalid is True and end_number >= 2 and eight_end_number >= 1 and db_eight_end_number >= 1):
            for valid_edge_node in valid_edge_cc:
                hx, hy, _ = valid_edge_node
                four_nes = [(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] \
                                if info_on_pix.get((x, y)) is not None and \
                                    mesh.has_edge(valid_edge_node, (x, y, info_on_pix[(x, y)][0]['depth'])) is False and \
                                    (edge_canvas[x, y] == -1 or edge_canvas[x, y] == valid_edge_id)]
                for ne in four_nes:
                    if invalid is True or (point_to_amount.get(ne) is None or point_to_amount[ne] < redundant_number) or \
                        point_to_id[ne] in point_to_adjoint.get(point_to_id[valid_edge_node], set()):
                        mesh.add_edge(valid_edge_node, ne)

    return mesh

def judge_dangle(mark, mesh, node):
    if not (1 <= node[0] < mesh.graph['H']-1) or not(1 <= node[1] < mesh.graph['W']-1):
        return mark
    mesh_neighbors = [*mesh.neighbors(node)]
    mesh_neighbors = [xx for xx in mesh_neighbors if 0 < xx[0] < mesh.graph['H'] - 1 and 0 < xx[1] < mesh.graph['W'] - 1]
    if len(mesh_neighbors) >= 3:
        return mark
    elif len(mesh_neighbors) <= 1:
        mark[node[0], node[1]] = (len(mesh_neighbors) + 1)
    else:
        dan_ne_node_a = mesh_neighbors[0]
        dan_ne_node_b = mesh_neighbors[1]
        if abs(dan_ne_node_a[0] - dan_ne_node_b[0]) > 1 or \
            abs(dan_ne_node_a[1] - dan_ne_node_b[1]) > 1:
            mark[node[0], node[1]] = 3

    return mark

def remove_dangling(mesh, edge_ccs, edge_mesh, info_on_pix, image, depth, config):

    tmp_edge_ccs = copy.deepcopy(edge_ccs)
    for edge_cc_id, valid_edge_cc in enumerate(tmp_edge_ccs):
        if len(valid_edge_cc) > 1 or len(valid_edge_cc) == 0:
            continue
        single_edge_node = [*valid_edge_cc][0]
        hx, hy, hz = single_edge_node
        eight_nes = set([(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1),
                         (hx + 1, hy + 1), (hx - 1, hy - 1), (hx - 1, hy + 1), (hx + 1, hy - 1)] \
                         if info_on_pix.get((x, y)) is not None])
        four_nes = [(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] \
                    if info_on_pix.get((x, y)) is not None]
        sub_mesh = mesh.subgraph(eight_nes).copy()
        ccs = netx.connected_components(sub_mesh)
        four_ccs = []
        for cc_id, _cc in enumerate(ccs):
            four_ccs.append(set())
            for cc_node in _cc:
                if abs(cc_node[0] - hx) + abs(cc_node[1] - hy) < 2:
                    four_ccs[cc_id].add(cc_node)
        largest_cc = sorted(four_ccs, key=lambda x: (len(x), -np.sum([abs(xx[2] - hz) for xx in x])))[-1]
        if len(largest_cc) < 2:
            for ne in four_nes:
                mesh.add_edge(single_edge_node, ne)
        else:
            mesh.remove_edges_from([(single_edge_node, ne) for ne in mesh.neighbors(single_edge_node)])
            new_depth = np.mean([xx[2] for xx in largest_cc])
            info_on_pix[(hx, hy)][0]['depth'] = new_depth
            info_on_pix[(hx, hy)][0]['disp'] = 1./new_depth
            new_node = (hx, hy, new_depth)
            mesh = refresh_node(single_edge_node, mesh.node[single_edge_node], new_node, dict(), mesh)
            edge_ccs[edge_cc_id] = set([new_node])
            for ne in largest_cc:
                mesh.add_edge(new_node, ne)

    mark = np.zeros((mesh.graph['H'], mesh.graph['W']))
    for edge_idx, edge_cc in enumerate(edge_ccs):
        for edge_node in edge_cc:
            if not (mesh.graph['bord_up'] <= edge_node[0] < mesh.graph['bord_down']-1) or \
               not (mesh.graph['bord_left'] <= edge_node[1] < mesh.graph['bord_right']-1):
                continue
            mesh_neighbors = [*mesh.neighbors(edge_node)]
            mesh_neighbors = [xx for xx in mesh_neighbors \
                                if mesh.graph['bord_up'] < xx[0] < mesh.graph['bord_down'] - 1 and \
                                   mesh.graph['bord_left'] < xx[1] < mesh.graph['bord_right'] - 1]
            if len([*mesh.neighbors(edge_node)]) >= 3:
                continue
            elif len([*mesh.neighbors(edge_node)]) <= 1:
                mark[edge_node[0], edge_node[1]] += (len([*mesh.neighbors(edge_node)]) + 1)
            else:
                dan_ne_node_a = [*mesh.neighbors(edge_node)][0]
                dan_ne_node_b = [*mesh.neighbors(edge_node)][1]
                if abs(dan_ne_node_a[0] - dan_ne_node_b[0]) > 1 or \
                    abs(dan_ne_node_a[1] - dan_ne_node_b[1]) > 1:
                    mark[edge_node[0], edge_node[1]] += 3
    mxs, mys = np.where(mark == 1)
    conn_0_nodes = [(x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth']) for x in zip(mxs, mys) \
                        if mesh.has_node((x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth']))]
    mxs, mys = np.where(mark == 2)
    conn_1_nodes = [(x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth']) for x in zip(mxs, mys) \
                        if mesh.has_node((x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth']))]
    for node in conn_0_nodes:
        hx, hy = node[0], node[1]
        four_nes = [(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] \
                     if info_on_pix.get((x, y)) is not None]
        re_depth = {'value' : 0, 'count': 0}
        for ne in four_nes:
            mesh.add_edge(node, ne)
            re_depth['value'] += cc_node[2]
            re_depth['count'] += 1.
        re_depth = re_depth['value'] / re_depth['count']
        mapping_dict = {node: (node[0], node[1], re_depth)}
        info_on_pix, mesh, edge_mesh = update_info(mapping_dict, info_on_pix, mesh, edge_mesh)
        depth[node[0], node[1]] = abs(re_depth)
        mark[node[0], node[1]] = 0
    for node in conn_1_nodes:
        hx, hy = node[0], node[1]
        eight_nes = set([(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1),
                                                                           (hx + 1, hy + 1), (hx - 1, hy - 1), (hx - 1, hy + 1), (hx + 1, hy - 1)] \
                        if info_on_pix.get((x, y)) is not None])
        self_nes = set([ne2 for ne1 in mesh.neighbors(node) for ne2 in mesh.neighbors(ne1) if ne2 in eight_nes])
        eight_nes = [*(eight_nes - self_nes)]
        sub_mesh = mesh.subgraph(eight_nes).copy()
        ccs = netx.connected_components(sub_mesh)
        largest_cc = sorted(ccs, key=lambda x: (len(x), -np.sum([abs(xx[0] - node[0]) + abs(xx[1] - node[1]) for xx in x])))[-1]

        mesh.remove_edges_from([(xx, node) for xx in mesh.neighbors(node)])
        re_depth = {'value' : 0, 'count': 0}
        for cc_node in largest_cc:
            if cc_node[0] == node[0] and cc_node[1] == node[1]:
                continue
            re_depth['value'] += cc_node[2]
            re_depth['count'] += 1.
            if abs(cc_node[0] - node[0]) + abs(cc_node[1] - node[1]) < 2:
                mesh.add_edge(cc_node, node)
        try:
            re_depth = re_depth['value'] / re_depth['count']
        except:
            re_depth = node[2]
        renode = (node[0], node[1], re_depth)
        mapping_dict = {node: renode}
        info_on_pix, mesh, edge_mesh = update_info(mapping_dict, info_on_pix, mesh, edge_mesh)
        depth[node[0], node[1]] = abs(re_depth)
        mark[node[0], node[1]] = 0
        edge_mesh, mesh, mark, info_on_pix = recursive_add_edge(edge_mesh, mesh, info_on_pix, renode, mark)
    mxs, mys = np.where(mark == 3)
    conn_2_nodes = [(x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth']) for x in zip(mxs, mys) \
                        if mesh.has_node((x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth'])) and \
                            mesh.degree((x[0], x[1], info_on_pix[(x[0], x[1])][0]['depth'])) == 2]
    sub_mesh = mesh.subgraph(conn_2_nodes).copy()
    ccs = netx.connected_components(sub_mesh)
    for cc in ccs:
        candidate_nodes = [xx for xx in cc if sub_mesh.degree(xx) == 1]
        for node in candidate_nodes:
            if mesh.has_node(node) is False:
                continue
            ne_node = [xx for xx in mesh.neighbors(node) if xx not in cc][0]
            hx, hy = node[0], node[1]
            eight_nes = set([(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1),
                                                                            (hx + 1, hy + 1), (hx - 1, hy - 1), (hx - 1, hy + 1), (hx + 1, hy - 1)] \
                              if info_on_pix.get((x, y)) is not None and (x, y, info_on_pix[(x, y)][0]['depth']) not in cc])
            ne_sub_mesh = mesh.subgraph(eight_nes).copy()
            ne_ccs = netx.connected_components(ne_sub_mesh)
            try:
                ne_cc = [ne_cc for ne_cc in ne_ccs if ne_node in ne_cc][0]
            except:
                import pdb; pdb.set_trace()
            largest_cc = [xx for xx in ne_cc if abs(xx[0] - node[0]) + abs(xx[1] - node[1]) == 1]
            mesh.remove_edges_from([(xx, node) for xx in mesh.neighbors(node)])
            re_depth = {'value' : 0, 'count': 0}
            for cc_node in largest_cc:
                re_depth['value'] += cc_node[2]
                re_depth['count'] += 1.
                mesh.add_edge(cc_node, node)
            try:
                re_depth = re_depth['value'] / re_depth['count']
            except:
                re_depth = node[2]
            renode = (node[0], node[1], re_depth)
            mapping_dict = {node: renode}
            info_on_pix, mesh, edge_mesh = update_info(mapping_dict, info_on_pix, mesh, edge_mesh)
            depth[node[0], node[1]] = abs(re_depth)
            mark[node[0], node[1]] = 0
            edge_mesh, mesh, mark, info_on_pix = recursive_add_edge(edge_mesh, mesh, info_on_pix, renode, mark)
            break
        if len(cc) == 1:
            node = [node for node in cc][0]
            hx, hy = node[0], node[1]
            nine_nes = set([(x, y, info_on_pix[(x, y)][0]['depth']) for x, y in [(hx, hy), (hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1),
                                                                                  (hx + 1, hy + 1), (hx - 1, hy - 1), (hx - 1, hy + 1), (hx + 1, hy - 1)] \
                                if info_on_pix.get((x, y)) is not None and mesh.has_node((x, y, info_on_pix[(x, y)][0]['depth']))])
            ne_sub_mesh = mesh.subgraph(nine_nes).copy()
            ne_ccs = netx.connected_components(ne_sub_mesh)
            for ne_cc in ne_ccs:
                if node in ne_cc:
                    re_depth = {'value' : 0, 'count': 0}
                    for ne in ne_cc:
                        if abs(ne[0] - node[0]) + abs(ne[1] - node[1]) == 1:
                            mesh.add_edge(node, ne)
                            re_depth['value'] += ne[2]
                            re_depth['count'] += 1.
                    re_depth = re_depth['value'] / re_depth['count']
                    mapping_dict = {node: (node[0], node[1], re_depth)}
                    info_on_pix, mesh, edge_mesh = update_info(mapping_dict, info_on_pix, mesh, edge_mesh)
                    depth[node[0], node[1]] = abs(re_depth)
                    mark[node[0], node[1]] = 0


    return mesh, info_on_pix, edge_mesh, depth, mark

def context_and_holes(mesh, edge_ccs, config, specific_edge_id, specific_edge_loc, depth_feat_model,
                      connect_points_ccs=None, inpaint_iter=0, filter_edge=False, vis_edge_id=None):
    edge_maps = np.zeros((mesh.graph['H'], mesh.graph['W'])) - 1
    mask_info = {}
    for edge_id, edge_cc in enumerate(edge_ccs):
        for edge_node in edge_cc:
            edge_maps[edge_node[0], edge_node[1]] = edge_id

    context_ccs = [set() for x in range(len(edge_ccs))]
    extend_context_ccs = [set() for x in range(len(edge_ccs))]
    extend_erode_context_ccs = [set() for x in range(len(edge_ccs))]
    extend_edge_ccs = [set() for x in range(len(edge_ccs))]
    accomp_extend_context_ccs = [set() for x in range(len(edge_ccs))]
    erode_context_ccs = [set() for x in range(len(edge_ccs))]
    broken_mask_ccs = [set() for x in range(len(edge_ccs))]
    invalid_extend_edge_ccs = [set() for x in range(len(edge_ccs))]
    intouched_ccs = [set() for x in range(len(edge_ccs))]
    redundant_ccs = [set() for x in range(len(edge_ccs))]
    if inpaint_iter == 0:
        background_thickness = config['background_thickness']
        context_thickness = config['context_thickness']
    else:
        background_thickness = config['background_thickness_2']
        context_thickness = config['context_thickness_2']

    mesh_nodes = mesh.nodes
    for edge_id, edge_cc in enumerate(edge_ccs):
        if context_thickness == 0 or (len(specific_edge_id) > 0 and edge_id not in specific_edge_id):
            continue
        edge_group = {}
        for edge_node in edge_cc:
            far_nodes = mesh_nodes[edge_node].get('far')
            if far_nodes is None:
                continue
            for far_node in far_nodes:
                if far_node in edge_cc:
                    continue
                context_ccs[edge_id].add(far_node)
                if mesh_nodes[far_node].get('edge_id') is not None:
                    if edge_group.get(mesh_nodes[far_node]['edge_id']) is None:
                        edge_group[mesh_nodes[far_node]['edge_id']] = set()
                    edge_group[mesh_nodes[far_node]['edge_id']].add(far_node)
        if len(edge_cc) > 2:
            for edge_key in [*edge_group.keys()]:
                if len(edge_group[edge_key]) == 1:
                    context_ccs[edge_id].remove([*edge_group[edge_key]][0])
    for edge_id, edge_cc in enumerate(edge_ccs):
        if inpaint_iter != 0:
            continue
        tmp_intouched_nodes = set()
        for edge_node in edge_cc:
            raw_intouched_nodes = set(mesh_nodes[edge_node].get('near')) if mesh_nodes[edge_node].get('near') is not None else set()
            tmp_intouched_nodes |= set([xx for xx in raw_intouched_nodes if mesh_nodes[xx].get('edge_id') is not None and \
                                                                         len(context_ccs[mesh_nodes[xx].get('edge_id')]) > 0])
        intouched_ccs[edge_id] |= tmp_intouched_nodes
        tmp_intouched_nodes = None
    mask_ccs = copy.deepcopy(edge_ccs)
    forbidden_len = 3
    forbidden_map = np.ones((mesh.graph['H'] - forbidden_len, mesh.graph['W'] - forbidden_len))
    forbidden_map = np.pad(forbidden_map, ((forbidden_len, forbidden_len), (forbidden_len, forbidden_len)), mode='constant').astype(np.bool)
    cur_tmp_mask_map = np.zeros_like(forbidden_map).astype(np.bool)
    passive_background = 10 if 10 is not None else background_thickness
    passive_context = 1 if 1 is not None else context_thickness

    for edge_id, edge_cc in enumerate(edge_ccs):
        cur_mask_cc = None; cur_mask_cc = []
        cur_context_cc = None; cur_context_cc = []
        cur_accomp_near_cc = None; cur_accomp_near_cc = []
        cur_invalid_extend_edge_cc = None; cur_invalid_extend_edge_cc = []
        cur_comp_far_cc = None; cur_comp_far_cc = []
        tmp_erode = []
        if len(context_ccs[edge_id]) == 0 or (len(specific_edge_id) > 0 and edge_id not in specific_edge_id):
            continue
        for i in range(max(background_thickness, context_thickness)):
            cur_tmp_mask_map.fill(False)
            if i == 0:
                tmp_mask_nodes = copy.deepcopy(mask_ccs[edge_id])
                tmp_intersect_nodes = []
                tmp_intersect_context_nodes = []
                mask_map = np.zeros((mesh.graph['H'], mesh.graph['W']), dtype=np.bool)
                context_depth = np.zeros((mesh.graph['H'], mesh.graph['W']))
                comp_cnt_depth = np.zeros((mesh.graph['H'], mesh.graph['W']))
                connect_map = np.zeros((mesh.graph['H'], mesh.graph['W']))
                for node in tmp_mask_nodes:
                    mask_map[node[0], node[1]] = True
                    depth_count = 0
                    if mesh_nodes[node].get('far') is not None:
                        for comp_cnt_node in mesh_nodes[node]['far']:
                            comp_cnt_depth[node[0], node[1]] += abs(comp_cnt_node[2])
                            depth_count += 1
                    if depth_count > 0:
                        comp_cnt_depth[node[0], node[1]] = comp_cnt_depth[node[0], node[1]] / depth_count
                    connect_node = []
                    if mesh_nodes[node].get('connect_point_id') is not None:
                        connect_node.append(mesh_nodes[node]['connect_point_id'])
                    connect_point_id = np.bincount(connect_node).argmax() if len(connect_node) > 0 else -1
                    if connect_point_id > -1 and connect_points_ccs is not None:
                        for xx in connect_points_ccs[connect_point_id]:
                            if connect_map[xx[0], xx[1]] == 0:
                                connect_map[xx[0], xx[1]] = xx[2]
                    if mesh_nodes[node].get('connect_point_exception') is not None:
                        for xx in mesh_nodes[node]['connect_point_exception']:
                            if connect_map[xx[0], xx[1]] == 0:
                                connect_map[xx[0], xx[1]] = xx[2]
                tmp_context_nodes = [*context_ccs[edge_id]]
                tmp_erode.append([*context_ccs[edge_id]])
                context_map = np.zeros((mesh.graph['H'], mesh.graph['W']), dtype=np.bool)
                if (context_map.astype(np.uint8) * mask_map.astype(np.uint8)).max() > 0:
                    import pdb; pdb.set_trace()
                for node in tmp_context_nodes:
                    context_map[node[0], node[1]] = True
                    context_depth[node[0], node[1]] = node[2]
                context_map[mask_map == True] = False
                if (context_map.astype(np.uint8) * mask_map.astype(np.uint8)).max() > 0:
                    import pdb; pdb.set_trace()
                tmp_intouched_nodes = [*intouched_ccs[edge_id]]
                intouched_map = np.zeros((mesh.graph['H'], mesh.graph['W']), dtype=np.bool)
                for node in tmp_intouched_nodes: intouched_map[node[0], node[1]] = True
                intouched_map[mask_map == True] = False
                tmp_redundant_nodes = set()
                tmp_noncont_nodes = set()
                noncont_map = np.zeros((mesh.graph['H'], mesh.graph['W']), dtype=np.bool)
                intersect_map = np.zeros((mesh.graph['H'], mesh.graph['W']), dtype=np.bool)
                intersect_context_map = np.zeros((mesh.graph['H'], mesh.graph['W']), dtype=np.bool)
            if i > passive_background and inpaint_iter == 0:
                new_tmp_intersect_nodes = None
                new_tmp_intersect_nodes = []
                for node in tmp_intersect_nodes:
                    nes = mesh.neighbors(node)
                    for ne in nes:
                        if bool(context_map[ne[0], ne[1]]) is False and \
                        bool(mask_map[ne[0], ne[1]]) is False and \
                        bool(forbidden_map[ne[0], ne[1]]) is True and \
                        bool(intouched_map[ne[0], ne[1]]) is False and\
                        bool(intersect_map[ne[0], ne[1]]) is False and\
                        bool(intersect_context_map[ne[0], ne[1]]) is False:
                            break_flag = False
                            if (i - passive_background) % 2 == 0 and (i - passive_background) % 8 != 0:
                                four_nes = [xx for xx in[[ne[0] - 1, ne[1]], [ne[0] + 1, ne[1]], [ne[0], ne[1] - 1], [ne[0], ne[1] + 1]] \
                                                if 0 <= xx[0] < mesh.graph['H'] and 0 <= xx[1] < mesh.graph['W']]
                                for fne in four_nes:
                                    if bool(mask_map[fne[0], fne[1]]) is True:
                                        break_flag = True
                                        break
                                if break_flag is True:
                                    continue
                            intersect_map[ne[0], ne[1]] = True
                            new_tmp_intersect_nodes.append(ne)
                tmp_intersect_nodes = None
                tmp_intersect_nodes = new_tmp_intersect_nodes

            if i > passive_context and inpaint_iter == 1:
                new_tmp_intersect_context_nodes = None
                new_tmp_intersect_context_nodes = []
                for node in tmp_intersect_context_nodes:
                    nes = mesh.neighbors(node)
                    for ne in nes:
                        if bool(context_map[ne[0], ne[1]]) is False and \
                        bool(mask_map[ne[0], ne[1]]) is False and \
                        bool(forbidden_map[ne[0], ne[1]]) is True and \
                        bool(intouched_map[ne[0], ne[1]]) is False and\
                        bool(intersect_map[ne[0], ne[1]]) is False and \
                        bool(intersect_context_map[ne[0], ne[1]]) is False:
                            intersect_context_map[ne[0], ne[1]] = True
                            new_tmp_intersect_context_nodes.append(ne)
                tmp_intersect_context_nodes = None
                tmp_intersect_context_nodes = new_tmp_intersect_context_nodes

            new_tmp_mask_nodes = None
            new_tmp_mask_nodes = []
            for node in tmp_mask_nodes:
                four_nes = {xx:[] for xx in [(node[0] - 1, node[1]), (node[0] + 1, node[1]), (node[0], node[1] - 1), (node[0], node[1] + 1)] if \
                            0 <= xx[0] < connect_map.shape[0] and 0 <= xx[1] < connect_map.shape[1]}
                if inpaint_iter > 0:
                    for ne in four_nes.keys():
                        if connect_map[ne[0], ne[1]] == True:
                            tmp_context_nodes.append((ne[0], ne[1], connect_map[ne[0], ne[1]]))
                            context_map[ne[0], ne[1]] = True
                nes = mesh.neighbors(node)
                if inpaint_iter > 0:
                    for ne in nes: four_nes[(ne[0], ne[1])].append(ne[2])
                    nes = []
                    for kfne, vfnes in four_nes.items(): vfnes.sort(key = lambda xx: abs(xx), reverse=True)
                    for kfne, vfnes in four_nes.items():
                        for vfne in vfnes: nes.append((kfne[0], kfne[1], vfne))
                for ne in nes:
                    if bool(context_map[ne[0], ne[1]]) is False and \
                       bool(mask_map[ne[0], ne[1]]) is False and \
                       bool(forbidden_map[ne[0], ne[1]]) is True and \
                       bool(intouched_map[ne[0], ne[1]]) is False and \
                       bool(intersect_map[ne[0], ne[1]]) is False and \
                       bool(intersect_context_map[ne[0], ne[1]]) is False:
                        if i == passive_background and inpaint_iter == 0:
                            if np.any(context_map[max(ne[0] - 1, 0):min(ne[0] + 2, mesh.graph['H']), max(ne[1] - 1, 0):min(ne[1] + 2, mesh.graph['W'])]) == True:
                                intersect_map[ne[0], ne[1]] = True
                                tmp_intersect_nodes.append(ne)
                                continue
                        if i < background_thickness:
                            if inpaint_iter == 0:
                                cur_mask_cc.append(ne)
                            elif mesh_nodes[ne].get('inpaint_id') == 1:
                                cur_mask_cc.append(ne)
                            else:
                                continue
                            mask_ccs[edge_id].add(ne)
                            if inpaint_iter == 0:
                                if comp_cnt_depth[node[0], node[1]] > 0 and comp_cnt_depth[ne[0], ne[1]] == 0:
                                    comp_cnt_depth[ne[0], ne[1]] = comp_cnt_depth[node[0], node[1]]
                                if mesh_nodes[ne].get('far') is not None:
                                    for comp_far_node in mesh_nodes[ne]['far']:
                                        cur_comp_far_cc.append(comp_far_node)
                                        cur_accomp_near_cc.append(ne)
                                        cur_invalid_extend_edge_cc.append(comp_far_node)
                                if mesh_nodes[ne].get('edge_id') is not None and \
                                    len(context_ccs[mesh_nodes[ne].get('edge_id')]) > 0:
                                    intouched_fars = set(mesh_nodes[ne].get('far')) if mesh_nodes[ne].get('far') is not None else set()
                                    accum_intouched_fars = set(intouched_fars)
                                    for intouched_far in intouched_fars:
                                        accum_intouched_fars |= set([*mesh.neighbors(intouched_far)])
                                    for intouched_far in accum_intouched_fars:
                                        if bool(mask_map[intouched_far[0], intouched_far[1]]) is True or \
                                        bool(context_map[intouched_far[0], intouched_far[1]]) is True:
                                            continue
                                        tmp_redundant_nodes.add(intouched_far)
                                        intouched_map[intouched_far[0], intouched_far[1]] = True
                                if mesh_nodes[ne].get('near') is not None:
                                    intouched_nears = set(mesh_nodes[ne].get('near'))
                                    for intouched_near in intouched_nears:
                                        if bool(mask_map[intouched_near[0], intouched_near[1]]) is True or \
                                        bool(context_map[intouched_near[0], intouched_near[1]]) is True:
                                            continue
                                        tmp_redundant_nodes.add(intouched_near)
                                        intouched_map[intouched_near[0], intouched_near[1]] = True
                        if not (mesh_nodes[ne].get('inpaint_id') != 1 and inpaint_iter == 1):
                            new_tmp_mask_nodes.append(ne)
                            mask_map[ne[0], ne[1]] = True
            tmp_mask_nodes = new_tmp_mask_nodes

            new_tmp_context_nodes = None
            new_tmp_context_nodes = []
            for node in tmp_context_nodes:
                nes = mesh.neighbors(node)
                if inpaint_iter > 0:
                    four_nes = {(node[0] - 1, node[1]):[], (node[0] + 1, node[1]):[], (node[0], node[1] - 1):[], (node[0], node[1] + 1):[]}
                    for ne in nes: four_nes[(ne[0], ne[1])].append(ne[2])
                    nes = []
                    for kfne, vfnes in four_nes.items(): vfnes.sort(key = lambda xx: abs(xx), reverse=True)
                    for kfne, vfnes in four_nes.items():
                        for vfne in vfnes: nes.append((kfne[0], kfne[1], vfne))
                for ne in nes:
                    mask_flag = (bool(mask_map[ne[0], ne[1]]) is False)
                    if bool(context_map[ne[0], ne[1]]) is False and mask_flag and \
                       bool(forbidden_map[ne[0], ne[1]]) is True and bool(noncont_map[ne[0], ne[1]]) is False and \
                       bool(intersect_context_map[ne[0], ne[1]]) is False:
                        if i == passive_context and inpaint_iter == 1:
                            mnes = mesh.neighbors(ne)
                            if any([mask_map[mne[0], mne[1]] == True for mne in mnes]) is True:
                                intersect_context_map[ne[0], ne[1]] = True
                                tmp_intersect_context_nodes.append(ne)
                                continue
                        if False and mesh_nodes[ne].get('near') is not None and mesh_nodes[ne].get('edge_id') != edge_id:
                            noncont_nears = set(mesh_nodes[ne].get('near'))
                            for noncont_near in noncont_nears:
                                if bool(context_map[noncont_near[0], noncont_near[1]]) is False:
                                    tmp_noncont_nodes.add(noncont_near)
                                    noncont_map[noncont_near[0], noncont_near[1]] = True
                        new_tmp_context_nodes.append(ne)
                        context_map[ne[0], ne[1]] = True
                        context_depth[ne[0], ne[1]] = ne[2]
            cur_context_cc.extend(new_tmp_context_nodes)
            tmp_erode.append(new_tmp_context_nodes)
            tmp_context_nodes = None
            tmp_context_nodes = new_tmp_context_nodes
            new_tmp_intouched_nodes = None; new_tmp_intouched_nodes = []

            for node in tmp_intouched_nodes:
                if bool(context_map[node[0], node[1]]) is True or bool(mask_map[node[0], node[1]]) is True:
                    continue
                nes = mesh.neighbors(node)

                for ne in nes:
                    if bool(context_map[ne[0], ne[1]]) is False and \
                       bool(mask_map[ne[0], ne[1]]) is False and \
                       bool(intouched_map[ne[0], ne[1]]) is False and \
                       bool(forbidden_map[ne[0], ne[1]]) is True:
                        new_tmp_intouched_nodes.append(ne)
                        intouched_map[ne[0], ne[1]] = True
            tmp_intouched_nodes = None
            tmp_intouched_nodes = set(new_tmp_intouched_nodes)
            new_tmp_redundant_nodes = None; new_tmp_redundant_nodes = []
            for node in tmp_redundant_nodes:
                if bool(context_map[node[0], node[1]]) is True or \
                   bool(mask_map[node[0], node[1]]) is True:
                    continue
                nes = mesh.neighbors(node)

                for ne in nes:
                    if bool(context_map[ne[0], ne[1]]) is False and \
                       bool(mask_map[ne[0], ne[1]]) is False and \
                       bool(intouched_map[ne[0], ne[1]]) is False and \
                       bool(forbidden_map[ne[0], ne[1]]) is True:
                        new_tmp_redundant_nodes.append(ne)
                        intouched_map[ne[0], ne[1]] = True
            tmp_redundant_nodes = None
            tmp_redundant_nodes = set(new_tmp_redundant_nodes)
            new_tmp_noncont_nodes = None; new_tmp_noncont_nodes = []
            for node in tmp_noncont_nodes:
                if bool(context_map[node[0], node[1]]) is True or \
                   bool(mask_map[node[0], node[1]]) is True:
                    continue
                nes = mesh.neighbors(node)
                rmv_flag = False
                for ne in nes:
                    if bool(context_map[ne[0], ne[1]]) is False and \
                       bool(mask_map[ne[0], ne[1]]) is False and \
                       bool(noncont_map[ne[0], ne[1]]) is False and \
                       bool(forbidden_map[ne[0], ne[1]]) is True:
                        patch_context_map = context_map[max(ne[0] - 1, 0):min(ne[0] + 2, context_map.shape[0]),
                                                        max(ne[1] - 1, 0):min(ne[1] + 2, context_map.shape[1])]
                        if bool(np.any(patch_context_map)) is True:
                            new_tmp_noncont_nodes.append(ne)
                            noncont_map[ne[0], ne[1]] = True
            tmp_noncont_nodes = None
            tmp_noncont_nodes = set(new_tmp_noncont_nodes)
        if inpaint_iter == 0:
            depth_dict = get_depth_from_maps(context_map, mask_map, context_depth, mesh.graph['H'], mesh.graph['W'], log_depth=config['log_depth'])
            mask_size = get_valid_size(depth_dict['mask'])
            mask_size = dilate_valid_size(mask_size, depth_dict['mask'], dilate=[20, 20])
            context_size = get_valid_size(depth_dict['context'])
            context_size = dilate_valid_size(context_size, depth_dict['context'], dilate=[20, 20])
            union_size = size_operation(mask_size, context_size, operation='+')
            depth_dict = depth_inpainting(None, None, None, None, mesh, config, union_size, depth_feat_model, None, given_depth_dict=depth_dict, spdb=False)
            near_depth_map, raw_near_depth_map = np.zeros((mesh.graph['H'], mesh.graph['W'])), np.zeros((mesh.graph['H'], mesh.graph['W']))
            filtered_comp_far_cc, filtered_accomp_near_cc = set(), set()
            for node in cur_accomp_near_cc:
                near_depth_map[node[0], node[1]] = depth_dict['output'][node[0], node[1]]
                raw_near_depth_map[node[0], node[1]] = node[2]
            for node in cur_comp_far_cc:
                four_nes = [xx for xx in [(node[0] - 1, node[1]), (node[0] + 1, node[1]), (node[0], node[1] - 1), (node[0], node[1] + 1)] \
                            if 0 <= xx[0] < mesh.graph['H'] and 0 <= xx[1] < mesh.graph['W'] and \
                            near_depth_map[xx[0], xx[1]] != 0 and \
                            abs(near_depth_map[xx[0], xx[1]]) < abs(node[2])]
                if len(four_nes) > 0:
                    filtered_comp_far_cc.add(node)
                for ne in four_nes:
                    filtered_accomp_near_cc.add((ne[0], ne[1], -abs(raw_near_depth_map[ne[0], ne[1]])))
            cur_comp_far_cc, cur_accomp_near_cc = filtered_comp_far_cc, filtered_accomp_near_cc
        mask_ccs[edge_id] |= set(cur_mask_cc)
        context_ccs[edge_id] |= set(cur_context_cc)
        accomp_extend_context_ccs[edge_id] |= set(cur_accomp_near_cc).intersection(cur_mask_cc)
        extend_edge_ccs[edge_id] |= set(cur_accomp_near_cc).intersection(cur_mask_cc)
        extend_context_ccs[edge_id] |= set(cur_comp_far_cc)
        invalid_extend_edge_ccs[edge_id] |= set(cur_invalid_extend_edge_cc)
        erode_size = [0]
        for tmp in tmp_erode:
            erode_size.append(len(tmp))
            if len(erode_size) > 1:
                erode_size[-1] += erode_size[-2]
        if inpaint_iter == 0:
            tmp_width = config['depth_edge_dilate']
        else:
            tmp_width = 0
        while float(erode_size[tmp_width]) / (erode_size[-1] + 1e-6) > 0.3:
            tmp_width = tmp_width - 1
        try:
            if tmp_width == 0:
                erode_context_ccs[edge_id] = set([])
            else:
                erode_context_ccs[edge_id] = set(reduce(lambda x, y : x + y, [] + tmp_erode[:tmp_width]))
        except:
            import pdb; pdb.set_trace()
        erode_context_cc = copy.deepcopy(erode_context_ccs[edge_id])
        for erode_context_node in erode_context_cc:
            if (inpaint_iter != 0 and (mesh_nodes[erode_context_node].get('inpaint_id') is None or
                                        mesh_nodes[erode_context_node].get('inpaint_id') == 0)):
                erode_context_ccs[edge_id].remove(erode_context_node)
            else:
                context_ccs[edge_id].remove(erode_context_node)
        context_map = np.zeros((mesh.graph['H'], mesh.graph['W']))
        for context_node in context_ccs[edge_id]:
            context_map[context_node[0], context_node[1]] = 1
        extend_context_ccs[edge_id] = extend_context_ccs[edge_id] - mask_ccs[edge_id] - accomp_extend_context_ccs[edge_id]
    if inpaint_iter == 0:
        all_ecnt_cc = set()
        for ecnt_id, ecnt_cc in enumerate(extend_context_ccs):
            constraint_context_ids = set()
            constraint_context_cc = set()
            constraint_erode_context_cc = set()
            tmp_mask_cc = set()
            accum_context_cc = None; accum_context_cc = []
            for ecnt_node in accomp_extend_context_ccs[ecnt_id]:
                if edge_maps[ecnt_node[0], ecnt_node[1]] > -1:
                    constraint_context_ids.add(int(round(edge_maps[ecnt_node[0], ecnt_node[1]])))
            constraint_erode_context_cc = erode_context_ccs[ecnt_id]
            for constraint_context_id in constraint_context_ids:
                constraint_context_cc = constraint_context_cc | context_ccs[constraint_context_id] | erode_context_ccs[constraint_context_id]
                constraint_erode_context_cc = constraint_erode_context_cc | erode_context_ccs[constraint_context_id]
            for i in range(background_thickness):
                if i == 0:
                    tmp_context_nodes = copy.deepcopy(ecnt_cc)
                    tmp_invalid_context_nodes = copy.deepcopy(invalid_extend_edge_ccs[ecnt_id])
                    tmp_mask_nodes = copy.deepcopy(accomp_extend_context_ccs[ecnt_id])
                    tmp_context_map = np.zeros((mesh.graph['H'], mesh.graph['W'])).astype(np.bool)
                    tmp_mask_map = np.zeros((mesh.graph['H'], mesh.graph['W'])).astype(np.bool)
                    tmp_invalid_context_map = np.zeros((mesh.graph['H'], mesh.graph['W'])).astype(np.bool)
                    for node in tmp_mask_nodes:
                        tmp_mask_map[node[0], node[1]] = True
                    for node in context_ccs[ecnt_id]:
                        tmp_context_map[node[0], node[1]] = True
                    for node in erode_context_ccs[ecnt_id]:
                        tmp_context_map[node[0], node[1]] = True
                    for node in extend_context_ccs[ecnt_id]:
                        tmp_context_map[node[0], node[1]] = True
                    for node in invalid_extend_edge_ccs[ecnt_id]:
                        tmp_invalid_context_map[node[0], node[1]] = True
                    init_invalid_context_map = tmp_invalid_context_map.copy()
                    init_context_map = tmp
                    if (tmp_mask_map.astype(np.uint8) * tmp_context_map.astype(np.uint8)).max() > 0:
                        import pdb; pdb.set_trace()
                    if vis_edge_id is not None and ecnt_id == vis_edge_id:
                        f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True)
                        ax1.imshow(tmp_context_map * 1); ax2.imshow(init_invalid_context_map * 1 + tmp_context_map * 2)
                        plt.show()
                        import pdb; pdb.set_trace()
                else:
                    tmp_context_nodes = new_tmp_context_nodes
                    new_tmp_context_nodes = None
                    tmp_mask_nodes = new_tmp_mask_nodes
                    new_tmp_mask_nodes = None
                    tmp_invalid_context_nodes = new_tmp_invalid_context_nodes
                    new_tmp_invalid_context_nodes = None
                new_tmp_context_nodes = None
                new_tmp_context_nodes = []
                new_tmp_invalid_context_nodes = None
                new_tmp_invalid_context_nodes = []
                new_tmp_mask_nodes = set([])
                for node in tmp_context_nodes:
                    for ne in mesh.neighbors(node):
                        if ne in constraint_context_cc and \
                            bool(tmp_mask_map[ne[0], ne[1]]) is False and \
                            bool(tmp_context_map[ne[0], ne[1]]) is False and \
                            bool(forbidden_map[ne[0], ne[1]]) is True:
                            new_tmp_context_nodes.append(ne)
                            tmp_context_map[ne[0], ne[1]] = True
                accum_context_cc.extend(new_tmp_context_nodes)
                for node in tmp_invalid_context_nodes:
                    for ne in mesh.neighbors(node):
                        if bool(tmp_mask_map[ne[0], ne[1]]) is False and \
                           bool(tmp_context_map[ne[0], ne[1]]) is False and \
                           bool(tmp_invalid_context_map[ne[0], ne[1]]) is False and \
                           bool(forbidden_map[ne[0], ne[1]]) is True:
                            tmp_invalid_context_map[ne[0], ne[1]] = True
                            new_tmp_invalid_context_nodes.append(ne)
                for node in tmp_mask_nodes:
                    for ne in mesh.neighbors(node):
                        if bool(tmp_mask_map[ne[0], ne[1]]) is False and \
                           bool(tmp_context_map[ne[0], ne[1]]) is False and \
                           bool(tmp_invalid_context_map[ne[0], ne[1]]) is False and \
                           bool(forbidden_map[ne[0], ne[1]]) is True:
                            new_tmp_mask_nodes.add(ne)
                            tmp_mask_map[ne[0], ne[1]] = True
            init_invalid_context_map[tmp_context_map] = False
            _, tmp_label_map = cv2.connectedComponents((init_invalid_context_map | tmp_context_map).astype(np.uint8), connectivity=8)
            tmp_label_ids = set(np.unique(tmp_label_map[init_invalid_context_map]))
            if (tmp_mask_map.astype(np.uint8) * tmp_context_map.astype(np.uint8)).max() > 0:
                import pdb; pdb.set_trace()
            if vis_edge_id is not None and ecnt_id == vis_edge_id:
                f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True)
                ax1.imshow(tmp_label_map); ax2.imshow(init_invalid_context_map * 1 + tmp_context_map * 2)
                plt.show()
                import pdb; pdb.set_trace()
            extend_context_ccs[ecnt_id] |= set(accum_context_cc)
            extend_context_ccs[ecnt_id] = extend_context_ccs[ecnt_id] - mask_ccs[ecnt_id]
            extend_erode_context_ccs[ecnt_id] = extend_context_ccs[ecnt_id] & constraint_erode_context_cc
            extend_context_ccs[ecnt_id] = extend_context_ccs[ecnt_id] - extend_erode_context_ccs[ecnt_id] - erode_context_ccs[ecnt_id]
            tmp_context_cc = context_ccs[ecnt_id] - extend_erode_context_ccs[ecnt_id] - erode_context_ccs[ecnt_id]
            if len(tmp_context_cc) > 0:
                context_ccs[ecnt_id] = tmp_context_cc
            tmp_mask_cc = tmp_mask_cc - context_ccs[ecnt_id] - erode_context_ccs[ecnt_id]
            mask_ccs[ecnt_id] = mask_ccs[ecnt_id] | tmp_mask_cc

    return context_ccs, mask_ccs, broken_mask_ccs, edge_ccs, erode_context_ccs, invalid_extend_edge_ccs, edge_maps, extend_context_ccs, extend_edge_ccs, extend_erode_context_ccs

def DL_inpaint_edge(mesh,
                    info_on_pix,
                    config,
                    image,
                    depth,
                    context_ccs,
                    erode_context_ccs,
                    extend_context_ccs,
                    extend_erode_context_ccs,
                    mask_ccs,
                    broken_mask_ccs,
                    edge_ccs,
                    extend_edge_ccs,
                    init_mask_connect,
                    edge_maps,
                    rgb_model=None,
                    depth_edge_model=None,
                    depth_edge_model_init=None,
                    depth_feat_model=None,
                    specific_edge_id=-1,
                    specific_edge_loc=None,
                    inpaint_iter=0):

    if isinstance(config["gpu_ids"], int) and (config["gpu_ids"] >= 0):
        device = config["gpu_ids"]
    else:
        device = "cpu"

    edge_map = np.zeros_like(depth)
    new_edge_ccs = [set() for _ in range(len(edge_ccs))]
    edge_maps_with_id = edge_maps
    edge_condition = lambda x, m: m.nodes[x].get('far') is not None and len(m.nodes[x].get('far')) > 0
    edge_map = get_map_from_ccs(edge_ccs, mesh.graph['H'], mesh.graph['W'], mesh, edge_condition)
    np_depth, np_image = depth.copy(), image.copy()
    image_c = image.shape[-1]
    image = torch.FloatTensor(image.transpose(2, 0, 1)).unsqueeze(0).to(device)
    if depth.ndim < 3:
        depth = depth[..., None]
    depth = torch.FloatTensor(depth.transpose(2, 0, 1)).unsqueeze(0).to(device)
    mesh.graph['max_edge_id'] = len(edge_ccs)
    connnect_points_ccs = [set() for _ in range(len(edge_ccs))]
    gp_time, tmp_mesh_time, bilateral_time = 0, 0, 0
    edges_infos = dict()
    edges_in_mask = [set() for _ in range(len(edge_ccs))]
    tmp_specific_edge_id = []
    for edge_id, (context_cc, mask_cc, erode_context_cc, extend_context_cc, edge_cc) in enumerate(zip(context_ccs, mask_ccs, erode_context_ccs, extend_context_ccs, edge_ccs)):
        if len(specific_edge_id) > 0:
            if edge_id not in specific_edge_id:
                continue
        if len(context_cc) < 1 or len(mask_cc) < 1:
            continue
        edge_dict = get_edge_from_nodes(context_cc | extend_context_cc, erode_context_cc | extend_erode_context_ccs[edge_id], mask_cc, edge_cc, extend_edge_ccs[edge_id],
                                        mesh.graph['H'], mesh.graph['W'], mesh)
        edge_dict['edge'], end_depth_maps, _ = \
            filter_irrelevant_edge_new(edge_dict['self_edge'], edge_dict['comp_edge'],
                                    edge_map,
                                    edge_maps_with_id,
                                    edge_id,
                                    edge_dict['context'],
                                    edge_dict['depth'], mesh, context_cc | erode_context_cc | extend_context_cc | extend_erode_context_ccs[edge_id], spdb=False)
        if specific_edge_loc is not None and \
            (specific_edge_loc is not None and edge_dict['mask'][specific_edge_loc[0], specific_edge_loc[1]] == 0):
            continue
        mask_size = get_valid_size(edge_dict['mask'])
        mask_size = dilate_valid_size(mask_size, edge_dict['mask'], dilate=[20, 20])
        context_size = get_valid_size(edge_dict['context'])
        context_size = dilate_valid_size(context_size, edge_dict['context'], dilate=[20, 20])
        union_size = size_operation(mask_size, context_size, operation='+')
        patch_edge_dict = dict()
        patch_edge_dict['mask'], patch_edge_dict['context'], patch_edge_dict['rgb'], \
            patch_edge_dict['disp'], patch_edge_dict['edge'] = \
            crop_maps_by_size(union_size, edge_dict['mask'], edge_dict['context'],
                                edge_dict['rgb'], edge_dict['disp'], edge_dict['edge'])
        x_anchor, y_anchor = [union_size['x_min'], union_size['x_max']], [union_size['y_min'], union_size['y_max']]
        tensor_edge_dict = convert2tensor(patch_edge_dict)
        input_edge_feat = torch.cat((tensor_edge_dict['rgb'],
                                        tensor_edge_dict['disp'],
                                        tensor_edge_dict['edge'],
                                        1 - tensor_edge_dict['context'],
                                        tensor_edge_dict['mask']), dim=1)
        if require_depth_edge(patch_edge_dict['edge'], patch_edge_dict['mask']) and inpaint_iter == 0:
            with torch.no_grad():
                depth_edge_output = depth_edge_model.forward_3P(tensor_edge_dict['mask'],
                                                                tensor_edge_dict['context'],
                                                                tensor_edge_dict['rgb'],
                                                                tensor_edge_dict['disp'],
                                                                tensor_edge_dict['edge'],
                                                                unit_length=128,
                                                                cuda=device)
                depth_edge_output = depth_edge_output.cpu()
            tensor_edge_dict['output'] = (depth_edge_output> config['ext_edge_threshold']).float() * tensor_edge_dict['mask'] + tensor_edge_dict['edge']
        else:
            tensor_edge_dict['output'] = tensor_edge_dict['edge']
            depth_edge_output = tensor_edge_dict['edge'] + 0
        patch_edge_dict['output'] = tensor_edge_dict['output'].squeeze().data.cpu().numpy()
        edge_dict['output'] = np.zeros((mesh.graph['H'], mesh.graph['W']))
        edge_dict['output'][union_size['x_min']:union_size['x_max'], union_size['y_min']:union_size['y_max']] = \
            patch_edge_dict['output']
        if require_depth_edge(patch_edge_dict['edge'], patch_edge_dict['mask']) and inpaint_iter == 0:
            if ((depth_edge_output> config['ext_edge_threshold']).float() * tensor_edge_dict['mask']).max() > 0:
                try:
                    edge_dict['fpath_map'], edge_dict['npath_map'], break_flag, npaths, fpaths, invalid_edge_id = \
                        clean_far_edge_new(edge_dict['output'], end_depth_maps, edge_dict['mask'], edge_dict['context'], mesh, info_on_pix, edge_dict['self_edge'], inpaint_iter, config)
                except:
                    import pdb; pdb.set_trace()
                pre_npath_map = edge_dict['npath_map'].copy()
                if config.get('repeat_inpaint_edge') is True:
                    for _ in range(2):
                        tmp_input_edge = ((edge_dict['npath_map'] > -1) + edge_dict['edge']).clip(0, 1)
                        patch_tmp_input_edge = crop_maps_by_size(union_size, tmp_input_edge)[0]
                        tensor_input_edge = torch.FloatTensor(patch_tmp_input_edge)[None, None, ...]
                        depth_edge_output = depth_edge_model.forward_3P(tensor_edge_dict['mask'],
                                                    tensor_edge_dict['context'],
                                                    tensor_edge_dict['rgb'],
                                                    tensor_edge_dict['disp'],
                                                    tensor_input_edge,
                                                    unit_length=128,
                                                    cuda=device)
                        depth_edge_output = depth_edge_output.cpu()
                        depth_edge_output = (depth_edge_output> config['ext_edge_threshold']).float() * tensor_edge_dict['mask'] + tensor_edge_dict['edge']
                        depth_edge_output = depth_edge_output.squeeze().data.cpu().numpy()
                        full_depth_edge_output = np.zeros((mesh.graph['H'], mesh.graph['W']))
                        full_depth_edge_output[union_size['x_min']:union_size['x_max'], union_size['y_min']:union_size['y_max']] = \
                            depth_edge_output
                        edge_dict['fpath_map'], edge_dict['npath_map'], break_flag, npaths, fpaths, invalid_edge_id = \
                            clean_far_edge_new(full_depth_edge_output, end_depth_maps, edge_dict['mask'], edge_dict['context'], mesh, info_on_pix, edge_dict['self_edge'], inpaint_iter, config)
                for nid in npaths.keys():
                    npath, fpath = npaths[nid], fpaths[nid]
                    start_mx, start_my, end_mx, end_my = -1, -1, -1, -1
                    if end_depth_maps[npath[0][0], npath[0][1]] != 0:
                        start_mx, start_my = npath[0][0], npath[0][1]
                    if end_depth_maps[npath[-1][0], npath[-1][1]] != 0:
                        end_mx, end_my = npath[-1][0], npath[-1][1]
                    if start_mx == -1:
                        import pdb; pdb.set_trace()
                    valid_end_pt = () if end_mx == -1 else (end_mx, end_my, info_on_pix[(end_mx, end_my)][0]['depth'])
                    new_edge_info = dict(fpath=fpath,
                                         npath=npath,
                                         cont_end_pts=valid_end_pt,
                                         mask_id=edge_id,
                                         comp_edge_id=nid,
                                         depth=end_depth_maps[start_mx, start_my])
                    if edges_infos.get((start_mx, start_my)) is None:
                        edges_infos[(start_mx, start_my)] = []
                    edges_infos[(start_mx, start_my)].append(new_edge_info)
                    edges_in_mask[edge_id].add((start_mx, start_my))
                    if len(valid_end_pt) > 0:
                        new_edge_info = dict(fpath=fpath[::-1],
                                             npath=npath[::-1],
                                             cont_end_pts=(start_mx, start_my, info_on_pix[(start_mx, start_my)][0]['depth']),
                                             mask_id=edge_id,
                                             comp_edge_id=nid,
                                             depth=end_depth_maps[end_mx, end_my])
                        if edges_infos.get((end_mx, end_my)) is None:
                            edges_infos[(end_mx, end_my)] = []
                        edges_infos[(end_mx, end_my)].append(new_edge_info)
                        edges_in_mask[edge_id].add((end_mx, end_my))
    for edge_id, (context_cc, mask_cc, erode_context_cc, extend_context_cc, edge_cc) in enumerate(zip(context_ccs, mask_ccs, erode_context_ccs, extend_context_ccs, edge_ccs)):
        if len(specific_edge_id) > 0:
            if edge_id not in specific_edge_id:
                continue
        if len(context_cc) < 1 or len(mask_cc) < 1:
            continue
        edge_dict = get_edge_from_nodes(context_cc | extend_context_cc, erode_context_cc | extend_erode_context_ccs[edge_id], mask_cc, edge_cc, extend_edge_ccs[edge_id],
                                        mesh.graph['H'], mesh.graph['W'], mesh)
        if specific_edge_loc is not None and \
            (specific_edge_loc is not None and edge_dict['mask'][specific_edge_loc[0], specific_edge_loc[1]] == 0):
            continue
        else:
            tmp_specific_edge_id.append(edge_id)
        edge_dict['edge'], end_depth_maps, _ = \
            filter_irrelevant_edge_new(edge_dict['self_edge'], edge_dict['comp_edge'],
                                    edge_map,
                                    edge_maps_with_id,
                                    edge_id,
                                    edge_dict['context'],
                                    edge_dict['depth'], mesh, context_cc | erode_context_cc | extend_context_cc | extend_erode_context_ccs[edge_id], spdb=False)
        discard_map = np.zeros_like(edge_dict['edge'])
        mask_size = get_valid_size(edge_dict['mask'])
        mask_size = dilate_valid_size(mask_size, edge_dict['mask'], dilate=[20, 20])
        context_size = get_valid_size(edge_dict['context'])
        context_size = dilate_valid_size(context_size, edge_dict['context'], dilate=[20, 20])
        union_size = size_operation(mask_size, context_size, operation='+')
        patch_edge_dict = dict()
        patch_edge_dict['mask'], patch_edge_dict['context'], patch_edge_dict['rgb'], \
            patch_edge_dict['disp'], patch_edge_dict['edge'] = \
            crop_maps_by_size(union_size, edge_dict['mask'], edge_dict['context'],
                                edge_dict['rgb'], edge_dict['disp'], edge_dict['edge'])
        x_anchor, y_anchor = [union_size['x_min'], union_size['x_max']], [union_size['y_min'], union_size['y_max']]
        tensor_edge_dict = convert2tensor(patch_edge_dict)
        input_edge_feat = torch.cat((tensor_edge_dict['rgb'],
                                        tensor_edge_dict['disp'],
                                        tensor_edge_dict['edge'],
                                        1 - tensor_edge_dict['context'],
                                        tensor_edge_dict['mask']), dim=1)
        edge_dict['output'] = edge_dict['edge'].copy()

        if require_depth_edge(patch_edge_dict['edge'], patch_edge_dict['mask']) and inpaint_iter == 0:
            edge_dict['fpath_map'], edge_dict['npath_map'] = edge_dict['fpath_map'] * 0 - 1, edge_dict['npath_map'] * 0 - 1
            end_pts = edges_in_mask[edge_id]
            for end_pt in end_pts:
                cur_edge_infos = edges_infos[(end_pt[0], end_pt[1])]
                cur_info = [xx for xx in cur_edge_infos if xx['mask_id'] == edge_id][0]
                other_infos = [xx for xx in cur_edge_infos if xx['mask_id'] != edge_id and len(xx['cont_end_pts']) > 0]
                if len(cur_info['cont_end_pts']) > 0 or (len(cur_info['cont_end_pts']) == 0 and len(other_infos) == 0):
                    for fnode in cur_info['fpath']:
                        edge_dict['fpath_map'][fnode[0], fnode[1]] = cur_info['comp_edge_id']
                    for fnode in cur_info['npath']:
                        edge_dict['npath_map'][fnode[0], fnode[1]] = cur_info['comp_edge_id']
            fnmap = edge_dict['fpath_map'] * 1
            fnmap[edge_dict['npath_map'] != -1] = edge_dict['npath_map'][edge_dict['npath_map'] != -1]
            for end_pt in end_pts:
                cur_edge_infos = edges_infos[(end_pt[0], end_pt[1])]
                cur_info = [xx for xx in cur_edge_infos if xx['mask_id'] == edge_id][0]
                cur_depth = cur_info['depth']
                other_infos = [xx for xx in cur_edge_infos if xx['mask_id'] != edge_id and len(xx['cont_end_pts']) > 0]
                comp_edge_id = cur_info['comp_edge_id']
                if len(cur_info['cont_end_pts']) == 0 and len(other_infos) > 0:
                    other_infos = sorted(other_infos, key=lambda aa: abs(abs(aa['cont_end_pts'][2]) - abs(cur_depth)))
                    for other_info in other_infos:
                        tmp_fmap, tmp_nmap = np.zeros((mesh.graph['H'], mesh.graph['W'])) - 1, np.zeros((mesh.graph['H'], mesh.graph['W'])) - 1
                        for fnode in other_info['fpath']:
                            if fnmap[fnode[0], fnode[1]] != -1:
                                tmp_fmap = tmp_fmap * 0 - 1
                                break
                            else:
                                tmp_fmap[fnode[0], fnode[1]] = comp_edge_id
                        if fnmap[fnode[0], fnode[1]] != -1:
                            continue
                        for fnode in other_info['npath']:
                            if fnmap[fnode[0], fnode[1]] != -1:
                                tmp_nmap = tmp_nmap * 0 - 1
                                break
                            else:
                                tmp_nmap[fnode[0], fnode[1]] = comp_edge_id
                        if fnmap[fnode[0], fnode[1]] != -1:
                            continue
                        break
                    if min(tmp_fmap.max(), tmp_nmap.max()) != -1:
                        edge_dict['fpath_map'] = tmp_fmap
                        edge_dict['fpath_map'][edge_dict['valid_area'] == 0] = -1
                        edge_dict['npath_map'] = tmp_nmap
                        edge_dict['npath_map'][edge_dict['valid_area'] == 0] = -1
                        discard_map = ((tmp_nmap != -1).astype(np.uint8) + (tmp_fmap != -1).astype(np.uint8)) * edge_dict['mask']
                    else:
                        for fnode in cur_info['fpath']:
                            edge_dict['fpath_map'][fnode[0], fnode[1]] = cur_info['comp_edge_id']
                        for fnode in cur_info['npath']:
                            edge_dict['npath_map'][fnode[0], fnode[1]] = cur_info['comp_edge_id']
            if edge_dict['npath_map'].min() == 0 or edge_dict['fpath_map'].min() == 0:
                import pdb; pdb.set_trace()
            edge_dict['output'] = (edge_dict['npath_map'] > -1) * edge_dict['mask'] + edge_dict['context'] * edge_dict['edge']
        mesh, _, _, _ = create_placeholder(edge_dict['context'], edge_dict['mask'],
                                  edge_dict['depth'], edge_dict['fpath_map'],
                                  edge_dict['npath_map'], mesh, inpaint_iter,
                                  edge_ccs,
                                  extend_edge_ccs[edge_id],
                                  edge_maps_with_id,
                                  edge_id)

        dxs, dys = np.where(discard_map != 0)
        for dx, dy in zip(dxs, dys):
            mesh.nodes[(dx, dy)]['inpaint_twice'] = False
        depth_dict = depth_inpainting(context_cc, extend_context_cc, erode_context_cc | extend_erode_context_ccs[edge_id], mask_cc, mesh, config, union_size, depth_feat_model, edge_dict['output'])
        refine_depth_output = depth_dict['output']*depth_dict['mask']
        for near_id in np.unique(edge_dict['npath_map'])[1:]:
            refine_depth_output = refine_depth_around_edge(refine_depth_output.copy(),
                                                            (edge_dict['fpath_map'] == near_id).astype(np.uint8) * edge_dict['mask'],
                                                            (edge_dict['fpath_map'] == near_id).astype(np.uint8),
                                                            (edge_dict['npath_map'] == near_id).astype(np.uint8) * edge_dict['mask'],
                                                            depth_dict['mask'].copy(),
                                                            depth_dict['output'] * depth_dict['context'],
                                                            config)
        depth_dict['output'][depth_dict['mask'] > 0] = refine_depth_output[depth_dict['mask'] > 0]
        rgb_dict = get_rgb_from_nodes(context_cc | extend_context_cc,
                                      erode_context_cc | extend_erode_context_ccs[edge_id], mask_cc, mesh.graph['H'], mesh.graph['W'], mesh)
        if np.all(rgb_dict['mask'] == edge_dict['mask']) is False:
            import pdb; pdb.set_trace()
        rgb_dict['edge'] = edge_dict['output']
        patch_rgb_dict = dict()
        patch_rgb_dict['mask'], patch_rgb_dict['context'], patch_rgb_dict['rgb'], \
            patch_rgb_dict['edge'] = crop_maps_by_size(union_size, rgb_dict['mask'],
                                                        rgb_dict['context'], rgb_dict['rgb'],
                                                        rgb_dict['edge'])
        tensor_rgb_dict = convert2tensor(patch_rgb_dict)
        resize_rgb_dict = {k: v.clone() for k, v in tensor_rgb_dict.items()}
        max_hw = np.array([*patch_rgb_dict['mask'].shape[-2:]]).max()
        init_frac = config['largest_size'] / (np.array([*patch_rgb_dict['mask'].shape[-2:]]).prod() ** 0.5)
        resize_hw = [patch_rgb_dict['mask'].shape[-2] * init_frac, patch_rgb_dict['mask'].shape[-1] * init_frac]
        resize_max_hw = max(resize_hw)
        frac = (np.floor(resize_max_hw / 128.) * 128.) / max_hw
        if frac < 1:
            resize_mark = torch.nn.functional.interpolate(torch.cat((resize_rgb_dict['mask'],
                                                            resize_rgb_dict['context']),
                                                            dim=1),
                                                            scale_factor=frac,
                                                            mode='area')
            resize_rgb_dict['mask'] = (resize_mark[:, 0:1] > 0).float()
            resize_rgb_dict['context'] = (resize_mark[:, 1:2] == 1).float()
            resize_rgb_dict['context'][resize_rgb_dict['mask'] > 0] = 0
            resize_rgb_dict['rgb'] = torch.nn.functional.interpolate(resize_rgb_dict['rgb'],
                                                                        scale_factor=frac,
                                                                        mode='area')
            resize_rgb_dict['rgb'] = resize_rgb_dict['rgb'] * resize_rgb_dict['context']
            resize_rgb_dict['edge'] = torch.nn.functional.interpolate(resize_rgb_dict['edge'],
                                                                        scale_factor=frac,
                                                                        mode='area')
            resize_rgb_dict['edge'] = (resize_rgb_dict['edge'] > 0).float() * 0
            resize_rgb_dict['edge'] = resize_rgb_dict['edge'] * (resize_rgb_dict['context'] + resize_rgb_dict['mask'])
        rgb_input_feat = torch.cat((resize_rgb_dict['rgb'], resize_rgb_dict['edge']), dim=1)
        rgb_input_feat[:, 3] = 1 - rgb_input_feat[:, 3]
        resize_mask = open_small_mask(resize_rgb_dict['mask'], resize_rgb_dict['context'], 3, 41)
        specified_hole = resize_mask
        with torch.no_grad():
            rgb_output = rgb_model.forward_3P(specified_hole,
                                            resize_rgb_dict['context'],
                                            resize_rgb_dict['rgb'],
                                            resize_rgb_dict['edge'],
                                            unit_length=128,
                                            cuda=device)
            rgb_output = rgb_output.cpu()
            if config.get('gray_image') is True:
                rgb_output = rgb_output.mean(1, keepdim=True).repeat((1,3,1,1))
            rgb_output = rgb_output.cpu()
        resize_rgb_dict['output'] = rgb_output * resize_rgb_dict['mask'] + resize_rgb_dict['rgb']
        tensor_rgb_dict['output'] = resize_rgb_dict['output']
        if frac < 1:
            tensor_rgb_dict['output'] = torch.nn.functional.interpolate(tensor_rgb_dict['output'],
                                                                        size=tensor_rgb_dict['mask'].shape[-2:],
                                                                        mode='bicubic')
            tensor_rgb_dict['output'] = tensor_rgb_dict['output'] * \
                                         tensor_rgb_dict['mask'] + (tensor_rgb_dict['rgb'] * tensor_rgb_dict['context'])
        patch_rgb_dict['output'] = tensor_rgb_dict['output'].data.cpu().numpy().squeeze().transpose(1,2,0)
        rgb_dict['output'] = np.zeros((mesh.graph['H'], mesh.graph['W'], 3))
        rgb_dict['output'][union_size['x_min']:union_size['x_max'], union_size['y_min']:union_size['y_max']] = \
            patch_rgb_dict['output']

        if require_depth_edge(patch_edge_dict['edge'], patch_edge_dict['mask']) or inpaint_iter > 0:
            edge_occlusion = True
        else:
            edge_occlusion = False
        for node in erode_context_cc:
            if rgb_dict['mask'][node[0], node[1]] > 0:
                for info in info_on_pix[(node[0], node[1])]:
                    if abs(info['depth']) == abs(node[2]):
                        info['update_color'] = (rgb_dict['output'][node[0], node[1]] * 255).astype(np.uint8)
        if frac < 1.:
            depth_edge_dilate_2_color_flag = False
        else:
            depth_edge_dilate_2_color_flag = True
        hxs, hys = np.where((rgb_dict['mask'] > 0) & (rgb_dict['erode'] == 0))
        for hx, hy in zip(hxs, hys):
            real_depth = None
            if abs(depth_dict['output'][hx, hy]) <= abs(np_depth[hx, hy]):
                depth_dict['output'][hx, hy] = np_depth[hx, hy] + 0.01
            node = (hx, hy, -depth_dict['output'][hx, hy])
            if info_on_pix.get((node[0], node[1])) is not None:
                for info in info_on_pix.get((node[0], node[1])):
                    if info.get('inpaint_id') is None or abs(info['inpaint_id'] < mesh.nodes[(hx, hy)]['inpaint_id']):
                        pre_depth = info['depth'] if info.get('real_depth') is None else info['real_depth']
                        if abs(node[2]) < abs(pre_depth):
                            node = (node[0], node[1], -(abs(pre_depth) + 0.001))
            if mesh.has_node(node):
                real_depth = node[2]
            while True:
                if mesh.has_node(node):
                    node = (node[0], node[1], -(abs(node[2]) + 0.001))
                else:
                    break
            if real_depth == node[2]:
                real_depth = None
            cur_disp = 1./node[2]
            if not(mesh.has_node(node)):
                if not mesh.has_node((node[0], node[1])):
                    print("2D node not found.")
                    import pdb; pdb.set_trace()
                if inpaint_iter == 1:
                    paint = (rgb_dict['output'][hx, hy] * 255).astype(np.uint8)
                else:
                    paint = (rgb_dict['output'][hx, hy] * 255).astype(np.uint8)
                ndict = dict(color=paint,
                                synthesis=True,
                                disp=cur_disp,
                                cc_id=set([edge_id]),
                                overlap_number=1.0,
                                refine_depth=False,
                                edge_occlusion=edge_occlusion,
                                depth_edge_dilate_2_color_flag=depth_edge_dilate_2_color_flag,
                                real_depth=real_depth)
                mesh, _, _ = refresh_node((node[0], node[1]), mesh.nodes[(node[0], node[1])], node, ndict, mesh, stime=True)
                if inpaint_iter == 0 and mesh.degree(node) < 4:
                    connnect_points_ccs[edge_id].add(node)
            if info_on_pix.get((hx, hy)) is None:
                info_on_pix[(hx, hy)] = []
            new_info = {'depth':node[2],
                        'color': paint,
                        'synthesis':True,
                        'disp':cur_disp,
                        'cc_id':set([edge_id]),
                        'inpaint_id':inpaint_iter + 1,
                        'edge_occlusion':edge_occlusion,
                        'overlap_number':1.0,
                        'real_depth': real_depth}
            info_on_pix[(hx, hy)].append(new_info)
    specific_edge_id = tmp_specific_edge_id
    for erode_id, erode_context_cc in enumerate(erode_context_ccs):
        if len(specific_edge_id) > 0 and erode_id not in specific_edge_id:
            continue
        for erode_node in erode_context_cc:
            for info in info_on_pix[(erode_node[0], erode_node[1])]:
                if info['depth'] == erode_node[2]:
                    info['color'] = info['update_color']
                    mesh.nodes[erode_node]['color'] = info['update_color']
                    np_image[(erode_node[0], erode_node[1])] = info['update_color']
    new_edge_ccs = [set() for _ in range(mesh.graph['max_edge_id'] + 1)]
    for node in mesh.nodes:
        if len(node) == 2:
            mesh.remove_node(node)
            continue
        if mesh.nodes[node].get('edge_id') is not None and mesh.nodes[node].get('inpaint_id') == inpaint_iter + 1:
            if mesh.nodes[node].get('inpaint_twice') is False:
                continue
            try:
                new_edge_ccs[mesh.nodes[node].get('edge_id')].add(node)
            except:
                import pdb; pdb.set_trace()
    specific_mask_nodes = None
    if inpaint_iter == 0:
        mesh, info_on_pix = refine_color_around_edge(mesh, info_on_pix, new_edge_ccs, config, False)

    return mesh, info_on_pix, specific_mask_nodes, new_edge_ccs, connnect_points_ccs, np_image


def write_ply(image,
              depth,
              int_mtx,
              ply_name,
              config,
              rgb_model,
              depth_edge_model,
              depth_edge_model_init,
              depth_feat_model):
    depth = depth.astype(np.float64)
    input_mesh, xy2depth, image, depth = create_mesh(depth, image, int_mtx, config)

    H, W = input_mesh.graph['H'], input_mesh.graph['W']
    input_mesh = tear_edges(input_mesh, config['depth_threshold'], xy2depth)
    input_mesh, info_on_pix = generate_init_node(input_mesh, config, min_node_in_cc=200)
    edge_ccs, input_mesh, edge_mesh = group_edges(input_mesh, config, image, remove_conflict_ordinal=False)
    edge_canvas = np.zeros((H, W)) - 1

    input_mesh, info_on_pix, depth = reassign_floating_island(input_mesh, info_on_pix, image, depth)
    input_mesh = update_status(input_mesh, info_on_pix)
    specific_edge_id = []
    edge_ccs, input_mesh, edge_mesh = group_edges(input_mesh, config, image, remove_conflict_ordinal=True)
    pre_depth = depth.copy()
    input_mesh, info_on_pix, edge_mesh, depth, aft_mark = remove_dangling(input_mesh, edge_ccs, edge_mesh, info_on_pix, image, depth, config)

    input_mesh, depth, info_on_pix = update_status(input_mesh, info_on_pix, depth)
    edge_ccs, input_mesh, edge_mesh = group_edges(input_mesh, config, image, remove_conflict_ordinal=True)
    edge_canvas = np.zeros((H, W)) - 1

    mesh, info_on_pix, depth = fill_missing_node(input_mesh, info_on_pix, image, depth)
    if config['extrapolate_border'] is True:
        pre_depth = depth.copy()
        input_mesh, info_on_pix, depth = refresh_bord_depth(input_mesh, info_on_pix, image, depth)
        input_mesh = remove_node_feat(input_mesh, 'edge_id')
        aft_depth = depth.copy()
        input_mesh, info_on_pix, depth, image = enlarge_border(input_mesh, info_on_pix, depth, image, config)
        noext_H, noext_W = H, W
        H, W = image.shape[:2]
        input_mesh, info_on_pix = fill_dummy_bord(input_mesh, info_on_pix, image, depth, config)
        edge_ccs, input_mesh, edge_mesh = \
            group_edges(input_mesh, config, image, remove_conflict_ordinal=True)
        input_mesh = combine_end_node(input_mesh, edge_mesh, edge_ccs, depth)
        input_mesh, depth, info_on_pix = update_status(input_mesh, info_on_pix, depth)
        edge_ccs, input_mesh, edge_mesh = \
            group_edges(input_mesh, config, image, remove_conflict_ordinal=True, spdb=False)
        input_mesh = remove_redundant_edge(input_mesh, edge_mesh, edge_ccs, info_on_pix, config, redundant_number=config['redundant_number'], spdb=False)
        input_mesh, depth, info_on_pix = update_status(input_mesh, info_on_pix, depth)
        edge_ccs, input_mesh, edge_mesh = group_edges(input_mesh, config, image, remove_conflict_ordinal=True)
        input_mesh = combine_end_node(input_mesh, edge_mesh, edge_ccs, depth)
        input_mesh = remove_redundant_edge(input_mesh, edge_mesh, edge_ccs, info_on_pix, config, redundant_number=config['redundant_number'], invalid=True, spdb=False)
        input_mesh, depth, info_on_pix = update_status(input_mesh, info_on_pix, depth)
        edge_ccs, input_mesh, edge_mesh = group_edges(input_mesh, config, image, remove_conflict_ordinal=True)
        input_mesh = combine_end_node(input_mesh, edge_mesh, edge_ccs, depth)
        input_mesh, depth, info_on_pix = update_status(input_mesh, info_on_pix, depth)
        edge_ccs, input_mesh, edge_mesh = group_edges(input_mesh, config, image, remove_conflict_ordinal=True)
        edge_condition = lambda x, m: m.nodes[x].get('far') is not None and len(m.nodes[x].get('far')) > 0
        edge_map = get_map_from_ccs(edge_ccs, input_mesh.graph['H'], input_mesh.graph['W'], input_mesh, edge_condition)
        other_edge_with_id = get_map_from_ccs(edge_ccs, input_mesh.graph['H'], input_mesh.graph['W'], real_id=True)
        info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs,
                                                depth_edge_model, depth_feat_model, rgb_model, config, direc="up")
        info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs,
                                                depth_edge_model, depth_feat_model, rgb_model, config, direc="left")
        info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs,
                                                depth_edge_model, depth_feat_model, rgb_model, config, direc="down")
        info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs,
                                                depth_edge_model, depth_feat_model, rgb_model, config, direc="right")
        info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs,
                                                depth_edge_model, depth_feat_model, rgb_model, config, direc="right-up")
        info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs,
                                                depth_edge_model, depth_feat_model, rgb_model, config, direc="right-down")
        info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs,
                                                depth_edge_model, depth_feat_model, rgb_model, config, direc="left-up")
        info_on_pix, input_mesh, image, depth, edge_ccs = extrapolate(input_mesh, info_on_pix, image, depth, other_edge_with_id, edge_map, edge_ccs,
                                                depth_edge_model, depth_feat_model, rgb_model, config, direc="left-down")
    specific_edge_loc = None
    specific_edge_id = []
    vis_edge_id = None
    context_ccs, mask_ccs, broken_mask_ccs, edge_ccs, erode_context_ccs, \
        init_mask_connect, edge_maps, extend_context_ccs, extend_edge_ccs, extend_erode_context_ccs = \
                                                                                context_and_holes(input_mesh,
                                                                                            edge_ccs,
                                                                                            config,
                                                                                            specific_edge_id,
                                                                                            specific_edge_loc,
                                                                                            depth_feat_model,
                                                                                            inpaint_iter=0,
                                                                                            vis_edge_id=vis_edge_id)
    edge_canvas = np.zeros((H, W))
    mask = np.zeros((H, W))
    context = np.zeros((H, W))
    vis_edge_ccs = filter_edge(input_mesh, edge_ccs, config)
    edge_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W'])) - 1
    specific_edge_loc = None
    FG_edge_maps = edge_maps.copy()
    edge_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W'])) - 1
    # for cc_id, cc in enumerate(edge_ccs):
    #     for node in cc:
    #         edge_canvas[node[0], node[1]] = cc_id
    # f, ((ax0, ax1, ax2)) = plt.subplots(1, 3, sharex=True, sharey=True); ax0.imshow(1./depth); ax1.imshow(image); ax2.imshow(edge_canvas); plt.show()
    input_mesh, info_on_pix, specific_edge_nodes, new_edge_ccs, connect_points_ccs, image = DL_inpaint_edge(input_mesh,
                                                                                                            info_on_pix,
                                                                                                            config,
                                                                                                            image,
                                                                                                            depth,
                                                                                                            context_ccs,
                                                                                                            erode_context_ccs,
                                                                                                            extend_context_ccs,
                                                                                                            extend_erode_context_ccs,
                                                                                                            mask_ccs,
                                                                                                            broken_mask_ccs,
                                                                                                            edge_ccs,
                                                                                                            extend_edge_ccs,
                                                                                                            init_mask_connect,
                                                                                                            edge_maps,
                                                                                                            rgb_model,
                                                                                                            depth_edge_model,
                                                                                                            depth_edge_model_init,
                                                                                                            depth_feat_model,
                                                                                                            specific_edge_id,
                                                                                                            specific_edge_loc,
                                                                                                            inpaint_iter=0)
    specific_edge_id = []
    edge_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W']))
    connect_points_ccs = [set() for _ in connect_points_ccs]
    context_ccs, mask_ccs, broken_mask_ccs, edge_ccs, erode_context_ccs, init_mask_connect, \
        edge_maps, extend_context_ccs, extend_edge_ccs, extend_erode_context_ccs = \
            context_and_holes(input_mesh, new_edge_ccs, config, specific_edge_id, specific_edge_loc, depth_feat_model, connect_points_ccs, inpaint_iter=1)
    mask_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W']))
    context_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W']))
    erode_context_ccs_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W']))
    edge_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W']))
    # edge_canvas = np.zeros((input_mesh.graph['H'], input_mesh.graph['W'])) - 1
    # for cc_id, cc in enumerate(edge_ccs):
    #     for node in cc:
    #         edge_canvas[node[0], node[1]] = cc_id
    specific_edge_id = []
    input_mesh, info_on_pix, specific_edge_nodes, new_edge_ccs, _, image = DL_inpaint_edge(input_mesh,
                                                                                    info_on_pix,
                                                                                    config,
                                                                                    image,
                                                                                    depth,
                                                                                    context_ccs,
                                                                                    erode_context_ccs,
                                                                                    extend_context_ccs,
                                                                                    extend_erode_context_ccs,
                                                                                    mask_ccs,
                                                                                    broken_mask_ccs,
                                                                                    edge_ccs,
                                                                                    extend_edge_ccs,
                                                                                    init_mask_connect,
                                                                                    edge_maps,
                                                                                    rgb_model,
                                                                                    depth_edge_model,
                                                                                    depth_edge_model_init,
                                                                                    depth_feat_model,
                                                                                    specific_edge_id,
                                                                                    specific_edge_loc,
                                                                                    inpaint_iter=1)
    vertex_id = 0
    input_mesh.graph['H'], input_mesh.graph['W'] = input_mesh.graph['noext_H'], input_mesh.graph['noext_W']
    background_canvas = np.zeros((input_mesh.graph['H'],
                                  input_mesh.graph['W'],
                                  3))
    ply_flag = config.get('save_ply')
    if ply_flag is True:
        node_str_list = []
    else:
        node_str_color = []
        node_str_point = []
    out_fmt = lambda x, x_flag: str(x) if x_flag is True else x
    point_time = 0
    hlight_time = 0
    cur_id_time = 0
    node_str_time = 0
    generate_face_time = 0
    point_list = []
    k_00, k_02, k_11, k_12 = \
        input_mesh.graph['cam_param_pix_inv'][0, 0], input_mesh.graph['cam_param_pix_inv'][0, 2], \
        input_mesh.graph['cam_param_pix_inv'][1, 1], input_mesh.graph['cam_param_pix_inv'][1, 2]
    w_offset = input_mesh.graph['woffset']
    h_offset = input_mesh.graph['hoffset']
    for pix_xy, pix_list in info_on_pix.items():
        for pix_idx, pix_info in enumerate(pix_list):
            pix_depth = pix_info['depth'] if pix_info.get('real_depth') is None else pix_info['real_depth']
            str_pt = [out_fmt(x, ply_flag) for x in reproject_3d_int_detail(pix_xy[0], pix_xy[1], pix_depth,
                      k_00, k_02, k_11, k_12, w_offset, h_offset)]
            if input_mesh.has_node((pix_xy[0], pix_xy[1], pix_info['depth'])) is False:
                return False
                continue
            if pix_info.get('overlap_number') is not None:
                str_color = [out_fmt(x, ply_flag) for x in (pix_info['color']/pix_info['overlap_number']).astype(np.uint8).tolist()]
            else:
                str_color = [out_fmt(x, ply_flag) for x in pix_info['color'].tolist()]
            if pix_info.get('edge_occlusion') is True:
                str_color.append(out_fmt(4, ply_flag))
            else:
                if pix_info.get('inpaint_id') is None:
                    str_color.append(out_fmt(1, ply_flag))
                else:
                    str_color.append(out_fmt(pix_info.get('inpaint_id') + 1, ply_flag))
            if pix_info.get('modified_border') is True or pix_info.get('ext_pixel') is True:
                if len(str_color) == 4:
                    str_color[-1] = out_fmt(5, ply_flag)
                else:
                    str_color.append(out_fmt(5, ply_flag))
            pix_info['cur_id'] = vertex_id
            input_mesh.nodes[(pix_xy[0], pix_xy[1], pix_info['depth'])]['cur_id'] = out_fmt(vertex_id, ply_flag)
            vertex_id += 1
            if ply_flag is True:
                node_str_list.append(' '.join(str_pt) + ' ' + ' '.join(str_color) + '\n')
            else:
                node_str_color.append(str_color)
                node_str_point.append(str_pt)
    str_faces = generate_face(input_mesh, info_on_pix, config)
    if config['save_ply'] is True:
        print("Writing mesh file %s ..." % ply_name)
        with open(ply_name, 'w') as ply_fi:
            ply_fi.write('ply\n' + 'format ascii 1.0\n')
            ply_fi.write('comment H ' + str(int(input_mesh.graph['H'])) + '\n')
            ply_fi.write('comment W ' + str(int(input_mesh.graph['W'])) + '\n')
            ply_fi.write('comment hFov ' + str(float(input_mesh.graph['hFov'])) + '\n')
            ply_fi.write('comment vFov ' + str(float(input_mesh.graph['vFov'])) + '\n')
            ply_fi.write('element vertex ' + str(len(node_str_list)) + '\n')
            ply_fi.write('property float x\n' + \
                         'property float y\n' + \
                         'property float z\n' + \
                         'property uchar red\n' + \
                         'property uchar green\n' + \
                         'property uchar blue\n' + \
                         'property uchar alpha\n')
            ply_fi.write('element face ' + str(len(str_faces)) + '\n')
            ply_fi.write('property list uchar int vertex_index\n')
            ply_fi.write('end_header\n')
            ply_fi.writelines(node_str_list)
            ply_fi.writelines(str_faces)
        ply_fi.close()
        return input_mesh
    else:
        H = int(input_mesh.graph['H'])
        W = int(input_mesh.graph['W'])
        hFov = input_mesh.graph['hFov']
        vFov = input_mesh.graph['vFov']
        node_str_color = np.array(node_str_color).astype(np.float32)
        node_str_color[..., :3] = node_str_color[..., :3] / 255.
        node_str_point = np.array(node_str_point)
        str_faces = np.array(str_faces)

        return node_str_point, node_str_color, str_faces, H, W, hFov, vFov

def read_ply(mesh_fi):
    ply_fi = open(mesh_fi, 'r')
    Height = None
    Width = None
    hFov = None
    vFov = None
    while True:
        line = ply_fi.readline().split('\n')[0]
        if line.startswith('element vertex'):
            num_vertex = int(line.split(' ')[-1])
        elif line.startswith('element face'):
            num_face = int(line.split(' ')[-1])
        elif line.startswith('comment'):
            if line.split(' ')[1] == 'H':
                Height = int(line.split(' ')[-1].split('\n')[0])
            if line.split(' ')[1] == 'W':
                Width = int(line.split(' ')[-1].split('\n')[0])
            if line.split(' ')[1] == 'hFov':
                hFov = float(line.split(' ')[-1].split('\n')[0])
            if line.split(' ')[1] == 'vFov':
                vFov = float(line.split(' ')[-1].split('\n')[0])
        elif line.startswith('end_header'):
            break
    contents = ply_fi.readlines()
    vertex_infos = contents[:num_vertex]
    face_infos = contents[num_vertex:]
    verts = []
    colors = []
    faces = []
    for v_info in vertex_infos:
        str_info = [float(v) for v in v_info.split('\n')[0].split(' ')]
        if len(str_info) == 6:
            vx, vy, vz, r, g, b = str_info
        else:
            vx, vy, vz, r, g, b, hi = str_info
        verts.append([vx, vy, vz])
        colors.append([r, g, b, hi])
    verts = np.array(verts)
    try:
        colors = np.array(colors)
        colors[..., :3] = colors[..., :3]/255.
    except:
        import pdb
        pdb.set_trace()

    for f_info in face_infos:
        _, v1, v2, v3 = [int(f) for f in f_info.split('\n')[0].split(' ')]
        faces.append([v1, v2, v3])
    faces = np.array(faces)


    return verts, colors, faces, Height, Width, hFov, vFov


class Canvas_view():
    def __init__(self,
                 fov,
                 verts,
                 faces,
                 colors,
                 canvas_size,
                 factor=1,
                 bgcolor='gray',
                 proj='perspective',
                 ):
        self.canvas = scene.SceneCanvas(bgcolor=bgcolor, size=(canvas_size*factor, canvas_size*factor))
        self.view = self.canvas.central_widget.add_view()
        self.view.camera = 'perspective'
        self.view.camera.fov = fov
        self.mesh = visuals.Mesh(shading=None)
        self.mesh.attach(Alpha(1.0))
        self.view.add(self.mesh)
        self.tr = self.view.camera.transform
        self.mesh.set_data(vertices=verts, faces=faces, vertex_colors=colors[:, :3])
        self.translate([0,0,0])
        self.rotate(axis=[1,0,0], angle=180)
        self.view_changed()

    def translate(self, trans=[0,0,0]):
        self.tr.translate(trans)

    def rotate(self, axis=[1,0,0], angle=0):
        self.tr.rotate(axis=axis, angle=angle)

    def view_changed(self):
        self.view.camera.view_changed()

    def render(self):
        return self.canvas.render()

    def reinit_mesh(self, verts, faces, colors):
        self.mesh.set_data(vertices=verts, faces=faces, vertex_colors=colors[:, :3])

    def reinit_camera(self, fov):
        self.view.camera.fov = fov
        self.view.camera.view_changed()


def output_3d_photo(verts, colors, faces, Height, Width, hFov, vFov, tgt_poses, video_traj_types, ref_pose,
                    output_dir, ref_image, int_mtx, config, image, videos_poses, video_basename, original_H=None, original_W=None,
                    border=None, depth=None, normal_canvas=None, all_canvas=None, mean_loc_depth=None):

    cam_mesh = netx.Graph()
    cam_mesh.graph['H'] = Height
    cam_mesh.graph['W'] = Width
    cam_mesh.graph['original_H'] = original_H
    cam_mesh.graph['original_W'] = original_W
    int_mtx_real_x = int_mtx[0] * Width
    int_mtx_real_y = int_mtx[1] * Height
    cam_mesh.graph['hFov'] = 2 * np.arctan((1. / 2.) * ((cam_mesh.graph['original_W']) / int_mtx_real_x[0]))
    cam_mesh.graph['vFov'] = 2 * np.arctan((1. / 2.) * ((cam_mesh.graph['original_H']) / int_mtx_real_y[1]))
    colors = colors[..., :3]

    fov_in_rad = max(cam_mesh.graph['vFov'], cam_mesh.graph['hFov'])
    fov = (fov_in_rad * 180 / np.pi)
    print("fov: " + str(fov))
    init_factor = 1
    if config.get('anti_flickering') is True:
        init_factor = 3
    if (cam_mesh.graph['original_H'] is not None) and (cam_mesh.graph['original_W'] is not None):
        canvas_w = cam_mesh.graph['original_W']
        canvas_h = cam_mesh.graph['original_H']
    else:
        canvas_w = cam_mesh.graph['W']
        canvas_h = cam_mesh.graph['H']
    canvas_size = max(canvas_h, canvas_w)
    if normal_canvas is None:
        normal_canvas = Canvas_view(fov,
                                    verts,
                                    faces,
                                    colors,
                                    canvas_size=canvas_size,
                                    factor=init_factor,
                                    bgcolor='gray',
                                    proj='perspective')
    else:
        normal_canvas.reinit_mesh(verts, faces, colors)
        normal_canvas.reinit_camera(fov)
    img = normal_canvas.render()
    backup_img, backup_all_img, all_img_wo_bound = img.copy(), img.copy() * 0, img.copy() * 0
    img = cv2.resize(img, (int(img.shape[1] / init_factor), int(img.shape[0] / init_factor)), interpolation=cv2.INTER_AREA)
    if border is None:
        border = [0, img.shape[0], 0, img.shape[1]]
    H, W = cam_mesh.graph['H'], cam_mesh.graph['W']
    if (cam_mesh.graph['original_H'] is not None) and (cam_mesh.graph['original_W'] is not None):
        aspect_ratio = cam_mesh.graph['original_H'] / cam_mesh.graph['original_W']
    else:
        aspect_ratio = cam_mesh.graph['H'] / cam_mesh.graph['W']
    if aspect_ratio > 1:
        img_h_len = cam_mesh.graph['H'] if cam_mesh.graph.get('original_H') is None else cam_mesh.graph['original_H']
        img_w_len = img_h_len / aspect_ratio
        anchor = [0,
                  img.shape[0],
                  int(max(0, int((img.shape[1])//2 - img_w_len//2))),
                  int(min(int((img.shape[1])//2 + img_w_len//2), (img.shape[1])-1))]
    elif aspect_ratio <= 1:
        img_w_len = cam_mesh.graph['W'] if cam_mesh.graph.get('original_W') is None else cam_mesh.graph['original_W']
        img_h_len = img_w_len * aspect_ratio
        anchor = [int(max(0, int((img.shape[0])//2 - img_h_len//2))),
                  int(min(int((img.shape[0])//2 + img_h_len//2), (img.shape[0])-1)),
                  0,
                  img.shape[1]]
    anchor = np.array(anchor)
    plane_width = np.tan(fov_in_rad/2.) * np.abs(mean_loc_depth)
    for video_pose, video_traj_type in zip(videos_poses, video_traj_types):
        stereos = []
        tops = []; buttoms = []; lefts = []; rights = []
        for tp_id, tp in enumerate(video_pose):
            rel_pose = np.linalg.inv(np.dot(tp, np.linalg.inv(ref_pose)))
            axis, angle = transforms3d.axangles.mat2axangle(rel_pose[0:3, 0:3])
            normal_canvas.rotate(axis=axis, angle=(angle*180)/np.pi)
            normal_canvas.translate(rel_pose[:3,3])
            new_mean_loc_depth = mean_loc_depth - float(rel_pose[2, 3])
            if 'dolly' in video_traj_type:
                new_fov = float((np.arctan2(plane_width, np.array([np.abs(new_mean_loc_depth)])) * 180. / np.pi) * 2)
                normal_canvas.reinit_camera(new_fov)
            else:
                normal_canvas.reinit_camera(fov)
            normal_canvas.view_changed()
            img = normal_canvas.render()
            img = cv2.GaussianBlur(img,(int(init_factor//2 * 2 + 1), int(init_factor//2 * 2 + 1)), 0)
            img = cv2.resize(img, (int(img.shape[1] / init_factor), int(img.shape[0] / init_factor)), interpolation=cv2.INTER_AREA)
            img = img[anchor[0]:anchor[1], anchor[2]:anchor[3]]
            img = img[int(border[0]):int(border[1]), int(border[2]):int(border[3])]

            if any(np.array(config['crop_border']) > 0.0):
                H_c, W_c, _ = img.shape
                o_t = int(H_c * config['crop_border'][0])
                o_l = int(W_c * config['crop_border'][1])
                o_b = int(H_c * config['crop_border'][2])
                o_r = int(W_c * config['crop_border'][3])
                img = img[o_t:H_c-o_b, o_l:W_c-o_r]
                img = cv2.resize(img, (W_c, H_c), interpolation=cv2.INTER_CUBIC)

            """
            img = cv2.resize(img, (int(img.shape[1] / init_factor), int(img.shape[0] / init_factor)), interpolation=cv2.INTER_CUBIC)
            img = img[anchor[0]:anchor[1], anchor[2]:anchor[3]]
            img = img[int(border[0]):int(border[1]), int(border[2]):int(border[3])]

            if config['crop_border'] is True:
                top, buttom, left, right = find_largest_rect(img, bg_color=(128, 128, 128))
                tops.append(top); buttoms.append(buttom); lefts.append(left); rights.append(right)
            """
            stereos.append(img[..., :3])
            normal_canvas.translate(-rel_pose[:3,3])
            normal_canvas.rotate(axis=axis, angle=-(angle*180)/np.pi)
            normal_canvas.view_changed()
        """
        if config['crop_border'] is True:
            atop, abuttom = min(max(tops), img.shape[0]//2 - 10), max(min(buttoms), img.shape[0]//2 + 10)
            aleft, aright = min(max(lefts), img.shape[1]//2 - 10), max(min(rights), img.shape[1]//2 + 10)
            atop -= atop % 2; abuttom -= abuttom % 2; aleft -= aleft % 2; aright -= aright % 2
        else:
            atop = 0; abuttom = img.shape[0] - img.shape[0] % 2; aleft = 0; aright = img.shape[1] - img.shape[1] % 2
        """
        atop = 0; abuttom = img.shape[0] - img.shape[0] % 2; aleft = 0; aright = img.shape[1] - img.shape[1] % 2
        crop_stereos = []
        for stereo in stereos:
            crop_stereos.append((stereo[atop:abuttom, aleft:aright, :3] * 1).astype(np.uint8))
            stereos = crop_stereos
        clip = ImageSequenceClip(stereos, fps=config['fps'])
        if isinstance(video_basename, list):
            video_basename = video_basename[0]
        clip.write_videofile(os.path.join(output_dir, video_basename + '_' + video_traj_type + '.mp4'), fps=config['fps'])



    return normal_canvas, all_canvas                                                                                                                                                                                                                                                                                                                                                                                                                                                              import os
import numpy as np
try:
    import cynetworkx as netx
except ImportError:
    import networkx as netx

import json
import scipy.misc as misc
#import OpenEXR
import scipy.signal as signal
import matplotlib.pyplot as plt
import cv2
import scipy.misc as misc
from skimage import io
from functools import partial
from vispy import scene, io
from vispy.scene import visuals
from functools import reduce
# from moviepy.editor import ImageSequenceClip
import scipy.misc as misc
from vispy.visuals.filters import Alpha
import cv2
from skimage.transform import resize
import copy
import torch
import os
from utils import refine_depth_around_edge, smooth_cntsyn_gap
from utils import require_depth_edge, filter_irrelevant_edge_new, open_small_mask
from skimage.feature import canny
from scipy import ndimage
import time
import transforms3d

def relabel_node(mesh, nodes, cur_node, new_node):
    if cur_node == new_node:
        return mesh
    mesh.add_node(new_node)
    for key, value in nodes[cur_node].items():
        nodes[new_node][key] = value
    for ne in mesh.neighbors(cur_node):
        mesh.add_edge(new_node, ne)
    mesh.remove_node(cur_node)

    return mesh

def filter_edge(mesh, edge_ccs, config, invalid=False):
    context_ccs = [set() for _ in edge_ccs]
    mesh_nodes = mesh.nodes
    for edge_id, edge_cc in enumerate(edge_ccs):
        if config['context_thickness'] == 0:
            continue
        edge_group = {}
        for edge_node in edge_cc:
            far_nodes = mesh_nodes[edge_node].get('far')
            if far_nodes is None:
                continue
            for far_node in far_nodes:
                context_ccs[edge_id].add(far_node)
                if mesh_nodes[far_node].get('edge_id') is not None:
                    if edge_group.get(mesh_nodes[far_node]['edge_id']) is None:
                        edge_group[mesh_nodes[far_node]['edge_id']] = set()
                    edge_group[mesh_nodes[far_node]['edge_id']].add(far_node)
        if len(edge_cc) > 2:
            for edge_key in [*edge_group.keys()]:
                if len(edge_group[edge_key]) == 1:
                    context_ccs[edge_id].remove([*edge_group[edge_key]][0])
    valid_edge_ccs = []
    for xidx, yy in enumerate(edge_ccs):
        if invalid is not True and len(context_ccs[xidx]) > 0:
            # if len(context_ccs[xidx]) > 0:
            valid_edge_ccs.append(yy)
        elif invalid is True and len(context_ccs[xidx]) == 0:
            valid_edge_ccs.append(yy)
        else:
            valid_edge_ccs.append(set())
    # valid_edge_ccs = [yy for xidx, yy in enumerate(edge_ccs) if len(context_ccs[xidx]) > 0]

    return valid_edge_ccs

def extrapolate(global_mesh,
                info_on_pix,
                image,
                depth,
                other_edge_with_id,
                edge_map,
                edge_ccs,
                depth_edge_model,
                depth_feat_model,
                rgb_feat_model,
                config,
                direc='right-up'):
    h_off, w_off = global_mesh.graph['hoffset'], global_mesh.graph['woffset']
    noext_H, noext_W = global_mesh.graph['noext_H'], global_mesh.graph['noext_W']

    if "up" in direc.lower() and "-" not in direc.lower():
        all_anchor = [0, h_off + config['context_thickness'], w_off, w_off + noext_W]
        global_shift = [all_anchor[0], all_anchor[2]]
        mask_anchor = [0, h_off, w_off, w_off + noext_W]
        context_anchor = [h_off, h_off + config['context_thickness'], w_off, w_off + noext_W]
        valid_line_anchor = [h_off, h_off + 1, w_off, w_off + noext_W]
        valid_anchor = [min(mask_anchor[0], context_anchor[0]), max(mask_anchor[1], context_anchor[1]),
                        min(mask_anchor[2], context_anchor[2]), max(mask_anchor[3], context_anchor[3])]
    elif "down" in direc.lower() and "-" not in direc.lower():
        all_anchor = [h_off + noext_H - config['context_thickness'], 2 * h_off + noext_H, w_off, w_off + noext_W]
        global_shift = [all_anchor[0], all_anchor[2]]
        mask_anchor = [h_off + noext_H, 2 * h_off + noext_H, w_off, w_off + noext_W]
        context_anchor = [h_off + noext_H - config['context_thickness'], h_off + noext_H, w_off, w_off + noext_W]
        valid_line_anchor = [h_off + noext_H - 1, h_off + noext_H, w_off, w_off + noext_W]
        valid_anchor = [min(mask_anchor[0], context_anchor[0]), max(mask_anchor[1], context_anchor[1]),
                        min(mask_anchor[2], context_anchor[2]), max(mask_anchor[3], context_anchor[3])]
    elif "left" in direc.lower() and "-" not in direc.lower():
        all_anchor = [h_off, h_off + noext_H, 0, w_off + config['context_thickness']]
        global_shift = [all_anchor[0], all_anchor[2]]
        mask_anchor = [h_off, h_off + noext_H, 0, w_off]
        context_anchor = [h_off, h_off + noext_H, w_off, w_off + config['context_thickness']]
        valid_line_anchor = [h_off, h_off + noext_H, w_off, w_off + 1]
        valid_anchor = [min(mask_anchor[0], context_anchor[0]), max(mask_anchor[1], context_anchor[1]),
                        min(mask_anchor[2], context_anchor[2]), max(mask_anchor[3], context_anchor[3])]
    elif "right" in direc.lower() and "-" not in direc.lower():
        all_anchor = [h_off, h_off + noext_H, w_off + noext_W - config['context_thickness'], 2 * w_off + noext_W]
        global_shift = [all_anchor[0], all_anchor[2]]
        mask_anchor = [h_off, h_off + noext_H, w_off + noext_W, 2 * w_off + noext_W]
        context_anchor = [h_off, h_off + noext_H, w_off + noext_W - config['context_thickness'], w_off + noext_W]
        valid_line_anchor = [h_off, h_off + noext_H, w_off + noext_W - 1, w_off + noext_W]
        valid_anchor = [min(mask_anchor[0], context_anchor[0]), max(mask_anchor[1], context_anchor[1]),
                        min(mask_anchor[2], context_anchor[2]), max(mask_anchor[3], context_anchor[3])]
    elif "left" in direc.lower() and "up" in direc.lower() and "-" in direc.lower():
        all_anchor = [0, h_off + config['context_thickness'], 0, w_off + config['context_thickness']]
        global_shift = [all_anchor[0], all_anchor[2]]
        mask_anchor = [0, h_off, 0, w_off]
        context_anchor = "inv-mask"
        valid_line_anchor = None
        valid_anchor = all_anchor
    elif "left" in direc.lower() and "down" in direc.lower() and "-" in direc.lower():
        all_anchor = [h_off + noext_H - config['context_thickness'], 2 * h_off + noext_H, 0, w_off + config['context_thickness']]
        global_shift = [all_anchor[0], all_anchor[2]]
        mask_anchor = [h_off + noext_H, 2 * h_off + noext_H, 0, w_off]
        context_anchor = "inv-mask"
        valid_line_anchor = None
        valid_anchor = all_anchor
    elif "right" in direc.lower() and "up" in direc.lower() and "-" in direc.lower():
        all_anchor = [0, h_off + config['context_thickness'], w_off + noext_W - config['context_thickness'], 2 * w_off + noext_W]
        global_shift = [all_anchor[0], all_anchor[2]]
        mask_anchor = [0, h_off, w_off + noext_W, 2 * w_off + noext_W]
        context_anchor = "inv-mask"
        valid_line_anchor = None
        valid_anchor = all_anchor
    elif "right" in direc.lower() and "down" in direc.lower() and "-" in direc.lower():
        all_anchor = [h_off + noext_H - config['context_thickness'], 2 * h_off + noext_H, w_off + noext_W - config['context_thickness'], 2 * w_off + noext_W]
        global_shift = [all_anchor[0], all_anchor[2]]
        mask_anchor = [h_off + noext_H, 2 * h_off + noext_H, w_off + noext_W, 2 * w_off + noext_W]
        context_anchor = "inv-mask"
        valid_line_anchor = None
        valid_anchor = all_anchor

    global_mask = np.zeros_like(depth)
    global_mask[mask_anchor[0]:mask_anchor[1],mask_anchor[2]:mask_anchor[3]] = 1
    mask = global_mask[valid_anchor[0]:valid_anchor[1], valid_anchor[2]:valid_anchor[3]] * 1
    context = 1 - mask
    global_context = np.zeros_like(depth)
    global_context[all_anchor[0]:all_anchor[1],all_anchor[2]:all_anchor[3]] = context
    # context = global_context[valid_anchor[0]:valid_anchor[1], valid_anchor[2]:valid_anchor[3]] * 1



    valid_area = mask + context
    input_rgb = image[valid_anchor[0]:valid_anchor[1], valid_anchor[2]:valid_anchor[3]] / 255. * context[..., None]
    input_depth = depth[valid_anchor[0]:valid_anchor[1], valid_anchor[2]:valid_anchor[3]] * context
    log_depth = np.log(input_depth + 1e-8)
    log_depth[mask > 0] = 0
    input_mean_depth = np.mean(log_depth[context > 0])
    input_zero_mean_depth = (log_depth - input_mean_depth) * context
    input_disp = 1./np.abs(input_depth)
    input_disp[mask > 0] = 0
    input_disp = input_disp / input_disp.max()
    valid_line = np.zeros_like(depth)
    if valid_line_anchor is not None:
        valid_line[valid_line_anchor[0]:valid_line_anchor[1], valid_line_anchor[2]:valid_line_anchor[3]] = 1
    valid_line = valid_line[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]]
    # f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True); ax1.imshow(global_context * 1 + global_mask * 2); ax2.imshow(image); plt.show()
    # f, ((ax1, ax2, ax3)) = plt.subplots(1, 3, sharex=True, sharey=True); ax1.imshow(context * 1 + mask * 2); ax2.imshow(input_rgb); ax3.imshow(valid_line); plt.show()
    # import pdb; pdb.set_trace()
    # return
    input_edge_map = edge_map[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]] * context
    input_other_edge_with_id = other_edge_with_id[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]]
    end_depth_maps = ((valid_line * input_edge_map) > 0) * input_depth


    if isinstance(config["gpu_ids"], int) and (config["gpu_ids"] >= 0):
        device = config["gpu_ids"]
    else:
        device = "cpu"

    valid_edge_ids = sorted(list(input_other_edge_with_id[(valid_line * input_edge_map) > 0]))
    valid_edge_ids = valid_edge_ids[1:] if (len(valid_edge_ids) > 0 and valid_edge_ids[0] == -1) else valid_edge_ids
    edge = reduce(lambda x, y: (x + (input_other_edge_with_id == y).astype(np.uint8)).clip(0, 1), [np.zeros_like(mask)] + list(valid_edge_ids))
    t_edge = torch.FloatTensor(edge).to(device)[None, None, ...]
    t_rgb = torch.FloatTensor(input_rgb).to(device).permute(2,0,1).unsqueeze(0)
    t_mask = torch.FloatTensor(mask).to(device)[None, None, ...]
    t_context = torch.FloatTensor(context).to(device)[None, None, ...]
    t_disp = torch.FloatTensor(input_disp).to(device)[None, None, ...]
    t_depth_zero_mean_depth = torch.FloatTensor(input_zero_mean_depth).to(device)[None, None, ...]

    depth_edge_output = depth_edge_model.forward_3P(t_mask, t_context, t_rgb, t_disp, t_edge, unit_length=128,
                                                    cuda=device)
    t_output_edge = (depth_edge_output> config['ext_edge_threshold']).float() * t_mask + t_edge
    output_raw_edge = t_output_edge.data.cpu().numpy().squeeze()
    # import pdb; pdb.set_trace()
    mesh = netx.Graph()
    hxs, hys = np.where(output_raw_edge * mask > 0)
    valid_map = mask + context
    for hx, hy in zip(hxs, hys):
        node = (hx, hy)
        mesh.add_node((hx, hy))
        eight_nes = [ne for ne in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1), \
                                   (hx + 1, hy + 1), (hx - 1, hy - 1), (hx - 1, hy + 1), (hx + 1, hy - 1)]\
                        if 0 <= ne[0] < output_raw_edge.shape[0] and 0 <= ne[1] < output_raw_edge.shape[1] and 0 < output_raw_edge[ne[0], ne[1]]]
        for ne in eight_nes:
            mesh.add_edge(node, ne, length=np.hypot(ne[0] - hx, ne[1] - hy))
            if end_depth_maps[ne[0], ne[1]] != 0:
                mesh.nodes[ne[0], ne[1]]['cnt'] = True
                mesh.nodes[ne[0], ne[1]]['depth'] = end_depth_maps[ne[0], ne[1]]
    ccs = [*netx.connected_components(mesh)]
    end_pts = []
    for cc in ccs:
        end_pts.append(set())
        for node in cc:
            if mesh.nodes[node].get('cnt') is not None:
                end_pts[-1].add((node[0], node[1], mesh.nodes[node]['depth']))
    fpath_map = np.zeros_like(output_raw_edge) - 1
    npath_map = np.zeros_like(output_raw_edge) - 1
    for end_pt, cc in zip(end_pts, ccs):
        sorted_end_pt = []
        if len(end_pt) >= 2:
            continue
        if len(end_pt) == 0:
            continue
        if len(end_pt) == 1:
            sub_mesh = mesh.subgraph(list(cc)).copy()
            pnodes = netx.periphery(sub_mesh)
            ends = [*end_pt]
            edge_id = global_mesh.nodes[(ends[0][0] + all_anchor[0], ends[0][1] + all_anchor[2], -ends[0][2])]['edge_id']
            pnodes = sorted(pnodes,
                            key=lambda x: np.hypot((x[0] - ends[0][0]), (x[1] - ends[0][1])),
                            reverse=True)[0]
            npath = [*netx.shortest_path(sub_mesh, (ends[0][0], ends[0][1]), pnodes, weight='length')]
            for np_node in npath:
                npath_map[np_node[0], np_node[1]] = edge_id
            fpath = []
            if global_mesh.nodes[(ends[0][0] + all_anchor[0], ends[0][1] + all_anchor[2], -ends[0][2])].get('far') is None:
                print("None far")
                import pdb; pdb.set_trace()
            else:
                fnodes = global_mesh.nodes[(ends[0][0] + all_anchor[0], ends[0][1] + all_anchor[2], -ends[0][2])].get('far')
                fnodes = [(xx[0] - all_anchor[0], xx[1] - all_anchor[2], xx[2]) for xx in fnodes]
                dmask = mask + 0
                did = 0
                while True:
                    did += 1
                    dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1)
                    if did > 3:
                        break
                    # ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0)]
                    ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0 and mask[fnode[0], fnode[1]] == 0)]
                    if len(ffnode) > 0:
                        fnode = ffnode[0]
                        break
                if len(ffnode) == 0:
                    continue
                fpath.append((fnode[0], fnode[1]))
                for step in range(0, len(npath) - 1):
                    parr = (npath[step + 1][0] - npath[step][0], npath[step + 1][1] - npath[step][1])
                    new_loc = (fpath[-1][0] + parr[0], fpath[-1][1] + parr[1])
                    new_loc_nes = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]),
                                                (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1)]\
                                        if xx[0] >= 0 and xx[0] < fpath_map.shape[0] and xx[1] >= 0 and xx[1] < fpath_map.shape[1]]
                    if np.sum([fpath_map[nlne[0], nlne[1]] for nlne in new_loc_nes]) != -4:
                        break
                    if npath_map[new_loc[0], new_loc[1]] != -1:
                        if npath_map[new_loc[0], new_loc[1]] != edge_id:
                            break
                        else:
                            continue
                    if valid_area[new_loc[0], new_loc[1]] == 0:
                        break
                    new_loc_nes_eight = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]),
                                                        (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1),
                                                        (new_loc[0] + 1, new_loc[1] + 1), (new_loc[0] + 1, new_loc[1] - 1),
                                                        (new_loc[0] - 1, new_loc[1] - 1), (new_loc[0] - 1, new_loc[1] + 1)]\
                                        if xx[0] >= 0 and xx[0] < fpath_map.shape[0] and xx[1] >= 0 and xx[1] < fpath_map.shape[1]]
                    if np.sum([int(npath_map[nlne[0], nlne[1]] == edge_id) for nlne in new_loc_nes_eight]) == 0:
                        break
                    fpath.append((fpath[-1][0] + parr[0], fpath[-1][1] + parr[1]))
                if step != len(npath) - 2:
                    for xx in npath[step+1:]:
                        if npath_map[xx[0], xx[1]] == edge_id:
                            npath_map[xx[0], xx[1]] = -1
            if len(fpath) > 0:
                for fp_node in fpath:
                    fpath_map[fp_node[0], fp_node[1]] = edge_id
    # import pdb; pdb.set_trace()
    far_edge = (fpath_map > -1).astype(np.uint8)
    update_edge = (npath_map > -1) * mask + edge
    t_update_edge = torch.FloatTensor(update_edge).to(device)[None, None, ...]
    depth_output = depth_feat_model.forward_3P(t_mask, t_context, t_depth_zero_mean_depth, t_update_edge, unit_length=128,
                                               cuda=device)
    depth_output = depth_output.cpu().data.numpy().squeeze()
    depth_output = np.exp(depth_output + input_mean_depth) * mask # + input_depth * context
    # if "right" in direc.lower() and "-" not in direc.lower():
    #     plt.imshow(depth_output); plt.show()
    #     import pdb; pdb.set_trace()
    #     f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True); ax1.imshow(depth_output); ax2.imshow(npath_map + fpath_map); plt.show()
    for near_id in np.unique(npath_map[npath_map > -1]):
        depth_output = refine_depth_around_edge(depth_output.copy(),
                                                (fpath_map == near_id).astype(np.uint8) * mask, # far_edge_map_in_mask,
                                                (fpath_map == near_id).astype(np.uint8), # far_edge_map,
                                                (npath_map == near_id).astype(np.uint8) * mask,
                                                mask.copy(),
                                                np.zeros_like(mask),
                                                config)
    # if "right" in direc.lower() and "-" not in direc.lower():
    #     plt.imshow(depth_output); plt.show()
    #     import pdb; pdb.set_trace()
    #     f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True); ax1.imshow(depth_output); ax2.imshow(npath_map + fpath_map); plt.show()
    rgb_output = rgb_feat_model.forward_3P(t_mask, t_context, t_rgb, t_update_edge, unit_length=128,
                                           cuda=device)

    # rgb_output = rgb_feat_model.forward_3P(t_mask, t_context, t_rgb, t_update_edge, unit_length=128, cuda=config['gpu_ids'])
    if config.get('gray_image') is True:
        rgb_output = rgb_output.mean(1, keepdim=True).repeat((1,3,1,1))
    rgb_output = ((rgb_output.squeeze().data.cpu().permute(1,2,0).numpy() * mask[..., None] + input_rgb) * 255).astype(np.uint8)
    image[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]][mask > 0] = rgb_output[mask > 0] # np.array([255,0,0]) # rgb_output[mask > 0]
    depth[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]][mask > 0] = depth_output[mask > 0]
    # nxs, nys = np.where(mask > -1)
    # for nx, ny in zip(nxs, nys):
    #     info_on_pix[(nx, ny)][0]['color'] = rgb_output[]


    nxs, nys = np.where((npath_map > -1))
    for nx, ny in zip(nxs, nys):
        n_id = npath_map[nx, ny]
        four_nes = [xx for xx in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)]\
                        if 0 <= xx[0] < fpath_map.shape[0] and 0 <= xx[1] < fpath_map.shape[1]]
        for nex, ney in four_nes:
            if fpath_map[nex, ney] == n_id:
                na, nb = (nx + all_anchor[0], ny + all_anchor[2], info_on_pix[(nx + all_anchor[0], ny + all_anchor[2])][0]['depth']), \
                        (nex + all_anchor[0], ney + all_anchor[2], info_on_pix[(nex + all_anchor[0], ney + all_anchor[2])][0]['depth'])
                if global_mesh.has_edge(na, nb):
                    global_mesh.remove_edge(na, nb)
    nxs, nys = np.where((fpath_map > -1))
    for nx, ny in zip(nxs, nys):
        n_id = fpath_map[nx, ny]
        four_nes = [xx for xx in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)]\
                        if 0 <= xx[0] < npath_map.shape[0] and 0 <= xx[1] < npath_map.shape[1]]
        for nex, ney in four_nes:
            if npath_map[nex, ney] == n_id:
                na, nb = (nx + all_anchor[0], ny + all_anchor[2], info_on_pix[(nx + all_anchor[0], ny + all_anchor[2])][0]['depth']), \
                        (nex + all_anchor[0], ney + all_anchor[2], info_on_pix[(nex + all_anchor[0], ney + all_anchor[2])][0]['depth'])
                if global_mesh.has_edge(na, nb):
                    global_mesh.remove_edge(na, nb)
    nxs, nys = np.where(mask > 0)
    for x, y in zip(nxs, nys):
        x = x + all_anchor[0]
        y = y + all_anchor[2]
        cur_node = (x, y, 0)
        new_node = (x, y, -abs(depth[x, y]))
        disp = 1. / -abs(depth[x, y])
        mapping_dict = {cur_node: new_node}
        info_on_pix, global_mesh = update_info(mapping_dict, info_on_pix, global_mesh)
        global_mesh.nodes[new_node]['color'] = image[x, y]
        global_mesh.nodes[new_node]['old_color'] = image[x, y]
        global_mesh.nodes[new_node]['disp'] = disp
        info_on_pix[(x, y)][0]['depth'] = -abs(depth[x, y])
        info_on_pix[(x, y)][0]['disp'] = disp
        info_on_pix[(x, y)][0]['color'] = image[x, y]


    nxs, nys = np.where((npath_map > -1))
    for nx, ny in zip(nxs, nys):
        self_node = (nx + all_anchor[0], ny + all_anchor[2], info_on_pix[(nx + all_anchor[0], ny + all_anchor[2])][0]['depth'])
        if global_mesh.has_node(self_node) is False:
            break
        n_id = int(round(npath_map[nx, ny]))
        four_nes = [xx for xx in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)]\
                        if 0 <= xx[0] < fpath_map.shape[0] and 0 <= xx[1] < fpath_map.shape[1]]
        for nex, ney in four_nes:
            ne_node = (nex + all_anchor[0], ney + all_anchor[2], info_on_pix[(nex + all_anchor[0], ney + all_anchor[2])][0]['depth'])
            if global_mesh.has_node(ne_node) is False:
                continue
            if fpath_map[nex, ney] == n_id:
                if global_mesh.nodes[self_node].get('edge_id') is None:
                    global_mesh.nodes[self_node]['edge_id'] = n_id
                    edge_ccs[n_id].add(self_node)
                    info_on_pix[(self_node[0], self_node[1])][0]['edge_id'] = n_id
                if global_mesh.has_edge(self_node, ne_node) is True:
                    global_mesh.remove_edge(self_node, ne_node)
                if global_mesh.nodes[self_node].get('far') is None:
                    global_mesh.nodes[self_node]['far'] = []
                global_mesh.nodes[self_node]['far'].append(ne_node)

    global_fpath_map = np.zeros_like(other_edge_with_id) - 1
    global_fpath_map[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]] = fpath_map
    fpath_ids = np.unique(global_fpath_map)
    fpath_ids = fpath_ids[1:] if fpath_ids.shape[0] > 0 and fpath_ids[0] == -1 else []
    fpath_real_id_map = np.zeros_like(global_fpath_map) - 1
    for fpath_id in fpath_ids:
        fpath_real_id = np.unique(((global_fpath_map == fpath_id).astype(np.int) * (other_edge_with_id + 1)) - 1)
        fpath_real_id = fpath_real_id[1:] if fpath_real_id.shape[0] > 0 and fpath_real_id[0] == -1 else []
        fpath_real_id = fpath_real_id.astype(np.int)
        fpath_real_id = np.bincount(fpath_real_id).argmax()
        fpath_real_id_map[global_fpath_map == fpath_id] = fpath_real_id
    nxs, nys = np.where((fpath_map > -1))
    for nx, ny in zip(nxs, nys):
        self_node = (nx + all_anchor[0], ny + all_anchor[2], info_on_pix[(nx + all_anchor[0], ny + all_anchor[2])][0]['depth'])
        n_id = fpath_map[nx, ny]
        four_nes = [xx for xx in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)]\
                        if 0 <= xx[0] < npath_map.shape[0] and 0 <= xx[1] < npath_map.shape[1]]
        for nex, ney in four_nes:
            ne_node = (nex + all_anchor[0], ney + all_anchor[2], info_on_pix[(nex + all_anchor[0], ney + all_anchor[2])][0]['depth'])
            if global_mesh.has_node(ne_node) is False:
                continue
            if npath_map[nex, ney] == n_id or global_mesh.nodes[ne_node].get('edge_id') == n_id:
                if global_mesh.has_edge(self_node, ne_node) is True:
                    global_mesh.remove_edge(self_node, ne_node)
                if global_mesh.nodes[self_node].get('near') is None:
                    global_mesh.nodes[self_node]['near'] = []
                if global_mesh.nodes[self_node].get('edge_id') is None:
                    f_id = int(round(fpath_real_id_map[self_node[0], self_node[1]]))
                    global_mesh.nodes[self_node]['edge_id'] = f_id
                    info_on_pix[(self_node[0], self_node[1])][0]['edge_id'] = f_id
                    edge_ccs[f_id].add(self_node)
                global_mesh.nodes[self_node]['near'].append(ne_node)

    return info_on_pix, global_mesh, image, depth, edge_ccs
    # for edge_cc in edge_ccs:
    #     for edge_node in edge_cc:
    #         edge_ccs
    # context_ccs, mask_ccs, broken_mask_ccs, edge_ccs, erode_context_ccs, init_mask_connect, edge_maps, extend_context_ccs, extend_edge_ccs

def get_valid_size(imap):
    x_max = np.where(imap.sum(1).squeeze() > 0)[0].max() + 1
    x_min = np.where(imap.sum(1).squeeze() > 0)[0].min()
    y_max = np.where(imap.sum(0).squeeze() > 0)[0].max() + 1
    y_min = np.where(imap.sum(0).squeeze() > 0)[0].min()
    size_dict = {'x_max':x_max, 'y_max':y_max, 'x_min':x_min, 'y_min':y_min}

    return size_dict

def dilate_valid_size(isize_dict, imap, dilate=[0, 0]):
    osize_dict = copy.deepcopy(isize_dict)
    osize_dict['x_min'] = max(0, osize_dict['x_min'] - dilate[0])
    osize_dict['x_max'] = min(imap.shape[0], osize_dict['x_max'] + dilate[0])
    osize_dict['y_min'] = max(0, osize_dict['y_min'] - dilate[0])
    osize_dict['y_max'] = min(imap.shape[1], osize_dict['y_max'] + dilate[1])

    return osize_dict

def size_operation(size_a, size_b, operation):
    assert operation == '+' or operation == '-', "Operation must be '+' (union) or '-' (exclude)"
    osize = {}
    if operation == '+':
        osize['x_min'] = min(size_a['x_min'], size_b['x_min'])
        osize['y_min'] = min(size_a['y_min'], size_b['y_min'])
        osize['x_max'] = max(size_a['x_max'], size_b['x_max'])
        osize['y_max'] = max(size_a['y_max'], size_b['y_max'])
    assert operation != '-', "Operation '-' is undefined !"

    return osize

def fill_dummy_bord(mesh, info_on_pix, image, depth, config):
    context = np.zeros_like(depth).astype(np.uint8)
    context[mesh.graph['hoffset']:mesh.graph['hoffset'] + mesh.graph['noext_H'],
            mesh.graph['woffset']:mesh.graph['woffset'] + mesh.graph['noext_W']] = 1
    mask = 1 - context
    xs, ys = np.where(mask > 0)
    depth = depth * context
    image = image * context[..., None]
    cur_depth = 0
    cur_disp = 0
    color = [0, 0, 0]
    for x, y in zip(xs, ys):
        cur_node = (x, y, cur_depth)
        mesh.add_node(cur_node, color=color,
                        synthesis=False,
                        disp=cur_disp,
                        cc_id=set(),
                        ext_pixel=True)
        info_on_pix[(x, y)] = [{'depth':cur_depth,
                    'color':mesh.nodes[(x, y, cur_depth)]['color'],
                    'synthesis':False,
                    'disp':mesh.nodes[cur_node]['disp'],
                    'ext_pixel':True}]
        # for x, y in zip(xs, ys):
        four_nes = [(xx, yy) for xx, yy in [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)] if\
                    0 <= x < mesh.graph['H'] and 0 <= y < mesh.graph['W'] and info_on_pix.get((xx, yy)) is not None]
        for ne in four_nes:
            # if (ne[0] - x) + (ne[1] - y) == 1 and info_on_pix.get((ne[0], ne[1])) is not None:
            mesh.add_edge(cur_node, (ne[0], ne[1], info_on_pix[(ne[0], ne[1])][0]['depth']))

    return mesh, info_on_pix


def enlarge_border(mesh, info_on_pix, depth, image, config):
    mesh.graph['hoffset'], mesh.graph['woffset'] = config['extrapolation_thickness'], config['extrapolation_thickness']
    mesh.graph['bord_up'], mesh.graph['bord_left'], mesh.graph['bord_down'], mesh.graph['bord_right'] = \
        0, 0, mesh.graph['H'], mesh.graph['W']
    # new_image = np.pad(image,
    #                    pad_width=((config['extrapolation_thickness'], config['extrapolation_thickness']),
    #                               (config['extrapolation_thickness'], config['extrapolation_thickness']), (0, 0)),
    #                    mode='constant')
    # new_depth = np.pad(depth,
    #                    pad_width=((config['extrapolation_thickness'], config['extrapolation_thickness']),
    #                               (config['extrapolation_thickness'], config['extrapolation_thickness'])),
    #                    mode='constant')

    return mesh, info_on_pix, depth, image

def fill_missing_node(mesh, info_on_pix, image, depth):
    for x in range(mesh.graph['bord_up'], mesh.graph['bord_down']):
        for y in range(mesh.graph['bord_left'], mesh.graph['bord_right']):
            if info_on_pix.get((x, y)) is None:
                print("fill missing node = ", x, y)
                import pdb; pdb.set_trace()
                re_depth, re_count = 0, 0
                for ne in [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)]:
                    if info_on_pix.get(ne) is not None:
                        re_depth += info_on_pix[ne][0]['depth']
                        re_count += 1
                if re_count == 0:
                    re_depth = -abs(depth[x, y])
                else:
                    re_depth = re_depth / re_count
                depth[x, y] = abs(re_depth)
                info_on_pix[(x, y)] = [{'depth':re_depth,
                                            'color':image[x, y],
                                            'synthesis':False,
                                            'disp':1./re_depth}]
                mesh.add_node((x, y, re_depth), color=image[x, y],
                                                synthesis=False,
                                                disp=1./re_depth,
                                                cc_id=set())
    return mesh, info_on_pix, depth



def refresh_bord_depth(mesh, info_on_pix, image, depth):
    H, W = mesh.graph['H'], mesh.graph['W']
    corner_nodes = [(mesh.graph['bord_up'], mesh.graph['bord_left']),
                    (mesh.graph['bord_up'], mesh.graph['bord_right'] - 1),
                    (mesh.graph['bord_down'] - 1, mesh.graph['bord_left']),
                    (mesh.graph['bord_down'] - 1, mesh.graph['bord_right'] - 1)]
                    # (0, W - 1), (H - 1, 0), (H - 1, W - 1)]
    bord_nodes = []
    bord_nodes += [(mesh.graph['bord_up'], xx) for xx in range(mesh.graph['bord_left'] + 1, mesh.graph['bord_right'] - 1)]
    bord_nodes += [(mesh.graph['bord_down'] - 1, xx) for xx in range(mesh.graph['bord_left'] + 1, mesh.graph['bord_right'] - 1)]
    bord_nodes += [(xx, mesh.graph['bord_left']) for xx in range(mesh.graph['bord_up'] + 1, mesh.graph['bord_down'] - 1)]
    bord_nodes += [(xx, mesh.graph['bord_right'] - 1) for xx in range(mesh.graph['bord_up'] + 1, mesh.graph['bord_down'] - 1)]
    for xy in bord_nodes:
        tgt_loc = None
        if xy[0] == mesh.graph['bord_up']:
            tgt_loc = (xy[0] + 1, xy[1])# (1, xy[1])
        elif xy[0] == mesh.graph['bord_down'] - 1:
            tgt_loc = (xy[0] - 1, xy[1]) # (H - 2, xy[1])
        elif xy[1] == mesh.graph['bord_left']:
            tgt_loc = (xy[0], xy[1] + 1)
        elif xy[1] == mesh.graph['bord_right'] - 1:
            tgt_loc = (xy[0], xy[1] - 1)
        if tgt_loc is not None:
            ne_infos = info_on_pix.get(tgt_loc)
            if ne_infos is None:
                import pdb; pdb.set_trace()
            # if ne_infos is not None and len(ne_infos) == 1:
            tgt_depth = ne_infos[0]['depth']
            tgt_disp = ne_infos[0]['disp']
            new_node = (xy[0], xy[1], tgt_depth)
            src_node = (tgt_loc[0], tgt_loc[1], tgt_depth)
            tgt_nes_loc = [(xx[0], xx[1]) \
                            for xx in mesh.neighbors(src_node)]
            tgt_nes_loc = [(xx[0] - tgt_loc[0] + xy[0], xx[1] - tgt_loc[1] + xy[1]) for xx in tgt_nes_loc \
                            if abs(xx[0] - xy[0]) == 1 and abs(xx[1] - xy[1]) == 1]
            tgt_nes_loc = [xx for xx in tgt_nes_loc if info_on_pix.get(xx) is not None]
            tgt_nes_loc.append(tgt_loc)
            # if (xy[0], xy[1]) == (559, 60):
            #     import pdb; pdb.set_trace()
            if info_on_pix.get(xy) is not None and len(info_on_pix.get(xy)) > 0:
                old_depth = info_on_pix[xy][0].get('depth')
                old_node = (xy[0], xy[1], old_depth)
                mesh.remove_edges_from([(old_ne, old_node) for old_ne in mesh.neighbors(old_node)])
                mesh.add_edges_from([((zz[0], zz[1], info_on_pix[zz][0]['depth']), old_node) for zz in tgt_nes_loc])
                mapping_dict = {old_node: new_node}
                # if old_node[2] == new_node[2]:
                #     print("mapping_dict = ", mapping_dict)
                info_on_pix, mesh = update_info(mapping_dict, info_on_pix, mesh)
            else:
                info_on_pix[xy] = []
                info_on_pix[xy][0] = info_on_pix[tgt_loc][0]
                info_on_pix['color'] = image[xy[0], xy[1]]
                info_on_pix['old_color'] = image[xy[0], xy[1]]
                mesh.add_node(new_node)
                mesh.add_edges_from([((zz[0], zz[1], info_on_pix[zz][0]['depth']), new_node) for zz in tgt_nes_loc])
            mesh.nodes[new_node]['far'] = None
            mesh.nodes[new_node]['near'] = None
            if mesh.nodes[src_node].get('far') is not None:
                redundant_nodes = [ne for ne in mesh.nodes[src_node]['far'] if (ne[0], ne[1]) == xy]
                [mesh.nodes[src_node]['far'].remove(aa) for aa in redundant_nodes]
            if mesh.nodes[src_node].get('near') is not None:
                redundant_nodes = [ne for ne in mesh.nodes[src_node]['near'] if (ne[0], ne[1]) == xy]
                [mesh.nodes[src_node]['near'].remove(aa) for aa in redundant_nodes]
    for xy in corner_nodes:
        hx, hy = xy
        four_nes = [xx for xx in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] if \
                        mesh.graph['bord_up'] <= xx[0] < mesh.graph['bord_down'] and \
                            mesh.graph['bord_left'] <= xx[1] < mesh.graph['bord_right']]
        ne_nodes = []
        ne_depths = []
        for ne_loc in four_nes:
            if info_on_pix.get(ne_loc) is not None:
                ne_depths.append(info_on_pix[ne_loc][0]['depth'])
                ne_nodes.append((ne_loc[0], ne_loc[1], info_on_pix[ne_loc][0]['depth']))
        new_node = (xy[0], xy[1], float(np.mean(ne_depths)))
        if info_on_pix.get(xy) is not None and len(info_on_pix.get(xy)) > 0:
            old_depth = info_on_pix[xy][0].get('depth')
            old_node = (xy[0], xy[1], old_depth)
            mesh.remove_edges_from([(old_ne, old_node) for old_ne in mesh.neighbors(old_node)])
            mesh.add_edges_from([(zz, old_node) for zz in ne_nodes])
            mapping_dict = {old_node: new_node}
            info_on_pix, mesh = update_info(mapping_dict, info_on_pix, mesh)
        else:
            info_on_pix[xy] = []
            info_on_pix[xy][0] = info_on_pix[ne_loc[-1]][0]
            info_on_pix['color'] = image[xy[0], xy[1]]
            info_on_pix['old_color'] = image[xy[0], xy[1]]
            mesh.add_node(new_node)
            mesh.add_edges_from([(zz, new_node) for zz in ne_nodes])
        mesh.nodes[new_node]['far'] = None
        mesh.nodes[new_node]['near'] = None
    for xy in bord_nodes + corner_nodes:
        # if (xy[0], xy[1]) == (559, 60):
        #     import pdb; pdb.set_trace()
        depth[xy[0], xy[1]] = abs(info_on_pix[xy][0]['depth'])
    for xy in bord_nodes:
        cur_node = (xy[0], xy[1], info_on_pix[xy][0]['depth'])
        nes = mesh.neighbors(cur_node)
        four_nes = set([(xy[0] + 1, xy[1]), (xy[0] - 1, xy[1]), (xy[0], xy[1] + 1), (xy[0], xy[1] - 1)]) - \
                   set([(ne[0], ne[1]) for ne in nes])
        four_nes = [ne for ne in four_nes if mesh.graph['bord_up'] <= ne[0] < mesh.graph['bord_down'] and \
                                             mesh.graph['bord_left'] <= ne[1] < mesh.graph['bord_right']]
        four_nes = [(ne[0], ne[1], info_on_pix[(ne[0], ne[1])][0]['depth']) for ne in four_nes]
        mesh.nodes[cur_node]['far'] = []
        mesh.nodes[cur_node]['near'] = []
        for ne in four_nes:
            if abs(ne[2]) >= abs(cur_node[2]):
                mesh.nodes[cur_node]['far'].append(ne)
            else:
                mesh.nodes[cur_node]['near'].append(ne)

    return mesh, info_on_pix, depth

def get_union_size(mesh, dilate, *alls_cc):
    all_cc = reduce(lambda x, y: x | y, [set()] + [*alls_cc])
    min_x, min_y, max_x, max_y = mesh.graph['H'], mesh.graph['W'], 0, 0
    H, W = mesh.graph['H'], mesh.graph['W']
    for node in all_cc:
        if node[0] < min_x:
            min_x = node[0]
        if node[0] > max_x:
            max_x = node[0]
        if node[1] < min_y:
            min_y = node[1]
        if node[1] > max_y:
            max_y = node[1]
    max_x = max_x + 1
    max_y = max_y + 1
    # mask_size = dilate_valid_size(mask_size, edge_dict['mask'], dilate=[20, 20])
    osize_dict = dict()
    osize_dict['x_min'] = max(0, min_x - dilate[0])
    osize_dict['x_max'] = min(H, max_x + dilate[0])
    osize_dict['y_min'] = max(0, min_y - dilate[1])
    osize_dict['y_max'] = min(W, max_y + dilate[1])

    return osize_dict

def incomplete_node(mesh, edge_maps, info_on_pix):
    vis_map = np.zeros((mesh.graph['H'], mesh.graph['W']))

    for node in mesh.nodes:
        if mesh.nodes[node].get('synthesis') is not True:
            connect_all_flag = False
            nes = [xx for xx in mesh.neighbors(node) if mesh.nodes[xx].get('synthesis') is not True]
            if len(nes) < 3 and 0 < node[0] < mesh.graph['H'] - 1 and 0 < node[1] < mesh.graph['W'] - 1:
                if len(nes) <= 1:
                    connect_all_flag = True
                else:
                    dan_ne_node_a = nes[0]
                    dan_ne_node_b = nes[1]
                    if abs(dan_ne_node_a[0] - dan_ne_node_b[0]) > 1 or \
                        abs(dan_ne_node_a[1] - dan_ne_node_b[1]) > 1:
                        connect_all_flag = True
            if connect_all_flag == True:
                vis_map[node[0], node[1]] = len(nes)
                four_nes = [(node[0] - 1, node[1]), (node[0] + 1, node[1]), (node[0], node[1] - 1), (node[0], node[1] + 1)]
                for ne in four_nes:
                    for info in info_on_pix[(ne[0], ne[1])]:
                        ne_node = (ne[0], ne[1], info['depth'])
                        if info.get('synthesis') is not True and mesh.has_node(ne_node):
                            mesh.add_edge(node, ne_node)
                            break

    return mesh

def edge_inpainting(edge_id, context_cc, erode_context_cc, mask_cc, edge_cc, extend_edge_cc,
                    mesh, edge_map, edge_maps_with_id, config, union_size, depth_edge_model, inpaint_iter):
    edge_dict = get_edge_from_nodes(context_cc, erode_context_cc, mask_cc, edge_cc, extend_edge_cc,
                                        mesh.graph['H'], mesh.graph['W'], mesh)
    edge_dict['edge'], end_depth_maps, _ = \
        filter_irrelevant_edge_new(edge_dict['self_edge'] + edge_dict['comp_edge'],
                                edge_map,
                                edge_maps_with_id,
                                edge_id,
                                edge_dict['context'],
                                edge_dict['depth'], mesh, context_cc | erode_context_cc, spdb=True)
    patch_edge_dict = dict()
    patch_edge_dict['mask'], patch_edge_dict['context'], patch_edge_dict['rgb'], \
        patch_edge_dict['disp'], patch_edge_dict['edge'] = \
        crop_maps_by_size(union_size, edge_dict['mask'], edge_dict['context'],
                            edge_dict['rgb'], edge_dict['disp'], edge_dict['edge'])
    tensor_edge_dict = convert2tensor(patch_edge_dict)
    if require_depth_edge(patch_edge_dict['edge'], patch_edge_dict['mask']) and inpaint_iter == 0:
        with torch.no_grad():
            device = config["gpu_ids"] if isinstance(config["gpu_ids"], int) and config["gpu_ids"] >= 0 else "cpu"
            depth_edge_output = depth_edge_model.forward_3P(tensor_edge_dict['mask'],
                                                            tensor_edge_dict['context'],
                                                            tensor_edge_dict['rgb'],
                                                            tensor_edge_dict['disp'],
                                                            tensor_edge_dict['edge'],
                                                            unit_length=128,
                                                            cuda=device)
            depth_edge_output = depth_edge_output.cpu()
        tensor_edge_dict['output'] = (depth_edge_output > config['ext_edge_threshold']).float() * tensor_edge_dict['mask'] + tensor_edge_dict['edge']
    else:
        tensor_edge_dict['output'] = tensor_edge_dict['edge']
        depth_edge_output = tensor_edge_dict['edge'] + 0
    patch_edge_dict['output'] = tensor_edge_dict['output'].squeeze().data.cpu().numpy()
    edge_dict['output'] = np.zeros((mesh.graph['H'], mesh.graph['W']))
    edge_dict['output'][union_size['x_min']:union_size['x_max'], union_size['y_min']:union_size['y_max']] = \
        patch_edge_dict['output']

    return edge_dict, end_depth_maps

def depth_inpainting(context_cc, extend_context_cc, erode_context_cc, mask_cc, mesh, config, union_size, depth_feat_model, edge_output, given_depth_dict=False, spdb=False):
    if given_depth_dict is False:
        depth_dict = get_depth_from_nodes(context_cc | extend_context_cc, erode_context_cc, mask_cc, mesh.graph['H'], mesh.graph['W'], mesh, config['log_depth'])
        if edge_output is not None:
            depth_dict['edge'] = edge_output
    else:
        depth_dict = given_depth_dict
    patch_depth_dict = dict()
    patch_depth_dict['mask'], patch_depth_dict['context'], patch_depth_dict['depth'], \
        patch_depth_dict['zero_mean_depth'], patch_depth_dict['edge'] = \
            crop_maps_by_size(union_size, depth_dict['mask'], depth_dict['context'],
                                depth_dict['real_depth'], depth_dict['zero_mean_depth'], depth_dict['edge'])
    tensor_depth_dict = convert2tensor(patch_depth_dict)
    resize_mask = open_small_mask(tensor_depth_dict['mask'], tensor_depth_dict['context'], 3, 41)
    with torch.no_grad():
        device = config["gpu_ids"] if isinstance(config["gpu_ids"], int) and config["gpu_ids"] >= 0 else "cpu"
        depth_output = depth_feat_model.forward_3P(resize_mask,
                                                    tensor_depth_dict['context'],
                                                    tensor_depth_dict['zero_mean_depth'],
                                                    tensor_depth_dict['edge'],
                                                    unit_length=128,
                                                    cuda=device)
        depth_output = depth_output.cpu()
    tensor_depth_dict['output'] = torch.exp(depth_output + depth_dict['mean_depth']) * \
                                            tensor_depth_dict['mask'] + tensor_depth_dict['depth']
    patch_depth_dict['output'] = tensor_depth_dict['output'].data.cpu().numpy().squeeze()
    depth_dict['output'] = np.zeros((mesh.graph['H'], mesh.graph['W']))
    depth_dict['output'][union_size['x_min']:union_size['x_max'], union_size['y_min']:union_size['y_max']] = \
        patch_depth_dict['output']
    depth_output = depth_dict['output'] * depth_dict['mask'] + depth_dict['depth'] * depth_dict['context']
    depth_output = smooth_cntsyn_gap(depth_dict['output'].copy() * depth_dict['mask'] + depth_dict['depth'] * depth_dict['context'],
                                    depth_dict['mask'], depth_dict['context'],
                                    init_mask_region=depth_dict['mask'])
    if spdb is True:
        f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True);
        ax1.imshow(depth_output * depth_dict['mask'] + depth_dict['depth']); ax2.imshow(depth_dict['output'] * depth_dict['mask'] + depth_dict['depth']); plt.show()
        import pdb; pdb.set_trace()
    depth_dict['output'] = depth_output * depth_dict['mask'] + depth_dict['depth'] * depth_dict['context']

    return depth_dict

def update_info(mapping_dict, info_on_pix, *meshes):
    rt_meshes = []
    for mesh in meshes:
        rt_meshes.append(relabel_node(mesh, mesh.nodes, [*mapping_dict.keys()][0], [*mapping_dict.values()][0]))
    x, y, _ = [*mapping_dict.keys()][0]
    info_on_pix[(x, y)][0]['depth'] = [*mapping_dict.values()][0][2]

    return [info_on_pix] + rt_meshes

def build_connection(mesh, cur_node, dst_node):
    if (abs(cur_node[0] - dst_node[0]) + abs(cur_node[1] - dst_node[1])) < 2:
        mesh.add_edge(cur_node, dst_node)
    if abs(cur_node[0] - dst_node[0]) > 1 or abs(cur_node[1] - dst_node[1]) > 1:
        return mesh
    ne_nodes = [*mesh.neighbors(cur_node)].copy()
    for ne_node in ne_nodes:
        if mesh.has_edge(ne_node, dst_node) or ne_node == dst_node:
            continue
        else:
            mesh = build_connection(mesh, ne_node, dst_node)

    return mesh

def recursive_add_edge(edge_mesh, mesh, info_on_pix, cur_node, mark):
    ne_nodes = [(x[0], x[1]) for x in edge_mesh.neighbors(cur_node)]
    for node_xy in ne_nodes:
        node = (node_xy[0], node_xy[1], info_on_pix[node_xy][0]['depth'])
        if mark[node[0], node[1]] != 3:
            continue
        else:
            mark[node[0], node[1]] = 0
            mesh.remove_edges_from([(xx, node) for xx in mesh.neighbors(node)])
            mesh = build_connection(mesh, cur_node, node)
            re_info = dict(depth=0, count=0)
            for re_ne in mesh.neighbors(node):
                re_info['depth'] += re_ne[2]
                re_info['count'] += 1.
            try:
                re_depth = re_info['depth'] / re_info['count']
            except:
                re_depth = node[2]
            re_node = (node_xy[0], node_xy[1], re_depth)
            mapping_dict = {node: re_node}
            info_on_pix, edge_mesh, mesh = update_info(mapping_dict, info_on_pix, edge_mesh, mesh)

            edge_mesh, mesh, mark, info_on_pix = recursive_add_edge(edge_mesh, mesh, info_on_pix, re_node, mark)

    return edge_mesh, mesh, mark, info_on_pix

def resize_for_edge(tensor_dict, largest_size):
    resize_dict = {k: v.clone() for k, v in tensor_dict.items()}
    frac = largest_size / np.array([*resize_dict['edge'].shape[-2:]]).max()
    if frac < 1:
        resize_mark = torch.nn.functional.interpolate(torch.cat((resize_dict['mask'],
                                                        resize_dict['context']),
                                                        dim=1),
                                                        scale_factor=frac,
                                                        mode='bilinear')
        resize_dict['mask'] = (resize_mark[:, 0:1] > 0).float()
        resize_dict['context'] = (resize_mark[:, 1:2] == 1).float()
        resize_dict['context'][resize_dict['mask'] > 0] = 0
        resize_dict['edge'] = torch.nn.functional.interpolate(resize_dict['edge'],
                                                                scale_factor=frac,
                                                                mode='bilinear')
        resize_dict['edge'] = (resize_dict['edge'] > 0).float()
        resize_dict['edge'] = resize_dict['edge'] * resize_dict['context']
        resize_dict['disp'] = torch.nn.functional.interpolate(resize_dict['disp'],
                                                                scale_factor=frac,
                                                                mode='nearest')
        resize_dict['disp'] = resize_dict['disp'] * resize_dict['context']
        resize_dict['rgb'] = torch.nn.functional.interpolate(resize_dict['rgb'],
                                                                    scale_factor=frac,
                                                                    mode='bilinear')
        resize_dict['rgb'] = resize_dict['rgb'] * resize_dict['context']
    return resize_dict

def get_map_from_nodes(nodes, height, width):
    omap = np.zeros((height, width))
    for n in nodes:
        omap[n[0], n[1]] = 1

    return omap

def get_map_from_ccs(ccs, height, width, condition_input=None, condition=None, real_id=False, id_shift=0):
    if condition is None:
        condition = lambda x, condition_input: True

    if real_id is True:
        omap = np.zeros((height, width)) + (-1) + id_shift
    else:
        omap = np.zeros((height, width))
    for cc_id, cc in enumerate(ccs):
        for n in cc:
            if condition(n, condition_input):
                if real_id is True:
                    omap[n[0], n[1]] = cc_id + id_shift
                else:
                    omap[n[0], n[1]] = 1
    return omap

def revise_map_by_nodes(nodes, imap, operation, limit_constr=None):
    assert operation == '+' or operation == '-', "Operation must be '+' (union) or '-' (exclude)"
    omap = copy.deepcopy(imap)
    revise_flag = True
    if operation == '+':
        for n in nodes:
            omap[n[0], n[1]] = 1
        if limit_constr is not None and omap.sum() > limit_constr:
            omap = imap
            revise_flag = False
    elif operation == '-':
        for n in nodes:
            omap[n[0], n[1]] = 0
        if limit_constr is not None and omap.sum() < limit_constr:
            omap = imap
            revise_flag = False

    return omap, revise_flag

def repaint_info(mesh, cc, x_anchor, y_anchor, source_type):
    if source_type == 'rgb':
        feat = np.zeros((3, x_anchor[1] - x_anchor[0], y_anchor[1] - y_anchor[0]))
    else:
        feat = np.zeros((1, x_anchor[1] - x_anchor[0], y_anchor[1] - y_anchor[0]))
    for node in cc:
        if source_type == 'rgb':
            feat[:, node[0] - x_anchor[0], node[1] - y_anchor[0]] = np.array(mesh.nodes[node]['color']) / 255.
        elif source_type == 'd':
            feat[:, node[0] - x_anchor[0], node[1] - y_anchor[0]] = abs(node[2])

    return feat

def get_context_from_nodes(mesh, cc, H, W, source_type=''):
    if 'rgb' in source_type or 'color' in source_type:
        feat = np.zeros((H, W, 3))
    else:
        feat = np.zeros((H, W))
    context = np.zeros((H, W))
    for node in cc:
        if 'rgb' in source_type or 'color' in source_type:
            feat[node[0], node[1]] = np.array(mesh.nodes[node]['color']) / 255.
            context[node[0], node[1]] = 1
        else:
            feat[node[0], node[1]] = abs(node[2])

    return feat, context

def get_mask_from_nodes(mesh, cc, H, W):
    mask = np.zeros((H, W))
    for node in cc:
        mask[node[0], node[1]] = abs(node[2])

    return mask


def get_edge_from_nodes(context_cc, erode_context_cc, mask_cc, edge_cc, extend_edge_cc, H, W, mesh):
    context = np.zeros((H, W))
    mask = np.zeros((H, W))
    rgb = np.zeros((H, W, 3))
    disp = np.zeros((H, W))
    depth = np.zeros((H, W))
    real_depth = np.zeros((H, W))
    edge = np.zeros((H, W))
    comp_edge = np.zeros((H, W))
    fpath_map = np.zeros((H, W)) - 1
    npath_map = np.zeros((H, W)) - 1
    near_depth = np.zeros((H, W))
    for node in context_cc:
        rgb[node[0], node[1]] = np.array(mesh.nodes[node]['color'])
        disp[node[0], node[1]] = mesh.nodes[node]['disp']
        depth[node[0], node[1]] = node[2]
        context[node[0], node[1]] = 1
    for node in erode_context_cc:
        rgb[node[0], node[1]] = np.array(mesh.nodes[node]['color'])
        disp[node[0], node[1]] = mesh.nodes[node]['disp']
        depth[node[0], node[1]] = node[2]
        context[node[0], node[1]] = 1
    rgb = rgb / 255.
    disp = np.abs(disp)
    disp = disp / disp.max()
    real_depth = depth.copy()
    for node in context_cc:
        if mesh.nodes[node].get('real_depth') is not None:
            real_depth[node[0], node[1]] = mesh.nodes[node]['real_depth']
    for node in erode_context_cc:
        if mesh.nodes[node].get('real_depth') is not None:
            real_depth[node[0], node[1]] = mesh.nodes[node]['real_depth']
    for node in mask_cc:
        mask[node[0], node[1]] = 1
        near_depth[node[0], node[1]] = node[2]
    for node in edge_cc:
        edge[node[0], node[1]] = 1
    for node in extend_edge_cc:
        comp_edge[node[0], node[1]] = 1
    rt_dict = {'rgb': rgb, 'disp': disp, 'depth': depth, 'real_depth': real_depth, 'self_edge': edge, 'context': context,
               'mask': mask, 'fpath_map': fpath_map, 'npath_map': npath_map, 'comp_edge': comp_edge, 'valid_area': context + mask,
               'near_depth': near_depth}

    return rt_dict

def get_depth_from_maps(context_map, mask_map, depth_map, H, W, log_depth=False):
    context = context_map.astype(np.uint8)
    mask = mask_map.astype(np.uint8).copy()
    depth = np.abs(depth_map)
    real_depth = depth.copy()
    zero_mean_depth = np.zeros((H, W))

    if log_depth is True:
        log_depth = np.log(real_depth + 1e-8) * context
        mean_depth = np.mean(log_depth[context > 0])
        zero_mean_depth = (log_depth - mean_depth) * context
    else:
        zero_mean_depth = real_depth
        mean_depth = 0
    edge = np.zeros_like(depth)

    rt_dict = {'depth': depth, 'real_depth': real_depth, 'context': context, 'mask': mask,
               'mean_depth': mean_depth, 'zero_mean_depth': zero_mean_depth, 'edge': edge}

    return rt_dict

def get_depth_from_nodes(context_cc, erode_context_cc, mask_cc, H, W, mesh, log_depth=False):
    context = np.zeros((H, W))
    mask = np.zeros((H, W))
    depth = np.zeros((H, W))
    real_depth = np.zeros((H, W))
    zero_mean_depth = np.zeros((H, W))
    for node in context_cc:
        depth[node[0], node[1]] = node[2]
        context[node[0], node[1]] = 1
    for node in erode_context_cc:
        depth[node[0], node[1]] = node[2]
        context[node[0], node[1]] = 1
    depth = np.abs(depth)
    real_depth = depth.copy()
    for node in context_cc:
        if mesh.nodes[node].get('real_depth') is not None:
            real_depth[node[0], node[1]] = mesh.nodes[node]['real_depth']
    for node in erode_context_cc:
        if mesh.nodes[node].get('real_depth') is not None:
            real_depth[node[0], node[1]] = mesh.nodes[node]['real_depth']
    real_depth = np.abs(real_depth)
    for node in mask_cc:
        mask[node[0], node[1]] = 1
    if log_depth is True:
        log_depth = np.log(real_depth + 1e-8) * context
        mean_depth = np.mean(log_depth[context > 0])
        zero_mean_depth = (log_depth - mean_depth) * context
    else:
        zero_mean_depth = real_depth
        mean_depth = 0

    rt_dict = {'depth': depth, 'real_depth': real_depth, 'context': context, 'mask': mask,
               'mean_depth': mean_depth, 'zero_mean_depth': zero_mean_depth}

    return rt_dict

def get_rgb_from_nodes(context_cc, erode_context_cc, mask_cc, H, W, mesh):
    context = np.zeros((H, W))
    mask = np.zeros((H, W))
    rgb = np.zeros((H, W, 3))
    erode_context = np.zeros((H, W))
    for node in context_cc:
        rgb[node[0], node[1]] = np.array(mesh.nodes[node]['color'])
        context[node[0], node[1]] = 1
    rgb = rgb / 255.
    for node in mask_cc:
        mask[node[0], node[1]] = 1
    for node in erode_context_cc:
        erode_context[node[0], node[1]] = 1
        mask[node[0], node[1]] = 1
    rt_dict = {'rgb': rgb, 'context': context, 'mask': mask,
               'erode': erode_context}

    return rt_dict

def crop_maps_by_size(size, *imaps):
    omaps = []
    for imap in imaps:
        omaps.append(imap[size['x_min']:size['x_max'], size['y_min']:size['y_max']].copy())

    return omaps

def convert2tensor(input_dict):
    rt_dict = {}
    for key, value in input_dict.items():
        if 'rgb' in key or 'color' in key:
            rt_dict[key] = torch.FloatTensor(value).permute(2, 0, 1)[None, ...]
        else:
            rt_dict[key] = torch.FloatTensor(value)[None, None, ...]

    return rt_dict                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                          import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import torch.nn.functional as F


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

    def init_weights(self, init_type='normal', gain=0.02):
        '''
        initialize network's weights
        init_type: normal | xavier | kaiming | orthogonal
        https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39
        '''

        def init_func(m):
            classname = m.__class__.__name__
            if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
                if init_type == 'normal':
                    nn.init.normal_(m.weight.data, 0.0, gain)
                elif init_type == 'xavier':
                    nn.init.xavier_normal_(m.weight.data, gain=gain)
                elif init_type == 'kaiming':
                    nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
                elif init_type == 'orthogonal':
                    nn.init.orthogonal_(m.weight.data, gain=gain)

                if hasattr(m, 'bias') and m.bias is not None:
                    nn.init.constant_(m.bias.data, 0.0)

            elif classname.find('BatchNorm2d') != -1:
                nn.init.normal_(m.weight.data, 1.0, gain)
                nn.init.constant_(m.bias.data, 0.0)

        self.apply(init_func)

def weights_init(init_type='gaussian'):
    def init_fun(m):
        classname = m.__class__.__name__
        if (classname.find('Conv') == 0 or classname.find(
                'Linear') == 0) and hasattr(m, 'weight'):
            if init_type == 'gaussian':
                nn.init.normal_(m.weight, 0.0, 0.02)
            elif init_type == 'xavier':
                nn.init.xavier_normal_(m.weight, gain=math.sqrt(2))
            elif init_type == 'kaiming':
                nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
            elif init_type == 'orthogonal':
                nn.init.orthogonal_(m.weight, gain=math.sqrt(2))
            elif init_type == 'default':
                pass
            else:
                assert 0, "Unsupported initialization: {}".format(init_type)
            if hasattr(m, 'bias') and m.bias is not None:
                nn.init.constant_(m.bias, 0.0)

    return init_fun

class PartialConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, groups=1, bias=True):
        super().__init__()
        self.input_conv = nn.Conv2d(in_channels, out_channels, kernel_size,
                                    stride, padding, dilation, groups, bias)
        self.mask_conv = nn.Conv2d(in_channels, out_channels, kernel_size,
                                   stride, padding, dilation, groups, False)
        self.input_conv.apply(weights_init('kaiming'))
        self.slide_winsize = in_channels * kernel_size * kernel_size

        torch.nn.init.constant_(self.mask_conv.weight, 1.0)

        # mask is not updated
        for param in self.mask_conv.parameters():
            param.requires_grad = False

    def forward(self, input, mask):
        # http://masc.cs.gmu.edu/wiki/partialconv
        # C(X) = W^T * X + b, C(0) = b, D(M) = 1 * M + 0 = sum(M)
        # W^T* (M .* X) / sum(M) + b = [C(M .* X) – C(0)] / D(M) + C(0)
        output = self.input_conv(input * mask)
        if self.input_conv.bias is not None:
            output_bias = self.input_conv.bias.view(1, -1, 1, 1).expand_as(
                output)
        else:
            output_bias = torch.zeros_like(output)

        with torch.no_grad():
            output_mask = self.mask_conv(mask)

        no_update_holes = output_mask == 0

        mask_sum = output_mask.masked_fill_(no_update_holes, 1.0)

        output_pre = ((output - output_bias) * self.slide_winsize) / mask_sum + output_bias
        output = output_pre.masked_fill_(no_update_holes, 0.0)

        new_mask = torch.ones_like(output)
        new_mask = new_mask.masked_fill_(no_update_holes, 0.0)

        return output, new_mask


class PCBActiv(nn.Module):
    def __init__(self, in_ch, out_ch, bn=True, sample='none-3', activ='relu',
                 conv_bias=False):
        super().__init__()
        if sample == 'down-5':
            self.conv = PartialConv(in_ch, out_ch, 5, 2, 2, bias=conv_bias)
        elif sample == 'down-7':
            self.conv = PartialConv(in_ch, out_ch, 7, 2, 3, bias=conv_bias)
        elif sample == 'down-3':
            self.conv = PartialConv(in_ch, out_ch, 3, 2, 1, bias=conv_bias)
        else:
            self.conv = PartialConv(in_ch, out_ch, 3, 1, 1, bias=conv_bias)

        if bn:
            self.bn = nn.BatchNorm2d(out_ch)
        if activ == 'relu':
            self.activation = nn.ReLU()
        elif activ == 'leaky':
            self.activation = nn.LeakyReLU(negative_slope=0.2)

    def forward(self, input, input_mask):
        h, h_mask = self.conv(input, input_mask)
        if hasattr(self, 'bn'):
            h = self.bn(h)
        if hasattr(self, 'activation'):
            h = self.activation(h)
        return h, h_mask

class Inpaint_Depth_Net(nn.Module):
    def __init__(self, layer_size=7, upsampling_mode='nearest'):
        super().__init__()
        in_channels = 4
        out_channels = 1
        self.freeze_enc_bn = False
        self.upsampling_mode = upsampling_mode
        self.layer_size = layer_size
        self.enc_1 = PCBActiv(in_channels, 64, bn=False, sample='down-7', conv_bias=True)
        self.enc_2 = PCBActiv(64, 128, sample='down-5', conv_bias=True)
        self.enc_3 = PCBActiv(128, 256, sample='down-5')
        self.enc_4 = PCBActiv(256, 512, sample='down-3')
        for i in range(4, self.layer_size):
            name = 'enc_{:d}'.format(i + 1)
            setattr(self, name, PCBActiv(512, 512, sample='down-3'))

        for i in range(4, self.layer_size):
            name = 'dec_{:d}'.format(i + 1)
            setattr(self, name, PCBActiv(512 + 512, 512, activ='leaky'))
        self.dec_4 = PCBActiv(512 + 256, 256, activ='leaky')
        self.dec_3 = PCBActiv(256 + 128, 128, activ='leaky')
        self.dec_2 = PCBActiv(128 + 64, 64, activ='leaky')
        self.dec_1 = PCBActiv(64 + in_channels, out_channels,
                              bn=False, activ=None, conv_bias=True)
    def add_border(self, input, mask_flag, PCONV=True):
        with torch.no_grad():
            h = input.shape[-2]
            w = input.shape[-1]
            require_len_unit = 2 ** self.layer_size
            residual_h = int(np.ceil(h / float(require_len_unit)) * require_len_unit - h) # + 2*require_len_unit
            residual_w = int(np.ceil(w / float(require_len_unit)) * require_len_unit - w) # + 2*require_len_unit
            enlarge_input = torch.zeros((input.shape[0], input.shape[1], h + residual_h, w + residual_w)).to(input.device)
            if mask_flag:
                if PCONV is False:
                    enlarge_input += 1.0
                enlarge_input = enlarge_input.clamp(0.0, 1.0)
            else:
                enlarge_input[:, 2, ...] = 0.0
            anchor_h = residual_h//2
            anchor_w = residual_w//2
            enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input

        return enlarge_input, [anchor_h, anchor_h+h, anchor_w, anchor_w+w]

    def forward_3P(self, mask, context, depth, edge, unit_length=128, cuda=None):
        with torch.no_grad():
            input = torch.cat((depth, edge, context, mask), dim=1)
            n, c, h, w = input.shape
            residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h)
            residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w)
            anchor_h = residual_h//2
            anchor_w = residual_w//2
            enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
            enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
            # enlarge_input[:, 3] = 1. - enlarge_input[:, 3]
            depth_output = self.forward(enlarge_input)
            depth_output = depth_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]
            # import pdb; pdb.set_trace()

        return depth_output

    def forward(self, input_feat, refine_border=False, sample=False, PCONV=True):
        input = input_feat
        input_mask = (input_feat[:, -2:-1] + input_feat[:, -1:]).clamp(0, 1).repeat(1, input.shape[1], 1, 1)

        vis_input = input.cpu().data.numpy()
        vis_input_mask = input_mask.cpu().data.numpy()
        H, W = input.shape[-2:]
        if refine_border is True:
            input, anchor = self.add_border(input, mask_flag=False)
            input_mask, anchor = self.add_border(input_mask, mask_flag=True, PCONV=PCONV)
        h_dict = {}  # for the output of enc_N
        h_mask_dict = {}  # for the output of enc_N
        h_dict['h_0'], h_mask_dict['h_0'] = input, input_mask

        h_key_prev = 'h_0'
        for i in range(1, self.layer_size + 1):
            l_key = 'enc_{:d}'.format(i)
            h_key = 'h_{:d}'.format(i)
            h_dict[h_key], h_mask_dict[h_key] = getattr(self, l_key)(
                h_dict[h_key_prev], h_mask_dict[h_key_prev])
            h_key_prev = h_key

        h_key = 'h_{:d}'.format(self.layer_size)
        h, h_mask = h_dict[h_key], h_mask_dict[h_key]

        for i in range(self.layer_size, 0, -1):
            enc_h_key = 'h_{:d}'.format(i - 1)
            dec_l_key = 'dec_{:d}'.format(i)

            h = F.interpolate(h, scale_factor=2, mode=self.upsampling_mode)
            h_mask = F.interpolate(h_mask, scale_factor=2, mode='nearest')

            h = torch.cat([h, h_dict[enc_h_key]], dim=1)
            h_mask = torch.cat([h_mask, h_mask_dict[enc_h_key]], dim=1)
            h, h_mask = getattr(self, dec_l_key)(h, h_mask)
        output = h
        if refine_border is True:
            h_mask = h_mask[..., anchor[0]:anchor[1], anchor[2]:anchor[3]]
            output = output[..., anchor[0]:anchor[1], anchor[2]:anchor[3]]

        return output

class Inpaint_Edge_Net(BaseNetwork):
    def __init__(self, residual_blocks=8, init_weights=True):
        super(Inpaint_Edge_Net, self).__init__()
        in_channels = 7
        out_channels = 1
        self.encoder = []
        # 0
        self.encoder_0 = nn.Sequential(
                            nn.ReflectionPad2d(3),
                            spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=7, padding=0), True),
                            nn.InstanceNorm2d(64, track_running_stats=False),
                            nn.ReLU(True))
        # 1
        self.encoder_1 = nn.Sequential(
                            spectral_norm(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding=1), True),
                            nn.InstanceNorm2d(128, track_running_stats=False),
                            nn.ReLU(True))
        # 2
        self.encoder_2 = nn.Sequential(
                            spectral_norm(nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1), True),
                            nn.InstanceNorm2d(256, track_running_stats=False),
                            nn.ReLU(True))
        # 3
        blocks = []
        for _ in range(residual_blocks):
            block = ResnetBlock(256, 2)
            blocks.append(block)

        self.middle = nn.Sequential(*blocks)
        # + 3
        self.decoder_0 = nn.Sequential(
                            spectral_norm(nn.ConvTranspose2d(in_channels=256+256, out_channels=128, kernel_size=4, stride=2, padding=1), True),
                            nn.InstanceNorm2d(128, track_running_stats=False),
                            nn.ReLU(True))
        # + 2
        self.decoder_1 = nn.Sequential(
                            spectral_norm(nn.ConvTranspose2d(in_channels=128+128, out_channels=64, kernel_size=4, stride=2, padding=1), True),
                            nn.InstanceNorm2d(64, track_running_stats=False),
                            nn.ReLU(True))
        # + 1
        self.decoder_2 = nn.Sequential(
                            nn.ReflectionPad2d(3),
                            nn.Conv2d(in_channels=64+64, out_channels=out_channels, kernel_size=7, padding=0),
                            )

        if init_weights:
            self.init_weights()

    def add_border(self, input, channel_pad_1=None):
        h = input.shape[-2]
        w = input.shape[-1]
        require_len_unit = 16
        residual_h = int(np.ceil(h / float(require_len_unit)) * require_len_unit - h) # + 2*require_len_unit
        residual_w = int(np.ceil(w / float(require_len_unit)) * require_len_unit - w) # + 2*require_len_unit
        enlarge_input = torch.zeros((input.shape[0], input.shape[1], h + residual_h, w + residual_w)).to(input.device)
        if channel_pad_1 is not None:
            for channel in channel_pad_1:
                enlarge_input[:, channel] = 1
        anchor_h = residual_h//2
        anchor_w = residual_w//2
        enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input

        return enlarge_input, [anchor_h, anchor_h+h, anchor_w, anchor_w+w]

    def forward_3P(self, mask, context, rgb, disp, edge, unit_length=128, cuda=None):
        with torch.no_grad():
            input = torch.cat((rgb, disp/disp.max(), edge, context, mask), dim=1)
            n, c, h, w = input.shape
            residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h)
            residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w)
            anchor_h = residual_h//2
            anchor_w = residual_w//2
            enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
            enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
            edge_output = self.forward(enlarge_input)
            edge_output = edge_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]

        return edge_output

    def forward(self, x, refine_border=False):
        if refine_border:
            x, anchor = self.add_border(x, [5])
        x1 = self.encoder_0(x)
        x2 = self.encoder_1(x1)
        x3 = self.encoder_2(x2)
        x4 = self.middle(x3)
        x5 = self.decoder_0(torch.cat((x4, x3), dim=1))
        x6 = self.decoder_1(torch.cat((x5, x2), dim=1))
        x7 = self.decoder_2(torch.cat((x6, x1), dim=1))
        x = torch.sigmoid(x7)
        if refine_border:
            x = x[..., anchor[0]:anchor[1], anchor[2]:anchor[3]]

        return x

class Inpaint_Color_Net(nn.Module):
    def __init__(self, layer_size=7, upsampling_mode='nearest', add_hole_mask=False, add_two_layer=False, add_border=False):
        super().__init__()
        self.freeze_enc_bn = False
        self.upsampling_mode = upsampling_mode
        self.layer_size = layer_size
        in_channels = 6
        self.enc_1 = PCBActiv(in_channels, 64, bn=False, sample='down-7')
        self.enc_2 = PCBActiv(64, 128, sample='down-5')
        self.enc_3 = PCBActiv(128, 256, sample='down-5')
        self.enc_4 = PCBActiv(256, 512, sample='down-3')
        self.enc_5 = PCBActiv(512, 512, sample='down-3')
        self.enc_6 = PCBActiv(512, 512, sample='down-3')
        self.enc_7 = PCBActiv(512, 512, sample='down-3')

        self.dec_7 = PCBActiv(512+512, 512, activ='leaky')
        self.dec_6 = PCBActiv(512+512, 512, activ='leaky')

        self.dec_5A = PCBActiv(512 + 512, 512, activ='leaky')
        self.dec_4A = PCBActiv(512 + 256, 256, activ='leaky')
        self.dec_3A = PCBActiv(256 + 128, 128, activ='leaky')
        self.dec_2A = PCBActiv(128 + 64, 64, activ='leaky')
        self.dec_1A = PCBActiv(64 + in_channels, 3, bn=False, activ=None, conv_bias=True)
        '''
        self.dec_5B = PCBActiv(512 + 512, 512, activ='leaky')
        self.dec_4B = PCBActiv(512 + 256, 256, activ='leaky')
        self.dec_3B = PCBActiv(256 + 128, 128, activ='leaky')
        self.dec_2B = PCBActiv(128 + 64, 64, activ='leaky')
        self.dec_1B = PCBActiv(64 + 4, 1, bn=False, activ=None, conv_bias=True)
        '''
    def cat(self, A, B):
        return torch.cat((A, B), dim=1)

    def upsample(self, feat, mask):
        feat = F.interpolate(feat, scale_factor=2, mode=self.upsampling_mode)
        mask = F.interpolate(mask, scale_factor=2, mode='nearest')

        return feat, mask

    def forward_3P(self, mask, context, rgb, edge, unit_length=128, cuda=None):
        with torch.no_grad():
            input = torch.cat((rgb, edge, context, mask), dim=1)
            n, c, h, w = input.shape
            residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h) # + 128
            residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w) # + 256
            anchor_h = residual_h//2
            anchor_w = residual_w//2
            enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
            enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
            # enlarge_input[:, 3] = 1. - enlarge_input[:, 3]
            enlarge_input = enlarge_input.to(cuda)
            rgb_output = self.forward(enlarge_input)
            rgb_output = rgb_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]

        return rgb_output

    def forward(self, input, add_border=False):
        input_mask = (input[:, -2:-1] + input[:, -1:]).clamp(0, 1)
        H, W = input.shape[-2:]
        f_0, h_0 = input, input_mask.repeat((1,input.shape[1],1,1))
        f_1, h_1 = self.enc_1(f_0, h_0)
        f_2, h_2 = self.enc_2(f_1, h_1)
        f_3, h_3 = self.enc_3(f_2, h_2)
        f_4, h_4 = self.enc_4(f_3, h_3)
        f_5, h_5 = self.enc_5(f_4, h_4)
        f_6, h_6 = self.enc_6(f_5, h_5)
        f_7, h_7 = self.enc_7(f_6, h_6)

        o_7, k_7 = self.upsample(f_7, h_7)
        o_6, k_6 = self.dec_7(self.cat(o_7, f_6), self.cat(k_7, h_6))
        o_6, k_6 = self.upsample(o_6, k_6)
        o_5, k_5 = self.dec_6(self.cat(o_6, f_5), self.cat(k_6, h_5))
        o_5, k_5 = self.upsample(o_5, k_5)
        o_5A, k_5A = o_5, k_5
        o_5B, k_5B = o_5, k_5
        ###############
        o_4A, k_4A = self.dec_5A(self.cat(o_5A, f_4), self.cat(k_5A, h_4))
        o_4A, k_4A = self.upsample(o_4A, k_4A)
        o_3A, k_3A = self.dec_4A(self.cat(o_4A, f_3), self.cat(k_4A, h_3))
        o_3A, k_3A = self.upsample(o_3A, k_3A)
        o_2A, k_2A = self.dec_3A(self.cat(o_3A, f_2), self.cat(k_3A, h_2))
        o_2A, k_2A = self.upsample(o_2A, k_2A)
        o_1A, k_1A = self.dec_2A(self.cat(o_2A, f_1), self.cat(k_2A, h_1))
        o_1A, k_1A = self.upsample(o_1A, k_1A)
        o_0A, k_0A = self.dec_1A(self.cat(o_1A, f_0), self.cat(k_1A, h_0))

        return torch.sigmoid(o_0A)

    def train(self, mode=True):
        """
        Override the default train() to freeze the BN parameters
        """
        super().train(mode)
        if self.freeze_enc_bn:
            for name, module in self.named_modules():
                if isinstance(module, nn.BatchNorm2d) and 'enc' in name:
                    module.eval()

class Discriminator(BaseNetwork):
    def __init__(self, use_sigmoid=True, use_spectral_norm=True, init_weights=True, in_channels=None):
        super(Discriminator, self).__init__()
        self.use_sigmoid = use_sigmoid
        self.conv1 = self.features = nn.Sequential(
            spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),
            nn.LeakyReLU(0.2, inplace=True),
        )

        self.conv2 = nn.Sequential(
            spectral_norm(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),
            nn.LeakyReLU(0.2, inplace=True),
        )

        self.conv3 = nn.Sequential(
            spectral_norm(nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),
            nn.LeakyReLU(0.2, inplace=True),
        )

        self.conv4 = nn.Sequential(
            spectral_norm(nn.Conv2d(in_channels=256, out_channels=512, kernel_size=4, stride=1, padding=1, bias=not use_spectral_norm), use_spectral_norm),
            nn.LeakyReLU(0.2, inplace=True),
        )

        self.conv5 = nn.Sequential(
            spectral_norm(nn.Conv2d(in_channels=512, out_channels=1, kernel_size=4, stride=1, padding=1, bias=not use_spectral_norm), use_spectral_norm),
        )

        if init_weights:
            self.init_weights()

    def forward(self, x):
        conv1 = self.conv1(x)
        conv2 = self.conv2(conv1)
        conv3 = self.conv3(conv2)
        conv4 = self.conv4(conv3)
        conv5 = self.conv5(conv4)

        outputs = conv5
        if self.use_sigmoid:
            outputs = torch.sigmoid(conv5)

        return outputs, [conv1, conv2, conv3, conv4, conv5]

class ResnetBlock(nn.Module):
    def __init__(self, dim, dilation=1):
        super(ResnetBlock, self).__init__()
        self.conv_block = nn.Sequential(
            nn.ReflectionPad2d(dilation),
            spectral_norm(nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3, padding=0, dilation=dilation, bias=not True), True),
            nn.InstanceNorm2d(dim, track_running_stats=False),
            nn.LeakyReLU(negative_slope=0.2),

            nn.ReflectionPad2d(1),
            spectral_norm(nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3, padding=0, dilation=1, bias=not True), True),
            nn.InstanceNorm2d(dim, track_running_stats=False),
        )

    def forward(self, x):
        out = x + self.conv_block(x)

        # Remove ReLU at the end of the residual block
        # http://torch.ch/blog/2016/02/04/resnets.html

        return out


def spectral_norm(module, mode=True):
    if mode:
        return nn.utils.spectral_norm(module)

    return module                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           [tool.black]
line-length = 120
target-version = ['py37']                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              opencv-python==4.2.0.32
vispy==0.6.4
moviepy==1.0.2
transforms3d==0.3.1
networkx==2.3
cynetworkx
scikit-image                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                  import os
import glob
import cv2
import scipy.misc as misc
from skimage.transform import resize
import numpy as np
from functools import reduce
from operator import mul
import torch
from torch import nn
import matplotlib.pyplot as plt
import re
try:
    import cynetworkx as netx
except ImportError:
    import networkx as netx
from scipy.ndimage import gaussian_filter
from skimage.feature import canny
import collections
import shutil
import imageio
import copy
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import time
from scipy.interpolate import interp1d
from collections import namedtuple

def path_planning(num_frames, x, y, z, path_type=''):
    if path_type == 'straight-line':
        corner_points = np.array([[0, 0, 0], [(0 + x) * 0.5, (0 + y) * 0.5, (0 + z) * 0.5], [x, y, z]])
        corner_t = np.linspace(0, 1, len(corner_points))
        t = np.linspace(0, 1, num_frames)
        cs = interp1d(corner_t, corner_points, axis=0, kind='quadratic')
        spline = cs(t)
        xs, ys, zs = [xx.squeeze() for xx in np.split(spline, 3, 1)]
    elif path_type == 'double-straight-line':
        corner_points = np.array([[-x, -y, -z], [0, 0, 0], [x, y, z]])
        corner_t = np.linspace(0, 1, len(corner_points))
        t = np.linspace(0, 1, num_frames)
        cs = interp1d(corner_t, corner_points, axis=0, kind='quadratic')
        spline = cs(t)
        xs, ys, zs = [xx.squeeze() for xx in np.split(spline, 3, 1)]        
    elif path_type == 'circle':
        xs, ys, zs = [], [], []
        for frame_id, bs_shift_val in enumerate(np.arange(-2.0, 2.0, (4./num_frames))):
            xs += [np.cos(bs_shift_val * np.pi) * 1 * x]
            ys += [np.sin(bs_shift_val * np.pi) * 1 * y]
            zs += [np.cos(bs_shift_val * np.pi/2.) * 1 * z]
        xs, ys, zs = np.array(xs), np.array(ys), np.array(zs)

    return xs, ys, zs

def open_small_mask(mask, context, open_iteration, kernel):
    np_mask = mask.cpu().data.numpy().squeeze().astype(np.uint8)
    raw_mask = np_mask.copy()
    np_context = context.cpu().data.numpy().squeeze().astype(np.uint8)
    np_input = np_mask + np_context
    for _ in range(open_iteration):
        np_input = cv2.erode(cv2.dilate(np_input, np.ones((kernel, kernel)), iterations=1), np.ones((kernel,kernel)), iterations=1)
    np_mask[(np_input - np_context) > 0] = 1
    out_mask = torch.FloatTensor(np_mask).to(mask)[None, None, ...]
    
    return out_mask

def filter_irrelevant_edge_new(self_edge, comp_edge, other_edges, other_edges_with_id, current_edge_id, context, depth, mesh, context_cc, spdb=False):
    other_edges = other_edges.squeeze().astype(np.uint8)
    other_edges_with_id = other_edges_with_id.squeeze()
    self_edge = self_edge.squeeze()
    dilate_bevel_self_edge = cv2.dilate((self_edge + comp_edge).astype(np.uint8), np.array([[1,1,1],[1,1,1],[1,1,1]]), iterations=1)
    dilate_cross_self_edge = cv2.dilate((self_edge + comp_edge).astype(np.uint8), np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8), iterations=1)
    edge_ids = np.unique(other_edges_with_id * context + (-1) * (1 - context)).astype(np.int)
    end_depth_maps = np.zeros_like(self_edge)
    self_edge_ids = np.sort(np.unique(other_edges_with_id[self_edge > 0]).astype(np.int))
    self_edge_ids = self_edge_ids[1:] if self_edge_ids.shape[0] > 0  and self_edge_ids[0] == -1 else self_edge_ids
    self_comp_ids = np.sort(np.unique(other_edges_with_id[comp_edge > 0]).astype(np.int))
    self_comp_ids = self_comp_ids[1:] if self_comp_ids.shape[0] > 0  and self_comp_ids[0] == -1 else self_comp_ids
    edge_ids = edge_ids[1:] if edge_ids[0] == -1 else edge_ids
    other_edges_info = []
    extend_other_edges = np.zeros_like(other_edges)
    if spdb is True:
        f, ((ax1, ax2, ax3)) = plt.subplots(1, 3, sharex=True, sharey=True); ax1.imshow(self_edge); ax2.imshow(context); ax3.imshow(other_edges_with_id * context + (-1) * (1 - context)); plt.show()
        import pdb; pdb.set_trace()
    filter_self_edge = np.zeros_like(self_edge)
    for self_edge_id in self_edge_ids:
        filter_self_edge[other_edges_with_id == self_edge_id] = 1
    dilate_self_comp_edge = cv2.dilate(comp_edge, kernel=np.ones((3, 3)), iterations=2)
    valid_self_comp_edge = np.zeros_like(comp_edge)
    for self_comp_id in self_comp_ids:
        valid_self_comp_edge[self_comp_id == other_edges_with_id] = 1
    self_comp_edge = dilate_self_comp_edge * valid_self_comp_edge
    filter_self_edge = (filter_self_edge + self_comp_edge).clip(0, 1)
    for edge_id in edge_ids:
        other_edge_locs = (other_edges_with_id == edge_id).astype(np.uint8)
        condition = (other_edge_locs * other_edges * context.astype(np.uint8))
        end_cross_point = dilate_cross_self_edge * condition * (1 - filter_self_edge)
        end_bevel_point = dilate_bevel_self_edge * condition * (1 - filter_self_edge)
        if end_bevel_point.max() != 0:
            end_depth_maps[end_bevel_point != 0] = depth[end_bevel_point != 0]
            if end_cross_point.max() == 0:
                nxs, nys = np.where(end_bevel_point != 0)
                for nx, ny in zip(nxs, nys):
                    bevel_node = [xx for xx in context_cc if xx[0] == nx and xx[1] == ny][0]
                for ne in mesh.neighbors(bevel_node):
                    if other_edges_with_id[ne[0], ne[1]] > -1 and dilate_cross_self_edge[ne[0], ne[1]] > 0:
                        extend_other_edges[ne[0], ne[1]] = 1
                        break
        else:
            other_edges[other_edges_with_id == edge_id] = 0
    other_edges = (other_edges + extend_other_edges).clip(0, 1) * context

    return other_edges, end_depth_maps, other_edges_info

def clean_far_edge_new(input_edge, end_depth_maps, mask, context, global_mesh, info_on_pix, self_edge, inpaint_id, config):
    mesh = netx.Graph()
    hxs, hys = np.where(input_edge * mask > 0)
    valid_near_edge = (input_edge != 0).astype(np.uint8) * context
    valid_map = mask + context
    invalid_edge_ids = []
    for hx, hy in zip(hxs, hys):
        node = (hx ,hy)
        mesh.add_node((hx, hy))
        eight_nes = [ne for ne in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1), \
                                   (hx + 1, hy + 1), (hx - 1, hy - 1), (hx - 1, hy + 1), (hx + 1, hy - 1)]\
                        if 0 <= ne[0] < input_edge.shape[0] and 0 <= ne[1] < input_edge.shape[1] and 0 < input_edge[ne[0], ne[1]]] # or end_depth_maps[ne[0], ne[1]] != 0]
        for ne in eight_nes:
            mesh.add_edge(node, ne, length=np.hypot(ne[0] - hx, ne[1] - hy))
            if end_depth_maps[ne[0], ne[1]] != 0:
                mesh.nodes[ne[0], ne[1]]['cnt'] = True
                if end_depth_maps[ne[0], ne[1]] == 0:
                    import pdb; pdb.set_trace()
                mesh.nodes[ne[0], ne[1]]['depth'] = end_depth_maps[ne[0], ne[1]]
            elif mask[ne[0], ne[1]] != 1:
                four_nes = [nne for nne in [(ne[0] + 1, ne[1]), (ne[0] - 1, ne[1]), (ne[0], ne[1] + 1), (ne[0], ne[1] - 1)]\
                                 if nne[0] < end_depth_maps.shape[0] and nne[0] >= 0 and nne[1] < end_depth_maps.shape[1] and nne[1] >= 0]
                for nne in four_nes:
                    if end_depth_maps[nne[0], nne[1]] != 0:
                        mesh.add_edge(nne, ne, length=np.hypot(nne[0] - ne[0], nne[1] - ne[1]))
                        mesh.nodes[nne[0], nne[1]]['cnt'] = True
                        mesh.nodes[nne[0], nne[1]]['depth'] = end_depth_maps[nne[0], nne[1]]
    ccs = [*netx.connected_components(mesh)]
    end_pts = []
    for cc in ccs:
        end_pts.append(set())
        for node in cc:
            if mesh.nodes[node].get('cnt') is not None:
                end_pts[-1].add((node[0], node[1], mesh.nodes[node]['depth']))    
    predef_npaths = [None for _ in range(len(ccs))]
    fpath_map = np.zeros_like(input_edge) - 1
    npath_map = np.zeros_like(input_edge) - 1
    npaths, fpaths = dict(), dict()
    break_flag = False
    end_idx = 0
    while end_idx < len(end_pts):
        end_pt, cc = [*zip(end_pts, ccs)][end_idx]
        end_idx += 1
        sorted_end_pt = []
        fpath = []
        iter_fpath = []
        if len(end_pt) > 2 or len(end_pt) == 0:
            if len(end_pt) > 2:
                continue
            continue
        if len(end_pt) == 2:
            ravel_end = [*end_pt]
            tmp_sub_mesh = mesh.subgraph(list(cc)).copy()
            tmp_npath = [*netx.shortest_path(tmp_sub_mesh, (ravel_end[0][0], ravel_end[0][1]), (ravel_end[1][0], ravel_end[1][1]), weight='length')]
            fpath_map1, npath_map1, disp_diff1 = plan_path(mesh, info_on_pix, cc, ravel_end[0:1], global_mesh, input_edge, mask, valid_map, inpaint_id, npath_map=None, fpath_map=None, npath=tmp_npath)
            fpath_map2, npath_map2, disp_diff2 = plan_path(mesh, info_on_pix, cc, ravel_end[1:2], global_mesh, input_edge, mask, valid_map, inpaint_id, npath_map=None, fpath_map=None, npath=tmp_npath)
            tmp_disp_diff = [disp_diff1, disp_diff2]
            self_end = []
            edge_len = []
            ds_edge = cv2.dilate(self_edge.astype(np.uint8), np.ones((3, 3)), iterations=1)
            if ds_edge[ravel_end[0][0], ravel_end[0][1]] > 0:
                self_end.append(1)
            else:
                self_end.append(0)
            if ds_edge[ravel_end[1][0], ravel_end[1][1]] > 0:
                self_end.append(1)
            else:
                self_end.append(0)
            edge_len = [np.count_nonzero(npath_map1), np.count_nonzero(npath_map2)]
            sorted_end_pts = [xx[0] for xx in sorted(zip(ravel_end, self_end, edge_len, [disp_diff1, disp_diff2]), key=lambda x: (x[1], x[2]), reverse=True)]
            re_npath_map1, re_fpath_map1 = (npath_map1 != -1).astype(np.uint8), (fpath_map1 != -1).astype(np.uint8)
            re_npath_map2, re_fpath_map2 = (npath_map2 != -1).astype(np.uint8), (fpath_map2 != -1).astype(np.uint8)
            if np.count_nonzero(re_npath_map1 * re_npath_map2 * mask) / \
                (np.count_nonzero((re_npath_map1 + re_npath_map2) * mask) + 1e-6) > 0.5\
                and np.count_nonzero(re_fpath_map1 * re_fpath_map2 * mask) / \
                     (np.count_nonzero((re_fpath_map1 + re_fpath_map2) * mask) + 1e-6) > 0.5\
                and tmp_disp_diff[0] != -1 and tmp_disp_diff[1] != -1:
                my_fpath_map, my_npath_map, npath, fpath = \
                    plan_path_e2e(mesh, cc, sorted_end_pts, global_mesh, input_edge, mask, valid_map, inpaint_id, npath_map=None, fpath_map=None)
                npath_map[my_npath_map != -1] = my_npath_map[my_npath_map != -1]
                fpath_map[my_fpath_map != -1] = my_fpath_map[my_fpath_map != -1]
                if len(fpath) > 0:
                    edge_id = global_mesh.nodes[[*sorted_end_pts][0]]['edge_id']
                    fpaths[edge_id] = fpath
                    npaths[edge_id] = npath
                invalid_edge_ids.append(edge_id)
            else:
                if tmp_disp_diff[0] != -1:
                    ratio_a = tmp_disp_diff[0] / (np.sum(tmp_disp_diff) + 1e-8)
                else:
                    ratio_a = 0
                if tmp_disp_diff[1] != -1:
                    ratio_b = tmp_disp_diff[1] / (np.sum(tmp_disp_diff) + 1e-8)
                else:
                    ratio_b = 0
                npath_len = len(tmp_npath)
                if npath_len > config['depth_edge_dilate_2'] * 2:
                    npath_len = npath_len - (config['depth_edge_dilate_2'] * 1)
                tmp_npath_a = tmp_npath[:int(np.floor(npath_len * ratio_a))]
                tmp_npath_b = tmp_npath[::-1][:int(np.floor(npath_len * ratio_b))]
                tmp_merge = []
                if len(tmp_npath_a) > 0 and sorted_end_pts[0][0] == tmp_npath_a[0][0] and sorted_end_pts[0][1] == tmp_npath_a[0][1]:
                    if len(tmp_npath_a) > 0 and mask[tmp_npath_a[-1][0], tmp_npath_a[-1][1]] > 0:
                        tmp_merge.append([sorted_end_pts[:1], tmp_npath_a])
                    if len(tmp_npath_b) > 0 and mask[tmp_npath_b[-1][0], tmp_npath_b[-1][1]] > 0:
                        tmp_merge.append([sorted_end_pts[1:2], tmp_npath_b])
                elif len(tmp_npath_b) > 0 and sorted_end_pts[0][0] == tmp_npath_b[0][0] and sorted_end_pts[0][1] == tmp_npath_b[0][1]:
                    if len(tmp_npath_b) > 0 and mask[tmp_npath_b[-1][0], tmp_npath_b[-1][1]] > 0:
                        tmp_merge.append([sorted_end_pts[:1], tmp_npath_b])
                    if len(tmp_npath_a) > 0 and mask[tmp_npath_a[-1][0], tmp_npath_a[-1][1]] > 0:
                        tmp_merge.append([sorted_end_pts[1:2], tmp_npath_a])
                for tmp_idx in range(len(tmp_merge)):
                    if len(tmp_merge[tmp_idx][1]) == 0:
                        continue
                    end_pts.append(tmp_merge[tmp_idx][0])
                    ccs.append(set(tmp_merge[tmp_idx][1]))
        if len(end_pt) == 1:
            sub_mesh = mesh.subgraph(list(cc)).copy()
            pnodes = netx.periphery(sub_mesh)
            if len(end_pt) == 1:
                ends = [*end_pt]
            elif len(sorted_end_pt) == 1:
                ends = [*sorted_end_pt]
            else:
                import pdb; pdb.set_trace()
            try:
                edge_id = global_mesh.nodes[ends[0]]['edge_id']
            except:
                import pdb; pdb.set_trace()
            pnodes = sorted(pnodes, 
                            key=lambda x: np.hypot((x[0] - ends[0][0]), (x[1] - ends[0][1])),
                            reverse=True)[0]
            npath = [*netx.shortest_path(sub_mesh, (ends[0][0], ends[0][1]), pnodes, weight='length')]
            for np_node in npath:
                npath_map[np_node[0], np_node[1]] = edge_id
            fpath = []
            if global_mesh.nodes[ends[0]].get('far') is None:
                print("None far")
            else:
                fnodes = global_mesh.nodes[ends[0]].get('far')
                dmask = mask + 0
                did = 0
                while True:
                    did += 1
                    dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1)
                    if did > 3:
                        break
                    ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0 and mask[fnode[0], fnode[1]] == 0 and\
                                                            global_mesh.nodes[fnode].get('inpaint_id') != inpaint_id + 1)]
                    if len(ffnode) > 0:
                        fnode = ffnode[0]
                        break
                if len(ffnode) == 0:
                    continue
                fpath.append((fnode[0], fnode[1]))
                barrel_dir = np.array([[1, 0], [1, 1], [0, 1], [-1, 1], [-1, 0], [-1, -1], [0, -1], [1, -1]])
                n2f_dir = (int(fnode[0] - npath[0][0]), int(fnode[1] - npath[0][1]))
                while True:
                    if barrel_dir[0, 0] == n2f_dir[0] and barrel_dir[0, 1] == n2f_dir[1]:
                        n2f_barrel = barrel_dir.copy()
                        break
                    barrel_dir = np.roll(barrel_dir, 1, axis=0)
                for step in range(0, len(npath)):
                    if step == 0:
                        continue
                    elif step == 1:
                        next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1])
                        while True:
                            if barrel_dir[0, 0] == next_dir[0] and barrel_dir[0, 1] == next_dir[1]:
                                next_barrel = barrel_dir.copy()
                                break
                            barrel_dir = np.roll(barrel_dir, 1, axis=0)
                        barrel_pair = np.stack((n2f_barrel, next_barrel), axis=0)
                        n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1])
                    elif step > 1:
                        next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1])
                        while True:
                            if barrel_pair[1, 0, 0] == next_dir[0] and barrel_pair[1, 0, 1] == next_dir[1]:
                                next_barrel = barrel_pair.copy()
                                break
                            barrel_pair = np.roll(barrel_pair, 1, axis=1)
                        n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1])
                    new_locs = []
                    if abs(n2f_dir[0]) == 1:
                        new_locs.append((npath[step][0] + n2f_dir[0], npath[step][1]))
                    if abs(n2f_dir[1]) == 1:
                        new_locs.append((npath[step][0], npath[step][1] + n2f_dir[1]))
                    if len(new_locs) > 1:
                        new_locs = sorted(new_locs, key=lambda xx: np.hypot((xx[0] - fpath[-1][0]), (xx[1] - fpath[-1][1])))
                    break_flag = False
                    for new_loc in new_locs:
                        new_loc_nes = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]), 
                                                    (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1)]\
                                            if xx[0] >= 0 and xx[0] < fpath_map.shape[0] and xx[1] >= 0 and xx[1] < fpath_map.shape[1]]
                        if np.all([(fpath_map[nlne[0], nlne[1]] == -1) for nlne in new_loc_nes]) != True:
                            break
                        if npath_map[new_loc[0], new_loc[1]] != -1:
                            if npath_map[new_loc[0], new_loc[1]] != edge_id:
                                break_flag = True
                                break
                            else:
                                continue
                        if valid_map[new_loc[0], new_loc[1]] == 0:
                            break_flag = True
                            break
                        fpath.append(new_loc)
                    if break_flag is True:
                        break
                if step != len(npath) - 1:
                    for xx in npath[step:]:
                        if npath_map[xx[0], xx[1]] == edge_id:
                            npath_map[xx[0], xx[1]] = -1
                    npath = npath[:step]
            if len(fpath) > 0:
                for fp_node in fpath:
                    fpath_map[fp_node[0], fp_node[1]] = edge_id
                fpaths[edge_id] = fpath
                npaths[edge_id] = npath
        fpath_map[valid_near_edge != 0] = -1
        if len(fpath) > 0:
            iter_fpath = copy.deepcopy(fpaths[edge_id])
        for node in iter_fpath:
            if valid_near_edge[node[0], node[1]] != 0:
                fpaths[edge_id].remove(node)

    return fpath_map, npath_map, False, npaths, fpaths, invalid_edge_ids

def plan_path_e2e(mesh, cc, end_pts, global_mesh, input_edge, mask, valid_map, inpaint_id, npath_map=None, fpath_map=None):
    my_npath_map = np.zeros_like(input_edge) - 1
    my_fpath_map = np.zeros_like(input_edge) - 1
    sub_mesh = mesh.subgraph(list(cc)).copy()
    ends_1, ends_2 = end_pts[0], end_pts[1]
    edge_id = global_mesh.nodes[ends_1]['edge_id']
    npath = [*netx.shortest_path(sub_mesh, (ends_1[0], ends_1[1]), (ends_2[0], ends_2[1]), weight='length')]
    for np_node in npath:
        my_npath_map[np_node[0], np_node[1]] = edge_id
    fpath = []
    if global_mesh.nodes[ends_1].get('far') is None:
        print("None far")
    else:
        fnodes = global_mesh.nodes[ends_1].get('far')
        dmask = mask + 0
        while True:
            dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1)
            ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0 and mask[fnode[0], fnode[1]] == 0 and\
                                                            global_mesh.nodes[fnode].get('inpaint_id') != inpaint_id + 1)]
            if len(ffnode) > 0:
                fnode = ffnode[0]
                break
        e_fnodes = global_mesh.nodes[ends_2].get('far')
        dmask = mask + 0
        while True:
            dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1)
            e_ffnode = [e_fnode for e_fnode in e_fnodes if (dmask[e_fnode[0], e_fnode[1]] > 0 and mask[e_fnode[0], e_fnode[1]] == 0 and\
                                                            global_mesh.nodes[e_fnode].get('inpaint_id') != inpaint_id + 1)]
            if len(e_ffnode) > 0:
                e_fnode = e_ffnode[0]
                break            
        fpath.append((fnode[0], fnode[1]))
        if len(e_ffnode) == 0 or len(ffnode) == 0:
            return my_npath_map, my_fpath_map, [], []
        barrel_dir = np.array([[1, 0], [1, 1], [0, 1], [-1, 1], [-1, 0], [-1, -1], [0, -1], [1, -1]])
        n2f_dir = (int(fnode[0] - npath[0][0]), int(fnode[1] - npath[0][1]))
        while True:
            if barrel_dir[0, 0] == n2f_dir[0] and barrel_dir[0, 1] == n2f_dir[1]:
                n2f_barrel = barrel_dir.copy()
                break
            barrel_dir = np.roll(barrel_dir, 1, axis=0)
        for step in range(0, len(npath)):
            if step == 0:
                continue
            elif step == 1:
                next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1])
                while True:
                    if barrel_dir[0, 0] == next_dir[0] and barrel_dir[0, 1] == next_dir[1]:
                        next_barrel = barrel_dir.copy()
                        break
                    barrel_dir = np.roll(barrel_dir, 1, axis=0)
                barrel_pair = np.stack((n2f_barrel, next_barrel), axis=0)
                n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1])
            elif step > 1:
                next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1])
                while True:
                    if barrel_pair[1, 0, 0] == next_dir[0] and barrel_pair[1, 0, 1] == next_dir[1]:
                        next_barrel = barrel_pair.copy()
                        break
                    barrel_pair = np.roll(barrel_pair, 1, axis=1)
                n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1])
            new_locs = []
            if abs(n2f_dir[0]) == 1:
                new_locs.append((npath[step][0] + n2f_dir[0], npath[step][1]))
            if abs(n2f_dir[1]) == 1:
                new_locs.append((npath[step][0], npath[step][1] + n2f_dir[1]))
            if len(new_locs) > 1:
                new_locs = sorted(new_locs, key=lambda xx: np.hypot((xx[0] - fpath[-1][0]), (xx[1] - fpath[-1][1])))
            break_flag = False
            for new_loc in new_locs:
                new_loc_nes = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]),
                                            (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1)]\
                                    if xx[0] >= 0 and xx[0] < my_fpath_map.shape[0] and xx[1] >= 0 and xx[1] < my_fpath_map.shape[1]]
                if fpath_map is not None and np.sum([fpath_map[nlne[0], nlne[1]] for nlne in new_loc_nes]) != 0:
                    break_flag = True
                    break
                if my_npath_map[new_loc[0], new_loc[1]] != -1:
                    continue
                if npath_map is not None and npath_map[new_loc[0], new_loc[1]] != edge_id:
                    break_flag = True
                    break
                fpath.append(new_loc)
            if break_flag is True:
                break
        if (e_fnode[0], e_fnode[1]) not in fpath:
            fpath.append((e_fnode[0], e_fnode[1]))
        if step != len(npath) - 1:
            for xx in npath[step:]:
                if my_npath_map[xx[0], xx[1]] == edge_id:
                    my_npath_map[xx[0], xx[1]] = -1
            npath = npath[:step]
        if len(fpath) > 0:
            for fp_node in fpath:
                my_fpath_map[fp_node[0], fp_node[1]] = edge_id
    
    return my_fpath_map, my_npath_map, npath, fpath

def plan_path(mesh, info_on_pix, cc, end_pt, global_mesh, input_edge, mask, valid_map, inpaint_id, npath_map=None, fpath_map=None, npath=None):
    my_npath_map = np.zeros_like(input_edge) - 1
    my_fpath_map = np.zeros_like(input_edge) - 1
    sub_mesh = mesh.subgraph(list(cc)).copy()
    pnodes = netx.periphery(sub_mesh)
    ends = [*end_pt]
    edge_id = global_mesh.nodes[ends[0]]['edge_id']
    pnodes = sorted(pnodes, 
                    key=lambda x: np.hypot((x[0] - ends[0][0]), (x[1] - ends[0][1])),
                    reverse=True)[0]
    if npath is None:
        npath = [*netx.shortest_path(sub_mesh, (ends[0][0], ends[0][1]), pnodes, weight='length')]
    else:
        if (ends[0][0], ends[0][1]) == npath[0]:
            npath = npath
        elif (ends[0][0], ends[0][1]) == npath[-1]:
            npath = npath[::-1]
        else:
            import pdb; pdb.set_trace()
    for np_node in npath:
        my_npath_map[np_node[0], np_node[1]] = edge_id
    fpath = []
    if global_mesh.nodes[ends[0]].get('far') is None:
        print("None far")
    else:
        fnodes = global_mesh.nodes[ends[0]].get('far')
        dmask = mask + 0
        did = 0
        while True:
            did += 1
            if did > 3:
                return my_fpath_map, my_npath_map, -1
            dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1)
            ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0 and mask[fnode[0], fnode[1]] == 0 and\
                                                            global_mesh.nodes[fnode].get('inpaint_id') != inpaint_id + 1)]
            if len(ffnode) > 0:
                fnode = ffnode[0]
                break
        
        fpath.append((fnode[0], fnode[1]))
        disp_diff = 0.
        for n_loc in npath:
            if mask[n_loc[0], n_loc[1]] != 0:
                disp_diff = abs(abs(1. / info_on_pix[(n_loc[0], n_loc[1])][0]['depth']) - abs(1. / ends[0][2]))
                break
        barrel_dir = np.array([[1, 0], [1, 1], [0, 1], [-1, 1], [-1, 0], [-1, -1], [0, -1], [1, -1]])
        n2f_dir = (int(fnode[0] - npath[0][0]), int(fnode[1] - npath[0][1]))
        while True:
            if barrel_dir[0, 0] == n2f_dir[0] and barrel_dir[0, 1] == n2f_dir[1]:
                n2f_barrel = barrel_dir.copy()
                break
            barrel_dir = np.roll(barrel_dir, 1, axis=0)
        for step in range(0, len(npath)):
            if step == 0:
                continue
            elif step == 1:
                next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1])
                while True:
                    if barrel_dir[0, 0] == next_dir[0] and barrel_dir[0, 1] == next_dir[1]:
                        next_barrel = barrel_dir.copy()
                        break
                    barrel_dir = np.roll(barrel_dir, 1, axis=0)
                barrel_pair = np.stack((n2f_barrel, next_barrel), axis=0)
                n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1])
            elif step > 1:
                next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1])
                while True:
                    if barrel_pair[1, 0, 0] == next_dir[0] and barrel_pair[1, 0, 1] == next_dir[1]:
                        next_barrel = barrel_pair.copy()
                        break
                    barrel_pair = np.roll(barrel_pair, 1, axis=1)
                n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1])
            new_locs = []
            if abs(n2f_dir[0]) == 1:
                new_locs.append((npath[step][0] + n2f_dir[0], npath[step][1]))
            if abs(n2f_dir[1]) == 1:
                new_locs.append((npath[step][0], npath[step][1] + n2f_dir[1]))
            if len(new_locs) > 1:
                new_locs = sorted(new_locs, key=lambda xx: np.hypot((xx[0] - fpath[-1][0]), (xx[1] - fpath[-1][1])))
            break_flag = False
            for new_loc in new_locs:
                new_loc_nes = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]),
                                        (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1)]\
                                if xx[0] >= 0 and xx[0] < my_fpath_map.shape[0] and xx[1] >= 0 and xx[1] < my_fpath_map.shape[1]]
                if fpath_map is not None and np.all([(fpath_map[nlne[0], nlne[1]] == -1) for nlne in new_loc_nes]) != True:
                    break_flag = True
                    break
                if np.all([(my_fpath_map[nlne[0], nlne[1]] == -1) for nlne in new_loc_nes]) != True:
                    break_flag = True
                    break 
                if my_npath_map[new_loc[0], new_loc[1]] != -1:
                    continue
                if npath_map is not None and npath_map[new_loc[0], new_loc[1]] != edge_id:
                    break_flag = True
                    break
                if valid_map[new_loc[0], new_loc[1]] == 0:
                    break_flag = True
                    break
                fpath.append(new_loc)
            if break_flag is True:
                break
        if step != len(npath) - 1:
            for xx in npath[step:]:
                if my_npath_map[xx[0], xx[1]] == edge_id:
                    my_npath_map[xx[0], xx[1]] = -1
            npath = npath[:step]
        if len(fpath) > 0:
            for fp_node in fpath:
                my_fpath_map[fp_node[0], fp_node[1]] = edge_id

    return my_fpath_map, my_npath_map, disp_diff

def refresh_node(old_node, old_feat, new_node, new_feat, mesh, stime=False):
    mesh.add_node(new_node)
    mesh.nodes[new_node].update(new_feat)
    mesh.nodes[new_node].update(old_feat)
    for ne in mesh.neighbors(old_node):
        mesh.add_edge(new_node, ne)
    if mesh.nodes[new_node].get('far') is not None:
        tmp_far_nodes = mesh.nodes[new_node]['far']
        for far_node in tmp_far_nodes:
            if mesh.has_node(far_node) is False:
                mesh.nodes[new_node]['far'].remove(far_node)
                continue
            if mesh.nodes[far_node].get('near') is not None:
                for idx in range(len(mesh.nodes[far_node].get('near'))):
                    if mesh.nodes[far_node]['near'][idx][0] == new_node[0] and mesh.nodes[far_node]['near'][idx][1] == new_node[1]:
                        if len(mesh.nodes[far_node]['near'][idx]) == len(old_node):
                            mesh.nodes[far_node]['near'][idx] = new_node
    if mesh.nodes[new_node].get('near') is not None:
        tmp_near_nodes = mesh.nodes[new_node]['near']
        for near_node in tmp_near_nodes:
            if mesh.has_node(near_node) is False:
                mesh.nodes[new_node]['near'].remove(near_node)
                continue        
            if mesh.nodes[near_node].get('far') is not None:
                for idx in range(len(mesh.nodes[near_node].get('far'))):
                    if mesh.nodes[near_node]['far'][idx][0] == new_node[0] and mesh.nodes[near_node]['far'][idx][1] == new_node[1]:
                        if len(mesh.nodes[near_node]['far'][idx]) == len(old_node):
                            mesh.nodes[near_node]['far'][idx] = new_node
    if new_node != old_node:
        mesh.remove_node(old_node)
    if stime is False:
        return mesh
    else:
        return mesh, None, None


def create_placeholder(context, mask, depth, fpath_map, npath_map, mesh, inpaint_id, edge_ccs, extend_edge_cc, all_edge_maps, self_edge_id):
    add_node_time = 0
    add_edge_time = 0
    add_far_near_time = 0
    valid_area = context + mask
    H, W = mesh.graph['H'], mesh.graph['W']
    edge_cc = edge_ccs[self_edge_id]
    num_com = len(edge_cc) + len(extend_edge_cc)
    hxs, hys = np.where(mask > 0)
    for hx, hy in zip(hxs, hys):
        mesh.add_node((hx, hy), inpaint_id=inpaint_id + 1, num_context=num_com)
    for hx, hy in zip(hxs, hys):
        four_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] if\
                        0 <= x < mesh.graph['H'] and 0 <= y < mesh.graph['W'] and valid_area[x, y] != 0]
        for ne in four_nes:
            if mask[ne[0], ne[1]] != 0:
                if not mesh.has_edge((hx, hy), ne):
                    mesh.add_edge((hx, hy), ne)
            elif depth[ne[0], ne[1]] != 0:
                if mesh.has_node((ne[0], ne[1], depth[ne[0], ne[1]])) and\
                    not mesh.has_edge((hx, hy), (ne[0], ne[1], depth[ne[0], ne[1]])):
                    mesh.add_edge((hx, hy), (ne[0], ne[1], depth[ne[0], ne[1]]))
                else:
                    print("Undefined context node.")
                    import pdb; pdb.set_trace()
    near_ids = np.unique(npath_map)
    if near_ids[0] == -1: near_ids = near_ids[1:]
    for near_id in near_ids:
        hxs, hys = np.where((fpath_map == near_id) & (mask > 0))
        if hxs.shape[0] > 0:
            mesh.graph['max_edge_id'] = mesh.graph['max_edge_id'] + 1
        else:
            break
        for hx, hy in zip(hxs, hys):
            mesh.nodes[(hx, hy)]['edge_id'] = int(round(mesh.graph['max_edge_id']))
            four_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] if\
                        x < mesh.graph['H'] and x >= 0 and y < mesh.graph['W'] and y >= 0 and npath_map[x, y] == near_id]
            for xx in four_nes:
                xx_n = copy.deepcopy(xx)
                if not mesh.has_node(xx_n):
                    if mesh.has_node((xx_n[0], xx_n[1], depth[xx_n[0], xx_n[1]])):
                        xx_n = (xx_n[0], xx_n[1], depth[xx_n[0], xx_n[1]])
                if mesh.has_edge((hx, hy), xx_n):
                    # pass
                    mesh.remove_edge((hx, hy), xx_n)
                if mesh.nodes[(hx, hy)].get('near') is None:
                    mesh.nodes[(hx, hy)]['near'] = []
                mesh.nodes[(hx, hy)]['near'].append(xx_n)
        connect_point_exception = set()
        hxs, hys = np.where((npath_map == near_id) & (all_edge_maps > -1))
        for hx, hy in zip(hxs, hys):
            unknown_id = int(round(all_edge_maps[hx, hy]))
            if unknown_id != near_id and unknown_id != self_edge_id:
                unknown_node = set([xx for xx in edge_ccs[unknown_id] if xx[0] == hx and xx[1] == hy])
                connect_point_exception |= unknown_node
        hxs, hys = np.where((npath_map == near_id) & (mask > 0))                
        if hxs.shape[0] > 0:
            mesh.graph['max_edge_id'] = mesh.graph['max_edge_id'] + 1
        else:
            break
        for hx, hy in zip(hxs, hys):
            mesh.nodes[(hx, hy)]['edge_id'] = int(round(mesh.graph['max_edge_id']))
            mesh.nodes[(hx, hy)]['connect_point_id'] = int(round(near_id)) 
            mesh.nodes[(hx, hy)]['connect_point_exception'] = connect_point_exception
            four_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] if\
                        x < mesh.graph['H'] and x >= 0 and y < mesh.graph['W'] and y >= 0 and fpath_map[x, y] == near_id]
            for xx in four_nes:
                xx_n = copy.deepcopy(xx)
                if not mesh.has_node(xx_n):
                    if mesh.has_node((xx_n[0], xx_n[1], depth[xx_n[0], xx_n[1]])):
                        xx_n = (xx_n[0], xx_n[1], depth[xx_n[0], xx_n[1]])
                if mesh.has_edge((hx, hy), xx_n):
                    mesh.remove_edge((hx, hy), xx_n)
                if mesh.nodes[(hx, hy)].get('far') is None:
                    mesh.nodes[(hx, hy)]['far'] = []
                mesh.nodes[(hx, hy)]['far'].append(xx_n)

    return mesh, add_node_time, add_edge_time, add_far_near_time

def clean_far_edge(mask_edge, mask_edge_with_id, context_edge, mask, info_on_pix, global_mesh, anchor):
    if isinstance(mask_edge, torch.Tensor):
        if mask_edge.is_cuda:
            mask_edge = mask_edge.cpu()
        mask_edge = mask_edge.data
        mask_edge = mask_edge.numpy()
    if isinstance(context_edge, torch.Tensor):
        if context_edge.is_cuda:
            context_edge = context_edge.cpu()
        context_edge = context_edge.data
        context_edge = context_edge.numpy()
    if isinstance(mask, torch.Tensor):
        if mask.is_cuda:
            mask = mask.cpu()
        mask = mask.data
        mask = mask.numpy()
    mask = mask.squeeze()
    mask_edge = mask_edge.squeeze()
    context_edge = context_edge.squeeze()
    valid_near_edge = np.zeros_like(mask_edge)
    far_edge = np.zeros_like(mask_edge)
    far_edge_with_id = np.ones_like(mask_edge) * -1
    near_edge_with_id = np.ones_like(mask_edge) * -1
    uncleaned_far_edge = np.zeros_like(mask_edge)
    # Detect if there is any valid pixel mask_edge, if not ==> return default value
    if mask_edge.sum() == 0:
        return far_edge, uncleaned_far_edge, far_edge_with_id, near_edge_with_id
    mask_edge_ids = dict(collections.Counter(mask_edge_with_id.flatten())).keys()
    for edge_id in mask_edge_ids:
        if edge_id < 0:
            continue
        specific_edge_map = (mask_edge_with_id == edge_id).astype(np.uint8)
        _, sub_specific_edge_maps = cv2.connectedComponents(specific_edge_map.astype(np.uint8), connectivity=8)
        for sub_edge_id in range(1, sub_specific_edge_maps.max() + 1):
            specific_edge_map = (sub_specific_edge_maps == sub_edge_id).astype(np.uint8)
            edge_pxs, edge_pys = np.where(specific_edge_map > 0)
            edge_mesh = netx.Graph()
            for edge_px, edge_py in zip(edge_pxs, edge_pys):
                edge_mesh.add_node((edge_px, edge_py))
                for ex in [edge_px-1, edge_px, edge_px+1]:
                    for ey in [edge_py-1, edge_py, edge_py+1]:
                        if edge_px == ex and edge_py == ey:
                            continue
                        if ex < 0 or ex >= specific_edge_map.shape[0] or ey < 0 or ey >= specific_edge_map.shape[1]:
                            continue
                        if specific_edge_map[ex, ey] == 1:
                            if edge_mesh.has_node((ex, ey)):
                                edge_mesh.add_edge((ex, ey), (edge_px, edge_py))
            periphery_nodes = netx.periphery(edge_mesh)
            path_diameter = netx.diameter(edge_mesh)
            start_near_node = None
            for node_s in periphery_nodes:
                for node_e in periphery_nodes:
                    if node_s != node_e:
                        if netx.shortest_path_length(edge_mesh, node_s, node_e) == path_diameter:
                            if np.any(context_edge[node_s[0]-1:node_s[0]+2, node_s[1]-1:node_s[1]+2].flatten()):
                                start_near_node = (node_s[0], node_s[1])
                                end_near_node = (node_e[0], node_e[1])
                                break
                            if np.any(context_edge[node_e[0]-1:node_e[0]+2, node_e[1]-1:node_e[1]+2].flatten()):
                                start_near_node = (node_e[0], node_e[1])
                                end_near_node = (node_s[0], node_s[1])
                                break
                if start_near_node is not None:
                    break
            if start_near_node is None:
                continue
            new_specific_edge_map = np.zeros_like(mask)
            for path_node in netx.shortest_path(edge_mesh, start_near_node, end_near_node):
                new_specific_edge_map[path_node[0], path_node[1]] = 1
            context_near_pxs, context_near_pys = np.where(context_edge[start_near_node[0]-1:start_near_node[0]+2, start_near_node[1]-1:start_near_node[1]+2] > 0)
            distance = np.abs((context_near_pxs - 1)) + np.abs((context_near_pys - 1))
            if (np.where(distance == distance.min())[0].shape[0]) > 1:
                closest_pxs = context_near_pxs[np.where(distance == distance.min())[0]]
                closest_pys = context_near_pys[np.where(distance == distance.min())[0]]
                closest_depths = []
                for closest_px, closest_py in zip(closest_pxs, closest_pys):
                    if info_on_pix.get((closest_px + start_near_node[0] - 1 + anchor[0], closest_py + start_near_node[1] - 1 + anchor[2])) is not None:
                        for info in info_on_pix.get((closest_px + start_near_node[0] - 1 + anchor[0], closest_py + start_near_node[1] - 1 + anchor[2])):
                            if info['synthesis'] is False:
                                closest_depths.append(abs(info['depth']))
                context_near_px, context_near_py = closest_pxs[np.array(closest_depths).argmax()], closest_pys[np.array(closest_depths).argmax()]
            else:
                context_near_px, context_near_py = context_near_pxs[distance.argmin()], context_near_pys[distance.argmin()]
            context_near_node = (start_near_node[0]-1 + context_near_px, start_near_node[1]-1 + context_near_py)
            far_node_list = []
            global_context_near_node = (context_near_node[0] + anchor[0], context_near_node[1] + anchor[2])
            if info_on_pix.get(global_context_near_node) is not None:
                for info in info_on_pix[global_context_near_node]:
                    if info['synthesis'] is False:
                        context_near_node_3d = (global_context_near_node[0], global_context_near_node[1], info['depth'])
                        if global_mesh.nodes[context_near_node_3d].get('far') is not None:
                            for far_node in global_mesh.nodes[context_near_node_3d].get('far'):
                                far_node = (far_node[0] - anchor[0], far_node[1] - anchor[2], far_node[2])
                                if mask[far_node[0], far_node[1]] == 0:
                                    far_node_list.append([far_node[0], far_node[1]])
            if len(far_node_list) > 0:
                far_nodes_dist = np.sum(np.abs(np.array(far_node_list) - np.array([[edge_px, edge_py]])), axis=1)
                context_far_node = tuple(far_node_list[far_nodes_dist.argmin()])
                corresponding_far_edge = np.zeros_like(mask_edge)
                corresponding_far_edge[context_far_node[0], context_far_node[1]] = 1
                surround_map = cv2.dilate(new_specific_edge_map.astype(np.uint8), 
                                            np.array([[1,1,1],[1,1,1],[1,1,1]]).astype(np.uint8), 
                                            iterations=1)
                specific_edge_map_wo_end_pt = new_specific_edge_map.copy()
                specific_edge_map_wo_end_pt[end_near_node[0], end_near_node[1]] = 0
                surround_map_wo_end_pt = cv2.dilate(specific_edge_map_wo_end_pt.astype(np.uint8), 
                                                    np.array([[1,1,1],[1,1,1],[1,1,1]]).astype(np.uint8), 
                                                    iterations=1)
                surround_map_wo_end_pt[new_specific_edge_map > 0] = 0
                surround_map_wo_end_pt[context_near_node[0], context_near_node[1]] = 0
                surround_map = surround_map_wo_end_pt.copy()
                _, far_edge_cc = cv2.connectedComponents(surround_map.astype(np.uint8), connectivity=4)
                start_far_node = None
                accompany_far_node = None
                if surround_map[context_far_node[0], context_far_node[1]] == 1:
                    start_far_node = context_far_node
                else:
                    four_nes = [(context_far_node[0] - 1, context_far_node[1]), 
                            (context_far_node[0] + 1, context_far_node[1]), 
                            (context_far_node[0], context_far_node[1] - 1), 
                            (context_far_node[0], context_far_node[1] + 1)]
                    candidate_bevel = []            
                    for ne in four_nes:
                        if surround_map[ne[0], ne[1]] == 1:
                            start_far_node = (ne[0], ne[1])
                            break
                        elif (ne[0] != context_near_node[0] or ne[1] != context_near_node[1]) and \
                                (ne[0] != start_near_node[0] or ne[1] != start_near_node[1]):
                            candidate_bevel.append((ne[0], ne[1]))
                    if start_far_node is None:
                        for ne in candidate_bevel:
                            if ne[0] == context_far_node[0]:
                                bevel_xys = [[ne[0] + 1, ne[1]], [ne[0] - 1, ne[1]]]
                            if ne[1] == context_far_node[1]:
                                bevel_xys = [[ne[0], ne[1] + 1], [ne[0], ne[1] - 1]]
                            for bevel_x, bevel_y in bevel_xys:
                                if surround_map[bevel_x, bevel_y] == 1:
                                    start_far_node = (bevel_x, bevel_y)
                                    accompany_far_node = (ne[0], ne[1])
                                    break
                            if start_far_node is not None:
                                break
                if start_far_node is not None:
                    for far_edge_id in range(1, far_edge_cc.max() + 1):
                        specific_far_edge = (far_edge_cc == far_edge_id).astype(np.uint8)
                        if specific_far_edge[start_far_node[0], start_far_node[1]] == 1:
                            if accompany_far_node is not None:
                                specific_far_edge[accompany_far_node] = 1
                            far_edge[specific_far_edge > 0] = 1
                            far_edge_with_id[specific_far_edge > 0] = edge_id
                            end_far_candidates = np.zeros_like(far_edge)
                            end_far_candidates[end_near_node[0], end_near_node[1]] = 1
                            end_far_candidates = cv2.dilate(end_far_candidates.astype(np.uint8), 
                                                            np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8), 
                                                            iterations=1)
                            end_far_candidates[end_near_node[0], end_near_node[1]] = 0
                            invalid_nodes = (((far_edge_cc != far_edge_id).astype(np.uint8) * \
                                              (far_edge_cc != 0).astype(np.uint8)).astype(np.uint8) + \
                                             (new_specific_edge_map).astype(np.uint8) + \
                                             (mask == 0).astype(np.uint8)).clip(0, 1)
                            end_far_candidates[invalid_nodes > 0] = 0
                            far_edge[end_far_candidates > 0] = 1
                            far_edge_with_id[end_far_candidates > 0] = edge_id
                            
                    far_edge[context_far_node[0], context_far_node[1]] = 1
                    far_edge_with_id[context_far_node[0], context_far_node[1]] = edge_id
                near_edge_with_id[(mask_edge_with_id == edge_id) > 0] = edge_id
    uncleaned_far_edge = far_edge.copy()
    far_edge[mask == 0] = 0

    return far_edge, uncleaned_far_edge, far_edge_with_id, near_edge_with_id

def get_MiDaS_samples(image_folder, depth_folder, config, specific=None, aft_certain=None):
    lines = [os.path.splitext(os.path.basename(xx))[0] for xx in glob.glob(os.path.join(image_folder, '*' + config['img_format']))]
    samples = []
    generic_pose = np.eye(4)
    assert len(config['traj_types']) == len(config['x_shift_range']) ==\
           len(config['y_shift_range']) == len(config['z_shift_range']) == len(config['video_postfix']), \
           "The number of elements in 'traj_types', 'x_shift_range', 'y_shift_range', 'z_shift_range' and \
               'video_postfix' should be equal."
    tgt_pose = [[generic_pose * 1]]
    tgts_poses = []
    for traj_idx in range(len(config['traj_types'])):
        tgt_poses = []
        sx, sy, sz = path_planning(config['num_frames'], config['x_shift_range'][traj_idx], config['y_shift_range'][traj_idx],
                                   config['z_shift_range'][traj_idx], path_type=config['traj_types'][traj_idx])
        for xx, yy, zz in zip(sx, sy, sz):
            tgt_poses.append(generic_pose * 1.)
            tgt_poses[-1][:3, -1] = np.array([xx, yy, zz])
        tgts_poses += [tgt_poses]    
    tgt_pose = generic_pose * 1
    
    aft_flag = True
    if aft_certain is not None and len(aft_certain) > 0:
        aft_flag = False
    for seq_dir in lines:
        if specific is not None and len(specific) > 0:
            if specific != seq_dir:
                continue
        if aft_certain is not None and len(aft_certain) > 0:
            if aft_certain == seq_dir:
                aft_flag = True
            if aft_flag is False:
                continue
        samples.append({})
        sdict = samples[-1]            
        sdict['depth_fi'] = os.path.join(depth_folder, seq_dir + config['depth_format'])
        sdict['ref_img_fi'] = os.path.join(image_folder, seq_dir + config['img_format'])
        H, W = imageio.imread(sdict['ref_img_fi']).shape[:2]
        sdict['int_mtx'] = np.array([[max(H, W), 0, W//2], [0, max(H, W), H//2], [0, 0, 1]]).astype(np.float32)
        if sdict['int_mtx'].max() > 1:
            sdict['int_mtx'][0, :] = sdict['int_mtx'][0, :] / float(W)
            sdict['int_mtx'][1, :] = sdict['int_mtx'][1, :] / float(H)
        sdict['ref_pose'] = np.eye(4)
        sdict['tgt_pose'] = tgt_pose
        sdict['tgts_poses'] = tgts_poses
        sdict['video_postfix'] = config['video_postfix']
        sdict['tgt_name'] = [os.path.splitext(os.path.basename(sdict['depth_fi']))[0]]
        sdict['src_pair_name'] = sdict['tgt_name'][0]

    return samples

def get_valid_size(imap):
    x_max = np.where(imap.sum(1).squeeze() > 0)[0].max() + 1
    x_min = np.where(imap.sum(1).squeeze() > 0)[0].min()
    y_max = np.where(imap.sum(0).squeeze() > 0)[0].max() + 1
    y_min = np.where(imap.sum(0).squeeze() > 0)[0].min()
    size_dict = {'x_max':x_max, 'y_max':y_max, 'x_min':x_min, 'y_min':y_min}
    
    return size_dict

def dilate_valid_size(isize_dict, imap, dilate=[0, 0]):
    osize_dict = copy.deepcopy(isize_dict)
    osize_dict['x_min'] = max(0, osize_dict['x_min'] - dilate[0])
    osize_dict['x_max'] = min(imap.shape[0], osize_dict['x_max'] + dilate[0])
    osize_dict['y_min'] = max(0, osize_dict['y_min'] - dilate[0])
    osize_dict['y_max'] = min(imap.shape[1], osize_dict['y_max'] + dilate[1])

    return osize_dict

def crop_maps_by_size(size, *imaps):
    omaps = []
    for imap in imaps:
        omaps.append(imap[size['x_min']:size['x_max'], size['y_min']:size['y_max']].copy())
    
    return omaps

def smooth_cntsyn_gap(init_depth_map, mask_region, context_region, init_mask_region=None):
    if init_mask_region is not None:
        curr_mask_region = init_mask_region * 1
    else:
        curr_mask_region = mask_region * 0
    depth_map = init_depth_map.copy()
    for _ in range(2):
        cm_mask = context_region + curr_mask_region
        depth_s1 = np.roll(depth_map, 1, 0)
        depth_s2 = np.roll(depth_map, -1, 0)
        depth_s3 = np.roll(depth_map, 1, 1)
        depth_s4 = np.roll(depth_map, -1, 1)
        mask_s1 = np.roll(cm_mask, 1, 0)
        mask_s2 = np.roll(cm_mask, -1, 0)
        mask_s3 = np.roll(cm_mask, 1, 1)
        mask_s4 = np.roll(cm_mask, -1, 1)
        fluxin_depths = (depth_s1 * mask_s1 + depth_s2 * mask_s2 + depth_s3 * mask_s3 + depth_s4 * mask_s4) / \
                        ((mask_s1 + mask_s2 + mask_s3 + mask_s4) + 1e-6)
        fluxin_mask = (fluxin_depths != 0) * mask_region
        init_mask = (fluxin_mask * (curr_mask_region >= 0).astype(np.float32) > 0).astype(np.uint8)
        depth_map[init_mask > 0] = fluxin_depths[init_mask > 0]
        if init_mask.shape[-1] > curr_mask_region.shape[-1]:
            curr_mask_region[init_mask.sum(-1, keepdims=True) > 0] = 1
        else:
            curr_mask_region[init_mask > 0] = 1
        depth_map[fluxin_mask > 0] = fluxin_depths[fluxin_mask > 0]

    return depth_map

def read_MiDaS_depth(disp_fi, disp_rescale=10., h=None, w=None):
    if 'npy' in os.path.splitext(disp_fi)[-1]:
        disp = np.load(disp_fi)
    else:
        disp = imageio.imread(disp_fi).astype(np.float32)
    disp = disp - disp.min()
    disp = cv2.blur(disp / disp.max(), ksize=(3, 3)) * disp.max()
    disp = (disp / disp.max()) * disp_rescale
    if h is not None and w is not None:
        disp = resize(disp / disp.max(), (h, w), order=1) * disp.max()
    depth = 1. / np.maximum(disp, 0.05)

    return depth

def follow_image_aspect_ratio(depth, image):
    H, W = image.shape[:2]
    image_aspect_ratio = H / W
    dH, dW = depth.shape[:2]
    depth_aspect_ratio = dH / dW
    if depth_aspect_ratio > image_aspect_ratio:
        resize_H = dH
        resize_W = dH / image_aspect_ratio
    else:
        resize_W = dW
        resize_H = dW * image_aspect_ratio
    depth = resize(depth / depth.max(), 
                    (int(resize_H), 
                    int(resize_W)), 
                    order=0) * depth.max()
    
    return depth

def depth_resize(depth, origin_size, image_size):
    if origin_size[0] is not 0:
        max_depth = depth.max()
        depth = depth / max_depth
        depth = resize(depth, origin_size, order=1, mode='edge')
        depth = depth * max_depth
    else:
        max_depth = depth.max()
        depth = depth / max_depth
        depth = resize(depth, image_size, order=1, mode='edge')
        depth = depth * max_depth

    return depth
    
def filter_irrelevant_edge(self_edge, other_edges, other_edges_with_id, current_edge_id, context, edge_ccs, mesh, anchor):
    other_edges = other_edges.squeeze()
    other_edges_with_id = other_edges_with_id.squeeze()
    
    self_edge = self_edge.squeeze()
    dilate_self_edge = cv2.dilate(self_edge.astype(np.uint8), np.array([[1,1,1],[1,1,1],[1,1,1]]).astype(np.uint8), iterations=1)
    edge_ids = collections.Counter(other_edges_with_id.flatten()).keys()
    other_edges_info = []
    # import ipdb
    # ipdb.set_trace()
    for edge_id in edge_ids:
        edge_id = int(edge_id)
        if edge_id >= 0:
            condition = ((other_edges_with_id == edge_id) * other_edges * context).astype(np.uint8)
            if dilate_self_edge[condition > 0].sum() == 0:
                other_edges[other_edges_with_id == edge_id] = 0
            else:
                num_condition, condition_labels = cv2.connectedComponents(condition, connectivity=8)
                for condition_id in range(1, num_condition):
                    isolate_condition = ((condition_labels == condition_id) > 0).astype(np.uint8)
                    num_end_group, end_group = cv2.connectedComponents(((dilate_self_edge * isolate_condition) > 0).astype(np.uint8), connectivity=8)
                    if num_end_group == 1:
                        continue
                    for end_id in range(1, num_end_group):
                        end_pxs, end_pys = np.where((end_group == end_id))
                        end_px, end_py = end_pxs[0], end_pys[0]
                        other_edges_info.append({})
                        other_edges_info[-1]['edge_id'] = edge_id
                        # other_edges_info[-1]['near_depth'] = None
                        other_edges_info[-1]['diff'] = None
                        other_edges_info[-1]['edge_map'] = np.zeros_like(self_edge)
                        other_edges_info[-1]['end_point_map'] = np.zeros_like(self_edge)
                        other_edges_info[-1]['end_point_map'][(end_group == end_id)] = 1
                        other_edges_info[-1]['forbidden_point_map'] = np.zeros_like(self_edge)
                        other_edges_info[-1]['forbidden_point_map'][(end_group != end_id) * (end_group != 0)] = 1
                        other_edges_info[-1]['forbidden_point_map'] = cv2.dilate(other_edges_info[-1]['forbidden_point_map'], kernel=np.array([[1,1,1],[1,1,1],[1,1,1]]), iterations=2)
                        for x in edge_ccs[edge_id]:
                            nx = x[0] - anchor[0]
                            ny = x[1] - anchor[1]
                            if nx == end_px and ny == end_py:
                                # other_edges_info[-1]['near_depth'] = abs(nx)
                                if mesh.nodes[x].get('far') is not None and len(mesh.nodes[x].get('far')) == 1:
                                    other_edges_info[-1]['diff'] = abs(1./abs([*mesh.nodes[x].get('far')][0][2]) - 1./abs(x[2]))
                                else:
                                    other_edges_info[-1]['diff'] = 0
                                # if end_group[nx, ny] != end_id and end_group[nx, ny] > 0:
                                #     continue
                            try:
                                if isolate_condition[nx, ny] == 1:
                                    other_edges_info[-1]['edge_map'][nx, ny] = 1
                            except:
                                pass
    try:
        other_edges_info = sorted(other_edges_info, key=lambda x : x['diff'], reverse=True)
    except:
        import pdb
        pdb.set_trace()
    # import pdb
    # pdb.set_trace()
    # other_edges = other_edges[..., None]
    for other_edge in other_edges_info:
        if other_edge['end_point_map'] is None:
            import pdb
            pdb.set_trace()

    other_edges = other_edges * context

    return other_edges, other_edges_info

def require_depth_edge(context_edge, mask):
    dilate_mask = cv2.dilate(mask, np.array([[1,1,1],[1,1,1],[1,1,1]]).astype(np.uint8), iterations=1)
    if (dilate_mask * context_edge).max() == 0:
        return False
    else:
        return True

def refine_color_around_edge(mesh, info_on_pix, edge_ccs, config, spdb=False):
    H, W = mesh.graph['H'], mesh.graph['W']
    tmp_edge_ccs = copy.deepcopy(edge_ccs)
    for edge_id, edge_cc in enumerate(edge_ccs):
        if len(edge_cc) == 0:
            continue
        near_maps = np.zeros((H, W)).astype(np.bool)
        far_maps = np.zeros((H, W)).astype(np.bool)
        tmp_far_nodes = set()
        far_nodes = set()
        near_nodes = set()
        end_nodes = set()        
        for i in range(5):
            if i == 0:
                for edge_node in edge_cc:
                    if mesh.nodes[edge_node].get('depth_edge_dilate_2_color_flag') is not True:
                        break
                    if mesh.nodes[edge_node].get('inpaint_id') == 1:
                        near_nodes.add(edge_node)
                        tmp_node = mesh.nodes[edge_node].get('far')
                        tmp_node = set(tmp_node) if tmp_node is not None else set()
                        tmp_far_nodes |= tmp_node
                rmv_tmp_far_nodes = set()
                for far_node in tmp_far_nodes:
                    if not(mesh.has_node(far_node) and mesh.nodes[far_node].get('inpaint_id') == 1):
                        rmv_tmp_far_nodes.add(far_node)
                if len(tmp_far_nodes - rmv_tmp_far_nodes) == 0:
                    break                        
                else:
                    for near_node in near_nodes:
                        near_maps[near_node[0], near_node[1]] = True
                        mesh.nodes[near_node]['refine_rgbd'] = True
                        mesh.nodes[near_node]['backup_depth'] = near_node[2] \
                                    if mesh.nodes[near_node].get('real_depth') is None else mesh.nodes[near_node]['real_depth']
                        mesh.nodes[near_node]['backup_color'] = mesh.nodes[near_node]['color']
                for far_node in tmp_far_nodes:
                    if mesh.has_node(far_node) and mesh.nodes[far_node].get('inpaint_id') == 1:
                        far_nodes.add(far_node)
                        far_maps[far_node[0], far_node[1]] = True
                        mesh.nodes[far_node]['refine_rgbd'] = True
                        mesh.nodes[far_node]['backup_depth'] = far_node[2] \
                                    if mesh.nodes[far_node].get('real_depth') is None else mesh.nodes[far_node]['real_depth']
                        mesh.nodes[far_node]['backup_color'] = mesh.nodes[far_node]['color']
                tmp_far_nodes = far_nodes
                tmp_near_nodes = near_nodes
            else:
                tmp_far_nodes = new_tmp_far_nodes
                tmp_near_nodes = new_tmp_near_nodes
                new_tmp_far_nodes = None
                new_tmp_near_nodes = None
            new_tmp_far_nodes = set()
            new_tmp_near_nodes = set()
            for node in tmp_near_nodes:
                for ne_node in mesh.neighbors(node):
                    if far_maps[ne_node[0], ne_node[1]] == False and \
                        near_maps[ne_node[0], ne_node[1]] == False:
                        if mesh.nodes[ne_node].get('inpaint_id') == 1:
                            new_tmp_near_nodes.add(ne_node)
                            near_maps[ne_node[0], ne_node[1]] = True
                            mesh.nodes[ne_node]['refine_rgbd'] = True
                            mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \
                                    if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth']
                            mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color']
                        else:
                            mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \
                                    if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth']
                            mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color']
                            end_nodes.add(node)
            near_nodes.update(new_tmp_near_nodes)
            for node in tmp_far_nodes:
                for ne_node in mesh.neighbors(node):
                    if far_maps[ne_node[0], ne_node[1]] == False and \
                        near_maps[ne_node[0], ne_node[1]] == False:
                        if mesh.nodes[ne_node].get('inpaint_id') == 1:
                            new_tmp_far_nodes.add(ne_node)
                            far_maps[ne_node[0], ne_node[1]] = True
                            mesh.nodes[ne_node]['refine_rgbd'] = True
                            mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \
                                    if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth']
                            mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color']
                        else:
                            mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \
                                    if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth']
                            mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color']
                            end_nodes.add(node)
            far_nodes.update(new_tmp_far_nodes)
        if len(far_nodes) == 0:
            tmp_edge_ccs[edge_id] = set()
            continue
        for node in new_tmp_far_nodes | new_tmp_near_nodes:
            for ne_node in mesh.neighbors(node):
                if far_maps[ne_node[0], ne_node[1]] == False and near_maps[ne_node[0], ne_node[1]] == False:
                    end_nodes.add(node)
                    mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \
                            if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth']
                    mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color']
        tmp_end_nodes = end_nodes
        
        refine_nodes = near_nodes | far_nodes
        remain_refine_nodes = copy.deepcopy(refine_nodes)
        accum_idx = 0
        while len(remain_refine_nodes) > 0:
            accum_idx += 1
            if accum_idx > 100:
                break
            new_tmp_end_nodes = None
            new_tmp_end_nodes = set()
            survive_tmp_end_nodes = set()
            for node in tmp_end_nodes:
                re_depth, re_color, re_count = 0, np.array([0., 0., 0.]), 0
                for ne_node in mesh.neighbors(node):
                    if mesh.nodes[ne_node].get('refine_rgbd') is True:
                        if ne_node not in tmp_end_nodes:
                            new_tmp_end_nodes.add(ne_node)
                    else:
                        try:
                            re_depth += mesh.nodes[ne_node]['backup_depth']
                            re_color += mesh.nodes[ne_node]['backup_color'].astype(np.float32)
                            re_count += 1.
                        except:
                            import pdb; pdb.set_trace()
                if re_count > 0:
                    re_depth = re_depth / re_count
                    re_color = re_color / re_count
                    mesh.nodes[node]['backup_depth'] = re_depth
                    mesh.nodes[node]['backup_color'] = re_color
                    mesh.nodes[node]['refine_rgbd'] = False
                else:
                    survive_tmp_end_nodes.add(node)
            for node in tmp_end_nodes - survive_tmp_end_nodes:
                if node in remain_refine_nodes:
                    remain_refine_nodes.remove(node)
            tmp_end_nodes = new_tmp_end_nodes
        if spdb == True:
            bfrd_canvas = np.zeros((H, W))
            bfrc_canvas = np.zeros((H, W, 3)).astype(np.uint8)
            aftd_canvas = np.zeros((H, W))
            aftc_canvas = np.zeros((H, W, 3)).astype(np.uint8)
            for node in refine_nodes:
                bfrd_canvas[node[0], node[1]] = abs(node[2])
                aftd_canvas[node[0], node[1]] = abs(mesh.nodes[node]['backup_depth'])
                bfrc_canvas[node[0], node[1]] = mesh.nodes[node]['color'].astype(np.uint8)
                aftc_canvas[node[0], node[1]] = mesh.nodes[node]['backup_color'].astype(np.uint8)
            f, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, sharex=True, sharey=True); 
            ax1.imshow(bfrd_canvas); 
            ax2.imshow(aftd_canvas); 
            ax3.imshow(bfrc_canvas); 
            ax4.imshow(aftc_canvas); 
            plt.show()
            import pdb; pdb.set_trace()
        for node in refine_nodes:
            if mesh.nodes[node].get('refine_rgbd') is not None:
                mesh.nodes[node].pop('refine_rgbd')
                mesh.nodes[node]['color'] = mesh.nodes[node]['backup_color']
                for info in info_on_pix[(node[0], node[1])]:
                    if info['depth'] == node[2]:
                        info['color'] = mesh.nodes[node]['backup_color']

    return mesh, info_on_pix

def refine_depth_around_edge(mask_depth, far_edge, uncleaned_far_edge, near_edge, mask, all_depth, config):
    if isinstance(mask_depth, torch.Tensor):
        if mask_depth.is_cuda:
            mask_depth = mask_depth.cpu()
        mask_depth = mask_depth.data
        mask_depth = mask_depth.numpy()
    if isinstance(far_edge, torch.Tensor):
        if far_edge.is_cuda:
            far_edge = far_edge.cpu()
        far_edge = far_edge.data
        far_edge = far_edge.numpy()
    if isinstance(uncleaned_far_edge, torch.Tensor):
        if uncleaned_far_edge.is_cuda:
            uncleaned_far_edge = uncleaned_far_edge.cpu()
        uncleaned_far_edge = uncleaned_far_edge.data
        uncleaned_far_edge = uncleaned_far_edge.numpy()
    if isinstance(near_edge, torch.Tensor):
        if near_edge.is_cuda:
            near_edge = near_edge.cpu()
        near_edge = near_edge.data
        near_edge = near_edge.numpy()
    if isinstance(mask, torch.Tensor):
        if mask.is_cuda:
            mask = mask.cpu()
        mask = mask.data
        mask = mask.numpy()
    mask = mask.squeeze()
    uncleaned_far_edge = uncleaned_far_edge.squeeze()
    far_edge = far_edge.squeeze()
    near_edge = near_edge.squeeze()
    mask_depth = mask_depth.squeeze()
    dilate_far_edge = cv2.dilate(uncleaned_far_edge.astype(np.uint8), kernel=np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8), iterations=1)
    near_edge[dilate_far_edge == 0] = 0
    dilate_near_edge = cv2.dilate(near_edge.astype(np.uint8), kernel=np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8), iterations=1)
    far_edge[dilate_near_edge == 0] = 0
    init_far_edge = far_edge.copy()
    init_near_edge = near_edge.copy()
    for i in range(config['depth_edge_dilate_2']):
        init_far_edge = cv2.dilate(init_far_edge, kernel=np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8), iterations=1)
        init_far_edge[init_near_edge == 1] = 0
        init_near_edge = cv2.dilate(init_near_edge, kernel=np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8), iterations=1)
        init_near_edge[init_far_edge == 1] = 0
    init_far_edge[mask == 0] = 0
    init_near_edge[mask == 0] = 0
    hole_far_edge = 1 - init_far_edge
    hole_near_edge = 1 - init_near_edge
    change = None
    while True:
        change = False
        hole_far_edge[init_near_edge == 1] = 0
        hole_near_edge[init_far_edge == 1] = 0
        far_pxs, far_pys = np.where((hole_far_edge == 0) * (init_far_edge == 1) > 0)
        current_hole_far_edge = hole_far_edge.copy()
        for far_px, far_py in zip(far_pxs, far_pys):
            min_px = max(far_px - 1, 0) 
            max_px = min(far_px + 2, mask.shape[0]-1)
            min_py = max(far_py - 1, 0) 
            max_py = min(far_py + 2, mask.shape[1]-1)
            hole_far = current_hole_far_edge[min_px: max_px, min_py: max_py]
            tmp_mask = mask[min_px: max_px, min_py: max_py]
            all_depth_patch = all_depth[min_px: max_px, min_py: max_py] * 0
            all_depth_mask = (all_depth_patch != 0).astype(np.uint8)
            cross_element = np.array([[0,1,0],[1,1,1],[0,1,0]])[min_px - (far_px - 1): max_px - (far_px - 1), min_py - (far_py - 1): max_py - (far_py - 1)]
            combine_mask = (tmp_mask + all_depth_mask).clip(0, 1) * hole_far * cross_element
            tmp_patch = combine_mask * (mask_depth[min_px: max_px, min_py: max_py] + all_depth_patch)
            number = np.count_nonzero(tmp_patch)
            if number > 0:
                mask_depth[far_px, far_py] = np.sum(tmp_patch).astype(np.float32) / max(number, 1e-6)
                hole_far_edge[far_px, far_py] = 1
                change = True
        near_pxs, near_pys = np.where((hole_near_edge == 0) * (init_near_edge == 1) > 0)
        current_hole_near_edge = hole_near_edge.copy()
        for near_px, near_py in zip(near_pxs, near_pys):
            min_px = max(near_px - 1, 0) 
            max_px = min(near_px + 2, mask.shape[0]-1)
            min_py = max(near_py - 1, 0) 
            max_py = min(near_py + 2, mask.shape[1]-1)
            hole_near = current_hole_near_edge[min_px: max_px, min_py: max_py]
            tmp_mask = mask[min_px: max_px, min_py: max_py]
            all_depth_patch = all_depth[min_px: max_px, min_py: max_py] * 0
            all_depth_mask = (all_depth_patch != 0).astype(np.uint8)            
            cross_element = np.array([[0,1,0],[1,1,1],[0,1,0]])[min_px - near_px + 1:max_px - near_px + 1, min_py - near_py + 1:max_py - near_py + 1]
            combine_mask = (tmp_mask + all_depth_mask).clip(0, 1) * hole_near * cross_element
            tmp_patch = combine_mask * (mask_depth[min_px: max_px, min_py: max_py] + all_depth_patch)
            number = np.count_nonzero(tmp_patch)
            if number > 0:                
                mask_depth[near_px, near_py] = np.sum(tmp_patch) / max(number, 1e-6)
                hole_near_edge[near_px, near_py] = 1
                change = True
        if change is False:
            break
        
    return mask_depth



def vis_depth_edge_connectivity(depth, config):
    disp = 1./depth
    u_diff = (disp[1:, :] - disp[:-1, :])[:-1, 1:-1]
    b_diff = (disp[:-1, :] - disp[1:, :])[1:, 1:-1]
    l_diff = (disp[:, 1:] - disp[:, :-1])[1:-1, :-1]
    r_diff = (disp[:, :-1] - disp[:, 1:])[1:-1, 1:]
    u_over = (np.abs(u_diff) > config['depth_threshold']).astype(np.float32)
    b_over = (np.abs(b_diff) > config['depth_threshold']).astype(np.float32)
    l_over = (np.abs(l_diff) > config['depth_threshold']).astype(np.float32)
    r_over = (np.abs(r_diff) > config['depth_threshold']).astype(np.float32)
    concat_diff = np.stack([u_diff, b_diff, r_diff, l_diff], axis=-1)
    concat_over = np.stack([u_over, b_over, r_over, l_over], axis=-1)
    over_diff = concat_diff * concat_over
    pos_over = (over_diff > 0).astype(np.float32).sum(-1).clip(0, 1)
    neg_over = (over_diff < 0).astype(np.float32).sum(-1).clip(0, 1)
    neg_over[(over_diff > 0).astype(np.float32).sum(-1) > 0] = 0
    _, edge_label = cv2.connectedComponents(pos_over.astype(np.uint8), connectivity=8)
    T_junction_maps = np.zeros_like(pos_over)
    for edge_id in range(1, edge_label.max() + 1):
        edge_map = (edge_label == edge_id).astype(np.uint8)
        edge_map = np.pad(edge_map, pad_width=((1,1),(1,1)), mode='constant')
        four_direc = np.roll(edge_map, 1, 1) + np.roll(edge_map, -1, 1) + np.roll(edge_map, 1, 0) + np.roll(edge_map, -1, 0)
        eight_direc = np.roll(np.roll(edge_map, 1, 1), 1, 0) + np.roll(np.roll(edge_map, 1, 1), -1, 0) + \
                      np.roll(np.roll(edge_map, -1, 1), 1, 0) + np.roll(np.roll(edge_map, -1, 1), -1, 0)
        eight_direc = (eight_direc + four_direc)[1:-1,1:-1]
        pos_over[eight_direc > 2] = 0
        T_junction_maps[eight_direc > 2] = 1
    _, edge_label = cv2.connectedComponents(pos_over.astype(np.uint8), connectivity=8)
    edge_label = np.pad(edge_label, 1, mode='constant')

    return edge_label



def max_size(mat, value=0):
    if not (mat and mat[0]): return (0, 0)
    it = iter(mat)
    prev = [(el==value) for el in next(it)]
    max_size = max_rectangle_size(prev)
    for row in it:
        hist = [(1+h) if el == value else 0 for h, el in zip(prev, row)]
        max_size = max(max_size, max_rectangle_size(hist), key=get_area)
        prev = hist                                               
    return max_size

def max_rectangle_size(histogram):
    Info = namedtuple('Info', 'start height')
    stack = []
    top = lambda: stack[-1]
    max_size = (0, 0) # height, width of the largest rectangle
    pos = 0 # current position in the histogram
    for pos, height in enumerate(histogram):
        start = pos # position where rectangle starts
        while True:
            if not stack or height > top().height:
                stack.append(Info(start, height)) # push
            if stack and height < top().height:
                max_size = max(max_size, (top().height, (pos-top().start)),
                               key=get_area)
                start, _ = stack.pop()
                continue
            break # height == top().height goes here
                
    pos += 1
    for start, height in stack:
        max_size = max(max_size, (height, (pos-start)),
                       key=get_area)

    return max_size

def get_area(size):
    return reduce(mul, size)

def find_anchors(matrix):
    matrix = [[*x] for x in matrix]
    mh, mw = max_size(matrix)
    matrix = np.array(matrix)
    # element = np.zeros((mh, mw))
    for i in range(matrix.shape[0] + 1 - mh):
        for j in range(matrix.shape[1] + 1 - mw):
            if matrix[i:i + mh, j:j + mw].max() == 0:
                return i, i + mh, j, j + mw

def find_largest_rect(dst_img, bg_color=(128, 128, 128)):
    valid = np.any(dst_img[..., :3] != bg_color, axis=-1) 
    dst_h, dst_w = dst_img.shape[:2]
    ret, labels = cv2.connectedComponents(np.uint8(valid == False)) 
    red_mat = np.zeros_like(labels) 
    # denoise 
    for i in range(1, np.max(labels)+1, 1): 
        x, y, w, h = cv2.boundingRect(np.uint8(labels==i)) 
        if x == 0 or (x+w) == dst_h or y == 0 or (y+h) == dst_w: 
            red_mat[labels==i] = 1 
    # crop 
    t, b, l, r = find_anchors(red_mat) 

    return t, b, l, r