Hugging Face's logo

Spaces:
zhigangjiang
/
3D-Room-Layout-Estimation_LGT-Net
Running

App Files Community
2
3D-Room-Layout-Estimation_LGT-Net / utils /boundary.py
zhigangjiang's picture
zhigangjiang
Update utils/boundary.py
c4adcb2
raw history blame contribute delete
No virus
14.4 kB
"""
@date: 2021/06/19
@description:
"""
import math
import functools
from scipy import stats
from scipy.ndimage.filters import maximum_filter
import numpy as np
from typing import List
from utils.conversion import uv2xyz, xyz2uv, depth2xyz, uv2pixel, depth2uv, pixel2uv, xyz2pixel, uv2lonlat
from utils.visibility_polygon import calc_visible_polygon
def connect_corners_uv(uv1: np.ndarray, uv2: np.ndarray, length=256) -> np.ndarray:
"""
:param uv1: [u, v]
:param uv2: [u, v]
:param length: Fix the total length in pixel coordinates
:return:
"""
# why -0.5? Check out the uv2Pixel function
p_u1 = uv1[0] * length - 0.5
p_u2 = uv2[0] * length - 0.5
if abs(p_u1 - p_u2) < length / 2:
start = np.ceil(min(p_u1, p_u2))
p = max(p_u1, p_u2)
end = np.floor(p)
if end == np.ceil(p):
end = end - 1
else:
start = np.ceil(max(p_u1, p_u2))
p = min(p_u1, p_u2) + length
end = np.floor(p)
if end == np.ceil(p):
end = end - 1
p_us = (np.arange(start, end + 1) % length).astype(np.float64)
if len(p_us) == 0:
return None
us = (p_us + 0.5) / length # why +0.5? Check out the uv2Pixel function
plan_y = boundary_type(np.array([uv1, uv2]))
xyz1 = uv2xyz(np.array(uv1), plan_y)
xyz2 = uv2xyz(np.array(uv2), plan_y)
x1 = xyz1[0]
z1 = xyz1[2]
x2 = xyz2[0]
z2 = xyz2[2]
d_x = x2 - x1
d_z = z2 - z1
lon_s = (us - 0.5) * 2 * np.pi
k = np.tan(lon_s)
ps = (k * z1 - x1) / (d_x - k * d_z)
cs = np.sqrt((z1 + ps * d_z) ** 2 + (x1 + ps * d_x) ** 2)
lats = np.arctan2(plan_y, cs)
vs = lats / np.pi + 0.5
uv = np.stack([us, vs], axis=-1)
if start == end:
return uv[0:1]
return uv
def connect_corners_xyz(uv1: np.ndarray, uv2: np.ndarray, step=0.01) -> np.ndarray:
"""
:param uv1: [u, v]
:param uv2: [u, v]
:param step: Fixed step size in xyz coordinates
:return:
"""
plan_y = boundary_type(np.array([uv1, uv2]))
xyz1 = uv2xyz(np.array(uv1), plan_y)
xyz2 = uv2xyz(np.array(uv2), plan_y)
vec = xyz2 - xyz1
norm = np.linalg.norm(vec, ord=2)
direct = vec / norm
xyz = np.array([xyz1 + direct * dis for dis in np.linspace(0, norm, int(norm / step))])
if len(xyz) == 0:
xyz = np.array([xyz2])
uv = xyz2uv(xyz)
return uv
def connect_corners(uv1: np.ndarray, uv2: np.ndarray, step=0.01, length=None) -> np.ndarray:
"""
:param uv1: [u, v]
:param uv2: [u, v]
:param step:
:param length:
:return: [[u1, v1], [u2, v2]....] if length!=None，length of return result = length
"""
if length is not None:
uv = connect_corners_uv(uv1, uv2, length)
elif step is not None:
uv = connect_corners_xyz(uv1, uv2, step)
else:
uv = np.array([uv1])
return uv
def visibility_corners(corners):
plan_y = boundary_type(corners)
xyz = uv2xyz(corners, plan_y)
xz = xyz[:, ::2]
xz = calc_visible_polygon(center=np.array([0, 0]), polygon=xz, show=False)
xyz = np.insert(xz, 1, plan_y, axis=1)
output = xyz2uv(xyz).astype(np.float32)
return output
def corners2boundary(corners: np.ndarray, step=0.01, length=None, visible=True) -> np.ndarray:
"""
When there is occlusion, even if the length is fixed, the final output length may be greater than the given length,
which is more defined as the fixed step size under UV
:param length:
:param step:
:param corners: [[u1, v1], [u2, v2]....]
:param visible:
:return: [[u1, v1], [u2, v2]....] if length!=None，length of return result = length
"""
assert step is not None or length is not None, "the step and length parameters cannot be null at the same time"
if len(corners) < 3:
return corners
if visible:
corners = visibility_corners(corners)
n_con = len(corners)
boundary = None
for j in range(n_con):
uv = connect_corners(corners[j], corners[(j + 1) % n_con], step, length)
if uv is None:
continue
if boundary is None:
boundary = uv
else:
boundary = np.concatenate((boundary, uv))
boundary = np.roll(boundary, -boundary.argmin(axis=0)[0], axis=0)
output_polygon = []
for i, p in enumerate(boundary):
q = boundary[(i + 1) % len(boundary)]
if int(p[0] * 10000) == int(q[0] * 10000):
continue
output_polygon.append(p)
output_polygon = np.array(output_polygon, dtype=np.float32)
return output_polygon
def corners2boundaries(ratio: float, corners_xyz: np.ndarray = None, corners_uv: np.ndarray = None, step=0.01,
length=None, visible=True):
"""
When both step and length are None, corners are also returned
:param ratio:
:param corners_xyz:
:param corners_uv:
:param step:
:param length:
:param visible:
:return: floor_boundary, ceil_boundary
"""
if corners_xyz is None:
plan_y = boundary_type(corners_uv)
xyz = uv2xyz(corners_uv, plan_y)
floor_xyz = xyz.copy()
ceil_xyz = xyz.copy()
if plan_y > 0:
ceil_xyz[:, 1] *= -ratio
else:
floor_xyz[:, 1] /= -ratio
else:
floor_xyz = corners_xyz.copy()
ceil_xyz = corners_xyz.copy()
if corners_xyz[0][1] > 0:
ceil_xyz[:, 1] *= -ratio
else:
floor_xyz[:, 1] /= -ratio
floor_uv = xyz2uv(floor_xyz)
ceil_uv = xyz2uv(ceil_xyz)
if step is None and length is None:
return floor_uv, ceil_uv
floor_boundary = corners2boundary(floor_uv, step, length, visible)
ceil_boundary = corners2boundary(ceil_uv, step, length, visible)
return floor_boundary, ceil_boundary
def depth2boundary(depth: np.array, step=0.01, length=None,):
xyz = depth2xyz(depth)
uv = xyz2uv(xyz)
return corners2boundary(uv, step, length, visible=False)
def depth2boundaries(ratio: float, depth: np.array, step=0.01, length=None,):
"""
:param ratio:
:param depth:
:param step:
:param length:
:return: floor_boundary, ceil_boundary
"""
xyz = depth2xyz(depth)
return corners2boundaries(ratio, corners_xyz=xyz, step=step, length=length, visible=False)
def boundary_type(corners: np.ndarray) -> int:
"""
Returns the boundary type that also represents the projection plane
:param corners:
:return:
"""
if is_ceil_boundary(corners):
plan_y = -1
elif is_floor_boundary(corners):
plan_y = 1
else:
# An intersection occurs and an exception is considered
assert False, 'corners error!'
return plan_y
def is_normal_layout(boundaries: List[np.array]):
if len(boundaries) != 2:
print("boundaries length must be 2!")
return False
if boundary_type(boundaries[0]) != -1:
print("ceil boundary error!")
return False
if boundary_type(boundaries[1]) != 1:
print("floor boundary error!")
return False
return True
def is_ceil_boundary(corners: np.ndarray) -> bool:
m = corners[..., 1].max()
return m < 0.5
def is_floor_boundary(corners: np.ndarray) -> bool:
m = corners[..., 1].min()
return m > 0.5
@functools.lru_cache()
def get_gauss_map(sigma=1.5, width=5):
x = np.arange(width*2 + 1) - width
y = stats.norm(0, sigma).pdf(x)
y = y / y.max()
return y
def get_heat_map(u_s, patch_num=256, sigma=2, window_width=15, show=False):
"""
:param window_width:
:param sigma:
:param u_s: [u1, u2, u3, ...]
:param patch_num
:param show
:return:
"""
pixel_us = uv2pixel(u_s, w=patch_num, axis=0)
gauss_map = get_gauss_map(sigma, window_width)
heat_map_all = []
for u in pixel_us:
heat_map = np.zeros(patch_num, dtype=np.float32)
left = u-window_width
right = u+window_width+1
offset = 0
if left < 0:
offset = left
elif right > patch_num:
offset = right - patch_num
left = left - offset
right = right - offset
heat_map[left:right] = gauss_map
if offset != 0:
heat_map = np.roll(heat_map, offset)
heat_map_all.append(heat_map)
heat_map_all = np.array(heat_map_all).max(axis=0)
if show:
import matplotlib.pyplot as plt
plt.imshow(heat_map_all[None].repeat(50, axis=0))
plt.show()
return heat_map_all
def find_peaks(signal, size=15*2+1, min_v=0.05, N=None):
# code from HorizonNet: https://github.com/sunset1995/HorizonNet/blob/master/inference.py
max_v = maximum_filter(signal, size=size, mode='wrap')
pk_loc = np.where(max_v == signal)[0]
pk_loc = pk_loc[signal[pk_loc] > min_v]
if N is not None:
order = np.argsort(-signal[pk_loc])
pk_loc = pk_loc[order[:N]]
pk_loc = pk_loc[np.argsort(pk_loc)]
return pk_loc, signal[pk_loc]
def get_object_cor(depth, size, center_u, patch_num=256):
width_u = size[0, center_u]
height_v = size[1, center_u]
boundary_v = size[2, center_u]
center_boundary_v = depth2uv(depth[center_u:center_u + 1])[0, 1]
center_bottom_v = center_boundary_v - boundary_v
center_top_v = center_bottom_v - height_v
base_v = center_boundary_v - 0.5
assert base_v > 0
center_u = pixel2uv(np.array([center_u]), w=patch_num, h=patch_num // 2, axis=0)[0]
center_boundary_uv = np.array([center_u, center_boundary_v])
center_bottom_uv = np.array([center_u, center_bottom_v])
center_top_uv = np.array([center_u, center_top_v])
left_u = center_u - width_u / 2
right_u = center_u + width_u / 2
left_u = 1 + left_u if left_u < 0 else left_u
right_u = right_u - 1 if right_u > 1 else right_u
pixel_u = uv2pixel(np.array([left_u, right_u]), w=patch_num, h=patch_num // 2, axis=0)
left_pixel_u = pixel_u[0]
right_pixel_u = pixel_u[1]
left_boundary_v = depth2uv(depth[left_pixel_u:left_pixel_u + 1])[0, 1]
right_boundary_v = depth2uv(depth[right_pixel_u:right_pixel_u + 1])[0, 1]
left_boundary_uv = np.array([left_u, left_boundary_v])
right_boundary_uv = np.array([right_u, right_boundary_v])
xyz = uv2xyz(np.array([left_boundary_uv, right_boundary_uv, center_boundary_uv]))
left_boundary_xyz = xyz[0]
right_boundary_xyz = xyz[1]
# need align
center_boundary_xyz = xyz[2]
center_bottom_xyz = uv2xyz(np.array([center_bottom_uv]))[0]
center_top_xyz = uv2xyz(np.array([center_top_uv]))[0]
center_boundary_norm = np.linalg.norm(center_boundary_xyz[::2])
center_bottom_norm = np.linalg.norm(center_bottom_xyz[::2])
center_top_norm = np.linalg.norm(center_top_xyz[::2])
center_bottom_xyz = center_bottom_xyz * center_boundary_norm / center_bottom_norm
center_top_xyz = center_top_xyz * center_boundary_norm / center_top_norm
left_bottom_xyz = left_boundary_xyz.copy()
left_bottom_xyz[1] = center_bottom_xyz[1]
right_bottom_xyz = right_boundary_xyz.copy()
right_bottom_xyz[1] = center_bottom_xyz[1]
left_top_xyz = left_boundary_xyz.copy()
left_top_xyz[1] = center_top_xyz[1]
right_top_xyz = right_boundary_xyz.copy()
right_top_xyz[1] = center_top_xyz[1]
uv = xyz2uv(np.array([left_bottom_xyz, right_bottom_xyz, left_top_xyz, right_top_xyz]))
left_bottom_uv = uv[0]
right_bottom_uv = uv[1]
left_top_uv = uv[2]
right_top_uv = uv[3]
return [left_bottom_uv, right_bottom_uv, left_top_uv, right_top_uv], \
[left_bottom_xyz, right_bottom_xyz, left_top_xyz, right_top_xyz]
def layout2depth(boundaries: List[np.array], return_mask=False, show=False, camera_height=1.6):
"""
:param camera_height:
:param boundaries: [[[u_f1, v_f2], [u_f2, v_f2],...], [[u_c1, v_c2], [u_c2, v_c2]]]
:param return_mask:
:param show:
:return:
"""
# code from HorizonNet: https://github.com/sunset1995/HorizonNet/blob/master/eval_general.py
w = len(boundaries[0])
h = w//2
# Convert corners to per-column boundary first
# Up -pi/2, Down pi/2
vf = uv2lonlat(boundaries[0])
vc = uv2lonlat(boundaries[1])
vc = vc[None, :, 1] # [1, w]
vf = vf[None, :, 1] # [1, w]
assert (vc > 0).sum() == 0
assert (vf < 0).sum() == 0
# Per-pixel v coordinate (vertical angle)
vs = ((np.arange(h) + 0.5) / h - 0.5) * np.pi
vs = np.repeat(vs[:, None], w, axis=1) # [h, w]
# Floor-plane to depth
floor_h = camera_height
floor_d = np.abs(floor_h / np.sin(vs))
# wall to camera distance on horizontal plane at cross camera center
cs = floor_h / np.tan(vf)
# Ceiling-plane to depth
ceil_h = np.abs(cs * np.tan(vc)) # [1, w]
ceil_d = np.abs(ceil_h / np.sin(vs)) # [h, w]
# Wall to depth
wall_d = np.abs(cs / np.cos(vs)) # [h, w]
# Recover layout depth
floor_mask = (vs > vf)
ceil_mask = (vs < vc)
wall_mask = (~floor_mask) & (~ceil_mask)
depth = np.zeros([h, w], np.float32) # [h, w]
depth[floor_mask] = floor_d[floor_mask]
depth[ceil_mask] = ceil_d[ceil_mask]
depth[wall_mask] = wall_d[wall_mask]
assert (depth == 0).sum() == 0
if return_mask:
return depth, floor_mask, ceil_mask, wall_mask
if show:
import matplotlib.pyplot as plt
plt.imshow(depth)
plt.show()
return depth
def calc_rotation(corners: np.ndarray):
xz = uv2xyz(corners)[..., 0::2]
max_norm = -1
max_v = None
for i in range(len(xz)):
p_c = xz[i]
p_n = xz[(i + 1) % len(xz)]
v_cn = p_n - p_c
v_norm = np.linalg.norm(v_cn)
if v_norm > max_norm:
max_norm = v_norm
max_v = v_cn
# v<-----------|o
# | | |
# | ----|----z |
# | | |
# | x \|/
# |------------u
# It is required that the vector be aligned on the x-axis, z equals y, and x is still x.
# In floorplan, x is displayed as the x-coordinate and z as the y-coordinate
rotation = np.arctan2(max_v[1], max_v[0])
return rotation
if __name__ == '__main__':
corners = np.array([[0.2, 0.7],
[0.4, 0.7],
[0.3, 0.6],
[0.6, 0.6],
[0.8, 0.7]])
get_heat_map(u=corners[..., 0], show=True, sigma=2, width=15)
pass