File size: 4,257 Bytes
88b0dcb |
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 |
"""
@Date: 2021/10/06
@description: Use the approach proposed by DuLa-Net
"""
import cv2
import numpy as np
import math
import matplotlib.pyplot as plt
from visualization.floorplan import draw_floorplan
def merge_near(lst, diag):
group = [[0, ]]
for i in range(1, len(lst)):
if lst[i] - np.mean(group[-1]) < diag * 0.02:
group[-1].append(lst[i])
else:
group.append([lst[i], ])
if len(group) == 1:
group = [lst[0], lst[-1]]
else:
group = [int(np.mean(x)) for x in group]
return group
def fit_layout_old(floor_xz, need_cube=False, show=False, block_eps=0.05):
show_radius = np.linalg.norm(floor_xz, axis=-1).max()
side_l = 512
floorplan = draw_floorplan(xz=floor_xz, show_radius=show_radius, show=show, scale=1, side_l=side_l).astype(np.uint8)
center = np.array([side_l / 2, side_l / 2])
polys = cv2.findContours(floorplan, 1, 2)
if isinstance(polys, tuple):
if len(polys) == 3:
# opencv 3
polys = list(polys[1])
else:
polys = list(polys[0])
polys.sort(key=lambda x: cv2.contourArea(x), reverse=True)
poly = polys[0]
sub_x, sub_y, w, h = cv2.boundingRect(poly)
floorplan_sub = floorplan[sub_y:sub_y + h, sub_x:sub_x + w]
sub_center = center - np.array([sub_x, sub_y])
polys = cv2.findContours(floorplan_sub, 1, 2)
if isinstance(polys, tuple):
if len(polys) == 3:
polys = polys[1]
else:
polys = polys[0]
poly = polys[0]
epsilon = 0.005 * cv2.arcLength(poly, True)
poly = cv2.approxPolyDP(poly, epsilon, True)
x_lst = [0, ]
y_lst = [0, ]
for i in range(len(poly)):
p1 = poly[i][0]
p2 = poly[(i + 1) % len(poly)][0]
if (p2[0] - p1[0]) == 0:
slope = 10
else:
slope = abs((p2[1] - p1[1]) / (p2[0] - p1[0]))
if slope <= 1:
s = int((p1[1] + p2[1]) / 2)
y_lst.append(s)
elif slope > 1:
s = int((p1[0] + p2[0]) / 2)
x_lst.append(s)
x_lst.append(floorplan_sub.shape[1])
y_lst.append(floorplan_sub.shape[0])
x_lst.sort()
y_lst.sort()
diag = math.sqrt(math.pow(floorplan_sub.shape[1], 2) + math.pow(floorplan_sub.shape[0], 2))
x_lst = merge_near(x_lst, diag)
y_lst = merge_near(y_lst, diag)
if need_cube and len(x_lst) > 2:
x_lst = [x_lst[0], x_lst[-1]]
if need_cube and len(y_lst) > 2:
y_lst = [y_lst[0], y_lst[-1]]
ans = np.zeros((floorplan_sub.shape[0], floorplan_sub.shape[1]))
for i in range(len(x_lst) - 1):
for j in range(len(y_lst) - 1):
sample = floorplan_sub[y_lst[j]:y_lst[j + 1], x_lst[i]:x_lst[i + 1]]
score = 0 if sample.size == 0 else sample.mean()
if score >= 0.3:
ans[y_lst[j]:y_lst[j + 1], x_lst[i]:x_lst[i + 1]] = 1
pred = np.uint8(ans)
pred_polys = cv2.findContours(pred, 1, 3)
if isinstance(pred_polys, tuple):
if len(pred_polys) == 3:
pred_polys = pred_polys[1]
else:
pred_polys = pred_polys[0]
polygon = [(p[0][1], p[0][0]) for p in pred_polys[0][::-1]]
v = np.array([p[0] + sub_y for p in polygon])
u = np.array([p[1] + sub_x for p in polygon])
# side_l
# v<-----------|o
# | | |
# | ----|----z | side_l
# | | |
# | x \|/
# |------------u
side_l = floorplan.shape[0]
pred_xz = np.concatenate((u[:, np.newaxis] - side_l // 2, side_l // 2 - v[:, np.newaxis]), axis=1)
pred_xz = pred_xz * show_radius / (side_l // 2)
if show:
draw_floorplan(pred_xz, show_radius=show_radius, show=show)
return pred_xz
if __name__ == '__main__':
from utils.conversion import uv2xyz
pano_img = np.zeros([512, 1024, 3])
corners = np.array([[0.1, 0.7],
[0.4, 0.7],
[0.3, 0.6],
[0.6, 0.6],
[0.8, 0.7]])
xz = uv2xyz(corners)[..., ::2]
draw_floorplan(xz, show=True, marker_color=None, center_color=0.8)
xz = fit_layout_old(xz)
draw_floorplan(xz, show=True, marker_color=None, center_color=0.8)
|