File size: 8,337 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 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 |
"""
@Date: 2021/07/28
@description:
"""
import os
import numpy as np
import cv2
import json
from PIL import Image
from utils.conversion import xyz2uv, pixel2uv
from utils.height import calc_ceil_ratio
def read_image(image_path, shape=None):
if shape is None:
shape = [512, 1024]
img = np.array(Image.open(image_path)).astype(np.float32) / 255
if img.shape[0] != shape[0] or img.shape[1] != shape[1]:
img = cv2.resize(img, dsize=tuple(shape[::-1]), interpolation=cv2.INTER_AREA)
return np.array(img)
def read_label(label_path, data_type='MP3D'):
if data_type == 'MP3D':
with open(label_path, 'r') as f:
label = json.load(f)
point_idx = [one['pointsIdx'][0] for one in label['layoutWalls']['walls']]
camera_height = label['cameraHeight']
room_height = label['layoutHeight']
camera_ceiling_height = room_height - camera_height
ratio = camera_ceiling_height / camera_height
xyz = [one['xyz'] for one in label['layoutPoints']['points']]
assert len(xyz) == len(point_idx), "len(xyz) != len(point_idx)"
xyz = [xyz[i] for i in point_idx]
xyz = np.asarray(xyz, dtype=np.float32)
xyz[:, 2] *= -1
xyz[:, 1] = camera_height
corners = xyz2uv(xyz)
elif data_type == 'Pano_S2D3D':
with open(label_path, 'r') as f:
lines = [line for line in f.readlines() if
len([c for c in line.split(' ') if c[0].isnumeric()]) > 1]
corners_list = np.array([line.strip().split() for line in lines], np.float32)
uv_list = pixel2uv(corners_list)
ceil_uv = uv_list[::2]
floor_uv = uv_list[1::2]
ratio = calc_ceil_ratio([ceil_uv, floor_uv], mode='mean')
corners = floor_uv
else:
return None
output = {
'ratio': np.array([ratio], dtype=np.float32),
'corners': corners,
'id': os.path.basename(label_path).split('.')[0]
}
return output
def move_not_simple_image(data_dir, simple_panos):
import shutil
for house_index in os.listdir(data_dir):
house_path = os.path.join(data_dir, house_index)
if not os.path.isdir(house_path) or house_index == 'visualization':
continue
floor_plan_path = os.path.join(house_path, 'floor_plans')
if os.path.exists(floor_plan_path):
print(f'move:{floor_plan_path}')
dst_floor_plan_path = floor_plan_path.replace('zind', 'zind2')
os.makedirs(dst_floor_plan_path, exist_ok=True)
shutil.move(floor_plan_path, dst_floor_plan_path)
panos_path = os.path.join(house_path, 'panos')
for pano in os.listdir(panos_path):
pano_path = os.path.join(panos_path, pano)
pano_index = '_'.join(pano.split('.')[0].split('_')[-2:])
if f'{house_index}_{pano_index}' not in simple_panos and os.path.exists(pano_path):
print(f'move:{pano_path}')
dst_pano_path = pano_path.replace('zind', 'zind2')
os.makedirs(os.path.dirname(dst_pano_path), exist_ok=True)
shutil.move(pano_path, dst_pano_path)
def read_zind(partition_path, simplicity_path, data_dir, mode, is_simple=True,
layout_type='layout_raw', is_ceiling_flat=False, plan_y=1):
with open(simplicity_path, 'r') as f:
simple_tag = json.load(f)
simple_panos = {}
for k in simple_tag.keys():
if not simple_tag[k]:
continue
split = k.split('_')
house_index = split[0]
pano_index = '_'.join(split[-2:])
simple_panos[f'{house_index}_{pano_index}'] = True
# move_not_simple_image(data_dir, simple_panos)
pano_list = []
with open(partition_path, 'r') as f1:
house_list = json.load(f1)[mode]
for house_index in house_list:
with open(os.path.join(data_dir, house_index, f"zind_data.json"), 'r') as f2:
data = json.load(f2)
panos = []
merger = data['merger']
for floor in merger.values():
for complete_room in floor.values():
for partial_room in complete_room.values():
for pano_index in partial_room:
pano = partial_room[pano_index]
pano['index'] = pano_index
panos.append(pano)
for pano in panos:
if layout_type not in pano:
continue
pano_index = pano['index']
if is_simple and f'{house_index}_{pano_index}' not in simple_panos.keys():
continue
if is_ceiling_flat and not pano['is_ceiling_flat']:
continue
layout = pano[layout_type]
# corners
corner_xz = np.array(layout['vertices'])
corner_xz[..., 0] = -corner_xz[..., 0]
corner_xyz = np.insert(corner_xz, 1, pano['camera_height'], axis=1)
corners = xyz2uv(corner_xyz).astype(np.float32)
# ratio
ratio = np.array([(pano['ceiling_height'] - pano['camera_height']) / pano['camera_height']], dtype=np.float32)
# Ours future work: detection window, door, opening
objects = {
'windows': [],
'doors': [],
'openings': [],
}
for label_index, wdo_type in enumerate(["windows", "doors", "openings"]):
if wdo_type not in layout:
continue
wdo_vertices = np.array(layout[wdo_type])
if len(wdo_vertices) == 0:
continue
assert len(wdo_vertices) % 3 == 0
for i in range(0, len(wdo_vertices), 3):
# In the Zind dataset, the camera height is 1, and the default camera height in our code is also 1,
# so the xyz coordinate here can be used directly
# Since we're taking the opposite z-axis, we're changing the order of left and right
left_bottom_xyz = np.array(
[-wdo_vertices[i + 1][0], -wdo_vertices[i + 2][0], wdo_vertices[i + 1][1]])
right_bottom_xyz = np.array(
[-wdo_vertices[i][0], -wdo_vertices[i + 2][0], wdo_vertices[i][1]])
center_bottom_xyz = (left_bottom_xyz + right_bottom_xyz) / 2
center_top_xyz = center_bottom_xyz.copy()
center_top_xyz[1] = -wdo_vertices[i + 2][1]
center_boundary_xyz = center_bottom_xyz.copy()
center_boundary_xyz[1] = plan_y
uv = xyz2uv(np.array([left_bottom_xyz, right_bottom_xyz,
center_bottom_xyz, center_top_xyz,
center_boundary_xyz]))
left_bottom_uv = uv[0]
right_bottom_uv = uv[1]
width_u = abs(right_bottom_uv[0] - left_bottom_uv[0])
width_u = 1 - width_u if width_u > 0.5 else width_u
assert width_u > 0, width_u
center_bottom_uv = uv[2]
center_top_uv = uv[3]
height_v = center_bottom_uv[1] - center_top_uv[1]
if height_v < 0:
continue
center_boundary_uv = uv[4]
boundary_v = center_boundary_uv[1] - center_bottom_uv[1] if wdo_type == 'windows' else 0
boundary_v = 0 if boundary_v < 0 else boundary_v
center_u = center_bottom_uv[0]
objects[wdo_type].append({
'width_u': width_u,
'height_v': height_v,
'boundary_v': boundary_v,
'center_u': center_u
})
pano_list.append({
'img_path': os.path.join(data_dir, house_index, pano['image_path']),
'corners': corners,
'objects': objects,
'ratio': ratio,
'id': f'{house_index}_{pano_index}',
'is_inside': pano['is_inside']
})
return pano_list
|