-
Notifications
You must be signed in to change notification settings - Fork 23
/
vis_layout.py
169 lines (154 loc) · 6.43 KB
/
vis_layout.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
import json
import numpy as np
import open3d as o3d
from PIL import Image
from scipy.signal import correlate2d
from scipy.ndimage import shift
from lib.misc.post_proc import np_coor2xy, np_coorx2u, np_coory2v
from eval_layout import layout_2_depth
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--img', required=True,
help='Image texture in equirectangular format')
parser.add_argument('--layout', required=True,
help='Txt or json file containing layout corners (cor_id)')
parser.add_argument('--out')
parser.add_argument('--no_vis', action='store_true')
parser.add_argument('--show_ceiling', action='store_true',
help='Rendering ceiling (skip by default)')
parser.add_argument('--ignore_floor', action='store_true',
help='Skip rendering floor')
parser.add_argument('--ignore_wall', action='store_true',
help='Skip rendering wall')
parser.add_argument('--ignore_wireframe', action='store_true',
help='Skip rendering wireframe')
args = parser.parse_args()
if not args.out and args.no_vis:
print('You may want to export (via --out) or visualize (without --vis)')
import sys; sys.exit()
# Reading source (texture img, cor_id txt)
equirect_texture = np.array(Image.open(args.img))
H, W = equirect_texture.shape[:2]
if args.layout.endswith('json'):
with open(args.layout) as f:
inferenced_result = json.load(f)
cor_id = np.array(inferenced_result['uv'], np.float32)
cor_id[:, 0] *= W
cor_id[:, 1] *= H
else:
cor_id = np.loadtxt(args.layout).astype(np.float32)
# Convert corners to layout
depth, floor_mask, ceil_mask, wall_mask = layout_2_depth(cor_id, H, W, return_mask=True)
coorx, coory = np.meshgrid(np.arange(W), np.arange(H))
us = np_coorx2u(coorx, W)
vs = np_coory2v(coory, H)
zs = depth * np.sin(vs)
cs = depth * np.cos(vs)
xs = cs * np.sin(us)
ys = -cs * np.cos(us)
# Aggregate mask
mask = np.ones_like(floor_mask)
if args.ignore_floor:
mask &= ~floor_mask
if not args.show_ceiling:
mask &= ~ceil_mask
if args.ignore_wall:
mask &= ~wall_mask
# Prepare ply's points and faces
xyzrgb = np.concatenate([
xs[...,None], ys[...,None], zs[...,None],
equirect_texture], -1)
xyzrgb = np.concatenate([xyzrgb, xyzrgb[:,[0]]], 1)
mask = np.concatenate([mask, mask[:,[0]]], 1)
lo_tri_template = np.array([
[0, 0, 0],
[0, 1, 0],
[0, 1, 1]])
up_tri_template = np.array([
[0, 0, 0],
[0, 1, 1],
[0, 0, 1]])
ma_tri_template = np.array([
[0, 0, 0],
[0, 1, 1],
[0, 1, 0]])
lo_mask = (correlate2d(mask, lo_tri_template, mode='same') == 3)
up_mask = (correlate2d(mask, up_tri_template, mode='same') == 3)
ma_mask = (correlate2d(mask, ma_tri_template, mode='same') == 3) & (~lo_mask) & (~up_mask)
ref_mask = (
lo_mask | (correlate2d(lo_mask, np.flip(lo_tri_template, (0,1)), mode='same') > 0) |\
up_mask | (correlate2d(up_mask, np.flip(up_tri_template, (0,1)), mode='same') > 0) |\
ma_mask | (correlate2d(ma_mask, np.flip(ma_tri_template, (0,1)), mode='same') > 0)
)
points = xyzrgb[ref_mask]
ref_id = np.full(ref_mask.shape, -1, np.int32)
ref_id[ref_mask] = np.arange(ref_mask.sum())
faces_lo_tri = np.stack([
ref_id[lo_mask],
ref_id[shift(lo_mask, [1, 0], cval=False, order=0)],
ref_id[shift(lo_mask, [1, 1], cval=False, order=0)],
], 1)
faces_up_tri = np.stack([
ref_id[up_mask],
ref_id[shift(up_mask, [1, 1], cval=False, order=0)],
ref_id[shift(up_mask, [0, 1], cval=False, order=0)],
], 1)
faces_ma_tri = np.stack([
ref_id[ma_mask],
ref_id[shift(ma_mask, [1, 0], cval=False, order=0)],
ref_id[shift(ma_mask, [0, 1], cval=False, order=0)],
], 1)
faces = np.concatenate([faces_lo_tri, faces_up_tri, faces_ma_tri])
# Dump results ply
if args.out:
ply_header = '\n'.join([
'ply',
'format ascii 1.0',
f'element vertex {len(points):d}',
'property float x',
'property float y',
'property float z',
'property uchar red',
'property uchar green',
'property uchar blue',
f'element face {len(faces):d}',
'property list uchar int vertex_indices',
'end_header',
])
with open(args.out, 'w') as f:
f.write(ply_header)
f.write('\n')
for x, y, z, r, g, b in points:
f.write(f'{x:.2f} {y:.2f} {z:.2f} {r:.0f} {g:.0f} {b:.0f}\n')
for i, j, k in faces:
f.write(f'3 {i:d} {j:d} {k:d}\n')
if not args.no_vis:
mesh = o3d.geometry.TriangleMesh()
mesh.vertices = o3d.utility.Vector3dVector(points[:, :3])
mesh.vertex_colors = o3d.utility.Vector3dVector(points[:, 3:] / 255.)
mesh.triangles = o3d.utility.Vector3iVector(faces)
draw_geometries = [mesh]
# Show wireframe
if not args.ignore_wireframe:
# Convert cor_id to 3d xyz
N = len(cor_id) // 2
floor_z = -1.6
floor_xy = np_coor2xy(cor_id[1::2], floor_z, W, H, floorW=1, floorH=1)
c = np.sqrt((floor_xy**2).sum(1))
v = np_coory2v(cor_id[0::2, 1], H)
ceil_z = (c * np.tan(v)).mean()
# Prepare wireframe in open3d
assert N == len(floor_xy)
wf_points = [[x, y, floor_z] for x, y in floor_xy] +\
[[x, y, ceil_z] for x, y in floor_xy]
wf_lines = [[i, (i+1)%N] for i in range(N)] +\
[[i+N, (i+1)%N+N] for i in range(N)] +\
[[i, i+N] for i in range(N)]
wf_colors = [[1, 0, 0] for i in range(len(wf_lines))]
wf_line_set = o3d.geometry.LineSet()
wf_line_set.points = o3d.utility.Vector3dVector(wf_points)
wf_line_set.lines = o3d.utility.Vector2iVector(wf_lines)
wf_line_set.colors = o3d.utility.Vector3dVector(wf_colors)
draw_geometries.append(wf_line_set)
o3d.visualization.draw_geometries(draw_geometries, mesh_show_back_face=True)