File size: 6,942 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 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 |
"""
@date: 2021/06/19
@description:
"""
import matplotlib.pyplot as plt
import cv2
import numpy as np
from utils.conversion import uv2pixel
from utils.boundary import corners2boundary, corners2boundaries, find_peaks, connect_corners_uv, get_object_cor, \
visibility_corners
def draw_boundary(pano_img, corners: np.ndarray = None, boundary: np.ndarray = None, draw_corners=True, show=False,
step=0.01, length=None, boundary_color=None, marker_color=None, title=None, visible=True):
if marker_color is None:
marker_color = [0, 0, 1]
if boundary_color is None:
boundary_color = [0, 1, 0]
assert corners is not None or boundary is not None, "corners or boundary error"
shape = sorted(pano_img.shape)
assert len(shape) > 1, "pano_img shape error"
w = shape[-1]
h = shape[-2]
pano_img = pano_img.copy()
if (corners is not None and len(corners) > 2) or \
(boundary is not None and len(boundary) > 2):
if isinstance(boundary_color, list) or isinstance(boundary_color, np.array):
if boundary is None:
boundary = corners2boundary(corners, step, length, visible)
boundary = uv2pixel(boundary, w, h)
pano_img[boundary[:, 1], boundary[:, 0]] = boundary_color
pano_img[np.clip(boundary[:, 1] + 1, 0, h - 1), boundary[:, 0]] = boundary_color
pano_img[np.clip(boundary[:, 1] - 1, 0, h - 1), boundary[:, 0]] = boundary_color
if pano_img.shape[1] > 512:
pano_img[np.clip(boundary[:, 1] + 1, 0, h - 1), np.clip(boundary[:, 0] + 1, 0, w - 1)] = boundary_color
pano_img[np.clip(boundary[:, 1] + 1, 0, h - 1), np.clip(boundary[:, 0] - 1, 0, w - 1)] = boundary_color
pano_img[np.clip(boundary[:, 1] - 1, 0, h - 1), np.clip(boundary[:, 0] + 1, 0, w - 1)] = boundary_color
pano_img[np.clip(boundary[:, 1] - 1, 0, h - 1), np.clip(boundary[:, 0] - 1, 0, w - 1)] = boundary_color
pano_img[boundary[:, 1], np.clip(boundary[:, 0] + 1, 0, w - 1)] = boundary_color
pano_img[boundary[:, 1], np.clip(boundary[:, 0] - 1, 0, w - 1)] = boundary_color
if corners is not None and draw_corners:
if visible:
corners = visibility_corners(corners)
corners = uv2pixel(corners, w, h)
for corner in corners:
cv2.drawMarker(pano_img, tuple(corner), marker_color, markerType=0, markerSize=10, thickness=2)
if show:
plt.figure(figsize=(10, 5))
if title is not None:
plt.title(title)
plt.axis('off')
plt.imshow(pano_img)
plt.show()
return pano_img
def draw_boundaries(pano_img, corners_list: list = None, boundary_list: list = None, draw_corners=True, show=False,
step=0.01, length=None, boundary_color=None, marker_color=None, title=None, ratio=None, visible=True):
"""
:param visible:
:param pano_img:
:param corners_list:
:param boundary_list:
:param draw_corners:
:param show:
:param step:
:param length:
:param boundary_color: RGB color
:param marker_color: RGB color
:param title:
:param ratio: ceil_height/camera_height
:return:
"""
assert corners_list is not None or boundary_list is not None, "corners_list or boundary_list error"
if corners_list is not None:
if ratio is not None and len(corners_list) == 1:
corners_list = corners2boundaries(ratio, corners_uv=corners_list[0], step=None, visible=visible)
for i, corners in enumerate(corners_list):
pano_img = draw_boundary(pano_img, corners=corners, draw_corners=draw_corners,
show=show if i == len(corners_list) - 1 else False,
step=step, length=length, boundary_color=boundary_color, marker_color=marker_color,
title=title, visible=visible)
elif boundary_list is not None:
if ratio is not None and len(boundary_list) == 1:
boundary_list = corners2boundaries(ratio, corners_uv=boundary_list[0], step=None, visible=visible)
for i, boundary in enumerate(boundary_list):
pano_img = draw_boundary(pano_img, boundary=boundary, draw_corners=draw_corners,
show=show if i == len(boundary_list) - 1 else False,
step=step, length=length, boundary_color=boundary_color, marker_color=marker_color,
title=title, visible=visible)
return pano_img
def draw_object(pano_img, heat_maps, size, depth, window_width=15, show=False):
# window, door, opening
colors = [[1, 0, 0], [1, 1, 0], [0, 0, 1]]
for i, heat_map in enumerate(heat_maps):
pk_u_s, _ = find_peaks(heat_map, size=window_width*2+1)
for pk_u in pk_u_s:
uv, xyz = get_object_cor(depth, size, center_u=pk_u, patch_num=len(heat_map))
bottom_poly = connect_corners_uv(uv[0], uv[1], length=pano_img.shape[1])
top_poly = connect_corners_uv(uv[2], uv[3], length=pano_img.shape[1])[::-1]
bottom_max_index = bottom_poly[..., 0].argmax()
if bottom_max_index != len(bottom_poly)-1:
top_max_index = top_poly[..., 0].argmax()
poly1 = np.concatenate([bottom_poly[:bottom_max_index+1], top_poly[top_max_index:]])
poly1 = uv2pixel(poly1, w=pano_img.shape[1], h=pano_img.shape[0])
poly1 = poly1[:, None, :]
poly2 = np.concatenate([bottom_poly[bottom_max_index+1:], top_poly[:top_max_index]])
poly2 = uv2pixel(poly2, w=pano_img.shape[1], h=pano_img.shape[0])
poly2 = poly2[:, None, :]
poly = [poly1, poly2]
else:
poly = np.concatenate([bottom_poly, top_poly])
poly = uv2pixel(poly, w=pano_img.shape[1], h=pano_img.shape[0])
poly = poly[:, None, :]
poly = [poly]
cv2.drawContours(pano_img, poly, -1, colors[i], 1)
#
# boundary_center_xyz = uv2xyz(np.array([pk_u, pk_v]))
#
# l_b_xyz =
if show:
plt.imshow(pano_img)
plt.show()
if __name__ == '__main__':
from visualization.floorplan import draw_floorplan
from utils.conversion import uv2xyz
pano_img = np.zeros([512, 1024, 3])
corners = np.array([[0.2, 0.7],
[0.4, 0.7],
[0.3, 0.6],
[0.6, 0.6],
[0.8, 0.7]])
# draw_boundary(pano_img, corners, show=True)
draw_boundaries(pano_img, corners_list=[corners], show=True, length=1024, ratio=1.2)
draw_floorplan(uv2xyz(corners)[..., ::2], show=True, marker_color=None, center_color=0.8) |