mesh_tools.py

import os
import numpy as np
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, cur_node, new_node):
    if cur_node == new_node:
        return mesh
    mesh.add_node(new_node)
    for key, value in mesh.nodes[cur_node].items():
        mesh.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]
    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]]] # 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
                mesh.nodes[ne[0], ne[1]]['depth'] = end_depth_maps[ne[0], ne[1]]
                # if (ne[0] + all_anchor[0], ne[1] + all_anchor[2]) == (559, 60):
                #     import pdb; pdb.set_trace()
    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, [*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))
    real_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