zhigangjiang's picture
no message
88b0dcb
"""
@Date: 2022/01/31
@description:
ZInd:
{'test': {'mw': 2789, 'aw': 381}, 'train': {'mw': 21228, 'aw': 3654}, 'val': {'mw': 2647, 'aw': 433}}
"""
import numpy as np
import matplotlib.pyplot as plt
import json
from tqdm import tqdm
from evaluation.iou import calc_IoU_2D
from visualization.floorplan import draw_floorplan
from visualization.boundary import draw_boundaries
from utils.conversion import depth2xyz, uv2xyz
def analyse_layout_type(dataset, show=False):
bar = tqdm(dataset, total=len(dataset), ncols=100)
manhattan = 0
atlanta = 0
corner_type = {}
for data in bar:
bar.set_description(f"Processing {data['id']}")
corners = data['corners']
corners = corners[corners[..., 0] + corners[..., 1] != 0] # Take effective corners
corners_count = str(len(corners)) if len(corners) < 10 else "10"
if corners_count not in corner_type:
corner_type[corners_count] = 0
corner_type[corners_count] += 1
all_xz = uv2xyz(corners)[..., ::2]
c = len(all_xz)
flag = False
for i in range(c - 1):
l1 = all_xz[i + 1] - all_xz[i]
l2 = all_xz[(i + 2) % c] - all_xz[i + 1]
a = (np.linalg.norm(l1)*np.linalg.norm(l2))
if a == 0:
continue
dot = np.dot(l1, l2)/a
if 0.9 > abs(dot) > 0.1:
# cos-1(0.1)=84.26 > angle > cos-1(0.9)=25.84 or
# cos-1(-0.9)=154.16 > angle > cos-1(-0.1)=95.74
flag = True
break
if flag:
atlanta += 1
else:
manhattan += 1
if flag and show:
draw_floorplan(all_xz, show=True)
draw_boundaries(data['image'].transpose(1, 2, 0), [corners], ratio=data['ratio'], show=True)
corner_type = dict(sorted(corner_type.items(), key=lambda item: int(item[0])))
return {'manhattan': manhattan, "atlanta": atlanta, "corner_type": corner_type}
def execute_analyse_layout_type(root_dir, dataset, modes=None):
if modes is None:
modes = ["train", "val", "test"]
iou2d_d = {}
for mode in modes:
print("mode: {}".format(mode))
types = analyse_layout_type(dataset(root_dir, mode), show=False)
iou2d_d[mode] = types
print(json.dumps(types, indent=4))
return iou2d_d
if __name__ == '__main__':
from dataset.zind_dataset import ZindDataset
from dataset.mp3d_dataset import MP3DDataset
iou2d_d = execute_analyse_layout_type(root_dir='../src/dataset/mp3d',
dataset=MP3DDataset)
# iou2d_d = execute_analyse_layout_type(root_dir='../src/dataset/zind',
# dataset=ZindDataset)
print(json.dumps(iou2d_d, indent=4))