|
""" |
|
@date: 2021/7/5 |
|
@description: |
|
""" |
|
import json |
|
import math |
|
import shutil |
|
|
|
import numpy as np |
|
from utils.boundary import * |
|
import dataset |
|
import os |
|
from tqdm import tqdm |
|
from PIL import Image |
|
from visualization.boundary import * |
|
from visualization.floorplan import * |
|
from shapely.geometry import Polygon, Point |
|
|
|
|
|
def filter_center(ceil_corners): |
|
xyz = uv2xyz(ceil_corners, plan_y=1.6) |
|
xz = xyz[:, ::2] |
|
poly = Polygon(xz).buffer(-0.01) |
|
return poly.contains(Point(0, 0)) |
|
|
|
|
|
def filter_boundary(corners): |
|
if is_ceil_boundary(corners): |
|
return True |
|
elif is_floor_boundary(corners): |
|
return True |
|
else: |
|
|
|
return False |
|
|
|
|
|
def filter_self_intersection(corners): |
|
xz = uv2xyz(corners)[:, ::2] |
|
poly = Polygon(xz) |
|
return poly.is_valid |
|
|
|
|
|
def filter_dataset(dataset, show=False, output_dir=None): |
|
if output_dir is None: |
|
output_dir = os.path.join(dataset.root_dir, dataset.mode) |
|
output_img_dir = os.path.join(output_dir, 'img_align') |
|
output_label_dir = os.path.join(output_dir, 'label_cor_align') |
|
else: |
|
output_dir = os.path.join(output_dir, dataset.mode) |
|
output_img_dir = os.path.join(output_dir, 'img') |
|
output_label_dir = os.path.join(output_dir, 'label_cor') |
|
|
|
if not os.path.exists(output_img_dir): |
|
os.makedirs(output_img_dir) |
|
|
|
if not os.path.exists(output_label_dir): |
|
os.makedirs(output_label_dir) |
|
|
|
bar = tqdm(dataset, total=len(dataset)) |
|
for data in bar: |
|
name = data['name'] |
|
bar.set_description(f"Processing {name}") |
|
img = data['img'] |
|
corners = data['corners'] |
|
|
|
if not filter_center(corners[1::2]): |
|
if show: |
|
draw_boundaries(img, corners_list=[corners[0::2], corners[1::2]], show=True) |
|
if not os.path.exists(data['img_path']): |
|
print("already remove") |
|
else: |
|
print(f"move {name}") |
|
shutil.move(data['img_path'], os.path.join(output_img_dir, os.path.basename(data['img_path']))) |
|
shutil.move(data['label_path'], os.path.join(output_label_dir, os.path.basename(data['label_path']))) |
|
|
|
|
|
def execute_filter_dataset(root_dir, dataset_name="PanoS2D3DDataset", modes=None, output_dir=None): |
|
if modes is None: |
|
modes = ["train", "test", "valid"] |
|
|
|
for mode in modes: |
|
print("mode: {}".format(mode)) |
|
|
|
filter_dataset(getattr(dataset, dataset_name)(root_dir, mode), show=False, output_dir=output_dir) |
|
|
|
|
|
if __name__ == '__main__': |
|
execute_filter_dataset(root_dir='/root/data/hd/hnet_dataset', |
|
dataset_name="PanoS2D3DDataset", modes=['train', "test", "valid"], |
|
output_dir='/root/data/hd/hnet_dataset_close') |
|
|
