zhigangjiang/3D-Room-Layout-Estimation_LGT-Net
160
1"""2@date: 2021/06/193@description:4"""5 6import matplotlib.pyplot as plt7import cv28import numpy as np9from utils.conversion import uv2pixel10from utils.boundary import corners2boundary, corners2boundaries, find_peaks, connect_corners_uv, get_object_cor, \11 visibility_corners12 13 14def draw_boundary(pano_img, corners: np.ndarray = None, boundary: np.ndarray = None, draw_corners=True, show=False,15 step=0.01, length=None, boundary_color=None, marker_color=None, title=None, visible=True):16 if marker_color is None:17 marker_color = [0, 0, 1]18 if boundary_color is None:19 boundary_color = [0, 1, 0]20 21 assert corners is not None or boundary is not None, "corners or boundary error"22 23 shape = sorted(pano_img.shape)24 assert len(shape) > 1, "pano_img shape error"25 w = shape[-1]26 h = shape[-2]27 28 pano_img = pano_img.copy()29 if (corners is not None and len(corners) > 2) or \30 (boundary is not None and len(boundary) > 2):31 if isinstance(boundary_color, list) or isinstance(boundary_color, np.array):32 if boundary is None:33 boundary = corners2boundary(corners, step, length, visible)34 35 boundary = uv2pixel(boundary, w, h)36 pano_img[boundary[:, 1], boundary[:, 0]] = boundary_color37 pano_img[np.clip(boundary[:, 1] + 1, 0, h - 1), boundary[:, 0]] = boundary_color38 pano_img[np.clip(boundary[:, 1] - 1, 0, h - 1), boundary[:, 0]] = boundary_color39 40 if pano_img.shape[1] > 512:41 pano_img[np.clip(boundary[:, 1] + 1, 0, h - 1), np.clip(boundary[:, 0] + 1, 0, w - 1)] = boundary_color42 pano_img[np.clip(boundary[:, 1] + 1, 0, h - 1), np.clip(boundary[:, 0] - 1, 0, w - 1)] = boundary_color43 pano_img[np.clip(boundary[:, 1] - 1, 0, h - 1), np.clip(boundary[:, 0] + 1, 0, w - 1)] = boundary_color44 pano_img[np.clip(boundary[:, 1] - 1, 0, h - 1), np.clip(boundary[:, 0] - 1, 0, w - 1)] = boundary_color45 46 pano_img[boundary[:, 1], np.clip(boundary[:, 0] + 1, 0, w - 1)] = boundary_color47 pano_img[boundary[:, 1], np.clip(boundary[:, 0] - 1, 0, w - 1)] = boundary_color48 49 if corners is not None and draw_corners:50 if visible:51 corners = visibility_corners(corners)52 corners = uv2pixel(corners, w, h)53 for corner in corners:54 cv2.drawMarker(pano_img, tuple(corner), marker_color, markerType=0, markerSize=10, thickness=2)55 56 if show:57 plt.figure(figsize=(10, 5))58 if title is not None:59 plt.title(title)60 61 plt.axis('off')62 plt.imshow(pano_img)63 plt.show()64 65 return pano_img66 67 68def draw_boundaries(pano_img, corners_list: list = None, boundary_list: list = None, draw_corners=True, show=False,69 step=0.01, length=None, boundary_color=None, marker_color=None, title=None, ratio=None, visible=True):70 """71 72 :param visible:73 :param pano_img:74 :param corners_list:75 :param boundary_list:76 :param draw_corners:77 :param show:78 :param step:79 :param length:80 :param boundary_color: RGB color81 :param marker_color: RGB color82 :param title:83 :param ratio: ceil_height/camera_height84 :return:85 """86 assert corners_list is not None or boundary_list is not None, "corners_list or boundary_list error"87 88 if corners_list is not None:89 if ratio is not None and len(corners_list) == 1:90 corners_list = corners2boundaries(ratio, corners_uv=corners_list[0], step=None, visible=visible)91 92 for i, corners in enumerate(corners_list):93 pano_img = draw_boundary(pano_img, corners=corners, draw_corners=draw_corners,94 show=show if i == len(corners_list) - 1 else False,95 step=step, length=length, boundary_color=boundary_color, marker_color=marker_color,96 title=title, visible=visible)97 elif boundary_list is not None:98 if ratio is not None and len(boundary_list) == 1:99 boundary_list = corners2boundaries(ratio, corners_uv=boundary_list[0], step=None, visible=visible)100 101 for i, boundary in enumerate(boundary_list):102 pano_img = draw_boundary(pano_img, boundary=boundary, draw_corners=draw_corners,103 show=show if i == len(boundary_list) - 1 else False,104 step=step, length=length, boundary_color=boundary_color, marker_color=marker_color,105 title=title, visible=visible)106 107 return pano_img108 109 110def draw_object(pano_img, heat_maps, size, depth, window_width=15, show=False):111 # window, door, opening112 colors = [[1, 0, 0], [1, 1, 0], [0, 0, 1]]113 for i, heat_map in enumerate(heat_maps):114 pk_u_s, _ = find_peaks(heat_map, size=window_width*2+1)115 for pk_u in pk_u_s:116 uv, xyz = get_object_cor(depth, size, center_u=pk_u, patch_num=len(heat_map))117 118 bottom_poly = connect_corners_uv(uv[0], uv[1], length=pano_img.shape[1])119 top_poly = connect_corners_uv(uv[2], uv[3], length=pano_img.shape[1])[::-1]120 121 bottom_max_index = bottom_poly[..., 0].argmax()122 if bottom_max_index != len(bottom_poly)-1:123 top_max_index = top_poly[..., 0].argmax()124 poly1 = np.concatenate([bottom_poly[:bottom_max_index+1], top_poly[top_max_index:]])125 poly1 = uv2pixel(poly1, w=pano_img.shape[1], h=pano_img.shape[0])126 poly1 = poly1[:, None, :]127 128 poly2 = np.concatenate([bottom_poly[bottom_max_index+1:], top_poly[:top_max_index]])129 poly2 = uv2pixel(poly2, w=pano_img.shape[1], h=pano_img.shape[0])130 poly2 = poly2[:, None, :]131 132 poly = [poly1, poly2]133 else:134 poly = np.concatenate([bottom_poly, top_poly])135 poly = uv2pixel(poly, w=pano_img.shape[1], h=pano_img.shape[0])136 poly = poly[:, None, :]137 poly = [poly]138 139 cv2.drawContours(pano_img, poly, -1, colors[i], 1)140 #141 # boundary_center_xyz = uv2xyz(np.array([pk_u, pk_v]))142 #143 # l_b_xyz =144 if show:145 plt.imshow(pano_img)146 plt.show()147 148 149if __name__ == '__main__':150 from visualization.floorplan import draw_floorplan151 from utils.conversion import uv2xyz152 153 pano_img = np.zeros([512, 1024, 3])154 corners = np.array([[0.2, 0.7],155 [0.4, 0.7],156 [0.3, 0.6],157 [0.6, 0.6],158 [0.8, 0.7]])159 # draw_boundary(pano_img, corners, show=True)160 draw_boundaries(pano_img, corners_list=[corners], show=True, length=1024, ratio=1.2)161 draw_floorplan(uv2xyz(corners)[..., ::2], show=True, marker_color=None, center_color=0.8)