|
""" |
|
@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: |
|
|
|
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 = 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) |
|
|
