File size: 2,833 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 |
"""
@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:
# An intersection occurs and an exception is considered
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')
|