utils.py

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