Saini commited on
Commit
0b9f920
1 Parent(s): a09d4b8
.gitignore ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #================================================
2
+ # User Specifics
3
+ #================================================
4
+ ## Don't push data up
5
+ image/*
6
+ !image/.empty
7
+ ## Ignore vim swap files
8
+ #*.swp
9
+ ## MAC why you do this?
10
+ #.DS_Store
11
+ ## Our own output files
12
+ #*.out
13
+
14
+ #================================================
15
+ # Python Specifics
16
+ #================================================
17
+
18
+ # Byte-compiled / optimized / DLL files
19
+ __pycache__/
20
+ *.py[cod]
21
+ *$py.class
22
+
23
+ # C extensions
24
+ *.so
25
+
26
+ # Distribution / packaging
27
+ .Python
28
+ env/
29
+ build/
30
+ develop-eggs/
31
+ dist/
32
+ downloads/
33
+ eggs/
34
+ .eggs/
35
+ lib/
36
+ lib64/
37
+ parts/
38
+ sdist/
39
+ var/
40
+ *.egg-info/
41
+ .installed.cfg
42
+ *.egg
43
+
44
+ # PyInstaller
45
+ # Usually these files are written by a python script from a template
46
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
47
+ *.manifest
48
+ *.spec
49
+
50
+ # Installer logs
51
+ pip-log.txt
52
+ pip-delete-this-directory.txt
53
+
54
+ # Unit test / coverage reports
55
+ htmlcov/
56
+ .tox/
57
+ .coverage
58
+ .coverage.*
59
+ .cache
60
+ nosetests.xml
61
+ coverage.xml
62
+ *,cover
63
+ .hypothesis/
64
+
65
+ # Translations
66
+ *.mo
67
+ *.pot
68
+
69
+ # Django stuff:
70
+ *.log
71
+ local_settings.py
72
+
73
+ # Flask stuff:
74
+ instance/
75
+ .webassets-cache
76
+
77
+ # Scrapy stuff:
78
+ .scrapy
79
+
80
+ # Sphinx documentation
81
+ docs/_build/
82
+
83
+ # PyBuilder
84
+ target/
85
+
86
+ # IPython Notebook
87
+ .ipynb_checkpoints
88
+
89
+ # pyenv
90
+ .python-version
91
+
92
+ # celery beat schedule file
93
+ celerybeat-schedule
94
+
95
+ # dotenv
96
+ .env
97
+
98
+ # virtualenv
99
+ venv/
100
+ ENV/
101
+
102
+ # Spyder project settings
103
+ .spyderproject
104
+
105
+ # Rope project settings
106
+ .ropeproject
DOCUMENTATION.md ADDED
@@ -0,0 +1,146 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Documentation
2
+
3
+ ## Python scripts
4
+
5
+ These files are for our monocular 3D Tracking pipeline:
6
+
7
+ `main.py` Execute 3D photo inpainting
8
+
9
+ `mesh.py` Functions about context-aware depth inpainting
10
+
11
+ `mesh_tools.py` Some common functions used in `mesh.py`
12
+
13
+ `utils.py` Some common functions used in image preprocessing, data loading
14
+
15
+ `networks.py` Network architectures of inpainting model
16
+
17
+
18
+ MiDaS/
19
+
20
+ `run.py` Execute depth estimation
21
+
22
+ `monodepth_net.py` Network architecture of depth estimation model
23
+
24
+ `MiDaS_utils.py` Some common functions in depth estimation
25
+
26
+
27
+ ## Configuration
28
+
29
+ ```bash
30
+ argument.yml
31
+ ```
32
+
33
+ - `depth_edge_model_ckpt: checkpoints/EdgeModel.pth`
34
+ - Pretrained model of depth-edge inpainting
35
+ - `depth_feat_model_ckpt: checkpoints/DepthModel.pth`
36
+ - Pretrained model of depth inpainting
37
+ - `rgb_feat_model_ckpt: checkpoints/ColorModel.pth`
38
+ - Pretrained model of color inpainting
39
+ - `MiDaS_model_ckpt: MiDaS/model.pt`
40
+ - Pretrained model of depth estimation
41
+ - `use_boostmonodepth: True`
42
+ - Use [BoostMonocularDepth](https://github.com/compphoto/BoostingMonocularDepth) to get sharper monocular depth estimation
43
+ - `fps: 40`
44
+ - Frame per second of output rendered video
45
+ - `num_frames: 240`
46
+ - Total number of frames in output rendered video
47
+ - `x_shift_range: [-0.03, -0.03, -0.03]`
48
+ - The translations on x-axis of output rendered videos.
49
+ - This parameter is a list. Each element corresponds to a specific camera motion.
50
+ - `y_shift_range: [-0.00, -0.00, -0.03]`
51
+ - The translations on y-axis of output rendered videos.
52
+ - This parameter is a list. Each element corresponds to a specific camera motion.
53
+ - `z_shift_range: [-0.07, -0.07, -0.07]`
54
+ - The translations on z-axis of output rendered videos.
55
+ - This parameter is a list. Each element corresponds to a specific camera motion.
56
+ - `traj_types: ['straight-line', 'circle', 'circle']`
57
+ - The type of camera trajectory.
58
+ - This parameter is a list.
59
+ - Currently, we only privode `straight-line` and `circle`.
60
+ - `video_postfix: ['zoom-in', 'swing', 'circle']`
61
+ - The postfix of video.
62
+ - This parameter is a list.
63
+ - Note that the number of elements in `x_shift_range`, `y_shift_range`, `z_shift_range`, `traj_types` and `video_postfix` should be equal.
64
+ - `specific: '' `
65
+ - The specific image name, use this to specify the image to be executed. By default, all the image in the folder will be executed.
66
+ - `longer_side_len: 960`
67
+ - The length of larger dimension in output resolution.
68
+ - `src_folder: image`
69
+ - Input image directory.
70
+ - `depth_folder: depth`
71
+ - Estimated depth directory.
72
+ - `mesh_folder: mesh`
73
+ - Output 3-D mesh directory.
74
+ - `video_folder: video`
75
+ - Output rendered video directory
76
+ - `load_ply: False`
77
+ - Action to load existed mesh (.ply) file
78
+ - `save_ply: True`
79
+ - Action to store the output mesh (.ply) file
80
+ - Disable this option `save_ply: False` to reduce the computational time.
81
+ - `inference_video: True`
82
+ - Action to rendered the output video
83
+ - `gpu_ids: 0`
84
+ - The ID of working GPU. Leave it blank or negative to use CPU.
85
+ - `offscreen_rendering: True`
86
+ - If you're executing the process in a remote server (via ssh), please switch on this flag.
87
+ - Sometimes, using off-screen rendering result in longer execution time.
88
+ - `img_format: '.jpg'`
89
+ - Input image format.
90
+ - `depth_format: '.npy'`
91
+ - Input depth (disparity) format. Use NumPy array file as default.
92
+ - If the user wants to edit the depth (disparity) map manually, we provide `.png` format depth (disparity) map.
93
+ - Remember to switch this parameter from `.npy` to `.png` when using depth (disparity) map with `.png` format.
94
+ - `require_midas: True`
95
+ - Set it to `True` if the user wants to use depth map estimated by `MiDaS`.
96
+ - Set it to `False` if the user wants to use manually edited depth map.
97
+ - If the user wants to edit the depth (disparity) map manually, we provide `.png` format depth (disparity) map.
98
+ - Remember to switch this parameter from `True` to `False` when using manually edited depth map.
99
+ - `depth_threshold: 0.04`
100
+ - A threshold in disparity, adjacent two pixels are discontinuity pixels
101
+ if the difference between them excceed this number.
102
+ - `ext_edge_threshold: 0.002`
103
+ - The threshold to define inpainted depth edge. A pixel in inpainted edge
104
+ map belongs to extended depth edge if the value of that pixel exceeds this number,
105
+ - `sparse_iter: 5`
106
+ - Total iteration numbers of bilateral median filter
107
+ - `filter_size: [7, 7, 5, 5, 5]`
108
+ - Window size of bilateral median filter in each iteration.
109
+ - `sigma_s: 4.0`
110
+ - Intensity term of bilateral median filter
111
+ - `sigma_r: 0.5`
112
+ - Spatial term of bilateral median filter
113
+ - `redundant_number: 12`
114
+ - The number defines short segments. If a depth edge is shorter than this number,
115
+ it is a short segment and removed.
116
+ - `background_thickness: 70`
117
+ - The thickness of synthesis area.
118
+ - `context_thickness: 140`
119
+ - The thickness of context area.
120
+ - `background_thickness_2: 70`
121
+ - The thickness of synthesis area when inpaint second time.
122
+ - `context_thickness_2: 70`
123
+ - The thickness of context area when inpaint second time.
124
+ - `discount_factor: 1.00`
125
+ - `log_depth: True`
126
+ - The scale of depth inpainting. If true, performing inpainting in log scale.
127
+ Otherwise, performing in linear scale.
128
+ - `largest_size: 512`
129
+ - The largest size of inpainted image patch.
130
+ - `depth_edge_dilate: 10`
131
+ - The thickness of dilated synthesis area.
132
+ - `depth_edge_dilate_2: 5`
133
+ - The thickness of dilated synthesis area when inpaint second time.
134
+ - `extrapolate_border: True`
135
+ - Action to extrapolate out-side the border.
136
+ - `extrapolation_thickness: 60`
137
+ - The thickness of extrapolated area.
138
+ - `repeat_inpaint_edge: True`
139
+ - Action to apply depth edge inpainting model repeatedly. Sometimes inpainting depth
140
+ edge once results in short inpinated edge, apply depth edge inpainting repeatedly
141
+ could help you prolong the inpainted depth edge.
142
+ - `crop_border: [0.03, 0.03, 0.05, 0.03]`
143
+ - The fraction of pixels to crop out around the borders `[top, left, bottom, right]`.
144
+ - `anti_flickering: True`
145
+ - Action to avoid flickering effect in the output video.
146
+ - This may result in longer computational time in rendering phase.
LICENSE ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ MIT License
3
+
4
+ Copyright (c) 2020 Virginia Tech Vision and Learning Lab
5
+
6
+ Permission is hereby granted, free of charge, to any person obtaining a copy
7
+ of this software and associated documentation files (the "Software"), to deal
8
+ in the Software without restriction, including without limitation the rights
9
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10
+ copies of the Software, and to permit persons to whom the Software is
11
+ furnished to do so, subject to the following conditions:
12
+
13
+ The above copyright notice and this permission notice shall be included in all
14
+ copies or substantial portions of the Software.
15
+
16
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22
+ SOFTWARE.
23
+
24
+ ------------------ LICENSE FOR MiDaS --------------------
25
+
26
+ MIT License
27
+
28
+ Copyright (c) 2019 Intel ISL (Intel Intelligent Systems Lab)
29
+
30
+ Permission is hereby granted, free of charge, to any person obtaining a copy
31
+ of this software and associated documentation files (the "Software"), to deal
32
+ in the Software without restriction, including without limitation the rights
33
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
34
+ copies of the Software, and to permit persons to whom the Software is
35
+ furnished to do so, subject to the following conditions:
36
+
37
+ The above copyright notice and this permission notice shall be included in all
38
+ copies or substantial portions of the Software.
39
+
40
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
41
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
42
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
43
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
44
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
45
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
46
+ SOFTWARE.
47
+
48
+ --------------------------- LICENSE FOR EdgeConnect --------------------------------
49
+
50
+ Attribution-NonCommercial 4.0 International
MiDaS/MiDaS_utils.py ADDED
@@ -0,0 +1,192 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Utils for monoDepth.
2
+ """
3
+ import sys
4
+ import re
5
+ import numpy as np
6
+ import cv2
7
+ import torch
8
+ import imageio
9
+
10
+
11
+ def read_pfm(path):
12
+ """Read pfm file.
13
+
14
+ Args:
15
+ path (str): path to file
16
+
17
+ Returns:
18
+ tuple: (data, scale)
19
+ """
20
+ with open(path, "rb") as file:
21
+
22
+ color = None
23
+ width = None
24
+ height = None
25
+ scale = None
26
+ endian = None
27
+
28
+ header = file.readline().rstrip()
29
+ if header.decode("ascii") == "PF":
30
+ color = True
31
+ elif header.decode("ascii") == "Pf":
32
+ color = False
33
+ else:
34
+ raise Exception("Not a PFM file: " + path)
35
+
36
+ dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("ascii"))
37
+ if dim_match:
38
+ width, height = list(map(int, dim_match.groups()))
39
+ else:
40
+ raise Exception("Malformed PFM header.")
41
+
42
+ scale = float(file.readline().decode("ascii").rstrip())
43
+ if scale < 0:
44
+ # little-endian
45
+ endian = "<"
46
+ scale = -scale
47
+ else:
48
+ # big-endian
49
+ endian = ">"
50
+
51
+ data = np.fromfile(file, endian + "f")
52
+ shape = (height, width, 3) if color else (height, width)
53
+
54
+ data = np.reshape(data, shape)
55
+ data = np.flipud(data)
56
+
57
+ return data, scale
58
+
59
+
60
+ def write_pfm(path, image, scale=1):
61
+ """Write pfm file.
62
+
63
+ Args:
64
+ path (str): pathto file
65
+ image (array): data
66
+ scale (int, optional): Scale. Defaults to 1.
67
+ """
68
+
69
+ with open(path, "wb") as file:
70
+ color = None
71
+
72
+ if image.dtype.name != "float32":
73
+ raise Exception("Image dtype must be float32.")
74
+
75
+ image = np.flipud(image)
76
+
77
+ if len(image.shape) == 3 and image.shape[2] == 3: # color image
78
+ color = True
79
+ elif (
80
+ len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1
81
+ ): # greyscale
82
+ color = False
83
+ else:
84
+ raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")
85
+
86
+ file.write("PF\n" if color else "Pf\n".encode())
87
+ file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))
88
+
89
+ endian = image.dtype.byteorder
90
+
91
+ if endian == "<" or endian == "=" and sys.byteorder == "little":
92
+ scale = -scale
93
+
94
+ file.write("%f\n".encode() % scale)
95
+
96
+ image.tofile(file)
97
+
98
+
99
+ def read_image(path):
100
+ """Read image and output RGB image (0-1).
101
+
102
+ Args:
103
+ path (str): path to file
104
+
105
+ Returns:
106
+ array: RGB image (0-1)
107
+ """
108
+ img = cv2.imread(path)
109
+
110
+ if img.ndim == 2:
111
+ img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
112
+
113
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
114
+
115
+ return img
116
+
117
+
118
+ def resize_image(img):
119
+ """Resize image and make it fit for network.
120
+
121
+ Args:
122
+ img (array): image
123
+
124
+ Returns:
125
+ tensor: data ready for network
126
+ """
127
+ height_orig = img.shape[0]
128
+ width_orig = img.shape[1]
129
+ unit_scale = 384.
130
+
131
+ if width_orig > height_orig:
132
+ scale = width_orig / unit_scale
133
+ else:
134
+ scale = height_orig / unit_scale
135
+
136
+ height = (np.ceil(height_orig / scale / 32) * 32).astype(int)
137
+ width = (np.ceil(width_orig / scale / 32) * 32).astype(int)
138
+
139
+ img_resized = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)
140
+
141
+ img_resized = (
142
+ torch.from_numpy(np.transpose(img_resized, (2, 0, 1))).contiguous().float()
143
+ )
144
+ img_resized = img_resized.unsqueeze(0)
145
+
146
+ return img_resized
147
+
148
+
149
+ def resize_depth(depth, width, height):
150
+ """Resize depth map and bring to CPU (numpy).
151
+
152
+ Args:
153
+ depth (tensor): depth
154
+ width (int): image width
155
+ height (int): image height
156
+
157
+ Returns:
158
+ array: processed depth
159
+ """
160
+ depth = torch.squeeze(depth[0, :, :, :]).to("cpu")
161
+ depth = cv2.blur(depth.numpy(), (3, 3))
162
+ depth_resized = cv2.resize(
163
+ depth, (width, height), interpolation=cv2.INTER_AREA
164
+ )
165
+
166
+ return depth_resized
167
+
168
+ def write_depth(path, depth, bits=1):
169
+ """Write depth map to pfm and png file.
170
+
171
+ Args:
172
+ path (str): filepath without extension
173
+ depth (array): depth
174
+ """
175
+ # write_pfm(path + ".pfm", depth.astype(np.float32))
176
+
177
+ depth_min = depth.min()
178
+ depth_max = depth.max()
179
+
180
+ max_val = (2**(8*bits))-1
181
+
182
+ if depth_max - depth_min > np.finfo("float").eps:
183
+ out = max_val * (depth - depth_min) / (depth_max - depth_min)
184
+ else:
185
+ out = 0
186
+
187
+ if bits == 1:
188
+ cv2.imwrite(path + ".png", out.astype("uint8"))
189
+ elif bits == 2:
190
+ cv2.imwrite(path + ".png", out.astype("uint16"))
191
+
192
+ return
MiDaS/model.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:617d916c0864b95880aed0b6be6d0629ce8b4c0d28361a559f8e5193a9bb554d
3
+ size 149751722
MiDaS/monodepth_net.py ADDED
@@ -0,0 +1,186 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """MonoDepthNet: Network for monocular depth estimation trained by mixing several datasets.
2
+ This file contains code that is adapted from
3
+ https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py
4
+ """
5
+ import torch
6
+ import torch.nn as nn
7
+ from torchvision import models
8
+
9
+
10
+ class MonoDepthNet(nn.Module):
11
+ """Network for monocular depth estimation.
12
+ """
13
+
14
+ def __init__(self, path=None, features=256):
15
+ """Init.
16
+
17
+ Args:
18
+ path (str, optional): Path to saved model. Defaults to None.
19
+ features (int, optional): Number of features. Defaults to 256.
20
+ """
21
+ super().__init__()
22
+
23
+ resnet = models.resnet50(pretrained=False)
24
+
25
+ self.pretrained = nn.Module()
26
+ self.scratch = nn.Module()
27
+ self.pretrained.layer1 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
28
+ resnet.maxpool, resnet.layer1)
29
+
30
+ self.pretrained.layer2 = resnet.layer2
31
+ self.pretrained.layer3 = resnet.layer3
32
+ self.pretrained.layer4 = resnet.layer4
33
+
34
+ # adjust channel number of feature maps
35
+ self.scratch.layer1_rn = nn.Conv2d(256, features, kernel_size=3, stride=1, padding=1, bias=False)
36
+ self.scratch.layer2_rn = nn.Conv2d(512, features, kernel_size=3, stride=1, padding=1, bias=False)
37
+ self.scratch.layer3_rn = nn.Conv2d(1024, features, kernel_size=3, stride=1, padding=1, bias=False)
38
+ self.scratch.layer4_rn = nn.Conv2d(2048, features, kernel_size=3, stride=1, padding=1, bias=False)
39
+
40
+ self.scratch.refinenet4 = FeatureFusionBlock(features)
41
+ self.scratch.refinenet3 = FeatureFusionBlock(features)
42
+ self.scratch.refinenet2 = FeatureFusionBlock(features)
43
+ self.scratch.refinenet1 = FeatureFusionBlock(features)
44
+
45
+ # adaptive output module: 2 convolutions and upsampling
46
+ self.scratch.output_conv = nn.Sequential(nn.Conv2d(features, 128, kernel_size=3, stride=1, padding=1),
47
+ nn.Conv2d(128, 1, kernel_size=3, stride=1, padding=1),
48
+ Interpolate(scale_factor=2, mode='bilinear'))
49
+
50
+ # load model
51
+ if path:
52
+ self.load(path)
53
+
54
+ def forward(self, x):
55
+ """Forward pass.
56
+
57
+ Args:
58
+ x (tensor): input data (image)
59
+
60
+ Returns:
61
+ tensor: depth
62
+ """
63
+ layer_1 = self.pretrained.layer1(x)
64
+ layer_2 = self.pretrained.layer2(layer_1)
65
+ layer_3 = self.pretrained.layer3(layer_2)
66
+ layer_4 = self.pretrained.layer4(layer_3)
67
+
68
+ layer_1_rn = self.scratch.layer1_rn(layer_1)
69
+ layer_2_rn = self.scratch.layer2_rn(layer_2)
70
+ layer_3_rn = self.scratch.layer3_rn(layer_3)
71
+ layer_4_rn = self.scratch.layer4_rn(layer_4)
72
+
73
+ path_4 = self.scratch.refinenet4(layer_4_rn)
74
+ path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
75
+ path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
76
+ path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
77
+
78
+ out = self.scratch.output_conv(path_1)
79
+
80
+ return out
81
+
82
+ def load(self, path):
83
+ """Load model from file.
84
+
85
+ Args:
86
+ path (str): file path
87
+ """
88
+ parameters = torch.load(path)
89
+
90
+ self.load_state_dict(parameters)
91
+
92
+
93
+ class Interpolate(nn.Module):
94
+ """Interpolation module.
95
+ """
96
+
97
+ def __init__(self, scale_factor, mode):
98
+ """Init.
99
+
100
+ Args:
101
+ scale_factor (float): scaling
102
+ mode (str): interpolation mode
103
+ """
104
+ super(Interpolate, self).__init__()
105
+
106
+ self.interp = nn.functional.interpolate
107
+ self.scale_factor = scale_factor
108
+ self.mode = mode
109
+
110
+ def forward(self, x):
111
+ """Forward pass.
112
+
113
+ Args:
114
+ x (tensor): input
115
+
116
+ Returns:
117
+ tensor: interpolated data
118
+ """
119
+ x = self.interp(x, scale_factor=self.scale_factor, mode=self.mode, align_corners=False)
120
+
121
+ return x
122
+
123
+
124
+ class ResidualConvUnit(nn.Module):
125
+ """Residual convolution module.
126
+ """
127
+
128
+ def __init__(self, features):
129
+ """Init.
130
+
131
+ Args:
132
+ features (int): number of features
133
+ """
134
+ super().__init__()
135
+
136
+ self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True)
137
+ self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=False)
138
+ self.relu = nn.ReLU(inplace=True)
139
+
140
+ def forward(self, x):
141
+ """Forward pass.
142
+
143
+ Args:
144
+ x (tensor): input
145
+
146
+ Returns:
147
+ tensor: output
148
+ """
149
+ out = self.relu(x)
150
+ out = self.conv1(out)
151
+ out = self.relu(out)
152
+ out = self.conv2(out)
153
+
154
+ return out + x
155
+
156
+
157
+ class FeatureFusionBlock(nn.Module):
158
+ """Feature fusion block.
159
+ """
160
+
161
+ def __init__(self, features):
162
+ """Init.
163
+
164
+ Args:
165
+ features (int): number of features
166
+ """
167
+ super().__init__()
168
+
169
+ self.resConfUnit = ResidualConvUnit(features)
170
+
171
+ def forward(self, *xs):
172
+ """Forward pass.
173
+
174
+ Returns:
175
+ tensor: output
176
+ """
177
+ output = xs[0]
178
+
179
+ if len(xs) == 2:
180
+ output += self.resConfUnit(xs[1])
181
+
182
+ output = self.resConfUnit(output)
183
+ output = nn.functional.interpolate(output, scale_factor=2,
184
+ mode='bilinear', align_corners=True)
185
+
186
+ return output
MiDaS/run.py ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Compute depth maps for images in the input folder.
2
+ """
3
+ import os
4
+ import glob
5
+ import torch
6
+ # from monodepth_net import MonoDepthNet
7
+ # import utils
8
+ import matplotlib.pyplot as plt
9
+ import numpy as np
10
+ import cv2
11
+ import imageio
12
+
13
+
14
+ def run_depth(img_names, input_path, output_path, model_path, Net, utils, target_w=None):
15
+ """Run MonoDepthNN to compute depth maps.
16
+
17
+ Args:
18
+ input_path (str): path to input folder
19
+ output_path (str): path to output folder
20
+ model_path (str): path to saved model
21
+ """
22
+ print("initialize")
23
+
24
+ # select device
25
+ device = torch.device("cpu")
26
+ print("device: %s" % device)
27
+
28
+ # load network
29
+ model = Net(model_path)
30
+ model.to(device)
31
+ model.eval()
32
+
33
+ # get input
34
+ # img_names = glob.glob(os.path.join(input_path, "*"))
35
+ num_images = len(img_names)
36
+
37
+ # create output folder
38
+ os.makedirs(output_path, exist_ok=True)
39
+
40
+ print("start processing")
41
+
42
+ for ind, img_name in enumerate(img_names):
43
+
44
+ print(" processing {} ({}/{})".format(img_name, ind + 1, num_images))
45
+
46
+ # input
47
+ img = utils.read_image(img_name)
48
+ w = img.shape[1]
49
+ scale = 640. / max(img.shape[0], img.shape[1])
50
+ target_height, target_width = int(round(img.shape[0] * scale)), int(round(img.shape[1] * scale))
51
+ img_input = utils.resize_image(img)
52
+ print(img_input.shape)
53
+ img_input = img_input.to(device)
54
+ # compute
55
+ with torch.no_grad():
56
+ out = model.forward(img_input)
57
+
58
+ depth = utils.resize_depth(out, target_width, target_height)
59
+ img = cv2.resize((img * 255).astype(np.uint8), (target_width, target_height), interpolation=cv2.INTER_AREA)
60
+
61
+ filename = os.path.join(
62
+ output_path, os.path.splitext(os.path.basename(img_name))[0]
63
+ )
64
+ np.save(filename + '.npy', depth)
65
+ utils.write_depth(filename, depth, bits=2)
66
+
67
+ print("finished")
68
+
69
+
70
+ # if __name__ == "__main__":
71
+ # # set paths
72
+ # INPUT_PATH = "image"
73
+ # OUTPUT_PATH = "output"
74
+ # MODEL_PATH = "model.pt"
75
+
76
+ # # set torch options
77
+ # torch.backends.cudnn.enabled = True
78
+ # torch.backends.cudnn.benchmark = True
79
+
80
+ # # compute depth maps
81
+ # run_depth(INPUT_PATH, OUTPUT_PATH, MODEL_PATH, Net, target_w=640)
README.md CHANGED
@@ -1,7 +1,7 @@
1
  ---
2
  title: 3D_Photo_Inpainting
3
- emoji: 🚀
4
- colorFrom: pink
5
  colorTo: red
6
  sdk: gradio
7
  app_file: app.py
1
  ---
2
  title: 3D_Photo_Inpainting
3
+ emoji: 👁
4
+ colorFrom: purple
5
  colorTo: red
6
  sdk: gradio
7
  app_file: app.py
app.py ADDED
@@ -0,0 +1,230 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Repo source: https://github.com/vt-vl-lab/3d-photo-inpainting
2
+
3
+ #import os
4
+ #os.environ['QT_DEBUG_PLUGINS'] = '1'
5
+
6
+ import subprocess
7
+ #subprocess.run('ldd /home/user/.local/lib/python3.8/site-packages/PyQt5/Qt/plugins/platforms/libqxcb.so', shell=True)
8
+ #subprocess.run('pip list', shell=True)
9
+ subprocess.run('nvidia-smi', shell=True)
10
+
11
+ from pyvirtualdisplay import Display
12
+ display = Display(visible=0, size=(1920, 1080)).start()
13
+ #subprocess.run('echo $DISPLAY', shell=True)
14
+
15
+ # 3d inpainting imports
16
+ import numpy as np
17
+ import argparse
18
+ import glob
19
+ import os
20
+ from functools import partial
21
+ import vispy
22
+ import scipy.misc as misc
23
+ from tqdm import tqdm
24
+ import yaml
25
+ import time
26
+ import sys
27
+ from mesh import write_ply, read_ply, output_3d_photo
28
+ from utils import get_MiDaS_samples, read_MiDaS_depth
29
+ import torch
30
+ import cv2
31
+ from skimage.transform import resize
32
+ import imageio
33
+ import copy
34
+ from networks import Inpaint_Color_Net, Inpaint_Depth_Net, Inpaint_Edge_Net
35
+ from MiDaS.run import run_depth
36
+ from boostmonodepth_utils import run_boostmonodepth
37
+ from MiDaS.monodepth_net import MonoDepthNet
38
+ import MiDaS.MiDaS_utils as MiDaS_utils
39
+ from bilateral_filtering import sparse_bilateral_filtering
40
+
41
+ import torch
42
+
43
+ # gradio imports
44
+ import gradio as gr
45
+ import uuid
46
+ from PIL import Image
47
+ from pathlib import Path
48
+ import shutil
49
+ from time import sleep
50
+
51
+ def inpaint(img_name, num_frames, fps, traj_type):
52
+
53
+ print(traj_type)
54
+
55
+ config = yaml.load(open('argument.yml', 'r'))
56
+
57
+ config['num_frames'] = num_frames
58
+ config['fps'] = fps
59
+
60
+ if torch.cuda.is_available():
61
+ config['gpu_ids'] = 0
62
+
63
+ if config['offscreen_rendering'] is True:
64
+ vispy.use(app='egl')
65
+
66
+ os.makedirs(config['mesh_folder'], exist_ok=True)
67
+ os.makedirs(config['video_folder'], exist_ok=True)
68
+ os.makedirs(config['depth_folder'], exist_ok=True)
69
+ sample_list = get_MiDaS_samples(config['src_folder'], config['depth_folder'], config, config['specific'], img_name.stem)
70
+ normal_canvas, all_canvas = None, None
71
+
72
+ if isinstance(config["gpu_ids"], int) and (config["gpu_ids"] >= 0):
73
+ device = config["gpu_ids"]
74
+ else:
75
+ device = "cpu"
76
+
77
+ print(f"running on device {device}")
78
+
79
+ for idx in tqdm(range(len(sample_list))):
80
+ depth = None
81
+ sample = sample_list[idx]
82
+ print("Current Source ==> ", sample['src_pair_name'])
83
+ mesh_fi = os.path.join(config['mesh_folder'], sample['src_pair_name'] +'.ply')
84
+ image = imageio.imread(sample['ref_img_fi'])
85
+
86
+ print(f"Running depth extraction at {time.time()}")
87
+ if config['use_boostmonodepth'] is True:
88
+ run_boostmonodepth(sample['ref_img_fi'], config['src_folder'], config['depth_folder'])
89
+ elif config['require_midas'] is True:
90
+ run_depth([sample['ref_img_fi']], config['src_folder'], config['depth_folder'],
91
+ config['MiDaS_model_ckpt'], MonoDepthNet, MiDaS_utils, target_w=640)
92
+
93
+ if 'npy' in config['depth_format']:
94
+ config['output_h'], config['output_w'] = np.load(sample['depth_fi']).shape[:2]
95
+ else:
96
+ config['output_h'], config['output_w'] = imageio.imread(sample['depth_fi']).shape[:2]
97
+ frac = config['longer_side_len'] / max(config['output_h'], config['output_w'])
98
+ config['output_h'], config['output_w'] = int(config['output_h'] * frac), int(config['output_w'] * frac)
99
+ config['original_h'], config['original_w'] = config['output_h'], config['output_w']
100
+ if image.ndim == 2:
101
+ image = image[..., None].repeat(3, -1)
102
+ if np.sum(np.abs(image[..., 0] - image[..., 1])) == 0 and np.sum(np.abs(image[..., 1] - image[..., 2])) == 0:
103
+ config['gray_image'] = True
104
+ else:
105
+ config['gray_image'] = False
106
+ image = cv2.resize(image, (config['output_w'], config['output_h']), interpolation=cv2.INTER_AREA)
107
+ depth = read_MiDaS_depth(sample['depth_fi'], 3.0, config['output_h'], config['output_w'])
108
+ mean_loc_depth = depth[depth.shape[0]//2, depth.shape[1]//2]
109
+ if not(config['load_ply'] is True and os.path.exists(mesh_fi)):
110
+ vis_photos, vis_depths = sparse_bilateral_filtering(depth.copy(), image.copy(), config, num_iter=config['sparse_iter'], spdb=False)
111
+ depth = vis_depths[-1]
112
+ model = None
113
+ torch.cuda.empty_cache()
114
+ print("Start Running 3D_Photo ...")
115
+ print(f"Loading edge model at {time.time()}")
116
+ depth_edge_model = Inpaint_Edge_Net(init_weights=True)
117
+ depth_edge_weight = torch.load(config['depth_edge_model_ckpt'],
118
+ map_location=torch.device(device))
119
+ depth_edge_model.load_state_dict(depth_edge_weight)
120
+ depth_edge_model = depth_edge_model.to(device)
121
+ depth_edge_model.eval()
122
+
123
+ print(f"Loading depth model at {time.time()}")
124
+ depth_feat_model = Inpaint_Depth_Net()
125
+ depth_feat_weight = torch.load(config['depth_feat_model_ckpt'],
126
+ map_location=torch.device(device))
127
+ depth_feat_model.load_state_dict(depth_feat_weight, strict=True)
128
+ depth_feat_model = depth_feat_model.to(device)
129
+ depth_feat_model.eval()
130
+ depth_feat_model = depth_feat_model.to(device)
131
+ print(f"Loading rgb model at {time.time()}")
132
+ rgb_model = Inpaint_Color_Net()
133
+ rgb_feat_weight = torch.load(config['rgb_feat_model_ckpt'],
134
+ map_location=torch.device(device))
135
+ rgb_model.load_state_dict(rgb_feat_weight)
136
+ rgb_model.eval()
137
+ rgb_model = rgb_model.to(device)
138
+ graph = None
139
+
140
+
141
+ print(f"Writing depth ply (and basically doing everything) at {time.time()}")
142
+ rt_info = write_ply(image,
143
+ depth,
144
+ sample['int_mtx'],
145
+ mesh_fi,
146
+ config,
147
+ rgb_model,
148
+ depth_edge_model,
149
+ depth_edge_model,
150
+ depth_feat_model)
151
+
152
+ if rt_info is False:
153
+ continue
154
+ rgb_model = None
155
+ color_feat_model = None
156
+ depth_edge_model = None
157
+ depth_feat_model = None
158
+ torch.cuda.empty_cache()
159
+ if config['save_ply'] is True or config['load_ply'] is True:
160
+ verts, colors, faces, Height, Width, hFov, vFov = read_ply(mesh_fi)
161
+ else:
162
+ verts, colors, faces, Height, Width, hFov, vFov = rt_info
163
+
164
+
165
+ print(f"Making video at {time.time()}")
166
+ videos_poses, video_basename = copy.deepcopy(sample['tgts_poses']), sample['tgt_name']
167
+ top = (config.get('original_h') // 2 - sample['int_mtx'][1, 2] * config['output_h'])
168
+ left = (config.get('original_w') // 2 - sample['int_mtx'][0, 2] * config['output_w'])
169
+ down, right = top + config['output_h'], left + config['output_w']
170
+ border = [int(xx) for xx in [top, down, left, right]]
171
+ normal_canvas, all_canvas = output_3d_photo(verts.copy(), colors.copy(), faces.copy(), copy.deepcopy(Height), copy.deepcopy(Width), copy.deepcopy(hFov), copy.deepcopy(vFov),
172
+ copy.deepcopy(sample['tgt_pose']), sample['video_postfix'], copy.deepcopy(sample['ref_pose']), copy.deepcopy(config['video_folder']),
173
+ image.copy(), copy.deepcopy(sample['int_mtx']), config, image,
174
+ videos_poses, video_basename, config.get('original_h'), config.get('original_w'), border=border, depth=depth, normal_canvas=normal_canvas, all_canvas=all_canvas,
175
+ mean_loc_depth=mean_loc_depth)
176
+
177
+ def resizer(input_img, max_img_size=512):
178
+ width, height = input_img.size
179
+ long_edge = height if height >= width else width
180
+ if long_edge > max_img_size:
181
+ ratio = max_img_size / long_edge
182
+ resized_width = int(ratio * width)
183
+ resized_height = int(ratio * height)
184
+ resized_input_img = input_img.resize((resized_width, resized_height), resample=2)
185
+ return resized_input_img
186
+
187
+ else:
188
+ return input_img
189
+
190
+ def main_app(input_img, num_frames, fps, traj_type):
191
+
192
+ # Save image in necessary folder for inpainting
193
+ img_name = Path(str(uuid.uuid4()) + '.jpg')
194
+ save_folder = Path('image')
195
+
196
+ input_img = resizer(input_img)
197
+ input_img.save(save_folder/img_name)
198
+
199
+ inpaint(img_name, num_frames, fps, traj_type)
200
+
201
+ #subprocess.run('ls -l', shell=True)
202
+ #subprocess.run('ls image -l', shell=True)
203
+ #subprocess.run('ls video/ -l', shell=True)
204
+
205
+ # Get output video path & return
206
+ input_img_path = str(save_folder/img_name)
207
+ out_vid_path = 'video/{0}_circle.mp4'.format(img_name.stem)
208
+
209
+ return out_vid_path
210
+
211
+ video_choices = ['dolly-zoom-in', 'zoom-in', 'circle', 'swing']
212
+ gradio_inputs = [gr.inputs.Image(type='pil', label='Input Image'),
213
+ gr.inputs.Slider(minimum=60, maximum=240, step=1, default=120, label="Number of Frames"),
214
+ gr.inputs.Slider(minimum=10, maximum=40, step=1, default=20, label="Frames per Second (FPS)"),
215
+ gr.inputs.Radio(choices=video_choices, default='circle', label='(Work-in-progress) What type of 3D video do you want?')]
216
+
217
+ gradio_outputs = [gr.outputs.Video(label='Output Video')]
218
+ examples = [ ['moon.jpg'], ['dog.jpg'] ]
219
+
220
+ description="Convert an image into a trajectory-following video. Images are automatically resized down to a max edge of 512. | NOTE: The current runtime for a sample is around 400-700 seconds. Running on a lower number of frames could help! Do be patient as this is on CPU-only, BUT if this space maybe gets a GPU one day, it's already configured to run with GPU-support :) If you have a GPU, feel free to use the author's original repo, or just `git clone https://huggingface.co/spaces/Classified/3D_Photo_Inpainting`, install packages and requirements, then `python app.py` to run the gradio GUI locally!"
221
+
222
+ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2004.04727' target='_blank'>3D Photography using Context-aware Layered Depth Inpainting</a> | <a href='https://shihmengli.github.io/3D-Photo-Inpainting/' target='_blank'>Github Project Page</a> | <a href='https://github.com/vt-vl-lab/3d-photo-inpainting' target='_blank'>Github Repo</a></p>"
223
+
224
+ iface = gr.Interface(fn=main_app, inputs=gradio_inputs , outputs=gradio_outputs, examples=examples,
225
+ title='3D Image Inpainting',
226
+ description=description,
227
+ article=article,
228
+ enable_queue=True)
229
+
230
+ iface.launch(debug=True)
argument.yml ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ depth_edge_model_ckpt: checkpoints/edge-model.pth
2
+ depth_feat_model_ckpt: checkpoints/depth-model.pth
3
+ rgb_feat_model_ckpt: checkpoints/color-model.pth
4
+ MiDaS_model_ckpt: MiDaS/model.pt
5
+ use_boostmonodepth: False
6
+ fps: 40
7
+ num_frames: 240
8
+ x_shift_range: [-0.015]
9
+ y_shift_range: [-0.015]
10
+ z_shift_range: [-0.05]
11
+ traj_types: ['circle']
12
+ video_postfix: ['circle']
13
+ #x_shift_range: [0.00, 0.00, -0.015, -0.015]
14
+ #y_shift_range: [0.00, 0.00, -0.015, -0.00]
15
+ #z_shift_range: [-0.05, -0.05, -0.05, -0.05]
16
+ #traj_types: ['double-straight-line', 'double-straight-line', 'circle', 'circle']
17
+ #video_postfix: ['dolly-zoom-in', 'zoom-in', 'circle', 'swing']
18
+ specific: ''
19
+ longer_side_len: 960
20
+ src_folder: image
21
+ depth_folder: depth
22
+ mesh_folder: mesh
23
+ video_folder: video
24
+ load_ply: False
25
+ save_ply: True
26
+ inference_video: True
27
+ gpu_ids: -1
28
+ offscreen_rendering: True
29
+ img_format: '.jpg'
30
+ depth_format: '.npy'
31
+ require_midas: True
32
+ depth_threshold: 0.04
33
+ ext_edge_threshold: 0.002
34
+ sparse_iter: 5
35
+ filter_size: [7, 7, 5, 5, 5]
36
+ sigma_s: 4.0
37
+ sigma_r: 0.5
38
+ redundant_number: 12
39
+ background_thickness: 70
40
+ context_thickness: 140
41
+ background_thickness_2: 70
42
+ context_thickness_2: 70
43
+ discount_factor: 1.00
44
+ log_depth: True
45
+ largest_size: 512
46
+ depth_edge_dilate: 10
47
+ depth_edge_dilate_2: 5
48
+ extrapolate_border: True
49
+ extrapolation_thickness: 60
50
+ repeat_inpaint_edge: True
51
+ crop_border: [0.03, 0.03, 0.05, 0.03]
52
+ anti_flickering: True
bilateral_filtering.py ADDED
@@ -0,0 +1,215 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ from functools import reduce
3
+
4
+ def sparse_bilateral_filtering(
5
+ depth, image, config, HR=False, mask=None, gsHR=True, edge_id=None, num_iter=None, num_gs_iter=None, spdb=False
6
+ ):
7
+ """
8
+ config:
9
+ - filter_size
10
+ """
11
+ import time
12
+
13
+ save_images = []
14
+ save_depths = []
15
+ save_discontinuities = []
16
+ vis_depth = depth.copy()
17
+ backup_vis_depth = vis_depth.copy()
18
+
19
+ depth_max = vis_depth.max()
20
+ depth_min = vis_depth.min()
21
+ vis_image = image.copy()
22
+ for i in range(num_iter):
23
+ if isinstance(config["filter_size"], list):
24
+ window_size = config["filter_size"][i]
25
+ else:
26
+ window_size = config["filter_size"]
27
+ vis_image = image.copy()
28
+ save_images.append(vis_image)
29
+ save_depths.append(vis_depth)
30
+ u_over, b_over, l_over, r_over = vis_depth_discontinuity(vis_depth, config, mask=mask)
31
+ vis_image[u_over > 0] = np.array([0, 0, 0])
32
+ vis_image[b_over > 0] = np.array([0, 0, 0])
33
+ vis_image[l_over > 0] = np.array([0, 0, 0])
34
+ vis_image[r_over > 0] = np.array([0, 0, 0])
35
+
36
+ discontinuity_map = (u_over + b_over + l_over + r_over).clip(0.0, 1.0)
37
+ discontinuity_map[depth == 0] = 1
38
+ save_discontinuities.append(discontinuity_map)
39
+ if mask is not None:
40
+ discontinuity_map[mask == 0] = 0
41
+ vis_depth = bilateral_filter(
42
+ vis_depth, config, discontinuity_map=discontinuity_map, HR=HR, mask=mask, window_size=window_size
43
+ )
44
+
45
+ return save_images, save_depths
46
+
47
+
48
+ def vis_depth_discontinuity(depth, config, vis_diff=False, label=False, mask=None):
49
+ """
50
+ config:
51
+ -
52
+ """
53
+ if label == False:
54
+ disp = 1./depth
55
+ u_diff = (disp[1:, :] - disp[:-1, :])[:-1, 1:-1]
56
+ b_diff = (disp[:-1, :] - disp[1:, :])[1:, 1:-1]
57
+ l_diff = (disp[:, 1:] - disp[:, :-1])[1:-1, :-1]
58
+ r_diff = (disp[:, :-1] - disp[:, 1:])[1:-1, 1:]
59
+ if mask is not None:
60
+ u_mask = (mask[1:, :] * mask[:-1, :])[:-1, 1:-1]
61
+ b_mask = (mask[:-1, :] * mask[1:, :])[1:, 1:-1]
62
+ l_mask = (mask[:, 1:] * mask[:, :-1])[1:-1, :-1]
63
+ r_mask = (mask[:, :-1] * mask[:, 1:])[1:-1, 1:]
64
+ u_diff = u_diff * u_mask
65
+ b_diff = b_diff * b_mask
66
+ l_diff = l_diff * l_mask
67
+ r_diff = r_diff * r_mask
68
+ u_over = (np.abs(u_diff) > config['depth_threshold']).astype(np.float32)
69
+ b_over = (np.abs(b_diff) > config['depth_threshold']).astype(np.float32)
70
+ l_over = (np.abs(l_diff) > config['depth_threshold']).astype(np.float32)
71
+ r_over = (np.abs(r_diff) > config['depth_threshold']).astype(np.float32)
72
+ else:
73
+ disp = depth
74
+ u_diff = (disp[1:, :] * disp[:-1, :])[:-1, 1:-1]
75
+ b_diff = (disp[:-1, :] * disp[1:, :])[1:, 1:-1]
76
+ l_diff = (disp[:, 1:] * disp[:, :-1])[1:-1, :-1]
77
+ r_diff = (disp[:, :-1] * disp[:, 1:])[1:-1, 1:]
78
+ if mask is not None:
79
+ u_mask = (mask[1:, :] * mask[:-1, :])[:-1, 1:-1]
80
+ b_mask = (mask[:-1, :] * mask[1:, :])[1:, 1:-1]
81
+ l_mask = (mask[:, 1:] * mask[:, :-1])[1:-1, :-1]
82
+ r_mask = (mask[:, :-1] * mask[:, 1:])[1:-1, 1:]
83
+ u_diff = u_diff * u_mask
84
+ b_diff = b_diff * b_mask
85
+ l_diff = l_diff * l_mask
86
+ r_diff = r_diff * r_mask
87
+ u_over = (np.abs(u_diff) > 0).astype(np.float32)
88
+ b_over = (np.abs(b_diff) > 0).astype(np.float32)
89
+ l_over = (np.abs(l_diff) > 0).astype(np.float32)
90
+ r_over = (np.abs(r_diff) > 0).astype(np.float32)
91
+ u_over = np.pad(u_over, 1, mode='constant')
92
+ b_over = np.pad(b_over, 1, mode='constant')
93
+ l_over = np.pad(l_over, 1, mode='constant')
94
+ r_over = np.pad(r_over, 1, mode='constant')
95
+ u_diff = np.pad(u_diff, 1, mode='constant')
96
+ b_diff = np.pad(b_diff, 1, mode='constant')
97
+ l_diff = np.pad(l_diff, 1, mode='constant')
98
+ r_diff = np.pad(r_diff, 1, mode='constant')
99
+
100
+ if vis_diff:
101
+ return [u_over, b_over, l_over, r_over], [u_diff, b_diff, l_diff, r_diff]
102
+ else:
103
+ return [u_over, b_over, l_over, r_over]
104
+
105
+ def bilateral_filter(depth, config, discontinuity_map=None, HR=False, mask=None, window_size=False):
106
+ sort_time = 0
107
+ replace_time = 0
108
+ filter_time = 0
109
+ init_time = 0
110
+ filtering_time = 0
111
+ sigma_s = config['sigma_s']
112
+ sigma_r = config['sigma_r']
113
+ if window_size == False:
114
+ window_size = config['filter_size']
115
+ midpt = window_size//2
116
+ ax = np.arange(-midpt, midpt+1.)
117
+ xx, yy = np.meshgrid(ax, ax)
118
+ if discontinuity_map is not None:
119
+ spatial_term = np.exp(-(xx**2 + yy**2) / (2. * sigma_s**2))
120
+
121
+ # padding
122
+ depth = depth[1:-1, 1:-1]
123
+ depth = np.pad(depth, ((1,1), (1,1)), 'edge')
124
+ pad_depth = np.pad(depth, (midpt,midpt), 'edge')
125
+ if discontinuity_map is not None:
126
+ discontinuity_map = discontinuity_map[1:-1, 1:-1]
127
+ discontinuity_map = np.pad(discontinuity_map, ((1,1), (1,1)), 'edge')
128
+ pad_discontinuity_map = np.pad(discontinuity_map, (midpt,midpt), 'edge')
129
+ pad_discontinuity_hole = 1 - pad_discontinuity_map
130
+ # filtering
131
+ output = depth.copy()
132
+ pad_depth_patches = rolling_window(pad_depth, [window_size, window_size], [1,1])
133
+ if discontinuity_map is not None:
134
+ pad_discontinuity_patches = rolling_window(pad_discontinuity_map, [window_size, window_size], [1,1])
135
+ pad_discontinuity_hole_patches = rolling_window(pad_discontinuity_hole, [window_size, window_size], [1,1])
136
+
137
+ if mask is not None:
138
+ pad_mask = np.pad(mask, (midpt,midpt), 'constant')
139
+ pad_mask_patches = rolling_window(pad_mask, [window_size, window_size], [1,1])
140
+ from itertools import product
141
+ if discontinuity_map is not None:
142
+ pH, pW = pad_depth_patches.shape[:2]
143
+ for pi in range(pH):
144
+ for pj in range(pW):
145
+ if mask is not None and mask[pi, pj] == 0:
146
+ continue
147
+ if discontinuity_map is not None:
148
+ if bool(pad_discontinuity_patches[pi, pj].any()) is False:
149
+ continue
150
+ discontinuity_patch = pad_discontinuity_patches[pi, pj]
151
+ discontinuity_holes = pad_discontinuity_hole_patches[pi, pj]
152
+ depth_patch = pad_depth_patches[pi, pj]
153
+ depth_order = depth_patch.ravel().argsort()
154
+ patch_midpt = depth_patch[window_size//2, window_size//2]
155
+ if discontinuity_map is not None:
156
+ coef = discontinuity_holes.astype(np.float32)
157
+ if mask is not None:
158
+ coef = coef * pad_mask_patches[pi, pj]
159
+ else:
160
+ range_term = np.exp(-(depth_patch-patch_midpt)**2 / (2. * sigma_r**2))
161
+ coef = spatial_term * range_term
162
+ if coef.max() == 0:
163
+ output[pi, pj] = patch_midpt
164
+ continue
165
+ if discontinuity_map is not None and (coef.max() == 0):
166
+ output[pi, pj] = patch_midpt
167
+ else:
168
+ coef = coef/(coef.sum())
169
+ coef_order = coef.ravel()[depth_order]
170
+ cum_coef = np.cumsum(coef_order)
171
+ ind = np.digitize(0.5, cum_coef)
172
+ output[pi, pj] = depth_patch.ravel()[depth_order][ind]
173
+ else:
174
+ pH, pW = pad_depth_patches.shape[:2]
175
+ for pi in range(pH):
176
+ for pj in range(pW):
177
+ if discontinuity_map is not None:
178
+ if pad_discontinuity_patches[pi, pj][window_size//2, window_size//2] == 1:
179
+ continue
180
+ discontinuity_patch = pad_discontinuity_patches[pi, pj]
181
+ discontinuity_holes = (1. - discontinuity_patch)
182
+ depth_patch = pad_depth_patches[pi, pj]
183
+ depth_order = depth_patch.ravel().argsort()
184
+ patch_midpt = depth_patch[window_size//2, window_size//2]
185
+ range_term = np.exp(-(depth_patch-patch_midpt)**2 / (2. * sigma_r**2))
186
+ if discontinuity_map is not None:
187
+ coef = spatial_term * range_term * discontinuity_holes
188
+ else:
189
+ coef = spatial_term * range_term
190
+ if coef.sum() == 0:
191
+ output[pi, pj] = patch_midpt
192
+ continue
193
+ if discontinuity_map is not None and (coef.sum() == 0):
194
+ output[pi, pj] = patch_midpt
195
+ else:
196
+ coef = coef/(coef.sum())
197
+ coef_order = coef.ravel()[depth_order]
198
+ cum_coef = np.cumsum(coef_order)
199
+ ind = np.digitize(0.5, cum_coef)
200
+ output[pi, pj] = depth_patch.ravel()[depth_order][ind]
201
+
202
+ return output
203
+
204
+ def rolling_window(a, window, strides):
205
+ assert len(a.shape)==len(window)==len(strides), "\'a\', \'window\', \'strides\' dimension mismatch"
206
+ shape_fn = lambda i,w,s: (a.shape[i]-w)//s + 1
207
+ shape = [shape_fn(i,w,s) for i,(w,s) in enumerate(zip(window, strides))] + list(window)
208
+ def acc_shape(i):
209
+ if i+1>=len(a.shape):
210
+ return 1
211
+ else:
212
+ return reduce(lambda x,y:x*y, a.shape[i+1:])
213
+ _strides = [acc_shape(i)*s*a.itemsize for i,s in enumerate(strides)] + list(a.strides)
214
+
215
+ return np.lib.stride_tricks.as_strided(a, shape=shape, strides=_strides)
boostmonodepth_utils.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import cv2
3
+ import glob
4
+ import numpy as np
5
+ import imageio
6
+ from MiDaS.MiDaS_utils import write_depth
7
+
8
+ BOOST_BASE = 'BoostingMonocularDepth'
9
+
10
+ BOOST_INPUTS = 'inputs'
11
+ BOOST_OUTPUTS = 'outputs'
12
+
13
+ def run_boostmonodepth(img_names, src_folder, depth_folder):
14
+
15
+ if not isinstance(img_names, list):
16
+ img_names = [img_names]
17
+
18
+ # remove irrelevant files first
19
+ clean_folder(os.path.join(BOOST_BASE, BOOST_INPUTS))
20
+ clean_folder(os.path.join(BOOST_BASE, BOOST_OUTPUTS))
21
+
22
+ tgt_names = []
23
+ for img_name in img_names:
24
+ base_name = os.path.basename(img_name)
25
+ tgt_name = os.path.join(BOOST_BASE, BOOST_INPUTS, base_name)
26
+ os.system(f'cp {img_name} {tgt_name}')
27
+
28
+ # keep only the file name here.
29
+ # they save all depth as .png file
30
+ tgt_names.append(os.path.basename(tgt_name).replace('.jpg', '.png'))
31
+
32
+ os.system(f'cd {BOOST_BASE} && python run.py --Final --data_dir {BOOST_INPUTS}/ --output_dir {BOOST_OUTPUTS} --depthNet 0')
33
+
34
+ for i, (img_name, tgt_name) in enumerate(zip(img_names, tgt_names)):
35
+ img = imageio.imread(img_name)
36
+ H, W = img.shape[:2]
37
+ scale = 640. / max(H, W)
38
+
39
+ # resize and save depth
40
+ target_height, target_width = int(round(H * scale)), int(round(W * scale))
41
+ depth = imageio.imread(os.path.join(BOOST_BASE, BOOST_OUTPUTS, tgt_name))
42
+ depth = np.array(depth).astype(np.float32)
43
+ depth = resize_depth(depth, target_width, target_height)
44
+ np.save(os.path.join(depth_folder, tgt_name.replace('.png', '.npy')), depth / 32768. - 1.)
45
+ write_depth(os.path.join(depth_folder, tgt_name.replace('.png', '')), depth)
46
+
47
+ def clean_folder(folder, img_exts=['.png', '.jpg', '.npy']):
48
+
49
+ for img_ext in img_exts:
50
+ paths_to_check = os.path.join(folder, f'*{img_ext}')
51
+ if len(glob.glob(paths_to_check)) == 0:
52
+ continue
53
+ print(paths_to_check)
54
+ os.system(f'rm {paths_to_check}')
55
+
56
+ def resize_depth(depth, width, height):
57
+ """Resize numpy (or image read by imageio) depth map
58
+
59
+ Args:
60
+ depth (numpy): depth
61
+ width (int): image width
62
+ height (int): image height
63
+
64
+ Returns:
65
+ array: processed depth
66
+ """
67
+ depth = cv2.blur(depth, (3, 3))
68
+ return cv2.resize(depth, (width, height), interpolation=cv2.INTER_AREA)
checkpoints/color-model.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:383c9b1db70097907a6f9c8abb0303e7056f50d5456a36f34ab784592b8b2c20
3
+ size 206331633
checkpoints/depth-model.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2d0e63e89a22762ddfa8bc8c9f8c992e5532b140123274ffc6e4171baa1b76f8
3
+ size 206272258
checkpoints/edge-model.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b1d768bd008ad5fe9f540004f870b8c3d355e4939b2009aa4db493fd313217c9
3
+ size 45974122
dog.jpg ADDED
download.sh ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/sh
2
+ fb_status=$(wget --spider -S https://filebox.ece.vt.edu/ 2>&1 | grep "HTTP/1.1 200 OK")
3
+
4
+ mkdir checkpoints
5
+
6
+ echo "downloading from filebox ..."
7
+ wget https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/color-model.pth
8
+ wget https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/depth-model.pth
9
+ wget https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/edge-model.pth
10
+ wget https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/model.pt
11
+
12
+ mv color-model.pth checkpoints/.
13
+ mv depth-model.pth checkpoints/.
14
+ mv edge-model.pth checkpoints/.
15
+ mv model.pt MiDaS/.
16
+
17
+ echo "cloning from BoostingMonocularDepth ..."
18
+ git clone https://github.com/compphoto/BoostingMonocularDepth.git
19
+ mkdir -p BoostingMonocularDepth/pix2pix/checkpoints/mergemodel/
20
+
21
+ echo "downloading mergenet weights ..."
22
+ wget https://filebox.ece.vt.edu/~jbhuang/project/3DPhoto/model/latest_net_G.pth
23
+ mv latest_net_G.pth BoostingMonocularDepth/pix2pix/checkpoints/mergemodel/
24
+ wget https://github.com/intel-isl/MiDaS/releases/download/v2/model-f46da743.pt
25
+ mv model-f46da743.pt BoostingMonocularDepth/midas/model.pt
gradio_queue.db ADDED
Binary file (16.4 kB). View file
image/.empty ADDED
File without changes
latest_net_G.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:50ec735d74ed6499562d898f41b49343e521808b8dae589aa3c2f5c9ac9f7462
3
+ size 318268048
main.py ADDED
@@ -0,0 +1,141 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import argparse
3
+ import glob
4
+ import os
5
+ from functools import partial
6
+ import vispy
7
+ import scipy.misc as misc
8
+ from tqdm import tqdm
9
+ import yaml
10
+ import time
11
+ import sys
12
+ from mesh import write_ply, read_ply, output_3d_photo
13
+ from utils import get_MiDaS_samples, read_MiDaS_depth
14
+ import torch
15
+ import cv2
16
+ from skimage.transform import resize
17
+ import imageio
18
+ import copy
19
+ from networks import Inpaint_Color_Net, Inpaint_Depth_Net, Inpaint_Edge_Net
20
+ from MiDaS.run import run_depth
21
+ from boostmonodepth_utils import run_boostmonodepth
22
+ from MiDaS.monodepth_net import MonoDepthNet
23
+ import MiDaS.MiDaS_utils as MiDaS_utils
24
+ from bilateral_filtering import sparse_bilateral_filtering
25
+
26
+ parser = argparse.ArgumentParser()
27
+ parser.add_argument('--config', type=str, default='argument.yml',help='Configure of post processing')
28
+ args = parser.parse_args()
29
+ config = yaml.load(open(args.config, 'r'))
30
+ if config['offscreen_rendering'] is True:
31
+ vispy.use(app='egl')
32
+ os.makedirs(config['mesh_folder'], exist_ok=True)
33
+ os.makedirs(config['video_folder'], exist_ok=True)
34
+ os.makedirs(config['depth_folder'], exist_ok=True)
35
+ sample_list = get_MiDaS_samples(config['src_folder'], config['depth_folder'], config, config['specific'])
36
+ normal_canvas, all_canvas = None, None
37
+
38
+ if isinstance(config["gpu_ids"], int) and (config["gpu_ids"] >= 0):
39
+ device = config["gpu_ids"]
40
+ else:
41
+ device = "cpu"
42
+
43
+ print(f"running on device {device}")
44
+
45
+ for idx in tqdm(range(len(sample_list))):
46
+ depth = None
47
+ sample = sample_list[idx]
48
+ print("Current Source ==> ", sample['src_pair_name'])
49
+ mesh_fi = os.path.join(config['mesh_folder'], sample['src_pair_name'] +'.ply')
50
+ image = imageio.imread(sample['ref_img_fi'])
51
+
52
+ print(f"Running depth extraction at {time.time()}")
53
+ if config['use_boostmonodepth'] is True:
54
+ run_boostmonodepth(sample['ref_img_fi'], config['src_folder'], config['depth_folder'])
55
+ elif config['require_midas'] is True:
56
+ run_depth([sample['ref_img_fi']], config['src_folder'], config['depth_folder'],
57
+ config['MiDaS_model_ckpt'], MonoDepthNet, MiDaS_utils, target_w=640)
58
+
59
+ if 'npy' in config['depth_format']:
60
+ config['output_h'], config['output_w'] = np.load(sample['depth_fi']).shape[:2]
61
+ else:
62
+ config['output_h'], config['output_w'] = imageio.imread(sample['depth_fi']).shape[:2]
63
+ frac = config['longer_side_len'] / max(config['output_h'], config['output_w'])
64
+ config['output_h'], config['output_w'] = int(config['output_h'] * frac), int(config['output_w'] * frac)
65
+ config['original_h'], config['original_w'] = config['output_h'], config['output_w']
66
+ if image.ndim == 2:
67
+ image = image[..., None].repeat(3, -1)
68
+ if np.sum(np.abs(image[..., 0] - image[..., 1])) == 0 and np.sum(np.abs(image[..., 1] - image[..., 2])) == 0:
69
+ config['gray_image'] = True
70
+ else:
71
+ config['gray_image'] = False
72
+ image = cv2.resize(image, (config['output_w'], config['output_h']), interpolation=cv2.INTER_AREA)
73
+ depth = read_MiDaS_depth(sample['depth_fi'], 3.0, config['output_h'], config['output_w'])
74
+ mean_loc_depth = depth[depth.shape[0]//2, depth.shape[1]//2]
75
+ if not(config['load_ply'] is True and os.path.exists(mesh_fi)):
76
+ vis_photos, vis_depths = sparse_bilateral_filtering(depth.copy(), image.copy(), config, num_iter=config['sparse_iter'], spdb=False)
77
+ depth = vis_depths[-1]
78
+ model = None
79
+ torch.cuda.empty_cache()
80
+ print("Start Running 3D_Photo ...")
81
+ print(f"Loading edge model at {time.time()}")
82
+ depth_edge_model = Inpaint_Edge_Net(init_weights=True)
83
+ depth_edge_weight = torch.load(config['depth_edge_model_ckpt'],
84
+ map_location=torch.device(device))
85
+ depth_edge_model.load_state_dict(depth_edge_weight)
86
+ depth_edge_model = depth_edge_model.to(device)
87
+ depth_edge_model.eval()
88
+
89
+ print(f"Loading depth model at {time.time()}")
90
+ depth_feat_model = Inpaint_Depth_Net()
91
+ depth_feat_weight = torch.load(config['depth_feat_model_ckpt'],
92
+ map_location=torch.device(device))
93
+ depth_feat_model.load_state_dict(depth_feat_weight, strict=True)
94
+ depth_feat_model = depth_feat_model.to(device)
95
+ depth_feat_model.eval()
96
+ depth_feat_model = depth_feat_model.to(device)
97
+ print(f"Loading rgb model at {time.time()}")
98
+ rgb_model = Inpaint_Color_Net()
99
+ rgb_feat_weight = torch.load(config['rgb_feat_model_ckpt'],
100
+ map_location=torch.device(device))
101
+ rgb_model.load_state_dict(rgb_feat_weight)
102
+ rgb_model.eval()
103
+ rgb_model = rgb_model.to(device)
104
+ graph = None
105
+
106
+
107
+ print(f"Writing depth ply (and basically doing everything) at {time.time()}")
108
+ rt_info = write_ply(image,
109
+ depth,
110
+ sample['int_mtx'],
111
+ mesh_fi,
112
+ config,
113
+ rgb_model,
114
+ depth_edge_model,
115
+ depth_edge_model,
116
+ depth_feat_model)
117
+
118
+ if rt_info is False:
119
+ continue
120
+ rgb_model = None
121
+ color_feat_model = None
122
+ depth_edge_model = None
123
+ depth_feat_model = None
124
+ torch.cuda.empty_cache()
125
+ if config['save_ply'] is True or config['load_ply'] is True:
126
+ verts, colors, faces, Height, Width, hFov, vFov = read_ply(mesh_fi)
127
+ else:
128
+ verts, colors, faces, Height, Width, hFov, vFov = rt_info
129
+
130
+
131
+ print(f"Making video at {time.time()}")
132
+ videos_poses, video_basename = copy.deepcopy(sample['tgts_poses']), sample['tgt_name']
133
+ top = (config.get('original_h') // 2 - sample['int_mtx'][1, 2] * config['output_h'])
134
+ left = (config.get('original_w') // 2 - sample['int_mtx'][0, 2] * config['output_w'])
135
+ down, right = top + config['output_h'], left + config['output_w']
136
+ border = [int(xx) for xx in [top, down, left, right]]
137
+ normal_canvas, all_canvas = output_3d_photo(verts.copy(), colors.copy(), faces.copy(), copy.deepcopy(Height), copy.deepcopy(Width), copy.deepcopy(hFov), copy.deepcopy(vFov),
138
+ copy.deepcopy(sample['tgt_pose']), sample['video_postfix'], copy.deepcopy(sample['ref_pose']), copy.deepcopy(config['video_folder']),
139
+ image.copy(), copy.deepcopy(sample['int_mtx']), config, image,
140
+ videos_poses, video_basename, config.get('original_h'), config.get('original_w'), border=border, depth=depth, normal_canvas=normal_canvas, all_canvas=all_canvas,
141
+ mean_loc_depth=mean_loc_depth)
mesh.py ADDED
The diff for this file is too large to render. See raw diff
mesh_tools.py ADDED
@@ -0,0 +1,1083 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import numpy as np
3
+ try:
4
+ import cynetworkx as netx
5
+ except ImportError:
6
+ import networkx as netx
7
+
8
+ import json
9
+ import scipy.misc as misc
10
+ #import OpenEXR
11
+ import scipy.signal as signal
12
+ import matplotlib.pyplot as plt
13
+ import cv2
14
+ import scipy.misc as misc
15
+ from skimage import io
16
+ from functools import partial
17
+ from vispy import scene, io
18
+ from vispy.scene import visuals
19
+ from functools import reduce
20
+ # from moviepy.editor import ImageSequenceClip
21
+ import scipy.misc as misc
22
+ from vispy.visuals.filters import Alpha
23
+ import cv2
24
+ from skimage.transform import resize
25
+ import copy
26
+ import torch
27
+ import os
28
+ from utils import refine_depth_around_edge, smooth_cntsyn_gap
29
+ from utils import require_depth_edge, filter_irrelevant_edge_new, open_small_mask
30
+ from skimage.feature import canny
31
+ from scipy import ndimage
32
+ import time
33
+ import transforms3d
34
+
35
+ def relabel_node(mesh, nodes, cur_node, new_node):
36
+ if cur_node == new_node:
37
+ return mesh
38
+ mesh.add_node(new_node)
39
+ for key, value in nodes[cur_node].items():
40
+ nodes[new_node][key] = value
41
+ for ne in mesh.neighbors(cur_node):
42
+ mesh.add_edge(new_node, ne)
43
+ mesh.remove_node(cur_node)
44
+
45
+ return mesh
46
+
47
+ def filter_edge(mesh, edge_ccs, config, invalid=False):
48
+ context_ccs = [set() for _ in edge_ccs]
49
+ mesh_nodes = mesh.nodes
50
+ for edge_id, edge_cc in enumerate(edge_ccs):
51
+ if config['context_thickness'] == 0:
52
+ continue
53
+ edge_group = {}
54
+ for edge_node in edge_cc:
55
+ far_nodes = mesh_nodes[edge_node].get('far')
56
+ if far_nodes is None:
57
+ continue
58
+ for far_node in far_nodes:
59
+ context_ccs[edge_id].add(far_node)
60
+ if mesh_nodes[far_node].get('edge_id') is not None:
61
+ if edge_group.get(mesh_nodes[far_node]['edge_id']) is None:
62
+ edge_group[mesh_nodes[far_node]['edge_id']] = set()
63
+ edge_group[mesh_nodes[far_node]['edge_id']].add(far_node)
64
+ if len(edge_cc) > 2:
65
+ for edge_key in [*edge_group.keys()]:
66
+ if len(edge_group[edge_key]) == 1:
67
+ context_ccs[edge_id].remove([*edge_group[edge_key]][0])
68
+ valid_edge_ccs = []
69
+ for xidx, yy in enumerate(edge_ccs):
70
+ if invalid is not True and len(context_ccs[xidx]) > 0:
71
+ # if len(context_ccs[xidx]) > 0:
72
+ valid_edge_ccs.append(yy)
73
+ elif invalid is True and len(context_ccs[xidx]) == 0:
74
+ valid_edge_ccs.append(yy)
75
+ else:
76
+ valid_edge_ccs.append(set())
77
+ # valid_edge_ccs = [yy for xidx, yy in enumerate(edge_ccs) if len(context_ccs[xidx]) > 0]
78
+
79
+ return valid_edge_ccs
80
+
81
+ def extrapolate(global_mesh,
82
+ info_on_pix,
83
+ image,
84
+ depth,
85
+ other_edge_with_id,
86
+ edge_map,
87
+ edge_ccs,
88
+ depth_edge_model,
89
+ depth_feat_model,
90
+ rgb_feat_model,
91
+ config,
92
+ direc='right-up'):
93
+ h_off, w_off = global_mesh.graph['hoffset'], global_mesh.graph['woffset']
94
+ noext_H, noext_W = global_mesh.graph['noext_H'], global_mesh.graph['noext_W']
95
+
96
+ if "up" in direc.lower() and "-" not in direc.lower():
97
+ all_anchor = [0, h_off + config['context_thickness'], w_off, w_off + noext_W]
98
+ global_shift = [all_anchor[0], all_anchor[2]]
99
+ mask_anchor = [0, h_off, w_off, w_off + noext_W]
100
+ context_anchor = [h_off, h_off + config['context_thickness'], w_off, w_off + noext_W]
101
+ valid_line_anchor = [h_off, h_off + 1, w_off, w_off + noext_W]
102
+ valid_anchor = [min(mask_anchor[0], context_anchor[0]), max(mask_anchor[1], context_anchor[1]),
103
+ min(mask_anchor[2], context_anchor[2]), max(mask_anchor[3], context_anchor[3])]
104
+ elif "down" in direc.lower() and "-" not in direc.lower():
105
+ all_anchor = [h_off + noext_H - config['context_thickness'], 2 * h_off + noext_H, w_off, w_off + noext_W]
106
+ global_shift = [all_anchor[0], all_anchor[2]]
107
+ mask_anchor = [h_off + noext_H, 2 * h_off + noext_H, w_off, w_off + noext_W]
108
+ context_anchor = [h_off + noext_H - config['context_thickness'], h_off + noext_H, w_off, w_off + noext_W]
109
+ valid_line_anchor = [h_off + noext_H - 1, h_off + noext_H, w_off, w_off + noext_W]
110
+ valid_anchor = [min(mask_anchor[0], context_anchor[0]), max(mask_anchor[1], context_anchor[1]),
111
+ min(mask_anchor[2], context_anchor[2]), max(mask_anchor[3], context_anchor[3])]
112
+ elif "left" in direc.lower() and "-" not in direc.lower():
113
+ all_anchor = [h_off, h_off + noext_H, 0, w_off + config['context_thickness']]
114
+ global_shift = [all_anchor[0], all_anchor[2]]
115
+ mask_anchor = [h_off, h_off + noext_H, 0, w_off]
116
+ context_anchor = [h_off, h_off + noext_H, w_off, w_off + config['context_thickness']]
117
+ valid_line_anchor = [h_off, h_off + noext_H, w_off, w_off + 1]
118
+ valid_anchor = [min(mask_anchor[0], context_anchor[0]), max(mask_anchor[1], context_anchor[1]),
119
+ min(mask_anchor[2], context_anchor[2]), max(mask_anchor[3], context_anchor[3])]
120
+ elif "right" in direc.lower() and "-" not in direc.lower():
121
+ all_anchor = [h_off, h_off + noext_H, w_off + noext_W - config['context_thickness'], 2 * w_off + noext_W]
122
+ global_shift = [all_anchor[0], all_anchor[2]]
123
+ mask_anchor = [h_off, h_off + noext_H, w_off + noext_W, 2 * w_off + noext_W]
124
+ context_anchor = [h_off, h_off + noext_H, w_off + noext_W - config['context_thickness'], w_off + noext_W]
125
+ valid_line_anchor = [h_off, h_off + noext_H, w_off + noext_W - 1, w_off + noext_W]
126
+ valid_anchor = [min(mask_anchor[0], context_anchor[0]), max(mask_anchor[1], context_anchor[1]),
127
+ min(mask_anchor[2], context_anchor[2]), max(mask_anchor[3], context_anchor[3])]
128
+ elif "left" in direc.lower() and "up" in direc.lower() and "-" in direc.lower():
129
+ all_anchor = [0, h_off + config['context_thickness'], 0, w_off + config['context_thickness']]
130
+ global_shift = [all_anchor[0], all_anchor[2]]
131
+ mask_anchor = [0, h_off, 0, w_off]
132
+ context_anchor = "inv-mask"
133
+ valid_line_anchor = None
134
+ valid_anchor = all_anchor
135
+ elif "left" in direc.lower() and "down" in direc.lower() and "-" in direc.lower():
136
+ all_anchor = [h_off + noext_H - config['context_thickness'], 2 * h_off + noext_H, 0, w_off + config['context_thickness']]
137
+ global_shift = [all_anchor[0], all_anchor[2]]
138
+ mask_anchor = [h_off + noext_H, 2 * h_off + noext_H, 0, w_off]
139
+ context_anchor = "inv-mask"
140
+ valid_line_anchor = None
141
+ valid_anchor = all_anchor
142
+ elif "right" in direc.lower() and "up" in direc.lower() and "-" in direc.lower():
143
+ all_anchor = [0, h_off + config['context_thickness'], w_off + noext_W - config['context_thickness'], 2 * w_off + noext_W]
144
+ global_shift = [all_anchor[0], all_anchor[2]]
145
+ mask_anchor = [0, h_off, w_off + noext_W, 2 * w_off + noext_W]
146
+ context_anchor = "inv-mask"
147
+ valid_line_anchor = None
148
+ valid_anchor = all_anchor
149
+ elif "right" in direc.lower() and "down" in direc.lower() and "-" in direc.lower():
150
+ 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]
151
+ global_shift = [all_anchor[0], all_anchor[2]]
152
+ mask_anchor = [h_off + noext_H, 2 * h_off + noext_H, w_off + noext_W, 2 * w_off + noext_W]
153
+ context_anchor = "inv-mask"
154
+ valid_line_anchor = None
155
+ valid_anchor = all_anchor
156
+
157
+ global_mask = np.zeros_like(depth)
158
+ global_mask[mask_anchor[0]:mask_anchor[1],mask_anchor[2]:mask_anchor[3]] = 1
159
+ mask = global_mask[valid_anchor[0]:valid_anchor[1], valid_anchor[2]:valid_anchor[3]] * 1
160
+ context = 1 - mask
161
+ global_context = np.zeros_like(depth)
162
+ global_context[all_anchor[0]:all_anchor[1],all_anchor[2]:all_anchor[3]] = context
163
+ # context = global_context[valid_anchor[0]:valid_anchor[1], valid_anchor[2]:valid_anchor[3]] * 1
164
+
165
+
166
+
167
+ valid_area = mask + context
168
+ input_rgb = image[valid_anchor[0]:valid_anchor[1], valid_anchor[2]:valid_anchor[3]] / 255. * context[..., None]
169
+ input_depth = depth[valid_anchor[0]:valid_anchor[1], valid_anchor[2]:valid_anchor[3]] * context
170
+ log_depth = np.log(input_depth + 1e-8)
171
+ log_depth[mask > 0] = 0
172
+ input_mean_depth = np.mean(log_depth[context > 0])
173
+ input_zero_mean_depth = (log_depth - input_mean_depth) * context
174
+ input_disp = 1./np.abs(input_depth)
175
+ input_disp[mask > 0] = 0
176
+ input_disp = input_disp / input_disp.max()
177
+ valid_line = np.zeros_like(depth)
178
+ if valid_line_anchor is not None:
179
+ valid_line[valid_line_anchor[0]:valid_line_anchor[1], valid_line_anchor[2]:valid_line_anchor[3]] = 1
180
+ valid_line = valid_line[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]]
181
+ # f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True); ax1.imshow(global_context * 1 + global_mask * 2); ax2.imshow(image); plt.show()
182
+ # 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()
183
+ # import pdb; pdb.set_trace()
184
+ # return
185
+ input_edge_map = edge_map[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]] * context
186
+ input_other_edge_with_id = other_edge_with_id[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]]
187
+ end_depth_maps = ((valid_line * input_edge_map) > 0) * input_depth
188
+
189
+
190
+ if isinstance(config["gpu_ids"], int) and (config["gpu_ids"] >= 0):
191
+ device = config["gpu_ids"]
192
+ else:
193
+ device = "cpu"
194
+
195
+ valid_edge_ids = sorted(list(input_other_edge_with_id[(valid_line * input_edge_map) > 0]))
196
+ valid_edge_ids = valid_edge_ids[1:] if (len(valid_edge_ids) > 0 and valid_edge_ids[0] == -1) else valid_edge_ids
197
+ 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))
198
+ t_edge = torch.FloatTensor(edge).to(device)[None, None, ...]
199
+ t_rgb = torch.FloatTensor(input_rgb).to(device).permute(2,0,1).unsqueeze(0)
200
+ t_mask = torch.FloatTensor(mask).to(device)[None, None, ...]
201
+ t_context = torch.FloatTensor(context).to(device)[None, None, ...]
202
+ t_disp = torch.FloatTensor(input_disp).to(device)[None, None, ...]
203
+ t_depth_zero_mean_depth = torch.FloatTensor(input_zero_mean_depth).to(device)[None, None, ...]
204
+
205
+ depth_edge_output = depth_edge_model.forward_3P(t_mask, t_context, t_rgb, t_disp, t_edge, unit_length=128,
206
+ cuda=device)
207
+ t_output_edge = (depth_edge_output> config['ext_edge_threshold']).float() * t_mask + t_edge
208
+ output_raw_edge = t_output_edge.data.cpu().numpy().squeeze()
209
+ # import pdb; pdb.set_trace()
210
+ mesh = netx.Graph()
211
+ hxs, hys = np.where(output_raw_edge * mask > 0)
212
+ valid_map = mask + context
213
+ for hx, hy in zip(hxs, hys):
214
+ node = (hx, hy)
215
+ mesh.add_node((hx, hy))
216
+ eight_nes = [ne for ne in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1), \
217
+ (hx + 1, hy + 1), (hx - 1, hy - 1), (hx - 1, hy + 1), (hx + 1, hy - 1)]\
218
+ 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]]]
219
+ for ne in eight_nes:
220
+ mesh.add_edge(node, ne, length=np.hypot(ne[0] - hx, ne[1] - hy))
221
+ if end_depth_maps[ne[0], ne[1]] != 0:
222
+ mesh.nodes[ne[0], ne[1]]['cnt'] = True
223
+ mesh.nodes[ne[0], ne[1]]['depth'] = end_depth_maps[ne[0], ne[1]]
224
+ ccs = [*netx.connected_components(mesh)]
225
+ end_pts = []
226
+ for cc in ccs:
227
+ end_pts.append(set())
228
+ for node in cc:
229
+ if mesh.nodes[node].get('cnt') is not None:
230
+ end_pts[-1].add((node[0], node[1], mesh.nodes[node]['depth']))
231
+ fpath_map = np.zeros_like(output_raw_edge) - 1
232
+ npath_map = np.zeros_like(output_raw_edge) - 1
233
+ for end_pt, cc in zip(end_pts, ccs):
234
+ sorted_end_pt = []
235
+ if len(end_pt) >= 2:
236
+ continue
237
+ if len(end_pt) == 0:
238
+ continue
239
+ if len(end_pt) == 1:
240
+ sub_mesh = mesh.subgraph(list(cc)).copy()
241
+ pnodes = netx.periphery(sub_mesh)
242
+ ends = [*end_pt]
243
+ edge_id = global_mesh.nodes[(ends[0][0] + all_anchor[0], ends[0][1] + all_anchor[2], -ends[0][2])]['edge_id']
244
+ pnodes = sorted(pnodes,
245
+ key=lambda x: np.hypot((x[0] - ends[0][0]), (x[1] - ends[0][1])),
246
+ reverse=True)[0]
247
+ npath = [*netx.shortest_path(sub_mesh, (ends[0][0], ends[0][1]), pnodes, weight='length')]
248
+ for np_node in npath:
249
+ npath_map[np_node[0], np_node[1]] = edge_id
250
+ fpath = []
251
+ if global_mesh.nodes[(ends[0][0] + all_anchor[0], ends[0][1] + all_anchor[2], -ends[0][2])].get('far') is None:
252
+ print("None far")
253
+ import pdb; pdb.set_trace()
254
+ else:
255
+ fnodes = global_mesh.nodes[(ends[0][0] + all_anchor[0], ends[0][1] + all_anchor[2], -ends[0][2])].get('far')
256
+ fnodes = [(xx[0] - all_anchor[0], xx[1] - all_anchor[2], xx[2]) for xx in fnodes]
257
+ dmask = mask + 0
258
+ did = 0
259
+ while True:
260
+ did += 1
261
+ dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1)
262
+ if did > 3:
263
+ break
264
+ # ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0)]
265
+ ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0 and mask[fnode[0], fnode[1]] == 0)]
266
+ if len(ffnode) > 0:
267
+ fnode = ffnode[0]
268
+ break
269
+ if len(ffnode) == 0:
270
+ continue
271
+ fpath.append((fnode[0], fnode[1]))
272
+ for step in range(0, len(npath) - 1):
273
+ parr = (npath[step + 1][0] - npath[step][0], npath[step + 1][1] - npath[step][1])
274
+ new_loc = (fpath[-1][0] + parr[0], fpath[-1][1] + parr[1])
275
+ new_loc_nes = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]),
276
+ (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1)]\
277
+ if xx[0] >= 0 and xx[0] < fpath_map.shape[0] and xx[1] >= 0 and xx[1] < fpath_map.shape[1]]
278
+ if np.sum([fpath_map[nlne[0], nlne[1]] for nlne in new_loc_nes]) != -4:
279
+ break
280
+ if npath_map[new_loc[0], new_loc[1]] != -1:
281
+ if npath_map[new_loc[0], new_loc[1]] != edge_id:
282
+ break
283
+ else:
284
+ continue
285
+ if valid_area[new_loc[0], new_loc[1]] == 0:
286
+ break
287
+ new_loc_nes_eight = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]),
288
+ (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1),
289
+ (new_loc[0] + 1, new_loc[1] + 1), (new_loc[0] + 1, new_loc[1] - 1),
290
+ (new_loc[0] - 1, new_loc[1] - 1), (new_loc[0] - 1, new_loc[1] + 1)]\
291
+ if xx[0] >= 0 and xx[0] < fpath_map.shape[0] and xx[1] >= 0 and xx[1] < fpath_map.shape[1]]
292
+ if np.sum([int(npath_map[nlne[0], nlne[1]] == edge_id) for nlne in new_loc_nes_eight]) == 0:
293
+ break
294
+ fpath.append((fpath[-1][0] + parr[0], fpath[-1][1] + parr[1]))
295
+ if step != len(npath) - 2:
296
+ for xx in npath[step+1:]:
297
+ if npath_map[xx[0], xx[1]] == edge_id:
298
+ npath_map[xx[0], xx[1]] = -1
299
+ if len(fpath) > 0:
300
+ for fp_node in fpath:
301
+ fpath_map[fp_node[0], fp_node[1]] = edge_id
302
+ # import pdb; pdb.set_trace()
303
+ far_edge = (fpath_map > -1).astype(np.uint8)
304
+ update_edge = (npath_map > -1) * mask + edge
305
+ t_update_edge = torch.FloatTensor(update_edge).to(device)[None, None, ...]
306
+ depth_output = depth_feat_model.forward_3P(t_mask, t_context, t_depth_zero_mean_depth, t_update_edge, unit_length=128,
307
+ cuda=device)
308
+ depth_output = depth_output.cpu().data.numpy().squeeze()
309
+ depth_output = np.exp(depth_output + input_mean_depth) * mask # + input_depth * context
310
+ # if "right" in direc.lower() and "-" not in direc.lower():
311
+ # plt.imshow(depth_output); plt.show()
312
+ # import pdb; pdb.set_trace()
313
+ # f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True); ax1.imshow(depth_output); ax2.imshow(npath_map + fpath_map); plt.show()
314
+ for near_id in np.unique(npath_map[npath_map > -1]):
315
+ depth_output = refine_depth_around_edge(depth_output.copy(),
316
+ (fpath_map == near_id).astype(np.uint8) * mask, # far_edge_map_in_mask,
317
+ (fpath_map == near_id).astype(np.uint8), # far_edge_map,
318
+ (npath_map == near_id).astype(np.uint8) * mask,
319
+ mask.copy(),
320
+ np.zeros_like(mask),
321
+ config)
322
+ # if "right" in direc.lower() and "-" not in direc.lower():
323
+ # plt.imshow(depth_output); plt.show()
324
+ # import pdb; pdb.set_trace()
325
+ # f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True); ax1.imshow(depth_output); ax2.imshow(npath_map + fpath_map); plt.show()
326
+ rgb_output = rgb_feat_model.forward_3P(t_mask, t_context, t_rgb, t_update_edge, unit_length=128,
327
+ cuda=device)
328
+
329
+ # rgb_output = rgb_feat_model.forward_3P(t_mask, t_context, t_rgb, t_update_edge, unit_length=128, cuda=config['gpu_ids'])
330
+ if config.get('gray_image') is True:
331
+ rgb_output = rgb_output.mean(1, keepdim=True).repeat((1,3,1,1))
332
+ rgb_output = ((rgb_output.squeeze().data.cpu().permute(1,2,0).numpy() * mask[..., None] + input_rgb) * 255).astype(np.uint8)
333
+ 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]
334
+ depth[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]][mask > 0] = depth_output[mask > 0]
335
+ # nxs, nys = np.where(mask > -1)
336
+ # for nx, ny in zip(nxs, nys):
337
+ # info_on_pix[(nx, ny)][0]['color'] = rgb_output[]
338
+
339
+
340
+ nxs, nys = np.where((npath_map > -1))
341
+ for nx, ny in zip(nxs, nys):
342
+ n_id = npath_map[nx, ny]
343
+ four_nes = [xx for xx in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)]\
344
+ if 0 <= xx[0] < fpath_map.shape[0] and 0 <= xx[1] < fpath_map.shape[1]]
345
+ for nex, ney in four_nes:
346
+ if fpath_map[nex, ney] == n_id:
347
+ na, nb = (nx + all_anchor[0], ny + all_anchor[2], info_on_pix[(nx + all_anchor[0], ny + all_anchor[2])][0]['depth']), \
348
+ (nex + all_anchor[0], ney + all_anchor[2], info_on_pix[(nex + all_anchor[0], ney + all_anchor[2])][0]['depth'])
349
+ if global_mesh.has_edge(na, nb):
350
+ global_mesh.remove_edge(na, nb)
351
+ nxs, nys = np.where((fpath_map > -1))
352
+ for nx, ny in zip(nxs, nys):
353
+ n_id = fpath_map[nx, ny]
354
+ four_nes = [xx for xx in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)]\
355
+ if 0 <= xx[0] < npath_map.shape[0] and 0 <= xx[1] < npath_map.shape[1]]
356
+ for nex, ney in four_nes:
357
+ if npath_map[nex, ney] == n_id:
358
+ na, nb = (nx + all_anchor[0], ny + all_anchor[2], info_on_pix[(nx + all_anchor[0], ny + all_anchor[2])][0]['depth']), \
359
+ (nex + all_anchor[0], ney + all_anchor[2], info_on_pix[(nex + all_anchor[0], ney + all_anchor[2])][0]['depth'])
360
+ if global_mesh.has_edge(na, nb):
361
+ global_mesh.remove_edge(na, nb)
362
+ nxs, nys = np.where(mask > 0)
363
+ for x, y in zip(nxs, nys):
364
+ x = x + all_anchor[0]
365
+ y = y + all_anchor[2]
366
+ cur_node = (x, y, 0)
367
+ new_node = (x, y, -abs(depth[x, y]))
368
+ disp = 1. / -abs(depth[x, y])
369
+ mapping_dict = {cur_node: new_node}
370
+ info_on_pix, global_mesh = update_info(mapping_dict, info_on_pix, global_mesh)
371
+ global_mesh.nodes[new_node]['color'] = image[x, y]
372
+ global_mesh.nodes[new_node]['old_color'] = image[x, y]
373
+ global_mesh.nodes[new_node]['disp'] = disp
374
+ info_on_pix[(x, y)][0]['depth'] = -abs(depth[x, y])
375
+ info_on_pix[(x, y)][0]['disp'] = disp
376
+ info_on_pix[(x, y)][0]['color'] = image[x, y]
377
+
378
+
379
+ nxs, nys = np.where((npath_map > -1))
380
+ for nx, ny in zip(nxs, nys):
381
+ self_node = (nx + all_anchor[0], ny + all_anchor[2], info_on_pix[(nx + all_anchor[0], ny + all_anchor[2])][0]['depth'])
382
+ if global_mesh.has_node(self_node) is False:
383
+ break
384
+ n_id = int(round(npath_map[nx, ny]))
385
+ four_nes = [xx for xx in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)]\
386
+ if 0 <= xx[0] < fpath_map.shape[0] and 0 <= xx[1] < fpath_map.shape[1]]
387
+ for nex, ney in four_nes:
388
+ ne_node = (nex + all_anchor[0], ney + all_anchor[2], info_on_pix[(nex + all_anchor[0], ney + all_anchor[2])][0]['depth'])
389
+ if global_mesh.has_node(ne_node) is False:
390
+ continue
391
+ if fpath_map[nex, ney] == n_id:
392
+ if global_mesh.nodes[self_node].get('edge_id') is None:
393
+ global_mesh.nodes[self_node]['edge_id'] = n_id
394
+ edge_ccs[n_id].add(self_node)
395
+ info_on_pix[(self_node[0], self_node[1])][0]['edge_id'] = n_id
396
+ if global_mesh.has_edge(self_node, ne_node) is True:
397
+ global_mesh.remove_edge(self_node, ne_node)
398
+ if global_mesh.nodes[self_node].get('far') is None:
399
+ global_mesh.nodes[self_node]['far'] = []
400
+ global_mesh.nodes[self_node]['far'].append(ne_node)
401
+
402
+ global_fpath_map = np.zeros_like(other_edge_with_id) - 1
403
+ global_fpath_map[all_anchor[0]:all_anchor[1], all_anchor[2]:all_anchor[3]] = fpath_map
404
+ fpath_ids = np.unique(global_fpath_map)
405
+ fpath_ids = fpath_ids[1:] if fpath_ids.shape[0] > 0 and fpath_ids[0] == -1 else []
406
+ fpath_real_id_map = np.zeros_like(global_fpath_map) - 1
407
+ for fpath_id in fpath_ids:
408
+ fpath_real_id = np.unique(((global_fpath_map == fpath_id).astype(np.int) * (other_edge_with_id + 1)) - 1)
409
+ fpath_real_id = fpath_real_id[1:] if fpath_real_id.shape[0] > 0 and fpath_real_id[0] == -1 else []
410
+ fpath_real_id = fpath_real_id.astype(np.int)
411
+ fpath_real_id = np.bincount(fpath_real_id).argmax()
412
+ fpath_real_id_map[global_fpath_map == fpath_id] = fpath_real_id
413
+ nxs, nys = np.where((fpath_map > -1))
414
+ for nx, ny in zip(nxs, nys):
415
+ self_node = (nx + all_anchor[0], ny + all_anchor[2], info_on_pix[(nx + all_anchor[0], ny + all_anchor[2])][0]['depth'])
416
+ n_id = fpath_map[nx, ny]
417
+ four_nes = [xx for xx in [(nx + 1, ny), (nx - 1, ny), (nx, ny + 1), (nx, ny - 1)]\
418
+ if 0 <= xx[0] < npath_map.shape[0] and 0 <= xx[1] < npath_map.shape[1]]
419
+ for nex, ney in four_nes:
420
+ ne_node = (nex + all_anchor[0], ney + all_anchor[2], info_on_pix[(nex + all_anchor[0], ney + all_anchor[2])][0]['depth'])
421
+ if global_mesh.has_node(ne_node) is False:
422
+ continue
423
+ if npath_map[nex, ney] == n_id or global_mesh.nodes[ne_node].get('edge_id') == n_id:
424
+ if global_mesh.has_edge(self_node, ne_node) is True:
425
+ global_mesh.remove_edge(self_node, ne_node)
426
+ if global_mesh.nodes[self_node].get('near') is None:
427
+ global_mesh.nodes[self_node]['near'] = []
428
+ if global_mesh.nodes[self_node].get('edge_id') is None:
429
+ f_id = int(round(fpath_real_id_map[self_node[0], self_node[1]]))
430
+ global_mesh.nodes[self_node]['edge_id'] = f_id
431
+ info_on_pix[(self_node[0], self_node[1])][0]['edge_id'] = f_id
432
+ edge_ccs[f_id].add(self_node)
433
+ global_mesh.nodes[self_node]['near'].append(ne_node)
434
+
435
+ return info_on_pix, global_mesh, image, depth, edge_ccs
436
+ # for edge_cc in edge_ccs:
437
+ # for edge_node in edge_cc:
438
+ # edge_ccs
439
+ # context_ccs, mask_ccs, broken_mask_ccs, edge_ccs, erode_context_ccs, init_mask_connect, edge_maps, extend_context_ccs, extend_edge_ccs
440
+
441
+ def get_valid_size(imap):
442
+ x_max = np.where(imap.sum(1).squeeze() > 0)[0].max() + 1
443
+ x_min = np.where(imap.sum(1).squeeze() > 0)[0].min()
444
+ y_max = np.where(imap.sum(0).squeeze() > 0)[0].max() + 1
445
+ y_min = np.where(imap.sum(0).squeeze() > 0)[0].min()
446
+ size_dict = {'x_max':x_max, 'y_max':y_max, 'x_min':x_min, 'y_min':y_min}
447
+
448
+ return size_dict
449
+
450
+ def dilate_valid_size(isize_dict, imap, dilate=[0, 0]):
451
+ osize_dict = copy.deepcopy(isize_dict)
452
+ osize_dict['x_min'] = max(0, osize_dict['x_min'] - dilate[0])
453
+ osize_dict['x_max'] = min(imap.shape[0], osize_dict['x_max'] + dilate[0])
454
+ osize_dict['y_min'] = max(0, osize_dict['y_min'] - dilate[0])
455
+ osize_dict['y_max'] = min(imap.shape[1], osize_dict['y_max'] + dilate[1])
456
+
457
+ return osize_dict
458
+
459
+ def size_operation(size_a, size_b, operation):
460
+ assert operation == '+' or operation == '-', "Operation must be '+' (union) or '-' (exclude)"
461
+ osize = {}
462
+ if operation == '+':
463
+ osize['x_min'] = min(size_a['x_min'], size_b['x_min'])
464
+ osize['y_min'] = min(size_a['y_min'], size_b['y_min'])
465
+ osize['x_max'] = max(size_a['x_max'], size_b['x_max'])
466
+ osize['y_max'] = max(size_a['y_max'], size_b['y_max'])
467
+ assert operation != '-', "Operation '-' is undefined !"
468
+
469
+ return osize
470
+
471
+ def fill_dummy_bord(mesh, info_on_pix, image, depth, config):
472
+ context = np.zeros_like(depth).astype(np.uint8)
473
+ context[mesh.graph['hoffset']:mesh.graph['hoffset'] + mesh.graph['noext_H'],
474
+ mesh.graph['woffset']:mesh.graph['woffset'] + mesh.graph['noext_W']] = 1
475
+ mask = 1 - context
476
+ xs, ys = np.where(mask > 0)
477
+ depth = depth * context
478
+ image = image * context[..., None]
479
+ cur_depth = 0
480
+ cur_disp = 0
481
+ color = [0, 0, 0]
482
+ for x, y in zip(xs, ys):
483
+ cur_node = (x, y, cur_depth)
484
+ mesh.add_node(cur_node, color=color,
485
+ synthesis=False,
486
+ disp=cur_disp,
487
+ cc_id=set(),
488
+ ext_pixel=True)
489
+ info_on_pix[(x, y)] = [{'depth':cur_depth,
490
+ 'color':mesh.nodes[(x, y, cur_depth)]['color'],
491
+ 'synthesis':False,
492
+ 'disp':mesh.nodes[cur_node]['disp'],
493
+ 'ext_pixel':True}]
494
+ # for x, y in zip(xs, ys):
495
+ four_nes = [(xx, yy) for xx, yy in [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)] if\
496
+ 0 <= x < mesh.graph['H'] and 0 <= y < mesh.graph['W'] and info_on_pix.get((xx, yy)) is not None]
497
+ for ne in four_nes:
498
+ # if (ne[0] - x) + (ne[1] - y) == 1 and info_on_pix.get((ne[0], ne[1])) is not None:
499
+ mesh.add_edge(cur_node, (ne[0], ne[1], info_on_pix[(ne[0], ne[1])][0]['depth']))
500
+
501
+ return mesh, info_on_pix
502
+
503
+
504
+ def enlarge_border(mesh, info_on_pix, depth, image, config):
505
+ mesh.graph['hoffset'], mesh.graph['woffset'] = config['extrapolation_thickness'], config['extrapolation_thickness']
506
+ mesh.graph['bord_up'], mesh.graph['bord_left'], mesh.graph['bord_down'], mesh.graph['bord_right'] = \
507
+ 0, 0, mesh.graph['H'], mesh.graph['W']
508
+ # new_image = np.pad(image,
509
+ # pad_width=((config['extrapolation_thickness'], config['extrapolation_thickness']),
510
+ # (config['extrapolation_thickness'], config['extrapolation_thickness']), (0, 0)),
511
+ # mode='constant')
512
+ # new_depth = np.pad(depth,
513
+ # pad_width=((config['extrapolation_thickness'], config['extrapolation_thickness']),
514
+ # (config['extrapolation_thickness'], config['extrapolation_thickness'])),
515
+ # mode='constant')
516
+
517
+ return mesh, info_on_pix, depth, image
518
+
519
+ def fill_missing_node(mesh, info_on_pix, image, depth):
520
+ for x in range(mesh.graph['bord_up'], mesh.graph['bord_down']):
521
+ for y in range(mesh.graph['bord_left'], mesh.graph['bord_right']):
522
+ if info_on_pix.get((x, y)) is None:
523
+ print("fill missing node = ", x, y)
524
+ import pdb; pdb.set_trace()
525
+ re_depth, re_count = 0, 0
526
+ for ne in [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)]:
527
+ if info_on_pix.get(ne) is not None:
528
+ re_depth += info_on_pix[ne][0]['depth']
529
+ re_count += 1
530
+ if re_count == 0:
531
+ re_depth = -abs(depth[x, y])
532
+ else:
533
+ re_depth = re_depth / re_count
534
+ depth[x, y] = abs(re_depth)
535
+ info_on_pix[(x, y)] = [{'depth':re_depth,
536
+ 'color':image[x, y],
537
+ 'synthesis':False,
538
+ 'disp':1./re_depth}]
539
+ mesh.add_node((x, y, re_depth), color=image[x, y],
540
+ synthesis=False,
541
+ disp=1./re_depth,
542
+ cc_id=set())
543
+ return mesh, info_on_pix, depth
544
+
545
+
546
+
547
+ def refresh_bord_depth(mesh, info_on_pix, image, depth):
548
+ H, W = mesh.graph['H'], mesh.graph['W']
549
+ corner_nodes = [(mesh.graph['bord_up'], mesh.graph['bord_left']),
550
+ (mesh.graph['bord_up'], mesh.graph['bord_right'] - 1),
551
+ (mesh.graph['bord_down'] - 1, mesh.graph['bord_left']),
552
+ (mesh.graph['bord_down'] - 1, mesh.graph['bord_right'] - 1)]
553
+ # (0, W - 1), (H - 1, 0), (H - 1, W - 1)]
554
+ bord_nodes = []
555
+ bord_nodes += [(mesh.graph['bord_up'], xx) for xx in range(mesh.graph['bord_left'] + 1, mesh.graph['bord_right'] - 1)]
556
+ bord_nodes += [(mesh.graph['bord_down'] - 1, xx) for xx in range(mesh.graph['bord_left'] + 1, mesh.graph['bord_right'] - 1)]
557
+ bord_nodes += [(xx, mesh.graph['bord_left']) for xx in range(mesh.graph['bord_up'] + 1, mesh.graph['bord_down'] - 1)]
558
+ bord_nodes += [(xx, mesh.graph['bord_right'] - 1) for xx in range(mesh.graph['bord_up'] + 1, mesh.graph['bord_down'] - 1)]
559
+ for xy in bord_nodes:
560
+ tgt_loc = None
561
+ if xy[0] == mesh.graph['bord_up']:
562
+ tgt_loc = (xy[0] + 1, xy[1])# (1, xy[1])
563
+ elif xy[0] == mesh.graph['bord_down'] - 1:
564
+ tgt_loc = (xy[0] - 1, xy[1]) # (H - 2, xy[1])
565
+ elif xy[1] == mesh.graph['bord_left']:
566
+ tgt_loc = (xy[0], xy[1] + 1)
567
+ elif xy[1] == mesh.graph['bord_right'] - 1:
568
+ tgt_loc = (xy[0], xy[1] - 1)
569
+ if tgt_loc is not None:
570
+ ne_infos = info_on_pix.get(tgt_loc)
571
+ if ne_infos is None:
572
+ import pdb; pdb.set_trace()
573
+ # if ne_infos is not None and len(ne_infos) == 1:
574
+ tgt_depth = ne_infos[0]['depth']
575
+ tgt_disp = ne_infos[0]['disp']
576
+ new_node = (xy[0], xy[1], tgt_depth)
577
+ src_node = (tgt_loc[0], tgt_loc[1], tgt_depth)
578
+ tgt_nes_loc = [(xx[0], xx[1]) \
579
+ for xx in mesh.neighbors(src_node)]
580
+ tgt_nes_loc = [(xx[0] - tgt_loc[0] + xy[0], xx[1] - tgt_loc[1] + xy[1]) for xx in tgt_nes_loc \
581
+ if abs(xx[0] - xy[0]) == 1 and abs(xx[1] - xy[1]) == 1]
582
+ tgt_nes_loc = [xx for xx in tgt_nes_loc if info_on_pix.get(xx) is not None]
583
+ tgt_nes_loc.append(tgt_loc)
584
+ # if (xy[0], xy[1]) == (559, 60):
585
+ # import pdb; pdb.set_trace()
586
+ if info_on_pix.get(xy) is not None and len(info_on_pix.get(xy)) > 0:
587
+ old_depth = info_on_pix[xy][0].get('depth')
588
+ old_node = (xy[0], xy[1], old_depth)
589
+ mesh.remove_edges_from([(old_ne, old_node) for old_ne in mesh.neighbors(old_node)])
590
+ mesh.add_edges_from([((zz[0], zz[1], info_on_pix[zz][0]['depth']), old_node) for zz in tgt_nes_loc])
591
+ mapping_dict = {old_node: new_node}
592
+ # if old_node[2] == new_node[2]:
593
+ # print("mapping_dict = ", mapping_dict)
594
+ info_on_pix, mesh = update_info(mapping_dict, info_on_pix, mesh)
595
+ else:
596
+ info_on_pix[xy] = []
597
+ info_on_pix[xy][0] = info_on_pix[tgt_loc][0]
598
+ info_on_pix['color'] = image[xy[0], xy[1]]
599
+ info_on_pix['old_color'] = image[xy[0], xy[1]]
600
+ mesh.add_node(new_node)
601
+ mesh.add_edges_from([((zz[0], zz[1], info_on_pix[zz][0]['depth']), new_node) for zz in tgt_nes_loc])
602
+ mesh.nodes[new_node]['far'] = None
603
+ mesh.nodes[new_node]['near'] = None
604
+ if mesh.nodes[src_node].get('far') is not None:
605
+ redundant_nodes = [ne for ne in mesh.nodes[src_node]['far'] if (ne[0], ne[1]) == xy]
606
+ [mesh.nodes[src_node]['far'].remove(aa) for aa in redundant_nodes]
607
+ if mesh.nodes[src_node].get('near') is not None:
608
+ redundant_nodes = [ne for ne in mesh.nodes[src_node]['near'] if (ne[0], ne[1]) == xy]
609
+ [mesh.nodes[src_node]['near'].remove(aa) for aa in redundant_nodes]
610
+ for xy in corner_nodes:
611
+ hx, hy = xy
612
+ four_nes = [xx for xx in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] if \
613
+ mesh.graph['bord_up'] <= xx[0] < mesh.graph['bord_down'] and \
614
+ mesh.graph['bord_left'] <= xx[1] < mesh.graph['bord_right']]
615
+ ne_nodes = []
616
+ ne_depths = []
617
+ for ne_loc in four_nes:
618
+ if info_on_pix.get(ne_loc) is not None:
619
+ ne_depths.append(info_on_pix[ne_loc][0]['depth'])
620
+ ne_nodes.append((ne_loc[0], ne_loc[1], info_on_pix[ne_loc][0]['depth']))
621
+ new_node = (xy[0], xy[1], float(np.mean(ne_depths)))
622
+ if info_on_pix.get(xy) is not None and len(info_on_pix.get(xy)) > 0:
623
+ old_depth = info_on_pix[xy][0].get('depth')
624
+ old_node = (xy[0], xy[1], old_depth)
625
+ mesh.remove_edges_from([(old_ne, old_node) for old_ne in mesh.neighbors(old_node)])
626
+ mesh.add_edges_from([(zz, old_node) for zz in ne_nodes])
627
+ mapping_dict = {old_node: new_node}
628
+ info_on_pix, mesh = update_info(mapping_dict, info_on_pix, mesh)
629
+ else:
630
+ info_on_pix[xy] = []
631
+ info_on_pix[xy][0] = info_on_pix[ne_loc[-1]][0]
632
+ info_on_pix['color'] = image[xy[0], xy[1]]
633
+ info_on_pix['old_color'] = image[xy[0], xy[1]]
634
+ mesh.add_node(new_node)
635
+ mesh.add_edges_from([(zz, new_node) for zz in ne_nodes])
636
+ mesh.nodes[new_node]['far'] = None
637
+ mesh.nodes[new_node]['near'] = None
638
+ for xy in bord_nodes + corner_nodes:
639
+ # if (xy[0], xy[1]) == (559, 60):
640
+ # import pdb; pdb.set_trace()
641
+ depth[xy[0], xy[1]] = abs(info_on_pix[xy][0]['depth'])
642
+ for xy in bord_nodes:
643
+ cur_node = (xy[0], xy[1], info_on_pix[xy][0]['depth'])
644
+ nes = mesh.neighbors(cur_node)
645
+ four_nes = set([(xy[0] + 1, xy[1]), (xy[0] - 1, xy[1]), (xy[0], xy[1] + 1), (xy[0], xy[1] - 1)]) - \
646
+ set([(ne[0], ne[1]) for ne in nes])
647
+ four_nes = [ne for ne in four_nes if mesh.graph['bord_up'] <= ne[0] < mesh.graph['bord_down'] and \
648
+ mesh.graph['bord_left'] <= ne[1] < mesh.graph['bord_right']]
649
+ four_nes = [(ne[0], ne[1], info_on_pix[(ne[0], ne[1])][0]['depth']) for ne in four_nes]
650
+ mesh.nodes[cur_node]['far'] = []
651
+ mesh.nodes[cur_node]['near'] = []
652
+ for ne in four_nes:
653
+ if abs(ne[2]) >= abs(cur_node[2]):
654
+ mesh.nodes[cur_node]['far'].append(ne)
655
+ else:
656
+ mesh.nodes[cur_node]['near'].append(ne)
657
+
658
+ return mesh, info_on_pix, depth
659
+
660
+ def get_union_size(mesh, dilate, *alls_cc):
661
+ all_cc = reduce(lambda x, y: x | y, [set()] + [*alls_cc])
662
+ min_x, min_y, max_x, max_y = mesh.graph['H'], mesh.graph['W'], 0, 0
663
+ H, W = mesh.graph['H'], mesh.graph['W']
664
+ for node in all_cc:
665
+ if node[0] < min_x:
666
+ min_x = node[0]
667
+ if node[0] > max_x:
668
+ max_x = node[0]
669
+ if node[1] < min_y:
670
+ min_y = node[1]
671
+ if node[1] > max_y:
672
+ max_y = node[1]
673
+ max_x = max_x + 1
674
+ max_y = max_y + 1
675
+ # mask_size = dilate_valid_size(mask_size, edge_dict['mask'], dilate=[20, 20])
676
+ osize_dict = dict()
677
+ osize_dict['x_min'] = max(0, min_x - dilate[0])
678
+ osize_dict['x_max'] = min(H, max_x + dilate[0])
679
+ osize_dict['y_min'] = max(0, min_y - dilate[1])
680
+ osize_dict['y_max'] = min(W, max_y + dilate[1])
681
+
682
+ return osize_dict
683
+
684
+ def incomplete_node(mesh, edge_maps, info_on_pix):
685
+ vis_map = np.zeros((mesh.graph['H'], mesh.graph['W']))
686
+
687
+ for node in mesh.nodes:
688
+ if mesh.nodes[node].get('synthesis') is not True:
689
+ connect_all_flag = False
690
+ nes = [xx for xx in mesh.neighbors(node) if mesh.nodes[xx].get('synthesis') is not True]
691
+ if len(nes) < 3 and 0 < node[0] < mesh.graph['H'] - 1 and 0 < node[1] < mesh.graph['W'] - 1:
692
+ if len(nes) <= 1:
693
+ connect_all_flag = True
694
+ else:
695
+ dan_ne_node_a = nes[0]
696
+ dan_ne_node_b = nes[1]
697
+ if abs(dan_ne_node_a[0] - dan_ne_node_b[0]) > 1 or \
698
+ abs(dan_ne_node_a[1] - dan_ne_node_b[1]) > 1:
699
+ connect_all_flag = True
700
+ if connect_all_flag == True:
701
+ vis_map[node[0], node[1]] = len(nes)
702
+ four_nes = [(node[0] - 1, node[1]), (node[0] + 1, node[1]), (node[0], node[1] - 1), (node[0], node[1] + 1)]
703
+ for ne in four_nes:
704
+ for info in info_on_pix[(ne[0], ne[1])]:
705
+ ne_node = (ne[0], ne[1], info['depth'])
706
+ if info.get('synthesis') is not True and mesh.has_node(ne_node):
707
+ mesh.add_edge(node, ne_node)
708
+ break
709
+
710
+ return mesh
711
+
712
+ def edge_inpainting(edge_id, context_cc, erode_context_cc, mask_cc, edge_cc, extend_edge_cc,
713
+ mesh, edge_map, edge_maps_with_id, config, union_size, depth_edge_model, inpaint_iter):
714
+ edge_dict = get_edge_from_nodes(context_cc, erode_context_cc, mask_cc, edge_cc, extend_edge_cc,
715
+ mesh.graph['H'], mesh.graph['W'], mesh)
716
+ edge_dict['edge'], end_depth_maps, _ = \
717
+ filter_irrelevant_edge_new(edge_dict['self_edge'] + edge_dict['comp_edge'],
718
+ edge_map,
719
+ edge_maps_with_id,
720
+ edge_id,
721
+ edge_dict['context'],
722
+ edge_dict['depth'], mesh, context_cc | erode_context_cc, spdb=True)
723
+ patch_edge_dict = dict()
724
+ patch_edge_dict['mask'], patch_edge_dict['context'], patch_edge_dict['rgb'], \
725
+ patch_edge_dict['disp'], patch_edge_dict['edge'] = \
726
+ crop_maps_by_size(union_size, edge_dict['mask'], edge_dict['context'],
727
+ edge_dict['rgb'], edge_dict['disp'], edge_dict['edge'])
728
+ tensor_edge_dict = convert2tensor(patch_edge_dict)
729
+ if require_depth_edge(patch_edge_dict['edge'], patch_edge_dict['mask']) and inpaint_iter == 0:
730
+ with torch.no_grad():
731
+ device = config["gpu_ids"] if isinstance(config["gpu_ids"], int) and config["gpu_ids"] >= 0 else "cpu"
732
+ depth_edge_output = depth_edge_model.forward_3P(tensor_edge_dict['mask'],
733
+ tensor_edge_dict['context'],
734
+ tensor_edge_dict['rgb'],
735
+ tensor_edge_dict['disp'],
736
+ tensor_edge_dict['edge'],
737
+ unit_length=128,
738
+ cuda=device)
739
+ depth_edge_output = depth_edge_output.cpu()
740
+ tensor_edge_dict['output'] = (depth_edge_output > config['ext_edge_threshold']).float() * tensor_edge_dict['mask'] + tensor_edge_dict['edge']
741
+ else:
742
+ tensor_edge_dict['output'] = tensor_edge_dict['edge']
743
+ depth_edge_output = tensor_edge_dict['edge'] + 0
744
+ patch_edge_dict['output'] = tensor_edge_dict['output'].squeeze().data.cpu().numpy()
745
+ edge_dict['output'] = np.zeros((mesh.graph['H'], mesh.graph['W']))
746
+ edge_dict['output'][union_size['x_min']:union_size['x_max'], union_size['y_min']:union_size['y_max']] = \
747
+ patch_edge_dict['output']
748
+
749
+ return edge_dict, end_depth_maps
750
+
751
+ 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):
752
+ if given_depth_dict is False:
753
+ 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'])
754
+ if edge_output is not None:
755
+ depth_dict['edge'] = edge_output
756
+ else:
757
+ depth_dict = given_depth_dict
758
+ patch_depth_dict = dict()
759
+ patch_depth_dict['mask'], patch_depth_dict['context'], patch_depth_dict['depth'], \
760
+ patch_depth_dict['zero_mean_depth'], patch_depth_dict['edge'] = \
761
+ crop_maps_by_size(union_size, depth_dict['mask'], depth_dict['context'],
762
+ depth_dict['real_depth'], depth_dict['zero_mean_depth'], depth_dict['edge'])
763
+ tensor_depth_dict = convert2tensor(patch_depth_dict)
764
+ resize_mask = open_small_mask(tensor_depth_dict['mask'], tensor_depth_dict['context'], 3, 41)
765
+ with torch.no_grad():
766
+ device = config["gpu_ids"] if isinstance(config["gpu_ids"], int) and config["gpu_ids"] >= 0 else "cpu"
767
+ depth_output = depth_feat_model.forward_3P(resize_mask,
768
+ tensor_depth_dict['context'],
769
+ tensor_depth_dict['zero_mean_depth'],
770
+ tensor_depth_dict['edge'],
771
+ unit_length=128,
772
+ cuda=device)
773
+ depth_output = depth_output.cpu()
774
+ tensor_depth_dict['output'] = torch.exp(depth_output + depth_dict['mean_depth']) * \
775
+ tensor_depth_dict['mask'] + tensor_depth_dict['depth']
776
+ patch_depth_dict['output'] = tensor_depth_dict['output'].data.cpu().numpy().squeeze()
777
+ depth_dict['output'] = np.zeros((mesh.graph['H'], mesh.graph['W']))
778
+ depth_dict['output'][union_size['x_min']:union_size['x_max'], union_size['y_min']:union_size['y_max']] = \
779
+ patch_depth_dict['output']
780
+ depth_output = depth_dict['output'] * depth_dict['mask'] + depth_dict['depth'] * depth_dict['context']
781
+ depth_output = smooth_cntsyn_gap(depth_dict['output'].copy() * depth_dict['mask'] + depth_dict['depth'] * depth_dict['context'],
782
+ depth_dict['mask'], depth_dict['context'],
783
+ init_mask_region=depth_dict['mask'])
784
+ if spdb is True:
785
+ f, ((ax1, ax2)) = plt.subplots(1, 2, sharex=True, sharey=True);
786
+ ax1.imshow(depth_output * depth_dict['mask'] + depth_dict['depth']); ax2.imshow(depth_dict['output'] * depth_dict['mask'] + depth_dict['depth']); plt.show()
787
+ import pdb; pdb.set_trace()
788
+ depth_dict['output'] = depth_output * depth_dict['mask'] + depth_dict['depth'] * depth_dict['context']
789
+
790
+ return depth_dict
791
+
792
+ def update_info(mapping_dict, info_on_pix, *meshes):
793
+ rt_meshes = []
794
+ for mesh in meshes:
795
+ rt_meshes.append(relabel_node(mesh, mesh.nodes, [*mapping_dict.keys()][0], [*mapping_dict.values()][0]))
796
+ x, y, _ = [*mapping_dict.keys()][0]
797
+ info_on_pix[(x, y)][0]['depth'] = [*mapping_dict.values()][0][2]
798
+
799
+ return [info_on_pix] + rt_meshes
800
+
801
+ def build_connection(mesh, cur_node, dst_node):
802
+ if (abs(cur_node[0] - dst_node[0]) + abs(cur_node[1] - dst_node[1])) < 2:
803
+ mesh.add_edge(cur_node, dst_node)
804
+ if abs(cur_node[0] - dst_node[0]) > 1 or abs(cur_node[1] - dst_node[1]) > 1:
805
+ return mesh
806
+ ne_nodes = [*mesh.neighbors(cur_node)].copy()
807
+ for ne_node in ne_nodes:
808
+ if mesh.has_edge(ne_node, dst_node) or ne_node == dst_node:
809
+ continue
810
+ else:
811
+ mesh = build_connection(mesh, ne_node, dst_node)
812
+
813
+ return mesh
814
+
815
+ def recursive_add_edge(edge_mesh, mesh, info_on_pix, cur_node, mark):
816
+ ne_nodes = [(x[0], x[1]) for x in edge_mesh.neighbors(cur_node)]
817
+ for node_xy in ne_nodes:
818
+ node = (node_xy[0], node_xy[1], info_on_pix[node_xy][0]['depth'])
819
+ if mark[node[0], node[1]] != 3:
820
+ continue
821
+ else:
822
+ mark[node[0], node[1]] = 0
823
+ mesh.remove_edges_from([(xx, node) for xx in mesh.neighbors(node)])
824
+ mesh = build_connection(mesh, cur_node, node)
825
+ re_info = dict(depth=0, count=0)
826
+ for re_ne in mesh.neighbors(node):
827
+ re_info['depth'] += re_ne[2]
828
+ re_info['count'] += 1.
829
+ try:
830
+ re_depth = re_info['depth'] / re_info['count']
831
+ except:
832
+ re_depth = node[2]
833
+ re_node = (node_xy[0], node_xy[1], re_depth)
834
+ mapping_dict = {node: re_node}
835
+ info_on_pix, edge_mesh, mesh = update_info(mapping_dict, info_on_pix, edge_mesh, mesh)
836
+
837
+ edge_mesh, mesh, mark, info_on_pix = recursive_add_edge(edge_mesh, mesh, info_on_pix, re_node, mark)
838
+
839
+ return edge_mesh, mesh, mark, info_on_pix
840
+
841
+ def resize_for_edge(tensor_dict, largest_size):
842
+ resize_dict = {k: v.clone() for k, v in tensor_dict.items()}
843
+ frac = largest_size / np.array([*resize_dict['edge'].shape[-2:]]).max()
844
+ if frac < 1:
845
+ resize_mark = torch.nn.functional.interpolate(torch.cat((resize_dict['mask'],
846
+ resize_dict['context']),
847
+ dim=1),
848
+ scale_factor=frac,
849
+ mode='bilinear')
850
+ resize_dict['mask'] = (resize_mark[:, 0:1] > 0).float()
851
+ resize_dict['context'] = (resize_mark[:, 1:2] == 1).float()
852
+ resize_dict['context'][resize_dict['mask'] > 0] = 0
853
+ resize_dict['edge'] = torch.nn.functional.interpolate(resize_dict['edge'],
854
+ scale_factor=frac,
855
+ mode='bilinear')
856
+ resize_dict['edge'] = (resize_dict['edge'] > 0).float()
857
+ resize_dict['edge'] = resize_dict['edge'] * resize_dict['context']
858
+ resize_dict['disp'] = torch.nn.functional.interpolate(resize_dict['disp'],
859
+ scale_factor=frac,
860
+ mode='nearest')
861
+ resize_dict['disp'] = resize_dict['disp'] * resize_dict['context']
862
+ resize_dict['rgb'] = torch.nn.functional.interpolate(resize_dict['rgb'],
863
+ scale_factor=frac,
864
+ mode='bilinear')
865
+ resize_dict['rgb'] = resize_dict['rgb'] * resize_dict['context']
866
+ return resize_dict
867
+
868
+ def get_map_from_nodes(nodes, height, width):
869
+ omap = np.zeros((height, width))
870
+ for n in nodes:
871
+ omap[n[0], n[1]] = 1
872
+
873
+ return omap
874
+
875
+ def get_map_from_ccs(ccs, height, width, condition_input=None, condition=None, real_id=False, id_shift=0):
876
+ if condition is None:
877
+ condition = lambda x, condition_input: True
878
+
879
+ if real_id is True:
880
+ omap = np.zeros((height, width)) + (-1) + id_shift
881
+ else:
882
+ omap = np.zeros((height, width))
883
+ for cc_id, cc in enumerate(ccs):
884
+ for n in cc:
885
+ if condition(n, condition_input):
886
+ if real_id is True:
887
+ omap[n[0], n[1]] = cc_id + id_shift
888
+ else:
889
+ omap[n[0], n[1]] = 1
890
+ return omap
891
+
892
+ def revise_map_by_nodes(nodes, imap, operation, limit_constr=None):
893
+ assert operation == '+' or operation == '-', "Operation must be '+' (union) or '-' (exclude)"
894
+ omap = copy.deepcopy(imap)
895
+ revise_flag = True
896
+ if operation == '+':
897
+ for n in nodes:
898
+ omap[n[0], n[1]] = 1
899
+ if limit_constr is not None and omap.sum() > limit_constr:
900
+ omap = imap
901
+ revise_flag = False
902
+ elif operation == '-':
903
+ for n in nodes:
904
+ omap[n[0], n[1]] = 0
905
+ if limit_constr is not None and omap.sum() < limit_constr:
906
+ omap = imap
907
+ revise_flag = False
908
+
909
+ return omap, revise_flag
910
+
911
+ def repaint_info(mesh, cc, x_anchor, y_anchor, source_type):
912
+ if source_type == 'rgb':
913
+ feat = np.zeros((3, x_anchor[1] - x_anchor[0], y_anchor[1] - y_anchor[0]))
914
+ else:
915
+ feat = np.zeros((1, x_anchor[1] - x_anchor[0], y_anchor[1] - y_anchor[0]))
916
+ for node in cc:
917
+ if source_type == 'rgb':
918
+ feat[:, node[0] - x_anchor[0], node[1] - y_anchor[0]] = np.array(mesh.nodes[node]['color']) / 255.
919
+ elif source_type == 'd':
920
+ feat[:, node[0] - x_anchor[0], node[1] - y_anchor[0]] = abs(node[2])
921
+
922
+ return feat
923
+
924
+ def get_context_from_nodes(mesh, cc, H, W, source_type=''):
925
+ if 'rgb' in source_type or 'color' in source_type:
926
+ feat = np.zeros((H, W, 3))
927
+ else:
928
+ feat = np.zeros((H, W))
929
+ context = np.zeros((H, W))
930
+ for node in cc:
931
+ if 'rgb' in source_type or 'color' in source_type:
932
+ feat[node[0], node[1]] = np.array(mesh.nodes[node]['color']) / 255.
933
+ context[node[0], node[1]] = 1
934
+ else:
935
+ feat[node[0], node[1]] = abs(node[2])
936
+
937
+ return feat, context
938
+
939
+ def get_mask_from_nodes(mesh, cc, H, W):
940
+ mask = np.zeros((H, W))
941
+ for node in cc:
942
+ mask[node[0], node[1]] = abs(node[2])
943
+
944
+ return mask
945
+
946
+
947
+ def get_edge_from_nodes(context_cc, erode_context_cc, mask_cc, edge_cc, extend_edge_cc, H, W, mesh):
948
+ context = np.zeros((H, W))
949
+ mask = np.zeros((H, W))
950
+ rgb = np.zeros((H, W, 3))
951
+ disp = np.zeros((H, W))
952
+ depth = np.zeros((H, W))
953
+ real_depth = np.zeros((H, W))
954
+ edge = np.zeros((H, W))
955
+ comp_edge = np.zeros((H, W))
956
+ fpath_map = np.zeros((H, W)) - 1
957
+ npath_map = np.zeros((H, W)) - 1
958
+ near_depth = np.zeros((H, W))
959
+ for node in context_cc:
960
+ rgb[node[0], node[1]] = np.array(mesh.nodes[node]['color'])
961
+ disp[node[0], node[1]] = mesh.nodes[node]['disp']
962
+ depth[node[0], node[1]] = node[2]
963
+ context[node[0], node[1]] = 1
964
+ for node in erode_context_cc:
965
+ rgb[node[0], node[1]] = np.array(mesh.nodes[node]['color'])
966
+ disp[node[0], node[1]] = mesh.nodes[node]['disp']
967
+ depth[node[0], node[1]] = node[2]
968
+ context[node[0], node[1]] = 1
969
+ rgb = rgb / 255.
970
+ disp = np.abs(disp)
971
+ disp = disp / disp.max()
972
+ real_depth = depth.copy()
973
+ for node in context_cc:
974
+ if mesh.nodes[node].get('real_depth') is not None:
975
+ real_depth[node[0], node[1]] = mesh.nodes[node]['real_depth']
976
+ for node in erode_context_cc:
977
+ if mesh.nodes[node].get('real_depth') is not None:
978
+ real_depth[node[0], node[1]] = mesh.nodes[node]['real_depth']
979
+ for node in mask_cc:
980
+ mask[node[0], node[1]] = 1
981
+ near_depth[node[0], node[1]] = node[2]
982
+ for node in edge_cc:
983
+ edge[node[0], node[1]] = 1
984
+ for node in extend_edge_cc:
985
+ comp_edge[node[0], node[1]] = 1
986
+ rt_dict = {'rgb': rgb, 'disp': disp, 'depth': depth, 'real_depth': real_depth, 'self_edge': edge, 'context': context,
987
+ 'mask': mask, 'fpath_map': fpath_map, 'npath_map': npath_map, 'comp_edge': comp_edge, 'valid_area': context + mask,
988
+ 'near_depth': near_depth}
989
+
990
+ return rt_dict
991
+
992
+ def get_depth_from_maps(context_map, mask_map, depth_map, H, W, log_depth=False):
993
+ context = context_map.astype(np.uint8)
994
+ mask = mask_map.astype(np.uint8).copy()
995
+ depth = np.abs(depth_map)
996
+ real_depth = depth.copy()
997
+ zero_mean_depth = np.zeros((H, W))
998
+
999
+ if log_depth is True:
1000
+ log_depth = np.log(real_depth + 1e-8) * context
1001
+ mean_depth = np.mean(log_depth[context > 0])
1002
+ zero_mean_depth = (log_depth - mean_depth) * context
1003
+ else:
1004
+ zero_mean_depth = real_depth
1005
+ mean_depth = 0
1006
+ edge = np.zeros_like(depth)
1007
+
1008
+ rt_dict = {'depth': depth, 'real_depth': real_depth, 'context': context, 'mask': mask,
1009
+ 'mean_depth': mean_depth, 'zero_mean_depth': zero_mean_depth, 'edge': edge}
1010
+
1011
+ return rt_dict
1012
+
1013
+ def get_depth_from_nodes(context_cc, erode_context_cc, mask_cc, H, W, mesh, log_depth=False):
1014
+ context = np.zeros((H, W))
1015
+ mask = np.zeros((H, W))
1016
+ depth = np.zeros((H, W))
1017
+ real_depth = np.zeros((H, W))
1018
+ zero_mean_depth = np.zeros((H, W))
1019
+ for node in context_cc:
1020
+ depth[node[0], node[1]] = node[2]
1021
+ context[node[0], node[1]] = 1
1022
+ for node in erode_context_cc:
1023
+ depth[node[0], node[1]] = node[2]
1024
+ context[node[0], node[1]] = 1
1025
+ depth = np.abs(depth)
1026
+ real_depth = depth.copy()
1027
+ for node in context_cc:
1028
+ if mesh.nodes[node].get('real_depth') is not None:
1029
+ real_depth[node[0], node[1]] = mesh.nodes[node]['real_depth']
1030
+ for node in erode_context_cc:
1031
+ if mesh.nodes[node].get('real_depth') is not None:
1032
+ real_depth[node[0], node[1]] = mesh.nodes[node]['real_depth']
1033
+ real_depth = np.abs(real_depth)
1034
+ for node in mask_cc:
1035
+ mask[node[0], node[1]] = 1
1036
+ if log_depth is True:
1037
+ log_depth = np.log(real_depth + 1e-8) * context
1038
+ mean_depth = np.mean(log_depth[context > 0])
1039
+ zero_mean_depth = (log_depth - mean_depth) * context
1040
+ else:
1041
+ zero_mean_depth = real_depth
1042
+ mean_depth = 0
1043
+
1044
+ rt_dict = {'depth': depth, 'real_depth': real_depth, 'context': context, 'mask': mask,
1045
+ 'mean_depth': mean_depth, 'zero_mean_depth': zero_mean_depth}
1046
+
1047
+ return rt_dict
1048
+
1049
+ def get_rgb_from_nodes(context_cc, erode_context_cc, mask_cc, H, W, mesh):
1050
+ context = np.zeros((H, W))
1051
+ mask = np.zeros((H, W))
1052
+ rgb = np.zeros((H, W, 3))
1053
+ erode_context = np.zeros((H, W))
1054
+ for node in context_cc:
1055
+ rgb[node[0], node[1]] = np.array(mesh.nodes[node]['color'])
1056
+ context[node[0], node[1]] = 1
1057
+ rgb = rgb / 255.
1058
+ for node in mask_cc:
1059
+ mask[node[0], node[1]] = 1
1060
+ for node in erode_context_cc:
1061
+ erode_context[node[0], node[1]] = 1
1062
+ mask[node[0], node[1]] = 1
1063
+ rt_dict = {'rgb': rgb, 'context': context, 'mask': mask,
1064
+ 'erode': erode_context}
1065
+
1066
+ return rt_dict
1067
+
1068
+ def crop_maps_by_size(size, *imaps):
1069
+ omaps = []
1070
+ for imap in imaps:
1071
+ omaps.append(imap[size['x_min']:size['x_max'], size['y_min']:size['y_max']].copy())
1072
+
1073
+ return omaps
1074
+
1075
+ def convert2tensor(input_dict):
1076
+ rt_dict = {}
1077
+ for key, value in input_dict.items():
1078
+ if 'rgb' in key or 'color' in key:
1079
+ rt_dict[key] = torch.FloatTensor(value).permute(2, 0, 1)[None, ...]
1080
+ else:
1081
+ rt_dict[key] = torch.FloatTensor(value)[None, None, ...]
1082
+
1083
+ return rt_dict
moon.jpg ADDED
networks.py ADDED
@@ -0,0 +1,501 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import numpy as np
4
+ import matplotlib.pyplot as plt
5
+ import torch.nn.functional as F
6
+
7
+
8
+ class BaseNetwork(nn.Module):
9
+ def __init__(self):
10
+ super(BaseNetwork, self).__init__()
11
+
12
+ def init_weights(self, init_type='normal', gain=0.02):
13
+ '''
14
+ initialize network's weights
15
+ init_type: normal | xavier | kaiming | orthogonal
16
+ https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39
17
+ '''
18
+
19
+ def init_func(m):
20
+ classname = m.__class__.__name__
21
+ if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
22
+ if init_type == 'normal':
23
+ nn.init.normal_(m.weight.data, 0.0, gain)
24
+ elif init_type == 'xavier':
25
+ nn.init.xavier_normal_(m.weight.data, gain=gain)
26
+ elif init_type == 'kaiming':
27
+ nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
28
+ elif init_type == 'orthogonal':
29
+ nn.init.orthogonal_(m.weight.data, gain=gain)
30
+
31
+ if hasattr(m, 'bias') and m.bias is not None:
32
+ nn.init.constant_(m.bias.data, 0.0)
33
+
34
+ elif classname.find('BatchNorm2d') != -1:
35
+ nn.init.normal_(m.weight.data, 1.0, gain)
36
+ nn.init.constant_(m.bias.data, 0.0)
37
+
38
+ self.apply(init_func)
39
+
40
+ def weights_init(init_type='gaussian'):
41
+ def init_fun(m):
42
+ classname = m.__class__.__name__
43
+ if (classname.find('Conv') == 0 or classname.find(
44
+ 'Linear') == 0) and hasattr(m, 'weight'):
45
+ if init_type == 'gaussian':
46
+ nn.init.normal_(m.weight, 0.0, 0.02)
47
+ elif init_type == 'xavier':
48
+ nn.init.xavier_normal_(m.weight, gain=math.sqrt(2))
49
+ elif init_type == 'kaiming':
50
+ nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
51
+ elif init_type == 'orthogonal':
52
+ nn.init.orthogonal_(m.weight, gain=math.sqrt(2))
53
+ elif init_type == 'default':
54
+ pass
55
+ else:
56
+ assert 0, "Unsupported initialization: {}".format(init_type)
57
+ if hasattr(m, 'bias') and m.bias is not None:
58
+ nn.init.constant_(m.bias, 0.0)
59
+
60
+ return init_fun
61
+
62
+ class PartialConv(nn.Module):
63
+ def __init__(self, in_channels, out_channels, kernel_size, stride=1,
64
+ padding=0, dilation=1, groups=1, bias=True):
65
+ super().__init__()
66
+ self.input_conv = nn.Conv2d(in_channels, out_channels, kernel_size,
67
+ stride, padding, dilation, groups, bias)
68
+ self.mask_conv = nn.Conv2d(in_channels, out_channels, kernel_size,
69
+ stride, padding, dilation, groups, False)
70
+ self.input_conv.apply(weights_init('kaiming'))
71
+ self.slide_winsize = in_channels * kernel_size * kernel_size
72
+
73
+ torch.nn.init.constant_(self.mask_conv.weight, 1.0)
74
+
75
+ # mask is not updated
76
+ for param in self.mask_conv.parameters():
77
+ param.requires_grad = False
78
+
79
+ def forward(self, input, mask):
80
+ # http://masc.cs.gmu.edu/wiki/partialconv
81
+ # C(X) = W^T * X + b, C(0) = b, D(M) = 1 * M + 0 = sum(M)
82
+ # W^T* (M .* X) / sum(M) + b = [C(M .* X) – C(0)] / D(M) + C(0)
83
+ output = self.input_conv(input * mask)
84
+ if self.input_conv.bias is not None:
85
+ output_bias = self.input_conv.bias.view(1, -1, 1, 1).expand_as(
86
+ output)
87
+ else:
88
+ output_bias = torch.zeros_like(output)
89
+
90
+ with torch.no_grad():
91
+ output_mask = self.mask_conv(mask)
92
+
93
+ no_update_holes = output_mask == 0
94
+
95
+ mask_sum = output_mask.masked_fill_(no_update_holes, 1.0)
96
+
97
+ output_pre = ((output - output_bias) * self.slide_winsize) / mask_sum + output_bias
98
+ output = output_pre.masked_fill_(no_update_holes, 0.0)
99
+
100
+ new_mask = torch.ones_like(output)
101
+ new_mask = new_mask.masked_fill_(no_update_holes, 0.0)
102
+
103
+ return output, new_mask
104
+
105
+
106
+ class PCBActiv(nn.Module):
107
+ def __init__(self, in_ch, out_ch, bn=True, sample='none-3', activ='relu',
108
+ conv_bias=False):
109
+ super().__init__()
110
+ if sample == 'down-5':
111
+ self.conv = PartialConv(in_ch, out_ch, 5, 2, 2, bias=conv_bias)
112
+ elif sample == 'down-7':
113
+ self.conv = PartialConv(in_ch, out_ch, 7, 2, 3, bias=conv_bias)
114
+ elif sample == 'down-3':
115
+ self.conv = PartialConv(in_ch, out_ch, 3, 2, 1, bias=conv_bias)
116
+ else:
117
+ self.conv = PartialConv(in_ch, out_ch, 3, 1, 1, bias=conv_bias)
118
+
119
+ if bn:
120
+ self.bn = nn.BatchNorm2d(out_ch)
121
+ if activ == 'relu':
122
+ self.activation = nn.ReLU()
123
+ elif activ == 'leaky':
124
+ self.activation = nn.LeakyReLU(negative_slope=0.2)
125
+
126
+ def forward(self, input, input_mask):
127
+ h, h_mask = self.conv(input, input_mask)
128
+ if hasattr(self, 'bn'):
129
+ h = self.bn(h)
130
+ if hasattr(self, 'activation'):
131
+ h = self.activation(h)
132
+ return h, h_mask
133
+
134
+ class Inpaint_Depth_Net(nn.Module):
135
+ def __init__(self, layer_size=7, upsampling_mode='nearest'):
136
+ super().__init__()
137
+ in_channels = 4
138
+ out_channels = 1
139
+ self.freeze_enc_bn = False
140
+ self.upsampling_mode = upsampling_mode
141
+ self.layer_size = layer_size
142
+ self.enc_1 = PCBActiv(in_channels, 64, bn=False, sample='down-7', conv_bias=True)
143
+ self.enc_2 = PCBActiv(64, 128, sample='down-5', conv_bias=True)
144
+ self.enc_3 = PCBActiv(128, 256, sample='down-5')
145
+ self.enc_4 = PCBActiv(256, 512, sample='down-3')
146
+ for i in range(4, self.layer_size):
147
+ name = 'enc_{:d}'.format(i + 1)
148
+ setattr(self, name, PCBActiv(512, 512, sample='down-3'))
149
+
150
+ for i in range(4, self.layer_size):
151
+ name = 'dec_{:d}'.format(i + 1)
152
+ setattr(self, name, PCBActiv(512 + 512, 512, activ='leaky'))
153
+ self.dec_4 = PCBActiv(512 + 256, 256, activ='leaky')
154
+ self.dec_3 = PCBActiv(256 + 128, 128, activ='leaky')
155
+ self.dec_2 = PCBActiv(128 + 64, 64, activ='leaky')
156
+ self.dec_1 = PCBActiv(64 + in_channels, out_channels,
157
+ bn=False, activ=None, conv_bias=True)
158
+ def add_border(self, input, mask_flag, PCONV=True):
159
+ with torch.no_grad():
160
+ h = input.shape[-2]
161
+ w = input.shape[-1]
162
+ require_len_unit = 2 ** self.layer_size
163
+ residual_h = int(np.ceil(h / float(require_len_unit)) * require_len_unit - h) # + 2*require_len_unit
164
+ residual_w = int(np.ceil(w / float(require_len_unit)) * require_len_unit - w) # + 2*require_len_unit
165
+ enlarge_input = torch.zeros((input.shape[0], input.shape[1], h + residual_h, w + residual_w)).to(input.device)
166
+ if mask_flag:
167
+ if PCONV is False:
168
+ enlarge_input += 1.0
169
+ enlarge_input = enlarge_input.clamp(0.0, 1.0)
170
+ else:
171
+ enlarge_input[:, 2, ...] = 0.0
172
+ anchor_h = residual_h//2
173
+ anchor_w = residual_w//2
174
+ enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
175
+
176
+ return enlarge_input, [anchor_h, anchor_h+h, anchor_w, anchor_w+w]
177
+
178
+ def forward_3P(self, mask, context, depth, edge, unit_length=128, cuda=None):
179
+ with torch.no_grad():
180
+ input = torch.cat((depth, edge, context, mask), dim=1)
181
+ n, c, h, w = input.shape
182
+ residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h)
183
+ residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w)
184
+ anchor_h = residual_h//2
185
+ anchor_w = residual_w//2
186
+ enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
187
+ enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
188
+ # enlarge_input[:, 3] = 1. - enlarge_input[:, 3]
189
+ depth_output = self.forward(enlarge_input)
190
+ depth_output = depth_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]
191
+ # import pdb; pdb.set_trace()
192
+
193
+ return depth_output
194
+
195
+ def forward(self, input_feat, refine_border=False, sample=False, PCONV=True):
196
+ input = input_feat
197
+ input_mask = (input_feat[:, -2:-1] + input_feat[:, -1:]).clamp(0, 1).repeat(1, input.shape[1], 1, 1)
198
+
199
+ vis_input = input.cpu().data.numpy()
200
+ vis_input_mask = input_mask.cpu().data.numpy()
201
+ H, W = input.shape[-2:]
202
+ if refine_border is True:
203
+ input, anchor = self.add_border(input, mask_flag=False)
204
+ input_mask, anchor = self.add_border(input_mask, mask_flag=True, PCONV=PCONV)
205
+ h_dict = {} # for the output of enc_N
206
+ h_mask_dict = {} # for the output of enc_N
207
+ h_dict['h_0'], h_mask_dict['h_0'] = input, input_mask
208
+
209
+ h_key_prev = 'h_0'
210
+ for i in range(1, self.layer_size + 1):
211
+ l_key = 'enc_{:d}'.format(i)
212
+ h_key = 'h_{:d}'.format(i)
213
+ h_dict[h_key], h_mask_dict[h_key] = getattr(self, l_key)(
214
+ h_dict[h_key_prev], h_mask_dict[h_key_prev])
215
+ h_key_prev = h_key
216
+
217
+ h_key = 'h_{:d}'.format(self.layer_size)
218
+ h, h_mask = h_dict[h_key], h_mask_dict[h_key]
219
+
220
+ for i in range(self.layer_size, 0, -1):
221
+ enc_h_key = 'h_{:d}'.format(i - 1)
222
+ dec_l_key = 'dec_{:d}'.format(i)
223
+
224
+ h = F.interpolate(h, scale_factor=2, mode=self.upsampling_mode)
225
+ h_mask = F.interpolate(h_mask, scale_factor=2, mode='nearest')
226
+
227
+ h = torch.cat([h, h_dict[enc_h_key]], dim=1)
228
+ h_mask = torch.cat([h_mask, h_mask_dict[enc_h_key]], dim=1)
229
+ h, h_mask = getattr(self, dec_l_key)(h, h_mask)
230
+ output = h
231
+ if refine_border is True:
232
+ h_mask = h_mask[..., anchor[0]:anchor[1], anchor[2]:anchor[3]]
233
+ output = output[..., anchor[0]:anchor[1], anchor[2]:anchor[3]]
234
+
235
+ return output
236
+
237
+ class Inpaint_Edge_Net(BaseNetwork):
238
+ def __init__(self, residual_blocks=8, init_weights=True):
239
+ super(Inpaint_Edge_Net, self).__init__()
240
+ in_channels = 7
241
+ out_channels = 1
242
+ self.encoder = []
243
+ # 0
244
+ self.encoder_0 = nn.Sequential(
245
+ nn.ReflectionPad2d(3),
246
+ spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=7, padding=0), True),
247
+ nn.InstanceNorm2d(64, track_running_stats=False),
248
+ nn.ReLU(True))
249
+ # 1
250
+ self.encoder_1 = nn.Sequential(
251
+ spectral_norm(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding=1), True),
252
+ nn.InstanceNorm2d(128, track_running_stats=False),
253
+ nn.ReLU(True))
254
+ # 2
255
+ self.encoder_2 = nn.Sequential(
256
+ spectral_norm(nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1), True),
257
+ nn.InstanceNorm2d(256, track_running_stats=False),
258
+ nn.ReLU(True))
259
+ # 3
260
+ blocks = []
261
+ for _ in range(residual_blocks):
262
+ block = ResnetBlock(256, 2)
263
+ blocks.append(block)
264
+
265
+ self.middle = nn.Sequential(*blocks)
266
+ # + 3
267
+ self.decoder_0 = nn.Sequential(
268
+ spectral_norm(nn.ConvTranspose2d(in_channels=256+256, out_channels=128, kernel_size=4, stride=2, padding=1), True),
269
+ nn.InstanceNorm2d(128, track_running_stats=False),
270
+ nn.ReLU(True))
271
+ # + 2
272
+ self.decoder_1 = nn.Sequential(
273
+ spectral_norm(nn.ConvTranspose2d(in_channels=128+128, out_channels=64, kernel_size=4, stride=2, padding=1), True),
274
+ nn.InstanceNorm2d(64, track_running_stats=False),
275
+ nn.ReLU(True))
276
+ # + 1
277
+ self.decoder_2 = nn.Sequential(
278
+ nn.ReflectionPad2d(3),
279
+ nn.Conv2d(in_channels=64+64, out_channels=out_channels, kernel_size=7, padding=0),
280
+ )
281
+
282
+ if init_weights:
283
+ self.init_weights()
284
+
285
+ def add_border(self, input, channel_pad_1=None):
286
+ h = input.shape[-2]
287
+ w = input.shape[-1]
288
+ require_len_unit = 16
289
+ residual_h = int(np.ceil(h / float(require_len_unit)) * require_len_unit - h) # + 2*require_len_unit
290
+ residual_w = int(np.ceil(w / float(require_len_unit)) * require_len_unit - w) # + 2*require_len_unit
291
+ enlarge_input = torch.zeros((input.shape[0], input.shape[1], h + residual_h, w + residual_w)).to(input.device)
292
+ if channel_pad_1 is not None:
293
+ for channel in channel_pad_1:
294
+ enlarge_input[:, channel] = 1
295
+ anchor_h = residual_h//2
296
+ anchor_w = residual_w//2
297
+ enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
298
+
299
+ return enlarge_input, [anchor_h, anchor_h+h, anchor_w, anchor_w+w]
300
+
301
+ def forward_3P(self, mask, context, rgb, disp, edge, unit_length=128, cuda=None):
302
+ with torch.no_grad():
303
+ input = torch.cat((rgb, disp/disp.max(), edge, context, mask), dim=1)
304
+ n, c, h, w = input.shape
305
+ residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h)
306
+ residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w)
307
+ anchor_h = residual_h//2
308
+ anchor_w = residual_w//2
309
+ enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
310
+ enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
311
+ edge_output = self.forward(enlarge_input)
312
+ edge_output = edge_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]
313
+
314
+ return edge_output
315
+
316
+ def forward(self, x, refine_border=False):
317
+ if refine_border:
318
+ x, anchor = self.add_border(x, [5])
319
+ x1 = self.encoder_0(x)
320
+ x2 = self.encoder_1(x1)
321
+ x3 = self.encoder_2(x2)
322
+ x4 = self.middle(x3)
323
+ x5 = self.decoder_0(torch.cat((x4, x3), dim=1))
324
+ x6 = self.decoder_1(torch.cat((x5, x2), dim=1))
325
+ x7 = self.decoder_2(torch.cat((x6, x1), dim=1))
326
+ x = torch.sigmoid(x7)
327
+ if refine_border:
328
+ x = x[..., anchor[0]:anchor[1], anchor[2]:anchor[3]]
329
+
330
+ return x
331
+
332
+ class Inpaint_Color_Net(nn.Module):
333
+ def __init__(self, layer_size=7, upsampling_mode='nearest', add_hole_mask=False, add_two_layer=False, add_border=False):
334
+ super().__init__()
335
+ self.freeze_enc_bn = False
336
+ self.upsampling_mode = upsampling_mode
337
+ self.layer_size = layer_size
338
+ in_channels = 6
339
+ self.enc_1 = PCBActiv(in_channels, 64, bn=False, sample='down-7')
340
+ self.enc_2 = PCBActiv(64, 128, sample='down-5')
341
+ self.enc_3 = PCBActiv(128, 256, sample='down-5')
342
+ self.enc_4 = PCBActiv(256, 512, sample='down-3')
343
+ self.enc_5 = PCBActiv(512, 512, sample='down-3')
344
+ self.enc_6 = PCBActiv(512, 512, sample='down-3')
345
+ self.enc_7 = PCBActiv(512, 512, sample='down-3')
346
+
347
+ self.dec_7 = PCBActiv(512+512, 512, activ='leaky')
348
+ self.dec_6 = PCBActiv(512+512, 512, activ='leaky')
349
+
350
+ self.dec_5A = PCBActiv(512 + 512, 512, activ='leaky')
351
+ self.dec_4A = PCBActiv(512 + 256, 256, activ='leaky')
352
+ self.dec_3A = PCBActiv(256 + 128, 128, activ='leaky')
353
+ self.dec_2A = PCBActiv(128 + 64, 64, activ='leaky')
354
+ self.dec_1A = PCBActiv(64 + in_channels, 3, bn=False, activ=None, conv_bias=True)
355
+ '''
356
+ self.dec_5B = PCBActiv(512 + 512, 512, activ='leaky')
357
+ self.dec_4B = PCBActiv(512 + 256, 256, activ='leaky')
358
+ self.dec_3B = PCBActiv(256 + 128, 128, activ='leaky')
359
+ self.dec_2B = PCBActiv(128 + 64, 64, activ='leaky')
360
+ self.dec_1B = PCBActiv(64 + 4, 1, bn=False, activ=None, conv_bias=True)
361
+ '''
362
+ def cat(self, A, B):
363
+ return torch.cat((A, B), dim=1)
364
+
365
+ def upsample(self, feat, mask):
366
+ feat = F.interpolate(feat, scale_factor=2, mode=self.upsampling_mode)
367
+ mask = F.interpolate(mask, scale_factor=2, mode='nearest')
368
+
369
+ return feat, mask
370
+
371
+ def forward_3P(self, mask, context, rgb, edge, unit_length=128, cuda=None):
372
+ with torch.no_grad():
373
+ input = torch.cat((rgb, edge, context, mask), dim=1)
374
+ n, c, h, w = input.shape
375
+ residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h) # + 128
376
+ residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w) # + 256
377
+ anchor_h = residual_h//2
378
+ anchor_w = residual_w//2
379
+ enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
380
+ enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
381
+ # enlarge_input[:, 3] = 1. - enlarge_input[:, 3]
382
+ enlarge_input = enlarge_input.to(cuda)
383
+ rgb_output = self.forward(enlarge_input)
384
+ rgb_output = rgb_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]
385
+
386
+ return rgb_output
387
+
388
+ def forward(self, input, add_border=False):
389
+ input_mask = (input[:, -2:-1] + input[:, -1:]).clamp(0, 1)
390
+ H, W = input.shape[-2:]
391
+ f_0, h_0 = input, input_mask.repeat((1,input.shape[1],1,1))
392
+ f_1, h_1 = self.enc_1(f_0, h_0)
393
+ f_2, h_2 = self.enc_2(f_1, h_1)
394
+ f_3, h_3 = self.enc_3(f_2, h_2)
395
+ f_4, h_4 = self.enc_4(f_3, h_3)
396
+ f_5, h_5 = self.enc_5(f_4, h_4)
397
+ f_6, h_6 = self.enc_6(f_5, h_5)
398
+ f_7, h_7 = self.enc_7(f_6, h_6)
399
+
400
+ o_7, k_7 = self.upsample(f_7, h_7)
401
+ o_6, k_6 = self.dec_7(self.cat(o_7, f_6), self.cat(k_7, h_6))
402
+ o_6, k_6 = self.upsample(o_6, k_6)
403
+ o_5, k_5 = self.dec_6(self.cat(o_6, f_5), self.cat(k_6, h_5))
404
+ o_5, k_5 = self.upsample(o_5, k_5)
405
+ o_5A, k_5A = o_5, k_5
406
+ o_5B, k_5B = o_5, k_5
407
+ ###############
408
+ o_4A, k_4A = self.dec_5A(self.cat(o_5A, f_4), self.cat(k_5A, h_4))
409
+ o_4A, k_4A = self.upsample(o_4A, k_4A)
410
+ o_3A, k_3A = self.dec_4A(self.cat(o_4A, f_3), self.cat(k_4A, h_3))
411
+ o_3A, k_3A = self.upsample(o_3A, k_3A)
412
+ o_2A, k_2A = self.dec_3A(self.cat(o_3A, f_2), self.cat(k_3A, h_2))
413
+ o_2A, k_2A = self.upsample(o_2A, k_2A)
414
+ o_1A, k_1A = self.dec_2A(self.cat(o_2A, f_1), self.cat(k_2A, h_1))
415
+ o_1A, k_1A = self.upsample(o_1A, k_1A)
416
+ o_0A, k_0A = self.dec_1A(self.cat(o_1A, f_0), self.cat(k_1A, h_0))
417
+
418
+ return torch.sigmoid(o_0A)
419
+
420
+ def train(self, mode=True):
421
+ """
422
+ Override the default train() to freeze the BN parameters
423
+ """
424
+ super().train(mode)
425
+ if self.freeze_enc_bn:
426
+ for name, module in self.named_modules():
427
+ if isinstance(module, nn.BatchNorm2d) and 'enc' in name:
428
+ module.eval()
429
+
430
+ class Discriminator(BaseNetwork):
431
+ def __init__(self, use_sigmoid=True, use_spectral_norm=True, init_weights=True, in_channels=None):
432
+ super(Discriminator, self).__init__()
433
+ self.use_sigmoid = use_sigmoid
434
+ self.conv1 = self.features = nn.Sequential(
435
+ spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),
436
+ nn.LeakyReLU(0.2, inplace=True),
437
+ )
438
+
439
+ self.conv2 = nn.Sequential(
440
+ spectral_norm(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),
441
+ nn.LeakyReLU(0.2, inplace=True),
442
+ )
443
+
444
+ self.conv3 = nn.Sequential(
445
+ spectral_norm(nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),
446
+ nn.LeakyReLU(0.2, inplace=True),
447
+ )
448
+
449
+ self.conv4 = nn.Sequential(
450
+ spectral_norm(nn.Conv2d(in_channels=256, out_channels=512, kernel_size=4, stride=1, padding=1, bias=not use_spectral_norm), use_spectral_norm),
451
+ nn.LeakyReLU(0.2, inplace=True),
452
+ )
453
+
454
+ self.conv5 = nn.Sequential(
455
+ spectral_norm(nn.Conv2d(in_channels=512, out_channels=1, kernel_size=4, stride=1, padding=1, bias=not use_spectral_norm), use_spectral_norm),
456
+ )
457
+
458
+ if init_weights:
459
+ self.init_weights()
460
+
461
+ def forward(self, x):
462
+ conv1 = self.conv1(x)
463
+ conv2 = self.conv2(conv1)
464
+ conv3 = self.conv3(conv2)
465
+ conv4 = self.conv4(conv3)
466
+ conv5 = self.conv5(conv4)
467
+
468
+ outputs = conv5
469
+ if self.use_sigmoid:
470
+ outputs = torch.sigmoid(conv5)
471
+
472
+ return outputs, [conv1, conv2, conv3, conv4, conv5]
473
+
474
+ class ResnetBlock(nn.Module):
475
+ def __init__(self, dim, dilation=1):
476
+ super(ResnetBlock, self).__init__()
477
+ self.conv_block = nn.Sequential(
478
+ nn.ReflectionPad2d(dilation),
479
+ spectral_norm(nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3, padding=0, dilation=dilation, bias=not True), True),
480
+ nn.InstanceNorm2d(dim, track_running_stats=False),
481
+ nn.LeakyReLU(negative_slope=0.2),
482
+
483
+ nn.ReflectionPad2d(1),
484
+ spectral_norm(nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3, padding=0, dilation=1, bias=not True), True),
485
+ nn.InstanceNorm2d(dim, track_running_stats=False),
486
+ )
487
+
488
+ def forward(self, x):
489
+ out = x + self.conv_block(x)
490
+
491
+ # Remove ReLU at the end of the residual block
492
+ # http://torch.ch/blog/2016/02/04/resnets.html
493
+
494
+ return out
495
+
496
+
497
+ def spectral_norm(module, mode=True):
498
+ if mode:
499
+ return nn.utils.spectral_norm(module)
500
+
501
+ return module
packages.txt ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
1
+ xvfb
2
+
3
+ libglfw3-dev
4
+ libgles2-mesa-dev
5
+
6
+ libxkbcommon-dev
7
+ libxkbcommon-x11-dev
8
+ libdbus-1-dev
pyproject.toml ADDED
@@ -0,0 +1,3 @@
 
 
 
1
+ [tool.black]
2
+ line-length = 120
3
+ target-version = ['py37']
requirements.txt ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch
2
+ torchvision
3
+ gradio
4
+ scikit-image
5
+ opencv-python==4.2.0.32
6
+ vispy
7
+ moviepy==1.0.2
8
+ transforms3d==0.3.1
9
+ networkx==2.3
10
+ pyyaml==5.4.1
11
+ PyQt5==5.13.2
12
+ pyvirtualdisplay
setup.py ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
1
+ from setuptools import setup
2
+
3
+ setup(
4
+ name='cynetworkx_workaround',
5
+ version='1.0',
6
+ description='A useful module',
7
+ install_requires=['cynetworkx'], #external packages as dependencies
8
+ )
utils.py ADDED
@@ -0,0 +1,1416 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import glob
3
+ import cv2
4
+ import scipy.misc as misc
5
+ from skimage.transform import resize
6
+ import numpy as np
7
+ from functools import reduce
8
+ from operator import mul
9
+ import torch
10
+ from torch import nn
11
+ import matplotlib.pyplot as plt
12
+ import re
13
+ try:
14
+ import cynetworkx as netx
15
+ except ImportError:
16
+ import networkx as netx
17
+ from scipy.ndimage import gaussian_filter
18
+ from skimage.feature import canny
19
+ import collections
20
+ import shutil
21
+ import imageio
22
+ import copy
23
+ from matplotlib import pyplot as plt
24
+ from mpl_toolkits.mplot3d import Axes3D
25
+ import time
26
+ from scipy.interpolate import interp1d
27
+ from collections import namedtuple
28
+
29
+ def path_planning(num_frames, x, y, z, path_type=''):
30
+ if path_type == 'straight-line':
31
+ corner_points = np.array([[0, 0, 0], [(0 + x) * 0.5, (0 + y) * 0.5, (0 + z) * 0.5], [x, y, z]])
32
+ corner_t = np.linspace(0, 1, len(corner_points))
33
+ t = np.linspace(0, 1, num_frames)
34
+ cs = interp1d(corner_t, corner_points, axis=0, kind='quadratic')
35
+ spline = cs(t)
36
+ xs, ys, zs = [xx.squeeze() for xx in np.split(spline, 3, 1)]
37
+ elif path_type == 'double-straight-line':
38
+ corner_points = np.array([[-x, -y, -z], [0, 0, 0], [x, y, z]])
39
+ corner_t = np.linspace(0, 1, len(corner_points))
40
+ t = np.linspace(0, 1, num_frames)
41
+ cs = interp1d(corner_t, corner_points, axis=0, kind='quadratic')
42
+ spline = cs(t)
43
+ xs, ys, zs = [xx.squeeze() for xx in np.split(spline, 3, 1)]
44
+ elif path_type == 'circle':
45
+ xs, ys, zs = [], [], []
46
+ for frame_id, bs_shift_val in enumerate(np.arange(-2.0, 2.0, (4./num_frames))):
47
+ xs += [np.cos(bs_shift_val * np.pi) * 1 * x]
48
+ ys += [np.sin(bs_shift_val * np.pi) * 1 * y]
49
+ zs += [np.cos(bs_shift_val * np.pi/2.) * 1 * z]
50
+ xs, ys, zs = np.array(xs), np.array(ys), np.array(zs)
51
+
52
+ return xs, ys, zs
53
+
54
+ def open_small_mask(mask, context, open_iteration, kernel):
55
+ np_mask = mask.cpu().data.numpy().squeeze().astype(np.uint8)
56
+ raw_mask = np_mask.copy()
57
+ np_context = context.cpu().data.numpy().squeeze().astype(np.uint8)
58
+ np_input = np_mask + np_context
59
+ for _ in range(open_iteration):
60
+ np_input = cv2.erode(cv2.dilate(np_input, np.ones((kernel, kernel)), iterations=1), np.ones((kernel,kernel)), iterations=1)
61
+ np_mask[(np_input - np_context) > 0] = 1
62
+ out_mask = torch.FloatTensor(np_mask).to(mask)[None, None, ...]
63
+
64
+ return out_mask
65
+
66
+ 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):
67
+ other_edges = other_edges.squeeze().astype(np.uint8)
68
+ other_edges_with_id = other_edges_with_id.squeeze()
69
+ self_edge = self_edge.squeeze()
70
+ 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)
71
+ 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)
72
+ edge_ids = np.unique(other_edges_with_id * context + (-1) * (1 - context)).astype(np.int)
73
+ end_depth_maps = np.zeros_like(self_edge)
74
+ self_edge_ids = np.sort(np.unique(other_edges_with_id[self_edge > 0]).astype(np.int))
75
+ self_edge_ids = self_edge_ids[1:] if self_edge_ids.shape[0] > 0 and self_edge_ids[0] == -1 else self_edge_ids
76
+ self_comp_ids = np.sort(np.unique(other_edges_with_id[comp_edge > 0]).astype(np.int))
77
+ self_comp_ids = self_comp_ids[1:] if self_comp_ids.shape[0] > 0 and self_comp_ids[0] == -1 else self_comp_ids
78
+ edge_ids = edge_ids[1:] if edge_ids[0] == -1 else edge_ids
79
+ other_edges_info = []
80
+ extend_other_edges = np.zeros_like(other_edges)
81
+ if spdb is True:
82
+ 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()
83
+ import pdb; pdb.set_trace()
84
+ filter_self_edge = np.zeros_like(self_edge)
85
+ for self_edge_id in self_edge_ids:
86
+ filter_self_edge[other_edges_with_id == self_edge_id] = 1
87
+ dilate_self_comp_edge = cv2.dilate(comp_edge, kernel=np.ones((3, 3)), iterations=2)
88
+ valid_self_comp_edge = np.zeros_like(comp_edge)
89
+ for self_comp_id in self_comp_ids:
90
+ valid_self_comp_edge[self_comp_id == other_edges_with_id] = 1
91
+ self_comp_edge = dilate_self_comp_edge * valid_self_comp_edge
92
+ filter_self_edge = (filter_self_edge + self_comp_edge).clip(0, 1)
93
+ for edge_id in edge_ids:
94
+ other_edge_locs = (other_edges_with_id == edge_id).astype(np.uint8)
95
+ condition = (other_edge_locs * other_edges * context.astype(np.uint8))
96
+ end_cross_point = dilate_cross_self_edge * condition * (1 - filter_self_edge)
97
+ end_bevel_point = dilate_bevel_self_edge * condition * (1 - filter_self_edge)
98
+ if end_bevel_point.max() != 0:
99
+ end_depth_maps[end_bevel_point != 0] = depth[end_bevel_point != 0]
100
+ if end_cross_point.max() == 0:
101
+ nxs, nys = np.where(end_bevel_point != 0)
102
+ for nx, ny in zip(nxs, nys):
103
+ bevel_node = [xx for xx in context_cc if xx[0] == nx and xx[1] == ny][0]
104
+ for ne in mesh.neighbors(bevel_node):
105
+ if other_edges_with_id[ne[0], ne[1]] > -1 and dilate_cross_self_edge[ne[0], ne[1]] > 0:
106
+ extend_other_edges[ne[0], ne[1]] = 1
107
+ break
108
+ else:
109
+ other_edges[other_edges_with_id == edge_id] = 0
110
+ other_edges = (other_edges + extend_other_edges).clip(0, 1) * context
111
+
112
+ return other_edges, end_depth_maps, other_edges_info
113
+
114
+ def clean_far_edge_new(input_edge, end_depth_maps, mask, context, global_mesh, info_on_pix, self_edge, inpaint_id, config):
115
+ mesh = netx.Graph()
116
+ hxs, hys = np.where(input_edge * mask > 0)
117
+ valid_near_edge = (input_edge != 0).astype(np.uint8) * context
118
+ valid_map = mask + context
119
+ invalid_edge_ids = []
120
+ for hx, hy in zip(hxs, hys):
121
+ node = (hx ,hy)
122
+ mesh.add_node((hx, hy))
123
+ eight_nes = [ne for ne in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1), \
124
+ (hx + 1, hy + 1), (hx - 1, hy - 1), (hx - 1, hy + 1), (hx + 1, hy - 1)]\
125
+ 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]
126
+ for ne in eight_nes:
127
+ mesh.add_edge(node, ne, length=np.hypot(ne[0] - hx, ne[1] - hy))
128
+ if end_depth_maps[ne[0], ne[1]] != 0:
129
+ mesh.nodes[ne[0], ne[1]]['cnt'] = True
130
+ if end_depth_maps[ne[0], ne[1]] == 0:
131
+ import pdb; pdb.set_trace()
132
+ mesh.nodes[ne[0], ne[1]]['depth'] = end_depth_maps[ne[0], ne[1]]
133
+ elif mask[ne[0], ne[1]] != 1:
134
+ 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)]\
135
+ if nne[0] < end_depth_maps.shape[0] and nne[0] >= 0 and nne[1] < end_depth_maps.shape[1] and nne[1] >= 0]
136
+ for nne in four_nes:
137
+ if end_depth_maps[nne[0], nne[1]] != 0:
138
+ mesh.add_edge(nne, ne, length=np.hypot(nne[0] - ne[0], nne[1] - ne[1]))
139
+ mesh.nodes[nne[0], nne[1]]['cnt'] = True
140
+ mesh.nodes[nne[0], nne[1]]['depth'] = end_depth_maps[nne[0], nne[1]]
141
+ ccs = [*netx.connected_components(mesh)]
142
+ end_pts = []
143
+ for cc in ccs:
144
+ end_pts.append(set())
145
+ for node in cc:
146
+ if mesh.nodes[node].get('cnt') is not None:
147
+ end_pts[-1].add((node[0], node[1], mesh.nodes[node]['depth']))
148
+ predef_npaths = [None for _ in range(len(ccs))]
149
+ fpath_map = np.zeros_like(input_edge) - 1
150
+ npath_map = np.zeros_like(input_edge) - 1
151
+ npaths, fpaths = dict(), dict()
152
+ break_flag = False
153
+ end_idx = 0
154
+ while end_idx < len(end_pts):
155
+ end_pt, cc = [*zip(end_pts, ccs)][end_idx]
156
+ end_idx += 1
157
+ sorted_end_pt = []
158
+ fpath = []
159
+ iter_fpath = []
160
+ if len(end_pt) > 2 or len(end_pt) == 0:
161
+ if len(end_pt) > 2:
162
+ continue
163
+ continue
164
+ if len(end_pt) == 2:
165
+ ravel_end = [*end_pt]
166
+ tmp_sub_mesh = mesh.subgraph(list(cc)).copy()
167
+ 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')]
168
+ 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)
169
+ 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)
170
+ tmp_disp_diff = [disp_diff1, disp_diff2]
171
+ self_end = []
172
+ edge_len = []
173
+ ds_edge = cv2.dilate(self_edge.astype(np.uint8), np.ones((3, 3)), iterations=1)
174
+ if ds_edge[ravel_end[0][0], ravel_end[0][1]] > 0:
175
+ self_end.append(1)
176
+ else:
177
+ self_end.append(0)
178
+ if ds_edge[ravel_end[1][0], ravel_end[1][1]] > 0:
179
+ self_end.append(1)
180
+ else:
181
+ self_end.append(0)
182
+ edge_len = [np.count_nonzero(npath_map1), np.count_nonzero(npath_map2)]
183
+ 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)]
184
+ re_npath_map1, re_fpath_map1 = (npath_map1 != -1).astype(np.uint8), (fpath_map1 != -1).astype(np.uint8)
185
+ re_npath_map2, re_fpath_map2 = (npath_map2 != -1).astype(np.uint8), (fpath_map2 != -1).astype(np.uint8)
186
+ if np.count_nonzero(re_npath_map1 * re_npath_map2 * mask) / \
187
+ (np.count_nonzero((re_npath_map1 + re_npath_map2) * mask) + 1e-6) > 0.5\
188
+ and np.count_nonzero(re_fpath_map1 * re_fpath_map2 * mask) / \
189
+ (np.count_nonzero((re_fpath_map1 + re_fpath_map2) * mask) + 1e-6) > 0.5\
190
+ and tmp_disp_diff[0] != -1 and tmp_disp_diff[1] != -1:
191
+ my_fpath_map, my_npath_map, npath, fpath = \
192
+ plan_path_e2e(mesh, cc, sorted_end_pts, global_mesh, input_edge, mask, valid_map, inpaint_id, npath_map=None, fpath_map=None)
193
+ npath_map[my_npath_map != -1] = my_npath_map[my_npath_map != -1]
194
+ fpath_map[my_fpath_map != -1] = my_fpath_map[my_fpath_map != -1]
195
+ if len(fpath) > 0:
196
+ edge_id = global_mesh.nodes[[*sorted_end_pts][0]]['edge_id']
197
+ fpaths[edge_id] = fpath
198
+ npaths[edge_id] = npath
199
+ invalid_edge_ids.append(edge_id)
200
+ else:
201
+ if tmp_disp_diff[0] != -1:
202
+ ratio_a = tmp_disp_diff[0] / (np.sum(tmp_disp_diff) + 1e-8)
203
+ else:
204
+ ratio_a = 0
205
+ if tmp_disp_diff[1] != -1:
206
+ ratio_b = tmp_disp_diff[1] / (np.sum(tmp_disp_diff) + 1e-8)
207
+ else:
208
+ ratio_b = 0
209
+ npath_len = len(tmp_npath)
210
+ if npath_len > config['depth_edge_dilate_2'] * 2:
211
+ npath_len = npath_len - (config['depth_edge_dilate_2'] * 1)
212
+ tmp_npath_a = tmp_npath[:int(np.floor(npath_len * ratio_a))]
213
+ tmp_npath_b = tmp_npath[::-1][:int(np.floor(npath_len * ratio_b))]
214
+ tmp_merge = []
215
+ 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]:
216
+ if len(tmp_npath_a) > 0 and mask[tmp_npath_a[-1][0], tmp_npath_a[-1][1]] > 0:
217
+ tmp_merge.append([sorted_end_pts[:1], tmp_npath_a])
218
+ if len(tmp_npath_b) > 0 and mask[tmp_npath_b[-1][0], tmp_npath_b[-1][1]] > 0:
219
+ tmp_merge.append([sorted_end_pts[1:2], tmp_npath_b])
220
+ 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]:
221
+ if len(tmp_npath_b) > 0 and mask[tmp_npath_b[-1][0], tmp_npath_b[-1][1]] > 0:
222
+ tmp_merge.append([sorted_end_pts[:1], tmp_npath_b])
223
+ if len(tmp_npath_a) > 0 and mask[tmp_npath_a[-1][0], tmp_npath_a[-1][1]] > 0:
224
+ tmp_merge.append([sorted_end_pts[1:2], tmp_npath_a])
225
+ for tmp_idx in range(len(tmp_merge)):
226
+ if len(tmp_merge[tmp_idx][1]) == 0:
227
+ continue
228
+ end_pts.append(tmp_merge[tmp_idx][0])
229
+ ccs.append(set(tmp_merge[tmp_idx][1]))
230
+ if len(end_pt) == 1:
231
+ sub_mesh = mesh.subgraph(list(cc)).copy()
232
+ pnodes = netx.periphery(sub_mesh)
233
+ if len(end_pt) == 1:
234
+ ends = [*end_pt]
235
+ elif len(sorted_end_pt) == 1:
236
+ ends = [*sorted_end_pt]
237
+ else:
238
+ import pdb; pdb.set_trace()
239
+ try:
240
+ edge_id = global_mesh.nodes[ends[0]]['edge_id']
241
+ except:
242
+ import pdb; pdb.set_trace()
243
+ pnodes = sorted(pnodes,
244
+ key=lambda x: np.hypot((x[0] - ends[0][0]), (x[1] - ends[0][1])),
245
+ reverse=True)[0]
246
+ npath = [*netx.shortest_path(sub_mesh, (ends[0][0], ends[0][1]), pnodes, weight='length')]
247
+ for np_node in npath:
248
+ npath_map[np_node[0], np_node[1]] = edge_id
249
+ fpath = []
250
+ if global_mesh.nodes[ends[0]].get('far') is None:
251
+ print("None far")
252
+ else:
253
+ fnodes = global_mesh.nodes[ends[0]].get('far')
254
+ dmask = mask + 0
255
+ did = 0
256
+ while True:
257
+ did += 1
258
+ dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1)
259
+ if did > 3:
260
+ break
261
+ ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0 and mask[fnode[0], fnode[1]] == 0 and\
262
+ global_mesh.nodes[fnode].get('inpaint_id') != inpaint_id + 1)]
263
+ if len(ffnode) > 0:
264
+ fnode = ffnode[0]
265
+ break
266
+ if len(ffnode) == 0:
267
+ continue
268
+ fpath.append((fnode[0], fnode[1]))
269
+ barrel_dir = np.array([[1, 0], [1, 1], [0, 1], [-1, 1], [-1, 0], [-1, -1], [0, -1], [1, -1]])
270
+ n2f_dir = (int(fnode[0] - npath[0][0]), int(fnode[1] - npath[0][1]))
271
+ while True:
272
+ if barrel_dir[0, 0] == n2f_dir[0] and barrel_dir[0, 1] == n2f_dir[1]:
273
+ n2f_barrel = barrel_dir.copy()
274
+ break
275
+ barrel_dir = np.roll(barrel_dir, 1, axis=0)
276
+ for step in range(0, len(npath)):
277
+ if step == 0:
278
+ continue
279
+ elif step == 1:
280
+ next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1])
281
+ while True:
282
+ if barrel_dir[0, 0] == next_dir[0] and barrel_dir[0, 1] == next_dir[1]:
283
+ next_barrel = barrel_dir.copy()
284
+ break
285
+ barrel_dir = np.roll(barrel_dir, 1, axis=0)
286
+ barrel_pair = np.stack((n2f_barrel, next_barrel), axis=0)
287
+ n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1])
288
+ elif step > 1:
289
+ next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1])
290
+ while True:
291
+ if barrel_pair[1, 0, 0] == next_dir[0] and barrel_pair[1, 0, 1] == next_dir[1]:
292
+ next_barrel = barrel_pair.copy()
293
+ break
294
+ barrel_pair = np.roll(barrel_pair, 1, axis=1)
295
+ n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1])
296
+ new_locs = []
297
+ if abs(n2f_dir[0]) == 1:
298
+ new_locs.append((npath[step][0] + n2f_dir[0], npath[step][1]))
299
+ if abs(n2f_dir[1]) == 1:
300
+ new_locs.append((npath[step][0], npath[step][1] + n2f_dir[1]))
301
+ if len(new_locs) > 1:
302
+ new_locs = sorted(new_locs, key=lambda xx: np.hypot((xx[0] - fpath[-1][0]), (xx[1] - fpath[-1][1])))
303
+ break_flag = False
304
+ for new_loc in new_locs:
305
+ new_loc_nes = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]),
306
+ (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1)]\
307
+ if xx[0] >= 0 and xx[0] < fpath_map.shape[0] and xx[1] >= 0 and xx[1] < fpath_map.shape[1]]
308
+ if np.all([(fpath_map[nlne[0], nlne[1]] == -1) for nlne in new_loc_nes]) != True:
309
+ break
310
+ if npath_map[new_loc[0], new_loc[1]] != -1:
311
+ if npath_map[new_loc[0], new_loc[1]] != edge_id:
312
+ break_flag = True
313
+ break
314
+ else:
315
+ continue
316
+ if valid_map[new_loc[0], new_loc[1]] == 0:
317
+ break_flag = True
318
+ break
319
+ fpath.append(new_loc)
320
+ if break_flag is True:
321
+ break
322
+ if step != len(npath) - 1:
323
+ for xx in npath[step:]:
324
+ if npath_map[xx[0], xx[1]] == edge_id:
325
+ npath_map[xx[0], xx[1]] = -1
326
+ npath = npath[:step]
327
+ if len(fpath) > 0:
328
+ for fp_node in fpath:
329
+ fpath_map[fp_node[0], fp_node[1]] = edge_id
330
+ fpaths[edge_id] = fpath
331
+ npaths[edge_id] = npath
332
+ fpath_map[valid_near_edge != 0] = -1
333
+ if len(fpath) > 0:
334
+ iter_fpath = copy.deepcopy(fpaths[edge_id])
335
+ for node in iter_fpath:
336
+ if valid_near_edge[node[0], node[1]] != 0:
337
+ fpaths[edge_id].remove(node)
338
+
339
+ return fpath_map, npath_map, False, npaths, fpaths, invalid_edge_ids
340
+
341
+ def plan_path_e2e(mesh, cc, end_pts, global_mesh, input_edge, mask, valid_map, inpaint_id, npath_map=None, fpath_map=None):
342
+ my_npath_map = np.zeros_like(input_edge) - 1
343
+ my_fpath_map = np.zeros_like(input_edge) - 1
344
+ sub_mesh = mesh.subgraph(list(cc)).copy()
345
+ ends_1, ends_2 = end_pts[0], end_pts[1]
346
+ edge_id = global_mesh.nodes[ends_1]['edge_id']
347
+ npath = [*netx.shortest_path(sub_mesh, (ends_1[0], ends_1[1]), (ends_2[0], ends_2[1]), weight='length')]
348
+ for np_node in npath:
349
+ my_npath_map[np_node[0], np_node[1]] = edge_id
350
+ fpath = []
351
+ if global_mesh.nodes[ends_1].get('far') is None:
352
+ print("None far")
353
+ else:
354
+ fnodes = global_mesh.nodes[ends_1].get('far')
355
+ dmask = mask + 0
356
+ while True:
357
+ dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1)
358
+ ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0 and mask[fnode[0], fnode[1]] == 0 and\
359
+ global_mesh.nodes[fnode].get('inpaint_id') != inpaint_id + 1)]
360
+ if len(ffnode) > 0:
361
+ fnode = ffnode[0]
362
+ break
363
+ e_fnodes = global_mesh.nodes[ends_2].get('far')
364
+ dmask = mask + 0
365
+ while True:
366
+ dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1)
367
+ 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\
368
+ global_mesh.nodes[e_fnode].get('inpaint_id') != inpaint_id + 1)]
369
+ if len(e_ffnode) > 0:
370
+ e_fnode = e_ffnode[0]
371
+ break
372
+ fpath.append((fnode[0], fnode[1]))
373
+ if len(e_ffnode) == 0 or len(ffnode) == 0:
374
+ return my_npath_map, my_fpath_map, [], []
375
+ barrel_dir = np.array([[1, 0], [1, 1], [0, 1], [-1, 1], [-1, 0], [-1, -1], [0, -1], [1, -1]])
376
+ n2f_dir = (int(fnode[0] - npath[0][0]), int(fnode[1] - npath[0][1]))
377
+ while True:
378
+ if barrel_dir[0, 0] == n2f_dir[0] and barrel_dir[0, 1] == n2f_dir[1]:
379
+ n2f_barrel = barrel_dir.copy()
380
+ break
381
+ barrel_dir = np.roll(barrel_dir, 1, axis=0)
382
+ for step in range(0, len(npath)):
383
+ if step == 0:
384
+ continue
385
+ elif step == 1:
386
+ next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1])
387
+ while True:
388
+ if barrel_dir[0, 0] == next_dir[0] and barrel_dir[0, 1] == next_dir[1]:
389
+ next_barrel = barrel_dir.copy()
390
+ break
391
+ barrel_dir = np.roll(barrel_dir, 1, axis=0)
392
+ barrel_pair = np.stack((n2f_barrel, next_barrel), axis=0)
393
+ n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1])
394
+ elif step > 1:
395
+ next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1])
396
+ while True:
397
+ if barrel_pair[1, 0, 0] == next_dir[0] and barrel_pair[1, 0, 1] == next_dir[1]:
398
+ next_barrel = barrel_pair.copy()
399
+ break
400
+ barrel_pair = np.roll(barrel_pair, 1, axis=1)
401
+ n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1])
402
+ new_locs = []
403
+ if abs(n2f_dir[0]) == 1:
404
+ new_locs.append((npath[step][0] + n2f_dir[0], npath[step][1]))
405
+ if abs(n2f_dir[1]) == 1:
406
+ new_locs.append((npath[step][0], npath[step][1] + n2f_dir[1]))
407
+ if len(new_locs) > 1:
408
+ new_locs = sorted(new_locs, key=lambda xx: np.hypot((xx[0] - fpath[-1][0]), (xx[1] - fpath[-1][1])))
409
+ break_flag = False
410
+ for new_loc in new_locs:
411
+ new_loc_nes = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]),
412
+ (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1)]\
413
+ if xx[0] >= 0 and xx[0] < my_fpath_map.shape[0] and xx[1] >= 0 and xx[1] < my_fpath_map.shape[1]]
414
+ if fpath_map is not None and np.sum([fpath_map[nlne[0], nlne[1]] for nlne in new_loc_nes]) != 0:
415
+ break_flag = True
416
+ break
417
+ if my_npath_map[new_loc[0], new_loc[1]] != -1:
418
+ continue
419
+ if npath_map is not None and npath_map[new_loc[0], new_loc[1]] != edge_id:
420
+ break_flag = True
421
+ break
422
+ fpath.append(new_loc)
423
+ if break_flag is True:
424
+ break
425
+ if (e_fnode[0], e_fnode[1]) not in fpath:
426
+ fpath.append((e_fnode[0], e_fnode[1]))
427
+ if step != len(npath) - 1:
428
+ for xx in npath[step:]:
429
+ if my_npath_map[xx[0], xx[1]] == edge_id:
430
+ my_npath_map[xx[0], xx[1]] = -1
431
+ npath = npath[:step]
432
+ if len(fpath) > 0:
433
+ for fp_node in fpath:
434
+ my_fpath_map[fp_node[0], fp_node[1]] = edge_id
435
+
436
+ return my_fpath_map, my_npath_map, npath, fpath
437
+
438
+ 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):
439
+ my_npath_map = np.zeros_like(input_edge) - 1
440
+ my_fpath_map = np.zeros_like(input_edge) - 1
441
+ sub_mesh = mesh.subgraph(list(cc)).copy()
442
+ pnodes = netx.periphery(sub_mesh)
443
+ ends = [*end_pt]
444
+ edge_id = global_mesh.nodes[ends[0]]['edge_id']
445
+ pnodes = sorted(pnodes,
446
+ key=lambda x: np.hypot((x[0] - ends[0][0]), (x[1] - ends[0][1])),
447
+ reverse=True)[0]
448
+ if npath is None:
449
+ npath = [*netx.shortest_path(sub_mesh, (ends[0][0], ends[0][1]), pnodes, weight='length')]
450
+ else:
451
+ if (ends[0][0], ends[0][1]) == npath[0]:
452
+ npath = npath
453
+ elif (ends[0][0], ends[0][1]) == npath[-1]:
454
+ npath = npath[::-1]
455
+ else:
456
+ import pdb; pdb.set_trace()
457
+ for np_node in npath:
458
+ my_npath_map[np_node[0], np_node[1]] = edge_id
459
+ fpath = []
460
+ if global_mesh.nodes[ends[0]].get('far') is None:
461
+ print("None far")
462
+ else:
463
+ fnodes = global_mesh.nodes[ends[0]].get('far')
464
+ dmask = mask + 0
465
+ did = 0
466
+ while True:
467
+ did += 1
468
+ if did > 3:
469
+ return my_fpath_map, my_npath_map, -1
470
+ dmask = cv2.dilate(dmask, np.ones((3, 3)), iterations=1)
471
+ ffnode = [fnode for fnode in fnodes if (dmask[fnode[0], fnode[1]] > 0 and mask[fnode[0], fnode[1]] == 0 and\
472
+ global_mesh.nodes[fnode].get('inpaint_id') != inpaint_id + 1)]
473
+ if len(ffnode) > 0:
474
+ fnode = ffnode[0]
475
+ break
476
+
477
+ fpath.append((fnode[0], fnode[1]))
478
+ disp_diff = 0.
479
+ for n_loc in npath:
480
+ if mask[n_loc[0], n_loc[1]] != 0:
481
+ disp_diff = abs(abs(1. / info_on_pix[(n_loc[0], n_loc[1])][0]['depth']) - abs(1. / ends[0][2]))
482
+ break
483
+ barrel_dir = np.array([[1, 0], [1, 1], [0, 1], [-1, 1], [-1, 0], [-1, -1], [0, -1], [1, -1]])
484
+ n2f_dir = (int(fnode[0] - npath[0][0]), int(fnode[1] - npath[0][1]))
485
+ while True:
486
+ if barrel_dir[0, 0] == n2f_dir[0] and barrel_dir[0, 1] == n2f_dir[1]:
487
+ n2f_barrel = barrel_dir.copy()
488
+ break
489
+ barrel_dir = np.roll(barrel_dir, 1, axis=0)
490
+ for step in range(0, len(npath)):
491
+ if step == 0:
492
+ continue
493
+ elif step == 1:
494
+ next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1])
495
+ while True:
496
+ if barrel_dir[0, 0] == next_dir[0] and barrel_dir[0, 1] == next_dir[1]:
497
+ next_barrel = barrel_dir.copy()
498
+ break
499
+ barrel_dir = np.roll(barrel_dir, 1, axis=0)
500
+ barrel_pair = np.stack((n2f_barrel, next_barrel), axis=0)
501
+ n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1])
502
+ elif step > 1:
503
+ next_dir = (npath[step][0] - npath[step - 1][0], npath[step][1] - npath[step - 1][1])
504
+ while True:
505
+ if barrel_pair[1, 0, 0] == next_dir[0] and barrel_pair[1, 0, 1] == next_dir[1]:
506
+ next_barrel = barrel_pair.copy()
507
+ break
508
+ barrel_pair = np.roll(barrel_pair, 1, axis=1)
509
+ n2f_dir = (barrel_pair[0, 0, 0], barrel_pair[0, 0, 1])
510
+ new_locs = []
511
+ if abs(n2f_dir[0]) == 1:
512
+ new_locs.append((npath[step][0] + n2f_dir[0], npath[step][1]))
513
+ if abs(n2f_dir[1]) == 1:
514
+ new_locs.append((npath[step][0], npath[step][1] + n2f_dir[1]))
515
+ if len(new_locs) > 1:
516
+ new_locs = sorted(new_locs, key=lambda xx: np.hypot((xx[0] - fpath[-1][0]), (xx[1] - fpath[-1][1])))
517
+ break_flag = False
518
+ for new_loc in new_locs:
519
+ new_loc_nes = [xx for xx in [(new_loc[0] + 1, new_loc[1]), (new_loc[0] - 1, new_loc[1]),
520
+ (new_loc[0], new_loc[1] + 1), (new_loc[0], new_loc[1] - 1)]\
521
+ if xx[0] >= 0 and xx[0] < my_fpath_map.shape[0] and xx[1] >= 0 and xx[1] < my_fpath_map.shape[1]]
522
+ if fpath_map is not None and np.all([(fpath_map[nlne[0], nlne[1]] == -1) for nlne in new_loc_nes]) != True:
523
+ break_flag = True
524
+ break
525
+ if np.all([(my_fpath_map[nlne[0], nlne[1]] == -1) for nlne in new_loc_nes]) != True:
526
+ break_flag = True
527
+ break
528
+ if my_npath_map[new_loc[0], new_loc[1]] != -1:
529
+ continue
530
+ if npath_map is not None and npath_map[new_loc[0], new_loc[1]] != edge_id:
531
+ break_flag = True
532
+ break
533
+ if valid_map[new_loc[0], new_loc[1]] == 0:
534
+ break_flag = True
535
+ break
536
+ fpath.append(new_loc)
537
+ if break_flag is True:
538
+ break
539
+ if step != len(npath) - 1:
540
+ for xx in npath[step:]:
541
+ if my_npath_map[xx[0], xx[1]] == edge_id:
542
+ my_npath_map[xx[0], xx[1]] = -1
543
+ npath = npath[:step]
544
+ if len(fpath) > 0:
545
+ for fp_node in fpath:
546
+ my_fpath_map[fp_node[0], fp_node[1]] = edge_id
547
+
548
+ return my_fpath_map, my_npath_map, disp_diff
549
+
550
+ def refresh_node(old_node, old_feat, new_node, new_feat, mesh, stime=False):
551
+ mesh.add_node(new_node)
552
+ mesh.nodes[new_node].update(new_feat)
553
+ mesh.nodes[new_node].update(old_feat)
554
+ for ne in mesh.neighbors(old_node):
555
+ mesh.add_edge(new_node, ne)
556
+ if mesh.nodes[new_node].get('far') is not None:
557
+ tmp_far_nodes = mesh.nodes[new_node]['far']
558
+ for far_node in tmp_far_nodes:
559
+ if mesh.has_node(far_node) is False:
560
+ mesh.nodes[new_node]['far'].remove(far_node)
561
+ continue
562
+ if mesh.nodes[far_node].get('near') is not None:
563
+ for idx in range(len(mesh.nodes[far_node].get('near'))):
564
+ if mesh.nodes[far_node]['near'][idx][0] == new_node[0] and mesh.nodes[far_node]['near'][idx][1] == new_node[1]:
565
+ if len(mesh.nodes[far_node]['near'][idx]) == len(old_node):
566
+ mesh.nodes[far_node]['near'][idx] = new_node
567
+ if mesh.nodes[new_node].get('near') is not None:
568
+ tmp_near_nodes = mesh.nodes[new_node]['near']
569
+ for near_node in tmp_near_nodes:
570
+ if mesh.has_node(near_node) is False:
571
+ mesh.nodes[new_node]['near'].remove(near_node)
572
+ continue
573
+ if mesh.nodes[near_node].get('far') is not None:
574
+ for idx in range(len(mesh.nodes[near_node].get('far'))):
575
+ if mesh.nodes[near_node]['far'][idx][0] == new_node[0] and mesh.nodes[near_node]['far'][idx][1] == new_node[1]:
576
+ if len(mesh.nodes[near_node]['far'][idx]) == len(old_node):
577
+ mesh.nodes[near_node]['far'][idx] = new_node
578
+ if new_node != old_node:
579
+ mesh.remove_node(old_node)
580
+ if stime is False:
581
+ return mesh
582
+ else:
583
+ return mesh, None, None
584
+
585
+
586
+ def create_placeholder(context, mask, depth, fpath_map, npath_map, mesh, inpaint_id, edge_ccs, extend_edge_cc, all_edge_maps, self_edge_id):
587
+ add_node_time = 0
588
+ add_edge_time = 0
589
+ add_far_near_time = 0
590
+ valid_area = context + mask
591
+ H, W = mesh.graph['H'], mesh.graph['W']
592
+ edge_cc = edge_ccs[self_edge_id]
593
+ num_com = len(edge_cc) + len(extend_edge_cc)
594
+ hxs, hys = np.where(mask > 0)
595
+ for hx, hy in zip(hxs, hys):
596
+ mesh.add_node((hx, hy), inpaint_id=inpaint_id + 1, num_context=num_com)
597
+ for hx, hy in zip(hxs, hys):
598
+ four_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] if\
599
+ 0 <= x < mesh.graph['H'] and 0 <= y < mesh.graph['W'] and valid_area[x, y] != 0]
600
+ for ne in four_nes:
601
+ if mask[ne[0], ne[1]] != 0:
602
+ if not mesh.has_edge((hx, hy), ne):
603
+ mesh.add_edge((hx, hy), ne)
604
+ elif depth[ne[0], ne[1]] != 0:
605
+ if mesh.has_node((ne[0], ne[1], depth[ne[0], ne[1]])) and\
606
+ not mesh.has_edge((hx, hy), (ne[0], ne[1], depth[ne[0], ne[1]])):
607
+ mesh.add_edge((hx, hy), (ne[0], ne[1], depth[ne[0], ne[1]]))
608
+ else:
609
+ print("Undefined context node.")
610
+ import pdb; pdb.set_trace()
611
+ near_ids = np.unique(npath_map)
612
+ if near_ids[0] == -1: near_ids = near_ids[1:]
613
+ for near_id in near_ids:
614
+ hxs, hys = np.where((fpath_map == near_id) & (mask > 0))
615
+ if hxs.shape[0] > 0:
616
+ mesh.graph['max_edge_id'] = mesh.graph['max_edge_id'] + 1
617
+ else:
618
+ break
619
+ for hx, hy in zip(hxs, hys):
620
+ mesh.nodes[(hx, hy)]['edge_id'] = int(round(mesh.graph['max_edge_id']))
621
+ four_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] if\
622
+ x < mesh.graph['H'] and x >= 0 and y < mesh.graph['W'] and y >= 0 and npath_map[x, y] == near_id]
623
+ for xx in four_nes:
624
+ xx_n = copy.deepcopy(xx)
625
+ if not mesh.has_node(xx_n):
626
+ if mesh.has_node((xx_n[0], xx_n[1], depth[xx_n[0], xx_n[1]])):
627
+ xx_n = (xx_n[0], xx_n[1], depth[xx_n[0], xx_n[1]])
628
+ if mesh.has_edge((hx, hy), xx_n):
629
+ # pass
630
+ mesh.remove_edge((hx, hy), xx_n)
631
+ if mesh.nodes[(hx, hy)].get('near') is None:
632
+ mesh.nodes[(hx, hy)]['near'] = []
633
+ mesh.nodes[(hx, hy)]['near'].append(xx_n)
634
+ connect_point_exception = set()
635
+ hxs, hys = np.where((npath_map == near_id) & (all_edge_maps > -1))
636
+ for hx, hy in zip(hxs, hys):
637
+ unknown_id = int(round(all_edge_maps[hx, hy]))
638
+ if unknown_id != near_id and unknown_id != self_edge_id:
639
+ unknown_node = set([xx for xx in edge_ccs[unknown_id] if xx[0] == hx and xx[1] == hy])
640
+ connect_point_exception |= unknown_node
641
+ hxs, hys = np.where((npath_map == near_id) & (mask > 0))
642
+ if hxs.shape[0] > 0:
643
+ mesh.graph['max_edge_id'] = mesh.graph['max_edge_id'] + 1
644
+ else:
645
+ break
646
+ for hx, hy in zip(hxs, hys):
647
+ mesh.nodes[(hx, hy)]['edge_id'] = int(round(mesh.graph['max_edge_id']))
648
+ mesh.nodes[(hx, hy)]['connect_point_id'] = int(round(near_id))
649
+ mesh.nodes[(hx, hy)]['connect_point_exception'] = connect_point_exception
650
+ four_nes = [(x, y) for x, y in [(hx + 1, hy), (hx - 1, hy), (hx, hy + 1), (hx, hy - 1)] if\
651
+ x < mesh.graph['H'] and x >= 0 and y < mesh.graph['W'] and y >= 0 and fpath_map[x, y] == near_id]
652
+ for xx in four_nes:
653
+ xx_n = copy.deepcopy(xx)
654
+ if not mesh.has_node(xx_n):
655
+ if mesh.has_node((xx_n[0], xx_n[1], depth[xx_n[0], xx_n[1]])):
656
+ xx_n = (xx_n[0], xx_n[1], depth[xx_n[0], xx_n[1]])
657
+ if mesh.has_edge((hx, hy), xx_n):
658
+ mesh.remove_edge((hx, hy), xx_n)
659
+ if mesh.nodes[(hx, hy)].get('far') is None:
660
+ mesh.nodes[(hx, hy)]['far'] = []
661
+ mesh.nodes[(hx, hy)]['far'].append(xx_n)
662
+
663
+ return mesh, add_node_time, add_edge_time, add_far_near_time
664
+
665
+ def clean_far_edge(mask_edge, mask_edge_with_id, context_edge, mask, info_on_pix, global_mesh, anchor):
666
+ if isinstance(mask_edge, torch.Tensor):
667
+ if mask_edge.is_cuda:
668
+ mask_edge = mask_edge.cpu()
669
+ mask_edge = mask_edge.data
670
+ mask_edge = mask_edge.numpy()
671
+ if isinstance(context_edge, torch.Tensor):
672
+ if context_edge.is_cuda:
673
+ context_edge = context_edge.cpu()
674
+ context_edge = context_edge.data
675
+ context_edge = context_edge.numpy()
676
+ if isinstance(mask, torch.Tensor):
677
+ if mask.is_cuda:
678
+ mask = mask.cpu()
679
+ mask = mask.data
680
+ mask = mask.numpy()
681
+ mask = mask.squeeze()
682
+ mask_edge = mask_edge.squeeze()
683
+ context_edge = context_edge.squeeze()
684
+ valid_near_edge = np.zeros_like(mask_edge)
685
+ far_edge = np.zeros_like(mask_edge)
686
+ far_edge_with_id = np.ones_like(mask_edge) * -1
687
+ near_edge_with_id = np.ones_like(mask_edge) * -1
688
+ uncleaned_far_edge = np.zeros_like(mask_edge)
689
+ # Detect if there is any valid pixel mask_edge, if not ==> return default value
690
+ if mask_edge.sum() == 0:
691
+ return far_edge, uncleaned_far_edge, far_edge_with_id, near_edge_with_id
692
+ mask_edge_ids = dict(collections.Counter(mask_edge_with_id.flatten())).keys()
693
+ for edge_id in mask_edge_ids:
694
+ if edge_id < 0:
695
+ continue
696
+ specific_edge_map = (mask_edge_with_id == edge_id).astype(np.uint8)
697
+ _, sub_specific_edge_maps = cv2.connectedComponents(specific_edge_map.astype(np.uint8), connectivity=8)
698
+ for sub_edge_id in range(1, sub_specific_edge_maps.max() + 1):
699
+ specific_edge_map = (sub_specific_edge_maps == sub_edge_id).astype(np.uint8)
700
+ edge_pxs, edge_pys = np.where(specific_edge_map > 0)
701
+ edge_mesh = netx.Graph()
702
+ for edge_px, edge_py in zip(edge_pxs, edge_pys):
703
+ edge_mesh.add_node((edge_px, edge_py))
704
+ for ex in [edge_px-1, edge_px, edge_px+1]:
705
+ for ey in [edge_py-1, edge_py, edge_py+1]:
706
+ if edge_px == ex and edge_py == ey:
707
+ continue
708
+ if ex < 0 or ex >= specific_edge_map.shape[0] or ey < 0 or ey >= specific_edge_map.shape[1]:
709
+ continue
710
+ if specific_edge_map[ex, ey] == 1:
711
+ if edge_mesh.has_node((ex, ey)):
712
+ edge_mesh.add_edge((ex, ey), (edge_px, edge_py))
713
+ periphery_nodes = netx.periphery(edge_mesh)
714
+ path_diameter = netx.diameter(edge_mesh)
715
+ start_near_node = None
716
+ for node_s in periphery_nodes:
717
+ for node_e in periphery_nodes:
718
+ if node_s != node_e:
719
+ if netx.shortest_path_length(edge_mesh, node_s, node_e) == path_diameter:
720
+ if np.any(context_edge[node_s[0]-1:node_s[0]+2, node_s[1]-1:node_s[1]+2].flatten()):
721
+ start_near_node = (node_s[0], node_s[1])
722
+ end_near_node = (node_e[0], node_e[1])
723
+ break
724
+ if np.any(context_edge[node_e[0]-1:node_e[0]+2, node_e[1]-1:node_e[1]+2].flatten()):
725
+ start_near_node = (node_e[0], node_e[1])
726
+ end_near_node = (node_s[0], node_s[1])
727
+ break
728
+ if start_near_node is not None:
729
+ break
730
+ if start_near_node is None:
731
+ continue
732
+ new_specific_edge_map = np.zeros_like(mask)
733
+ for path_node in netx.shortest_path(edge_mesh, start_near_node, end_near_node):
734
+ new_specific_edge_map[path_node[0], path_node[1]] = 1
735
+ 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)
736
+ distance = np.abs((context_near_pxs - 1)) + np.abs((context_near_pys - 1))
737
+ if (np.where(distance == distance.min())[0].shape[0]) > 1:
738
+ closest_pxs = context_near_pxs[np.where(distance == distance.min())[0]]
739
+ closest_pys = context_near_pys[np.where(distance == distance.min())[0]]
740
+ closest_depths = []
741
+ for closest_px, closest_py in zip(closest_pxs, closest_pys):
742
+ 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:
743
+ 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])):
744
+ if info['synthesis'] is False:
745
+ closest_depths.append(abs(info['depth']))
746
+ context_near_px, context_near_py = closest_pxs[np.array(closest_depths).argmax()], closest_pys[np.array(closest_depths).argmax()]
747
+ else:
748
+ context_near_px, context_near_py = context_near_pxs[distance.argmin()], context_near_pys[distance.argmin()]
749
+ context_near_node = (start_near_node[0]-1 + context_near_px, start_near_node[1]-1 + context_near_py)
750
+ far_node_list = []
751
+ global_context_near_node = (context_near_node[0] + anchor[0], context_near_node[1] + anchor[2])
752
+ if info_on_pix.get(global_context_near_node) is not None:
753
+ for info in info_on_pix[global_context_near_node]:
754
+ if info['synthesis'] is False:
755
+ context_near_node_3d = (global_context_near_node[0], global_context_near_node[1], info['depth'])
756
+ if global_mesh.nodes[context_near_node_3d].get('far') is not None:
757
+ for far_node in global_mesh.nodes[context_near_node_3d].get('far'):
758
+ far_node = (far_node[0] - anchor[0], far_node[1] - anchor[2], far_node[2])
759
+ if mask[far_node[0], far_node[1]] == 0:
760
+ far_node_list.append([far_node[0], far_node[1]])
761
+ if len(far_node_list) > 0:
762
+ far_nodes_dist = np.sum(np.abs(np.array(far_node_list) - np.array([[edge_px, edge_py]])), axis=1)
763
+ context_far_node = tuple(far_node_list[far_nodes_dist.argmin()])
764
+ corresponding_far_edge = np.zeros_like(mask_edge)
765
+ corresponding_far_edge[context_far_node[0], context_far_node[1]] = 1
766
+ surround_map = cv2.dilate(new_specific_edge_map.astype(np.uint8),
767
+ np.array([[1,1,1],[1,1,1],[1,1,1]]).astype(np.uint8),
768
+ iterations=1)
769
+ specific_edge_map_wo_end_pt = new_specific_edge_map.copy()
770
+ specific_edge_map_wo_end_pt[end_near_node[0], end_near_node[1]] = 0
771
+ surround_map_wo_end_pt = cv2.dilate(specific_edge_map_wo_end_pt.astype(np.uint8),
772
+ np.array([[1,1,1],[1,1,1],[1,1,1]]).astype(np.uint8),
773
+ iterations=1)
774
+ surround_map_wo_end_pt[new_specific_edge_map > 0] = 0
775
+ surround_map_wo_end_pt[context_near_node[0], context_near_node[1]] = 0
776
+ surround_map = surround_map_wo_end_pt.copy()
777
+ _, far_edge_cc = cv2.connectedComponents(surround_map.astype(np.uint8), connectivity=4)
778
+ start_far_node = None
779
+ accompany_far_node = None
780
+ if surround_map[context_far_node[0], context_far_node[1]] == 1:
781
+ start_far_node = context_far_node
782
+ else:
783
+ four_nes = [(context_far_node[0] - 1, context_far_node[1]),
784
+ (context_far_node[0] + 1, context_far_node[1]),
785
+ (context_far_node[0], context_far_node[1] - 1),
786
+ (context_far_node[0], context_far_node[1] + 1)]
787
+ candidate_bevel = []
788
+ for ne in four_nes:
789
+ if surround_map[ne[0], ne[1]] == 1:
790
+ start_far_node = (ne[0], ne[1])
791
+ break
792
+ elif (ne[0] != context_near_node[0] or ne[1] != context_near_node[1]) and \
793
+ (ne[0] != start_near_node[0] or ne[1] != start_near_node[1]):
794
+ candidate_bevel.append((ne[0], ne[1]))
795
+ if start_far_node is None:
796
+ for ne in candidate_bevel:
797
+ if ne[0] == context_far_node[0]:
798
+ bevel_xys = [[ne[0] + 1, ne[1]], [ne[0] - 1, ne[1]]]
799
+ if ne[1] == context_far_node[1]:
800
+ bevel_xys = [[ne[0], ne[1] + 1], [ne[0], ne[1] - 1]]
801
+ for bevel_x, bevel_y in bevel_xys:
802
+ if surround_map[bevel_x, bevel_y] == 1:
803
+ start_far_node = (bevel_x, bevel_y)
804
+ accompany_far_node = (ne[0], ne[1])
805
+ break
806
+ if start_far_node is not None:
807
+ break
808
+ if start_far_node is not None:
809
+ for far_edge_id in range(1, far_edge_cc.max() + 1):
810
+ specific_far_edge = (far_edge_cc == far_edge_id).astype(np.uint8)
811
+ if specific_far_edge[start_far_node[0], start_far_node[1]] == 1:
812
+ if accompany_far_node is not None:
813
+ specific_far_edge[accompany_far_node] = 1
814
+ far_edge[specific_far_edge > 0] = 1
815
+ far_edge_with_id[specific_far_edge > 0] = edge_id
816
+ end_far_candidates = np.zeros_like(far_edge)
817
+ end_far_candidates[end_near_node[0], end_near_node[1]] = 1
818
+ end_far_candidates = cv2.dilate(end_far_candidates.astype(np.uint8),
819
+ np.array([[0,1,0],[1,1,1],[0,1,0]]).astype(np.uint8),
820
+ iterations=1)
821
+ end_far_candidates[end_near_node[0], end_near_node[1]] = 0
822
+ invalid_nodes = (((far_edge_cc != far_edge_id).astype(np.uint8) * \
823
+ (far_edge_cc != 0).astype(np.uint8)).astype(np.uint8) + \
824
+ (new_specific_edge_map).astype(np.uint8) + \
825
+ (mask == 0).astype(np.uint8)).clip(0, 1)
826
+ end_far_candidates[invalid_nodes > 0] = 0
827
+ far_edge[end_far_candidates > 0] = 1
828
+ far_edge_with_id[end_far_candidates > 0] = edge_id
829
+
830
+ far_edge[context_far_node[0], context_far_node[1]] = 1
831
+ far_edge_with_id[context_far_node[0], context_far_node[1]] = edge_id
832
+ near_edge_with_id[(mask_edge_with_id == edge_id) > 0] = edge_id
833
+ uncleaned_far_edge = far_edge.copy()
834
+ far_edge[mask == 0] = 0
835
+
836
+ return far_edge, uncleaned_far_edge, far_edge_with_id, near_edge_with_id
837
+
838
+ def get_MiDaS_samples(image_folder, depth_folder, config, specific=None, aft_certain=None):
839
+ lines = [os.path.splitext(os.path.basename(xx))[0] for xx in glob.glob(os.path.join(image_folder, '*' + config['img_format']))]
840
+ samples = []
841
+ generic_pose = np.eye(4)
842
+ assert len(config['traj_types']) == len(config['x_shift_range']) ==\
843
+ len(config['y_shift_range']) == len(config['z_shift_range']) == len(config['video_postfix']), \
844
+ "The number of elements in 'traj_types', 'x_shift_range', 'y_shift_range', 'z_shift_range' and \
845
+ 'video_postfix' should be equal."
846
+ tgt_pose = [[generic_pose * 1]]
847
+ tgts_poses = []
848
+ for traj_idx in range(len(config['traj_types'])):
849
+ tgt_poses = []
850
+ sx, sy, sz = path_planning(config['num_frames'], config['x_shift_range'][traj_idx], config['y_shift_range'][traj_idx],
851
+ config['z_shift_range'][traj_idx], path_type=config['traj_types'][traj_idx])
852
+ for xx, yy, zz in zip(sx, sy, sz):
853
+ tgt_poses.append(generic_pose * 1.)
854
+ tgt_poses[-1][:3, -1] = np.array([xx, yy, zz])
855
+ tgts_poses += [tgt_poses]
856
+ tgt_pose = generic_pose * 1
857
+
858
+ aft_flag = True
859
+ if aft_certain is not None and len(aft_certain) > 0:
860
+ aft_flag = False
861
+ for seq_dir in lines:
862
+ if specific is not None and len(specific) > 0:
863
+ if specific != seq_dir:
864
+ continue
865
+ if aft_certain is not None and len(aft_certain) > 0:
866
+ if aft_certain == seq_dir:
867
+ aft_flag = True
868
+ if aft_flag is False:
869
+ continue
870
+ samples.append({})
871
+ sdict = samples[-1]
872
+ sdict['depth_fi'] = os.path.join(depth_folder, seq_dir + config['depth_format'])
873
+ sdict['ref_img_fi'] = os.path.join(image_folder, seq_dir + config['img_format'])
874
+ H, W = imageio.imread(sdict['ref_img_fi']).shape[:2]
875
+ sdict['int_mtx'] = np.array([[max(H, W), 0, W//2], [0, max(H, W), H//2], [0, 0, 1]]).astype(np.float32)
876
+ if sdict['int_mtx'].max() > 1:
877
+ sdict['int_mtx'][0, :] = sdict['int_mtx'][0, :] / float(W)
878
+ sdict['int_mtx'][1, :] = sdict['int_mtx'][1, :] / float(H)
879
+ sdict['ref_pose'] = np.eye(4)
880
+ sdict['tgt_pose'] = tgt_pose
881
+ sdict['tgts_poses'] = tgts_poses
882
+ sdict['video_postfix'] = config['video_postfix']
883
+ sdict['tgt_name'] = [os.path.splitext(os.path.basename(sdict['depth_fi']))[0]]
884
+ sdict['src_pair_name'] = sdict['tgt_name'][0]
885
+
886
+ return samples
887
+
888
+ def get_valid_size(imap):
889
+ x_max = np.where(imap.sum(1).squeeze() > 0)[0].max() + 1
890
+ x_min = np.where(imap.sum(1).squeeze() > 0)[0].min()
891
+ y_max = np.where(imap.sum(0).squeeze() > 0)[0].max() + 1
892
+ y_min = np.where(imap.sum(0).squeeze() > 0)[0].min()
893
+ size_dict = {'x_max':x_max, 'y_max':y_max, 'x_min':x_min, 'y_min':y_min}
894
+
895
+ return size_dict
896
+
897
+ def dilate_valid_size(isize_dict, imap, dilate=[0, 0]):
898
+ osize_dict = copy.deepcopy(isize_dict)
899
+ osize_dict['x_min'] = max(0, osize_dict['x_min'] - dilate[0])
900
+ osize_dict['x_max'] = min(imap.shape[0], osize_dict['x_max'] + dilate[0])
901
+ osize_dict['y_min'] = max(0, osize_dict['y_min'] - dilate[0])
902
+ osize_dict['y_max'] = min(imap.shape[1], osize_dict['y_max'] + dilate[1])
903
+
904
+ return osize_dict
905
+
906
+ def crop_maps_by_size(size, *imaps):
907
+ omaps = []
908
+ for imap in imaps:
909
+ omaps.append(imap[size['x_min']:size['x_max'], size['y_min']:size['y_max']].copy())
910
+
911
+ return omaps
912
+
913
+ def smooth_cntsyn_gap(init_depth_map, mask_region, context_region, init_mask_region=None):
914
+ if init_mask_region is not None:
915
+ curr_mask_region = init_mask_region * 1
916
+ else:
917
+ curr_mask_region = mask_region * 0
918
+ depth_map = init_depth_map.copy()
919
+ for _ in range(2):
920
+ cm_mask = context_region + curr_mask_region
921
+ depth_s1 = np.roll(depth_map, 1, 0)
922
+ depth_s2 = np.roll(depth_map, -1, 0)
923
+ depth_s3 = np.roll(depth_map, 1, 1)
924
+ depth_s4 = np.roll(depth_map, -1, 1)
925
+ mask_s1 = np.roll(cm_mask, 1, 0)
926
+ mask_s2 = np.roll(cm_mask, -1, 0)
927
+ mask_s3 = np.roll(cm_mask, 1, 1)
928
+ mask_s4 = np.roll(cm_mask, -1, 1)
929
+ fluxin_depths = (depth_s1 * mask_s1 + depth_s2 * mask_s2 + depth_s3 * mask_s3 + depth_s4 * mask_s4) / \
930
+ ((mask_s1 + mask_s2 + mask_s3 + mask_s4) + 1e-6)
931
+ fluxin_mask = (fluxin_depths != 0) * mask_region
932
+ init_mask = (fluxin_mask * (curr_mask_region >= 0).astype(np.float32) > 0).astype(np.uint8)
933
+ depth_map[init_mask > 0] = fluxin_depths[init_mask > 0]
934
+ if init_mask.shape[-1] > curr_mask_region.shape[-1]:
935
+ curr_mask_region[init_mask.sum(-1, keepdims=True) > 0] = 1
936
+ else:
937
+ curr_mask_region[init_mask > 0] = 1
938
+ depth_map[fluxin_mask > 0] = fluxin_depths[fluxin_mask > 0]
939
+
940
+ return depth_map
941
+
942
+ def read_MiDaS_depth(disp_fi, disp_rescale=10., h=None, w=None):
943
+ if 'npy' in os.path.splitext(disp_fi)[-1]:
944
+ disp = np.load(disp_fi)
945
+ else:
946
+ disp = imageio.imread(disp_fi).astype(np.float32)
947
+ disp = disp - disp.min()
948
+ disp = cv2.blur(disp / disp.max(), ksize=(3, 3)) * disp.max()
949
+ disp = (disp / disp.max()) * disp_rescale
950
+ if h is not None and w is not None:
951
+ disp = resize(disp / disp.max(), (h, w), order=1) * disp.max()
952
+ depth = 1. / np.maximum(disp, 0.05)
953
+
954
+ return depth
955
+
956
+ def follow_image_aspect_ratio(depth, image):
957
+ H, W = image.shape[:2]
958
+ image_aspect_ratio = H / W
959
+ dH, dW = depth.shape[:2]
960
+ depth_aspect_ratio = dH / dW
961
+ if depth_aspect_ratio > image_aspect_ratio:
962
+ resize_H = dH
963
+ resize_W = dH / image_aspect_ratio
964
+ else:
965
+ resize_W = dW
966
+ resize_H = dW * image_aspect_ratio
967
+ depth = resize(depth / depth.max(),
968
+ (int(resize_H),
969
+ int(resize_W)),
970
+ order=0) * depth.max()
971
+
972
+ return depth
973
+
974
+ def depth_resize(depth, origin_size, image_size):
975
+ if origin_size[0] is not 0:
976
+ max_depth = depth.max()
977
+ depth = depth / max_depth
978
+ depth = resize(depth, origin_size, order=1, mode='edge')
979
+ depth = depth * max_depth
980
+ else:
981
+ max_depth = depth.max()
982
+ depth = depth / max_depth
983
+ depth = resize(depth, image_size, order=1, mode='edge')
984
+ depth = depth * max_depth
985
+
986
+ return depth
987
+
988
+ def filter_irrelevant_edge(self_edge, other_edges, other_edges_with_id, current_edge_id, context, edge_ccs, mesh, anchor):
989
+ other_edges = other_edges.squeeze()
990
+ other_edges_with_id = other_edges_with_id.squeeze()
991
+
992
+ self_edge = self_edge.squeeze()
993
+ 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)
994
+ edge_ids = collections.Counter(other_edges_with_id.flatten()).keys()
995
+ other_edges_info = []
996
+ # import ipdb
997
+ # ipdb.set_trace()
998
+ for edge_id in edge_ids:
999
+ edge_id = int(edge_id)
1000
+ if edge_id >= 0:
1001
+ condition = ((other_edges_with_id == edge_id) * other_edges * context).astype(np.uint8)
1002
+ if dilate_self_edge[condition > 0].sum() == 0:
1003
+ other_edges[other_edges_with_id == edge_id] = 0
1004
+ else:
1005
+ num_condition, condition_labels = cv2.connectedComponents(condition, connectivity=8)
1006
+ for condition_id in range(1, num_condition):
1007
+ isolate_condition = ((condition_labels == condition_id) > 0).astype(np.uint8)
1008
+ num_end_group, end_group = cv2.connectedComponents(((dilate_self_edge * isolate_condition) > 0).astype(np.uint8), connectivity=8)
1009
+ if num_end_group == 1:
1010
+ continue
1011
+ for end_id in range(1, num_end_group):
1012
+ end_pxs, end_pys = np.where((end_group == end_id))
1013
+ end_px, end_py = end_pxs[0], end_pys[0]
1014
+ other_edges_info.append({})
1015
+ other_edges_info[-1]['edge_id'] = edge_id
1016
+ # other_edges_info[-1]['near_depth'] = None
1017
+ other_edges_info[-1]['diff'] = None
1018
+ other_edges_info[-1]['edge_map'] = np.zeros_like(self_edge)
1019
+ other_edges_info[-1]['end_point_map'] = np.zeros_like(self_edge)
1020
+ other_edges_info[-1]['end_point_map'][(end_group == end_id)] = 1
1021
+ other_edges_info[-1]['forbidden_point_map'] = np.zeros_like(self_edge)
1022
+ other_edges_info[-1]['forbidden_point_map'][(end_group != end_id) * (end_group != 0)] = 1
1023
+ 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)
1024
+ for x in edge_ccs[edge_id]:
1025
+ nx = x[0] - anchor[0]
1026
+ ny = x[1] - anchor[1]
1027
+ if nx == end_px and ny == end_py:
1028
+ # other_edges_info[-1]['near_depth'] = abs(nx)
1029
+ if mesh.nodes[x].get('far') is not None and len(mesh.nodes[x].get('far')) == 1:
1030
+ other_edges_info[-1]['diff'] = abs(1./abs([*mesh.nodes[x].get('far')][0][2]) - 1./abs(x[2]))
1031
+ else:
1032
+ other_edges_info[-1]['diff'] = 0
1033
+ # if end_group[nx, ny] != end_id and end_group[nx, ny] > 0:
1034
+ # continue
1035
+ try:
1036
+ if isolate_condition[nx, ny] == 1:
1037
+ other_edges_info[-1]['edge_map'][nx, ny] = 1
1038
+ except:
1039
+ pass
1040
+ try:
1041
+ other_edges_info = sorted(other_edges_info, key=lambda x : x['diff'], reverse=True)
1042
+ except:
1043
+ import pdb
1044
+ pdb.set_trace()
1045
+ # import pdb
1046
+ # pdb.set_trace()
1047
+ # other_edges = other_edges[..., None]
1048
+ for other_edge in other_edges_info:
1049
+ if other_edge['end_point_map'] is None:
1050
+ import pdb
1051
+ pdb.set_trace()
1052
+
1053
+ other_edges = other_edges * context
1054
+
1055
+ return other_edges, other_edges_info
1056
+
1057
+ def require_depth_edge(context_edge, mask):
1058
+ dilate_mask = cv2.dilate(mask, np.array([[1,1,1],[1,1,1],[1,1,1]]).astype(np.uint8), iterations=1)
1059
+ if (dilate_mask * context_edge).max() == 0:
1060
+ return False
1061
+ else:
1062
+ return True
1063
+
1064
+ def refine_color_around_edge(mesh, info_on_pix, edge_ccs, config, spdb=False):
1065
+ H, W = mesh.graph['H'], mesh.graph['W']
1066
+ tmp_edge_ccs = copy.deepcopy(edge_ccs)
1067
+ for edge_id, edge_cc in enumerate(edge_ccs):
1068
+ if len(edge_cc) == 0:
1069
+ continue
1070
+ near_maps = np.zeros((H, W)).astype(np.bool)
1071
+ far_maps = np.zeros((H, W)).astype(np.bool)
1072
+ tmp_far_nodes = set()
1073
+ far_nodes = set()
1074
+ near_nodes = set()
1075
+ end_nodes = set()
1076
+ for i in range(5):
1077
+ if i == 0:
1078
+ for edge_node in edge_cc:
1079
+ if mesh.nodes[edge_node].get('depth_edge_dilate_2_color_flag') is not True:
1080
+ break
1081
+ if mesh.nodes[edge_node].get('inpaint_id') == 1:
1082
+ near_nodes.add(edge_node)
1083
+ tmp_node = mesh.nodes[edge_node].get('far')
1084
+ tmp_node = set(tmp_node) if tmp_node is not None else set()
1085
+ tmp_far_nodes |= tmp_node
1086
+ rmv_tmp_far_nodes = set()
1087
+ for far_node in tmp_far_nodes:
1088
+ if not(mesh.has_node(far_node) and mesh.nodes[far_node].get('inpaint_id') == 1):
1089
+ rmv_tmp_far_nodes.add(far_node)
1090
+ if len(tmp_far_nodes - rmv_tmp_far_nodes) == 0:
1091
+ break
1092
+ else:
1093
+ for near_node in near_nodes:
1094
+ near_maps[near_node[0], near_node[1]] = True
1095
+ mesh.nodes[near_node]['refine_rgbd'] = True
1096
+ mesh.nodes[near_node]['backup_depth'] = near_node[2] \
1097
+ if mesh.nodes[near_node].get('real_depth') is None else mesh.nodes[near_node]['real_depth']
1098
+ mesh.nodes[near_node]['backup_color'] = mesh.nodes[near_node]['color']
1099
+ for far_node in tmp_far_nodes:
1100
+ if mesh.has_node(far_node) and mesh.nodes[far_node].get('inpaint_id') == 1:
1101
+ far_nodes.add(far_node)
1102
+ far_maps[far_node[0], far_node[1]] = True
1103
+ mesh.nodes[far_node]['refine_rgbd'] = True
1104
+ mesh.nodes[far_node]['backup_depth'] = far_node[2] \
1105
+ if mesh.nodes[far_node].get('real_depth') is None else mesh.nodes[far_node]['real_depth']
1106
+ mesh.nodes[far_node]['backup_color'] = mesh.nodes[far_node]['color']
1107
+ tmp_far_nodes = far_nodes
1108
+ tmp_near_nodes = near_nodes
1109
+ else:
1110
+ tmp_far_nodes = new_tmp_far_nodes
1111
+ tmp_near_nodes = new_tmp_near_nodes
1112
+ new_tmp_far_nodes = None
1113
+ new_tmp_near_nodes = None
1114
+ new_tmp_far_nodes = set()
1115
+ new_tmp_near_nodes = set()
1116
+ for node in tmp_near_nodes:
1117
+ for ne_node in mesh.neighbors(node):
1118
+ if far_maps[ne_node[0], ne_node[1]] == False and \
1119
+ near_maps[ne_node[0], ne_node[1]] == False:
1120
+ if mesh.nodes[ne_node].get('inpaint_id') == 1:
1121
+ new_tmp_near_nodes.add(ne_node)
1122
+ near_maps[ne_node[0], ne_node[1]] = True
1123
+ mesh.nodes[ne_node]['refine_rgbd'] = True
1124
+ mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \
1125
+ if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth']
1126
+ mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color']
1127
+ else:
1128
+ mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \
1129
+ if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth']
1130
+ mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color']
1131
+ end_nodes.add(node)
1132
+ near_nodes.update(new_tmp_near_nodes)
1133
+ for node in tmp_far_nodes:
1134
+ for ne_node in mesh.neighbors(node):
1135
+ if far_maps[ne_node[0], ne_node[1]] == False and \
1136
+ near_maps[ne_node[0], ne_node[1]] == False:
1137
+ if mesh.nodes[ne_node].get('inpaint_id') == 1:
1138
+ new_tmp_far_nodes.add(ne_node)
1139
+ far_maps[ne_node[0], ne_node[1]] = True
1140
+ mesh.nodes[ne_node]['refine_rgbd'] = True
1141
+ mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \
1142
+ if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth']
1143
+ mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color']
1144
+ else:
1145
+ mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \
1146
+ if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth']
1147
+ mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color']
1148
+ end_nodes.add(node)
1149
+ far_nodes.update(new_tmp_far_nodes)
1150
+ if len(far_nodes) == 0:
1151
+ tmp_edge_ccs[edge_id] = set()
1152
+ continue
1153
+ for node in new_tmp_far_nodes | new_tmp_near_nodes:
1154
+ for ne_node in mesh.neighbors(node):
1155
+ if far_maps[ne_node[0], ne_node[1]] == False and near_maps[ne_node[0], ne_node[1]] == False:
1156
+ end_nodes.add(node)
1157
+ mesh.nodes[ne_node]['backup_depth'] = ne_node[2] \
1158
+ if mesh.nodes[ne_node].get('real_depth') is None else mesh.nodes[ne_node]['real_depth']
1159
+ mesh.nodes[ne_node]['backup_color'] = mesh.nodes[ne_node]['color']
1160
+ tmp_end_nodes = end_nodes
1161
+
1162
+ refine_nodes = near_nodes | far_nodes
1163
+ remain_refine_nodes = copy.deepcopy(refine_nodes)
1164
+ accum_idx = 0
1165
+ while len(remain_refine_nodes) > 0:
1166
+ accum_idx += 1
1167
+ if accum_idx > 100:
1168
+ break
1169
+ new_tmp_end_nodes = None
1170
+ new_tmp_end_nodes = set()
1171
+ survive_tmp_end_nodes = set()
1172
+ for node in tmp_end_nodes:
1173
+ re_depth, re_color, re_count = 0, np.array([0., 0., 0.]), 0
1174
+ for ne_node in mesh.neighbors(node):
1175
+ if mesh.nodes[ne_node].get('refine_rgbd') is True:
1176
+ if ne_node not in tmp_end_nodes:
1177
+ new_tmp_end_nodes.add(ne_node)
1178
+ else:
1179
+ try:
1180
+ re_depth += mesh.nodes[ne_node]['backup_depth']
1181
+ re_color += mesh.nodes[ne_node]['backup_color'].astype(np.float32)
1182
+ re_count += 1.
1183
+ except:
1184
+ import pdb; pdb.set_trace()
1185
+ if re_count > 0:
1186
+ re_depth = re_depth / re_count
1187
+ re_color = re_color / re_count
1188
+ mesh.nodes[node]['backup_depth'] = re_depth
1189
+ mesh.nodes[node]['backup_color'] = re_color
1190
+ mesh.nodes[node]['refine_rgbd'] = False
1191
+ else:
1192
+ survive_tmp_end_nodes.add(node)
1193
+ for node in tmp_end_nodes - survive_tmp_end_nodes:
1194
+ if node in remain_refine_nodes:
1195
+ remain_refine_nodes.remove(node)
1196
+ tmp_end_nodes = new_tmp_end_nodes
1197
+ if spdb == True:
1198
+ bfrd_canvas = np.zeros((H, W))
1199
+ bfrc_canvas = np.zeros((H, W, 3)).astype(np.uint8)
1200
+ aftd_canvas = np.zeros((H, W))
1201
+ aftc_canvas = np.zeros((H, W, 3)).astype(np.uint8)
1202
+ for node in refine_nodes:
1203
+ bfrd_canvas[node[0], node[1]] = abs(node[2])
1204
+ aftd_canvas[node[0], node[1]] = abs(mesh.nodes[node]['backup_depth'])
1205
+ bfrc_canvas[node[0], node[1]] = mesh.nodes[node]['color'].astype(np.uint8)
1206
+ aftc_canvas[node[0], node[1]] = mesh.nodes[node]['backup_color'].astype(np.uint8)
1207
+ f, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, sharex=True, sharey=True);
1208
+ ax1.imshow(bfrd_canvas);
1209
+ ax2.imshow(aftd_canvas);
1210
+ ax3.imshow(bfrc_canvas);
1211
+ ax4.imshow(aftc_canvas);
1212
+ plt.show()
1213
+ import pdb; pdb.set_trace()
1214
+ for node in refine_nodes:
1215
+ if mesh.nodes[node].get('refine_rgbd') is not None:
1216
+ mesh.nodes[node].pop('refine_rgbd')
1217
+ mesh.nodes[node]['color'] = mesh.nodes[node]['backup_color']
1218
+ for info in info_on_pix[(node[0], node[1])]:
1219
+ if info['depth'] == node[2]:
1220
+ info['color'] = mesh.nodes[node]['backup_color']
1221
+
1222
+ return mesh, info_on_pix
1223
+
1224
+ def refine_depth_around_edge(mask_depth, far_edge, uncleaned_far_edge, near_edge, mask, all_depth, config):
1225
+ if isinstance(mask_depth, torch.Tensor):
1226
+ if mask_depth.is_cuda:
1227
+ mask_depth = mask_depth.cpu()
1228
+ mask_depth = mask_depth.data
1229
+ mask_depth = mask_depth.numpy()
1230
+ if isinstance(far_edge, torch.Tensor):
1231
+ if far_edge.is_cuda:
1232
+ far_edge = far_edge.cpu()
1233
+ far_edge = far_edge.data
1234
+ far_edge = far_edge.numpy()
1235
+ if isinstance(uncleaned_far_edge, torch.Tensor):
1236
+ if uncleaned_far_edge.is_cuda:
1237
+ uncleaned_far_edge = uncleaned_far_edge.cpu()
1238
+ uncleaned_far_edge = uncleaned_far_edge.data
1239
+ uncleaned_far_edge = uncleaned_far_edge.numpy()
1240
+ if isinstance(near_edge, torch.Tensor):
1241
+ if near_edge.is_cuda:
1242
+ near_edge = near_edge.cpu()
1243
+ near_edge = near_edge.data
1244
+ near_edge = near_edge.numpy()
1245
+ if isinstance(mask, torch.Tensor):
1246
+ if mask.is_cuda:
1247
+ mask = mask.cpu()
1248
+ mask = mask.data
1249
+ mask = mask.numpy()
1250
+ mask = mask.squeeze()
1251
+ uncleaned_far_edge = uncleaned_far_edge.squeeze()
1252
+ far_edge = far_edge.squeeze()
1253
+ near_edge = near_edge.squeeze()
1254
+ mask_depth = mask_depth.squeeze()
1255
+ 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)
1256
+ near_edge[dilate_far_edge == 0] = 0
1257
+ 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)
1258
+ far_edge[dilate_near_edge == 0] = 0
1259
+ init_far_edge = far_edge.copy()
1260
+ init_near_edge = near_edge.copy()
1261
+ for i in range(config['depth_edge_dilate_2']):
1262
+ 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)
1263
+ init_far_edge[init_near_edge == 1] = 0
1264
+ 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)
1265
+ init_near_edge[init_far_edge == 1] = 0
1266
+ init_far_edge[mask == 0] = 0
1267
+ init_near_edge[mask == 0] = 0
1268
+ hole_far_edge = 1 - init_far_edge
1269
+ hole_near_edge = 1 - init_near_edge
1270
+ change = None
1271
+ while True:
1272
+ change = False
1273
+ hole_far_edge[init_near_edge == 1] = 0
1274
+ hole_near_edge[init_far_edge == 1] = 0
1275
+ far_pxs, far_pys = np.where((hole_far_edge == 0) * (init_far_edge == 1) > 0)
1276
+ current_hole_far_edge = hole_far_edge.copy()
1277
+ for far_px, far_py in zip(far_pxs, far_pys):
1278
+ min_px = max(far_px - 1, 0)
1279
+ max_px = min(far_px + 2, mask.shape[0]-1)
1280
+ min_py = max(far_py - 1, 0)
1281
+ max_py = min(far_py + 2, mask.shape[1]-1)
1282
+ hole_far = current_hole_far_edge[min_px: max_px, min_py: max_py]
1283
+ tmp_mask = mask[min_px: max_px, min_py: max_py]
1284
+ all_depth_patch = all_depth[min_px: max_px, min_py: max_py] * 0
1285
+ all_depth_mask = (all_depth_patch != 0).astype(np.uint8)
1286
+ 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)]
1287
+ combine_mask = (tmp_mask + all_depth_mask).clip(0, 1) * hole_far * cross_element
1288
+ tmp_patch = combine_mask * (mask_depth[min_px: max_px, min_py: max_py] + all_depth_patch)
1289
+ number = np.count_nonzero(tmp_patch)
1290
+ if number > 0:
1291
+ mask_depth[far_px, far_py] = np.sum(tmp_patch).astype(np.float32) / max(number, 1e-6)
1292
+ hole_far_edge[far_px, far_py] = 1
1293
+ change = True
1294
+ near_pxs, near_pys = np.where((hole_near_edge == 0) * (init_near_edge == 1) > 0)
1295
+ current_hole_near_edge = hole_near_edge.copy()
1296
+ for near_px, near_py in zip(near_pxs, near_pys):
1297
+ min_px = max(near_px - 1, 0)
1298
+ max_px = min(near_px + 2, mask.shape[0]-1)
1299
+ min_py = max(near_py - 1, 0)
1300
+ max_py = min(near_py + 2, mask.shape[1]-1)
1301
+ hole_near = current_hole_near_edge[min_px: max_px, min_py: max_py]
1302
+ tmp_mask = mask[min_px: max_px, min_py: max_py]
1303
+ all_depth_patch = all_depth[min_px: max_px, min_py: max_py] * 0
1304
+ all_depth_mask = (all_depth_patch != 0).astype(np.uint8)
1305
+ 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]
1306
+ combine_mask = (tmp_mask + all_depth_mask).clip(0, 1) * hole_near * cross_element
1307
+ tmp_patch = combine_mask * (mask_depth[min_px: max_px, min_py: max_py] + all_depth_patch)
1308
+ number = np.count_nonzero(tmp_patch)
1309
+ if number > 0:
1310
+ mask_depth[near_px, near_py] = np.sum(tmp_patch) / max(number, 1e-6)
1311
+ hole_near_edge[near_px, near_py] = 1
1312
+ change = True
1313
+ if change is False:
1314
+ break
1315
+
1316
+ return mask_depth
1317
+
1318
+
1319
+
1320
+ def vis_depth_edge_connectivity(depth, config):
1321
+ disp = 1./depth
1322
+ u_diff = (disp[1:, :] - disp[:-1, :])[:-1, 1:-1]
1323
+ b_diff = (disp[:-1, :] - disp[1:, :])[1:, 1:-1]
1324
+ l_diff = (disp[:, 1:] - disp[:, :-1])[1:-1, :-1]
1325
+ r_diff = (disp[:, :-1] - disp[:, 1:])[1:-1, 1:]
1326
+ u_over = (np.abs(u_diff) > config['depth_threshold']).astype(np.float32)
1327
+ b_over = (np.abs(b_diff) > config['depth_threshold']).astype(np.float32)
1328
+ l_over = (np.abs(l_diff) > config['depth_threshold']).astype(np.float32)
1329
+ r_over = (np.abs(r_diff) > config['depth_threshold']).astype(np.float32)
1330
+ concat_diff = np.stack([u_diff, b_diff, r_diff, l_diff], axis=-1)
1331
+ concat_over = np.stack([u_over, b_over, r_over, l_over], axis=-1)
1332
+ over_diff = concat_diff * concat_over
1333
+ pos_over = (over_diff > 0).astype(np.float32).sum(-1).clip(0, 1)
1334
+ neg_over = (over_diff < 0).astype(np.float32).sum(-1).clip(0, 1)
1335
+ neg_over[(over_diff > 0).astype(np.float32).sum(-1) > 0] = 0
1336
+ _, edge_label = cv2.connectedComponents(pos_over.astype(np.uint8), connectivity=8)
1337
+ T_junction_maps = np.zeros_like(pos_over)
1338
+ for edge_id in range(1, edge_label.max() + 1):
1339
+ edge_map = (edge_label == edge_id).astype(np.uint8)
1340
+ edge_map = np.pad(edge_map, pad_width=((1,1),(1,1)), mode='constant')
1341
+ 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)
1342
+ eight_direc = np.roll(np.roll(edge_map, 1, 1), 1, 0) + np.roll(np.roll(edge_map, 1, 1), -1, 0) + \
1343
+ np.roll(np.roll(edge_map, -1, 1), 1, 0) + np.roll(np.roll(edge_map, -1, 1), -1, 0)
1344
+ eight_direc = (eight_direc + four_direc)[1:-1,1:-1]
1345
+ pos_over[eight_direc > 2] = 0
1346
+ T_junction_maps[eight_direc > 2] = 1
1347
+ _, edge_label = cv2.connectedComponents(pos_over.astype(np.uint8), connectivity=8)
1348
+ edge_label = np.pad(edge_label, 1, mode='constant')
1349
+
1350
+ return edge_label
1351
+
1352
+
1353
+
1354
+ def max_size(mat, value=0):
1355
+ if not (mat and mat[0]): return (0, 0)
1356
+ it = iter(mat)
1357
+ prev = [(el==value) for el in next(it)]
1358
+ max_size = max_rectangle_size(prev)
1359
+ for row in it:
1360
+ hist = [(1+h) if el == value else 0 for h, el in zip(prev, row)]
1361
+ max_size = max(max_size, max_rectangle_size(hist), key=get_area)
1362
+ prev = hist
1363
+ return max_size
1364
+
1365
+ def max_rectangle_size(histogram):
1366
+ Info = namedtuple('Info', 'start height')
1367
+ stack = []
1368
+ top = lambda: stack[-1]
1369
+ max_size = (0, 0) # height, width of the largest rectangle
1370
+ pos = 0 # current position in the histogram
1371
+ for pos, height in enumerate(histogram):
1372
+ start = pos # position where rectangle starts
1373
+ while True:
1374
+ if not stack or height > top().height:
1375
+ stack.append(Info(start, height)) # push
1376
+ if stack and height < top().height:
1377
+ max_size = max(max_size, (top().height, (pos-top().start)),
1378
+ key=get_area)
1379
+ start, _ = stack.pop()
1380
+ continue
1381
+ break # height == top().height goes here
1382
+
1383
+ pos += 1
1384
+ for start, height in stack:
1385
+ max_size = max(max_size, (height, (pos-start)),
1386
+ key=get_area)
1387
+
1388
+ return max_size
1389
+
1390
+ def get_area(size):
1391
+ return reduce(mul, size)
1392
+
1393
+ def find_anchors(matrix):
1394
+ matrix = [[*x] for x in matrix]
1395
+ mh, mw = max_size(matrix)
1396
+ matrix = np.array(matrix)
1397
+ # element = np.zeros((mh, mw))
1398
+ for i in range(matrix.shape[0] + 1 - mh):
1399
+ for j in range(matrix.shape[1] + 1 - mw):
1400
+ if matrix[i:i + mh, j:j + mw].max() == 0:
1401
+ return i, i + mh, j, j + mw
1402
+
1403
+ def find_largest_rect(dst_img, bg_color=(128, 128, 128)):
1404
+ valid = np.any(dst_img[..., :3] != bg_color, axis=-1)
1405
+ dst_h, dst_w = dst_img.shape[:2]
1406
+ ret, labels = cv2.connectedComponents(np.uint8(valid == False))
1407
+ red_mat = np.zeros_like(labels)
1408
+ # denoise
1409
+ for i in range(1, np.max(labels)+1, 1):
1410
+ x, y, w, h = cv2.boundingRect(np.uint8(labels==i))
1411
+ if x == 0 or (x+w) == dst_h or y == 0 or (y+h) == dst_w:
1412
+ red_mat[labels==i] = 1
1413
+ # crop
1414
+ t, b, l, r = find_anchors(red_mat)
1415
+
1416
+ return t, b, l, r