├── IP.py ├── LICENSE ├── QP.py ├── README.md ├── RecordReader.py ├── RecordWriterCustom.py ├── RecordWriterTango.py ├── augmentation_tf.py ├── checkpoint └── floornet_hybrid1_branch0123_wsf │ ├── checkpoint │ ├── checkpoint.ckpt.data-00000-of-00001 │ ├── checkpoint.ckpt.index │ └── checkpoint.ckpt.meta ├── evaluate.py ├── floorplan_utils.py ├── icons ├── bathtub.jpg ├── bed.jpg ├── cabinet.jpg ├── cooking_counter.jpg ├── refrigerator.jpg ├── sofa.jpg ├── table.jpg ├── toilet.jpg └── washing_basin.jpg ├── train.py └── utils.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Chen Liu 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # FloorNet: A Unified Framework for Floorplan Reconstruction from 3D Scans 2 | By Chen Liu^\*, Jiaye Wu^\*, and Yasutaka Furukawa (^\* indicates equal contribution) 3 | 4 | ## Introduction 5 | 6 | This paper proposes FloorNet, a novel neural network, to turn RGBD videos of indoor spaces into vector-graphics floorplans. FloorNet consists of three branches, PointNet branch, Floorplan branch, and Image branch. For more details, please refer to our ECCV 2018 [paper](https://arxiv.org/abs/1804.00090) or visit our [project website](http://art-programmer.github.io/floornet.html). This is a follow-up work of our floorplan transformation project which you can find [here](https://github.com/art-programmer/FloorplanTransformation). 7 | 8 | ## Updates 9 | [12/22/2018] We now provide a free IP solver (not relying on Gurobi) at IP.py. The functionality of IP.py should be similar to QP.py which uses Gurobi to solve the IP problem. You might want to consider the free solver if you don't have a Gurobi license. 10 | 11 | ## Dependencies 12 | Python 2.7, TensorFlow (>= 1.3), numpy, opencv 3, CUDA (>= 8.0), Gurobi (free only for academic usages). 13 | 14 | ## Data 15 | 16 | ### Dataset used in the paper 17 | 18 | We collect 155 scans of residential units and annotated corresponding floorplan information. Among 155 scans, 135 are used for training and 20 are for testing. We convert data to tfrecords files which can be downloaded [here](https://drive.google.com/open?id=16lyX_xTiALUzKyst86WJHlhpTDr8XPF_) (or [here](https://mega.nz/#F!5yQy0b5T!ykkR4dqwGO9J5EwnKT_GBw) if you cannot access the previous one). Please put the downloaded files under folder *data/*. 19 | 20 | Here are the links to the raw [point clouds](https://drive.google.com/open?id=1JJlD0qsgMpiU5Jq9TNm3uDPjvqi88aZn), [annotations](https://drive.google.com/open?id=1hYDE2SXLA8Cq7LEK67xO-UMeTSPJ5rcB), and their [associations](https://drive.google.com/open?id=125TAmYWk22EyzCdlbGIfX4Z4DRMhru_V). Please refer to **RecordWriterTango.py** to see how to convert the raw data and annotations to tfrecords files. 21 | 22 | 23 | ### Using custom data 24 | 25 | To generate training/testing data from other data source, the data should be converted to tfrecords as what we did in **RecordWriterTango.py** (an example of our raw data before processed by **RecordWriterTango.py** is provided [here](https://mega.nz/#!dnohjKZa!I3NJZ806vNK-UYp-ap7OynGnS5E-E5AK_z5WsX8n1Ls)). Please refer to [this guide](http://warmspringwinds.github.io/tensorflow/tf-slim/2016/12/21/tfrecords-guide/) for how to generate and read tfrecords. 26 | 27 | Basically, every data sample(tf.train.Example) should at least contain the following components: 28 | 29 | 30 | 1. Inputs: 31 | 32 | - a point cloud (randomly sampled 50,000 points) 33 | - a mapping from point cloud's 3D space to 2D space of the 256x256 top-view density image. 34 | - It contains 50,000 indices, one for each point. 35 | - For point (x, y, z), index = round((y - min(Y) + padding) / (maxRange + 2 * padding) * 256) * 256 + round((x - min(X) + padding) / (maxRange + 2 * padding) * 256). 36 | - maxRange = max(max(X) - min(X), max(Y) - min(Y)) 37 | - padding could be any small value, say 0.05 maxRange 38 | - optional: image features of the RGB video stream, if the image branch is enabled 39 | 40 | 2. Labels: 41 | 42 | - Corners and their corresponding types 43 | - Total number of corners 44 | - A ground-truth icon segmentation map 45 | - A ground-truth room segmentation map 46 | 47 | Again, please refer to **RecordWriterTango.py** for exact details. 48 | 49 | **NEW:** We added a template file, **RecordWriterCustom.py** for using custom data. 50 | 51 | 52 | ## Annotator 53 | For reference, a similar (but not the same) annotator written in Python is [here](https://github.com/art-programmer/FloorplanAnnotator). You need to make some changes to annotate your own data. 54 | 55 | ## Training 56 | To train the network from scratch, please run: 57 | ```bash 58 | python train.py --restore=0 59 | ``` 60 | 61 | ## Evaluation 62 | To evaluate the performance of our trained model, please run: 63 | ```bash 64 | python train.py --task=evaluate --separateIconLoss 65 | ``` 66 | 67 | ## Generate 3D models 68 | We can popup the reconstructed floorplan to generate 3D models. Please refer to our previous project, [FloorplanTransformation](https://github.com/art-programmer/FloorplanTransformation), for more details. 69 | 70 | ## Contact 71 | 72 | If you have any questions, please contact me at chenliu@wustl.edu. 73 | -------------------------------------------------------------------------------- /RecordReader.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import threading 4 | import PIL.Image as Image 5 | from functools import partial 6 | from multiprocessing import Pool 7 | import cv2 8 | import sys 9 | import os 10 | import functools 11 | from floorplan_utils import * 12 | from augmentation_tf import * 13 | from utils import * 14 | #sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 15 | 16 | NUM_THREADS = 8 17 | 18 | def parse_fn(example, augmentation, readImageFeatures = False, convertCorners=False, kernelSize=11): 19 | if readImageFeatures: 20 | features = tf.parse_single_example( 21 | example, 22 | # Defaults are not specified since both keys are required. 23 | features={ 24 | 'image_path': tf.FixedLenFeature([], tf.string), 25 | 'points': tf.FixedLenFeature([NUM_POINTS * (NUM_INPUT_CHANNELS - 1)], tf.float32), 26 | 'point_indices': tf.FixedLenFeature([NUM_POINTS], tf.int64), 27 | 'corner': tf.FixedLenFeature([MAX_NUM_CORNERS * 3], tf.int64), 28 | 'num_corners': tf.FixedLenFeature([], tf.int64), 29 | 'icon': tf.FixedLenFeature([], tf.string), 30 | 'room': tf.FixedLenFeature([], tf.string), 31 | 'image': tf.FixedLenFeature(sum([size * size * numChannels for size, numChannels in zip(SIZES, NUM_CHANNELS)[1:]]), tf.float32), 32 | 'flags': tf.FixedLenFeature([2], tf.int64), 33 | }) 34 | else: 35 | features = tf.parse_single_example( 36 | example, 37 | # Defaults are not specified since both keys are required. 38 | features={ 39 | 'image_path': tf.FixedLenFeature([], tf.string), 40 | 'points': tf.FixedLenFeature([NUM_POINTS * (NUM_INPUT_CHANNELS - 1)], tf.float32), 41 | 'point_indices': tf.FixedLenFeature([NUM_POINTS], tf.int64), 42 | 'corner': tf.FixedLenFeature([MAX_NUM_CORNERS * 3], tf.int64), 43 | 'num_corners': tf.FixedLenFeature([], tf.int64), 44 | 'icon': tf.FixedLenFeature([], tf.string), 45 | 'room': tf.FixedLenFeature([], tf.string), 46 | 'flags': tf.FixedLenFeature([2], tf.int64), 47 | }) 48 | pass 49 | 50 | points = tf.reshape(features['points'], (NUM_POINTS, NUM_INPUT_CHANNELS - 1)) 51 | point_indices = tf.cast(features['point_indices'], tf.int32) 52 | #point_indices = features['point_indices'] 53 | 54 | corners = tf.cast(tf.reshape(features['corner'], [MAX_NUM_CORNERS, 3]), tf.int32) 55 | numCorners = features['num_corners'] 56 | corners = corners[:numCorners] 57 | iconSegmentation = tf.reshape(tf.decode_raw(features['icon'], tf.uint8), (HEIGHT, WIDTH, 1)) 58 | roomSegmentation = tf.reshape(tf.decode_raw(features['room'], tf.uint8), (HEIGHT, WIDTH, 1)) 59 | heatmaps = tf.stack([iconSegmentation, roomSegmentation], axis=0) 60 | 61 | if readImageFeatures: 62 | imageFeature = features['image'] 63 | imageFeatures = {} 64 | offset = 0 65 | for index, (size, numChannels) in enumerate(zip(SIZES, NUM_CHANNELS)[1:]): 66 | imageFeatures[index] = tf.reshape(imageFeature[offset:offset + size * size * numChannels], (size, size, numChannels)) 67 | offset += size * size * numChannels 68 | continue 69 | else: 70 | imageFeatures = {} 71 | pass 72 | 73 | flags = features['flags'] 74 | if 'w' in augmentation: 75 | point_indices, corners, heatmaps = tf.cond(tf.logical_or(tf.equal(flags[0], 0), tf.equal(flags[0], 4)), lambda: augmentWarping(point_indices, corners, heatmaps, gridStride=32, randomScale=2), lambda: (point_indices, corners, heatmaps)) 76 | #point_indices, corners, heatmaps = augmentWarping(point_indices, corners, heatmaps, gridStride=32, randomScale=4) 77 | pass 78 | if 's' in augmentation: 79 | points, point_indices, corners, heatmaps, imageFeatures = augmentScaling(points, point_indices, corners, heatmaps, imageFeatures) 80 | pass 81 | if 'f' in augmentation: 82 | points, point_indices, corners, heatmaps, imageFeatures = augmentFlipping(points, point_indices, corners, heatmaps, imageFeatures) 83 | pass 84 | 85 | iconSegmentation = tf.cast(tf.squeeze(heatmaps[0]), tf.int32) 86 | roomSegmentation = tf.cast(tf.squeeze(heatmaps[1]), tf.int32) 87 | 88 | roomSegmentation = tf.minimum(roomSegmentation, NUM_ROOMS - 1) 89 | 90 | # point_indices_stack = getCoarseIndicesMaps(point_indices, WIDTH, HEIGHT, 0) 91 | 92 | corners = tf.reshape(tf.concat([corners, tf.zeros((MAX_NUM_CORNERS - tf.shape(corners)[0], 3), dtype=tf.int32)], axis=0), (MAX_NUM_CORNERS, 3)) 93 | if convertCorners: 94 | cornerSegmentation = tf.stack([tf.sparse_to_dense(tf.stack([corners[:, 1], corners[:, 0]], axis=1), (HEIGHT, WIDTH), corners[:, 2], validate_indices=False)], axis=0) 95 | cornerHeatmaps = tf.one_hot(cornerSegmentation, depth=NUM_CORNERS + 1, axis=-1)[:, :, :, 1:] 96 | #kernel_size = kernelSize 97 | #neighbor_kernel_array = disk(kernel_size) 98 | #neighbor_kernel = tf.constant(neighbor_kernel_array.reshape(-1), shape=neighbor_kernel_array.shape, dtype=tf.float32) 99 | #neighbor_kernel = tf.reshape(neighbor_kernel, [kernel_size, kernel_size, 1, 1]) 100 | #cornerHeatmaps = tf.nn.depthwise_conv2d(cornerHeatmaps, tf.tile(neighbor_kernel, [1, 1, NUM_CORNERS, 1]), strides=[1, 1, 1, 1], padding='SAME') 101 | corners = tf.cast(cornerHeatmaps > 0.5, tf.float32) 102 | 103 | # cornerSegmentation = tf.sparse_to_dense(tf.stack([corners[:, 1], corners[:, 0]], axis=1), (HEIGHT, WIDTH), corners[:, 2], validate_indices=False) 104 | # cornerHeatmaps = tf.one_hot(cornerSegmentation, depth=NUM_CORNERS, axis=-1) 105 | # kernel = tf.tile(tf.expand_dims(tf.constant(disk(11)), -1), [1, 1, NUM_CORNERS]) 106 | # cornerHeatmaps = tf.nn.dilation2d(tf.expand_dims(cornerHeatmaps, 0), kernel, [1, 1, 1, 1], [1, 1, 1, 1], 'SAME')[0] 107 | 108 | imagePath = features['image_path'] 109 | 110 | points = tf.concat([points, tf.ones((NUM_POINTS, 1))], axis=1) 111 | 112 | if readImageFeatures: 113 | input_dict = {'points': points, 'point_indices': point_indices, 'image_features': imageFeatures, 'image_path': imagePath, 'flags': flags} 114 | else: 115 | input_dict = {'points': points, 'point_indices': point_indices, 'image_path': imagePath, 'flags': flags} 116 | pass 117 | gt_dict = {'corner': corners, 'icon': iconSegmentation, 'room': roomSegmentation, 'num_corners': numCorners} 118 | return input_dict, gt_dict 119 | 120 | # class RecordReader(): 121 | # def __init__(self): 122 | # return 123 | 124 | # def getBatch(self, filename_queues, batchSize, augmentation, min_after_dequeue = 1000, random=True, getLocal=False, getSegmentation=False, test=True): 125 | # reader = tf.TFRecordReader() 126 | # queueIndex = tf.cast(tf.random_uniform([len(filename_queues)], maxval=5), tf.int32)[0] 127 | # #filename_queue = tf.cond(tf.equal(queueIndex, 0), lambda: filename_queues[0], lambda: tf.cond(tf.equal(queueIndex, 1), lambda: filename_queues[1], lambda: tf.cond(tf.equal(queueIndex, 2), lambda: filename_queues[2], lambda: tf.cond(tf.equal(queueIndex, 3), lambda: filename_queues[3], lambda: filename_queues[4])))) 128 | 129 | # _, serialized_example = reader.read(filename_queues[queueIndex]) 130 | 131 | # features = tf.parse_single_example( 132 | # serialized_example, 133 | # # Defaults are not specified since both keys are required. 134 | # features={ 135 | # 'image_path': tf.FixedLenFeature([], tf.string), 136 | # 'points': tf.FixedLenFeature([NUM_POINTS * 7], tf.float32), 137 | # 'point_indices': tf.FixedLenFeature([NUM_POINTS], tf.int32), 138 | # 'corner': tf.FixedLenFeature([MAX_NUM_CORNERS * 3], tf.int32), 139 | # 'num_corners': tf.FixedLenFeature([], tf.int32), 140 | # 'icon': tf.FixedLenFeature([], tf.string), 141 | # 'room': tf.FixedLenFeature([], tf.string), 142 | # 'image': tf.FixedLenFeature(sum([size * size * numChannels for size, numChannels in zip(SIZES, NUM_CHANNELS)[1:]]), tf.float32), 143 | # 'flags': tf.FixedLenFeature([2], tf.int32), 144 | # }) 145 | 146 | 147 | # points = tf.reshape(features['points'], (NUM_POINTS, NUM_INPUT_CHANNELS)) 148 | # point_indices = tf.reshape(features['point_indices'], (NUM_POINTS)) 149 | # corners = tf.reshape(features['corner'], [MAX_NUM_CORNERS, 3]) 150 | # numCorners = tf.cast(features['num_corners'], tf.int32) 151 | # corners = corners[:numCorners] 152 | # iconSegmentation = tf.reshape(tf.decode_raw(features['icon'], tf.uint8), (HEIGHT, WIDTH, 1)) 153 | # roomSegmentation = tf.reshape(tf.decode_raw(features['room'], tf.uint8), (HEIGHT, WIDTH, 1)) 154 | # heatmaps = tf.stack([icon, room], axis=0) 155 | 156 | # flags = features['flags'] 157 | 158 | # if 'w' in augmentation: 159 | # point_indices, corners, heatmaps = tf.cond(flags[0] == 0 or flags[0] == 4, lambda: augmentWarping(point_indices, corners, heatmaps, stride=32, randomScale=4), lambda: pointcloudIndices, corners, heatmaps) 160 | # pass 161 | # if 's' in augmentation: 162 | # points, point_indices, corners, heatmaps = augmentScaling(points, point_indices, corners, heatmaps) 163 | # pass 164 | # if 'f' in augmentation: 165 | # points, point_indices, corners, heatmaps = augmentFlipping(points, point_indices, corners, heatmaps) 166 | # pass 167 | 168 | # iconSegmentation = tf.squeeze(heatmaps[0]) 169 | # roomSegmentation = tf.squeeze(heatmaps[1]) 170 | 171 | # point_indices_stack = getCoarseIndicesMaps(point_indices, WIDTH, HEIGHT, 0) 172 | 173 | # cornerSegmentation = tf.sparse_to_dense(tf.stack([corners[:, 1], corners[:, 0]], axis=1), (HEIGHT, WIDTH), corners[:, 2]) 174 | # cornerHeatmaps = tf.one_hot(cornerSegmentation, depth=NUM_CORNERS, axis=-1) 175 | 176 | # kernel = tf.tile(tf.expand_dims(tf.constant(disk(11)), -1), [1, 1, NUM_CORNERS]) 177 | # cornerHeatmaps = nn.diltion2d(cornerHeatmaps, kernel, [1, 1, 1, 1], [1, 1, 1, 1], 'VALID') 178 | 179 | # imagePath = features['image_path'] 180 | 181 | # imageFeatures = [] 182 | # imageFeature = tf.reduce_sum(tf.reshape(features['image'], [2, -1]), axis=0) 183 | # offset = 0 184 | # for size, numChannels in zip(SIZES, NUM_CHANNELS)[1:]: 185 | # imageFeatures.append(tf.reshape(imageFeature[offset:offset + size * size * numChannels], (size[1], size, numChannels))) 186 | # offset += size * size * numChannels 187 | # continue 188 | 189 | # if random: 190 | # points_inp, point_indices_inp, corner_gt, icon_gt, room_gt, image_features_inp, image_path_inp, flags_inp = tf.train.shuffle_batch([points, point_indices_stack, cornerHeatmaps, iconSegmentation, roomSegmentation, imageFeatures, imagePath, flags], batch_size=batchSize, capacity=min_after_dequeue + (NUM_THREADS + 2) * batchSize, num_threads=NUM_THREADS, min_after_dequeue=min_after_dequeue) 191 | # else: 192 | # points_inp, point_indices_inp, corner_gt, icon_gt, room_gt, image_features_inp, image_path_inp, flags_inp = tf.train.batch([points, point_indices_stack, cornerHeatmaps, iconSegmentation, roomSegmentation, imageFeatures, imagePath, flags], batch_size=batchSize, capacity=min_after_dequeue + (NUM_THREADS + 2) * batchSize, num_threads=1) 193 | # pass 194 | 195 | # input_dict = {'points': points_inp, 'point_indices': point_indices_inp, 'image_features': image_features_inp, 'image_path': image_path_inp, 'flags': flags_inp} 196 | # gt_dict = {'corner': corner_gt, 'icon': icon_gt, 'room': room_gt} 197 | # return input_dict, gt_dict 198 | 199 | 200 | def getDatasetTrain(filenames, augmentation, readImageFeatures, batchSize): 201 | if len(filenames) == 1: 202 | return tf.data.TFRecordDataset(filenames[0]).repeat().map(functools.partial(parse_fn, augmentation=augmentation, readImageFeatures=readImageFeatures), num_parallel_calls=NUM_THREADS).batch(batchSize).prefetch(1) 203 | else: 204 | return tf.data.Dataset.from_tensor_slices(filenames).interleave(lambda x: tf.data.TFRecordDataset(x).repeat().map(functools.partial(parse_fn, augmentation=augmentation), num_parallel_calls=NUM_THREADS), cycle_length=len(filenames), block_length=batchSize).batch(batchSize).prefetch(100 / len(filenames)).shuffle(100 / len(filenames)) 205 | 206 | 207 | def getDatasetVal(filenames, augmentation, readImageFeatures, batchSize): 208 | #return tf.data.Dataset.from_tensor_slices(filenames).interleave(lambda x: tf.data.TFRecordDataset(x).map(functools.partial(parse_fn, augmentation=augmentation, readImageFeatures=readImageFeatures), num_parallel_calls=NUM_THREADS), cycle_length=len(filenames), block_length=batchSize).apply(tf.contrib.data.batch_and_drop_remainder(batchSize)).prefetch(1) 209 | if len(filenames) == 1: 210 | return tf.data.TFRecordDataset(filenames[0]).repeat().map(functools.partial(parse_fn, augmentation=augmentation, readImageFeatures=readImageFeatures), num_parallel_calls=NUM_THREADS).batch(batchSize).prefetch(1) 211 | else: 212 | return tf.data.Dataset.from_tensor_slices(filenames).interleave(lambda x: tf.data.TFRecordDataset(x).map(functools.partial(parse_fn, augmentation=augmentation, readImageFeatures=readImageFeatures), num_parallel_calls=NUM_THREADS), cycle_length=len(filenames), block_length=batchSize).apply(tf.contrib.data.batch_and_drop_remainder(batchSize)).prefetch(1) 213 | -------------------------------------------------------------------------------- /RecordWriterCustom.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import numpy as np 3 | import cv2 4 | import os 5 | import tensorflow as tf 6 | 7 | HEIGHT = 256 8 | WIDTH = 256 9 | NUM_POINTS = 50000 10 | NUM_CHANNELS = [7, 64, 64, 64, 128, 256] 11 | SIZES = [WIDTH, WIDTH // 2, WIDTH // 4, WIDTH // 8, WIDTH // 16, WIDTH // 32] 12 | 13 | def _bytes_feature(value): 14 | return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) 15 | 16 | def _int64_feature(value): 17 | return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) 18 | 19 | def _float_feature(value): 20 | return tf.train.Feature(float_list=tf.train.FloatList(value=value)) 21 | 22 | 23 | def fitPlane(points): 24 | if points.shape[0] == points.shape[1]: 25 | return np.linalg.solve(points, np.ones(points.shape[0])) 26 | else: 27 | return np.linalg.lstsq(points, np.ones(points.shape[0]))[0] 28 | 29 | class RecordWriterCustom(): 30 | def __init__(self, filename_list, numPoints=50000, numInputChannels=6): 31 | """filename_list: a list of filename dictionaries Each element is a dict of filenames for data entries, which must contains 'point_cloud' 32 | numPoints: the number of sample points 33 | numInputChannels: [3: XYZ, 6:XYZ + RGB, 9: XYZ + RGB + Normals] 34 | """ 35 | self.numPoints = numPoints 36 | self.numInputChannels = numInputChannels 37 | #super(RecordWriterTango, self).__init__() 38 | self.writer = tf.python_io.TFRecordWriter('data/' + filename + '_' + options.task + '.tfrecords') 39 | for filename_dict in filename_list: 40 | self.writeExample(filename_dict) 41 | continue 42 | self.writer.close() 43 | return 44 | 45 | 46 | def loadPoints(self, filename): 47 | """Function to load the point cloud""" 48 | assert(False, "Not implemented") 49 | return 50 | 51 | def rotatePoints(self, points, rotation_matrix=None): 52 | """Function to rotate the point cloud to be axis-aligned. 53 | Recommand annotate the rotation by adapting the annotator here: https://github.com/art-programmer/FloorplanAnnotator. 54 | If the annotation is not available we use heuristics to rotate the point cloud 55 | """ 56 | #assert(False, "Not implemented") 57 | if transformation: 58 | rotated_points = points.copy() 59 | rotated_points[:, :3] = np.matmul(points[:, :3], rotation_matrix.transpose()) 60 | if points.shape[-1] >= 9: 61 | rotated_points[:, 6:9] = np.matmul(points[:, 6:9], rotation_matrix.transpose()) 62 | pass 63 | return rotated_points 64 | 65 | if points.shape[-1] >= 9: 66 | ## Point normal exists 67 | normals = points[:, 6:9] 68 | else: 69 | ## Compute the point normal by ourselves. Consider using PCL if available. If not we sample some points and compute their normals 70 | sampled_points = points[np.random.choice(np.arange(len(points), dtype=np.int32), 100)] 71 | scale = (points[:, :3].max(0) - points[:, :3].min(0)).max() 72 | normals = [] 73 | for sampled_point in sampled_points: 74 | neighbor_points = points[np.linalg.norm(points - sampled_point, axis=1) < scale * 0.02] 75 | if len(neighbor_points) >= 3: 76 | ## Consider to use try catch for fitPlane as it might encounter singular cases 77 | plane = fitPlane(neighbor_points) 78 | normal = plane / max(np.linalg.norm(plane), 1e-4) 79 | normals.append(normal) 80 | pass 81 | continue 82 | normals = np.stack(normals, axis=0) 83 | pass 84 | 85 | 86 | polarAngles = np.arange(16) * np.pi / 2 / 16 87 | azimuthalAngles = np.arange(64) * np.pi * 2 / 64 88 | polarAngles = np.expand_dims(polarAngles, -1) 89 | azimuthalAngles = np.expand_dims(azimuthalAngles, 0) 90 | 91 | normalBins = np.stack([np.sin(polarAngles) * np.cos(azimuthalAngles), np.tile(np.cos(polarAngles), [1, azimuthalAngles.shape[1]]), -np.sin(polarAngles) * np.sin(azimuthalAngles)], axis=2) 92 | normalBins = np.reshape(normalBins, [-1, 3]) 93 | numBins = normalBins.shape[0] 94 | 95 | 96 | normalDiff = np.tensordot(normals, normalBins, axes=([1], [1])) 97 | normalDiffSign = np.sign(normalDiff) 98 | normalDiff = np.maximum(normalDiff, -normalDiff) 99 | normalMask = one_hot(np.argmax(normalDiff, axis=-1), numBins) 100 | bins = normalMask.sum(0) 101 | maxNormals = np.expand_dims(normals, 1) * np.expand_dims(normalMask, -1) 102 | maxNormals *= np.expand_dims(normalDiffSign, -1) 103 | averageNormals = maxNormals.sum(0) / np.maximum(np.expand_dims(bins, -1), 1e-4) 104 | averageNormals /= np.maximum(np.linalg.norm(averageNormals, axis=-1, keepdims=True), 1e-4) 105 | dominantNormal_1 = averageNormals[np.argmax(bins)] 106 | 107 | dotThreshold_1 = np.cos(np.deg2rad(100)) 108 | dotThreshold_2 = np.cos(np.deg2rad(80)) 109 | 110 | dot_1 = np.tensordot(normalBins, dominantNormal_1, axes=([1], [0])) 111 | bins[np.logical_or(dot_1 < dotThreshold_1, dot_1 > dotThreshold_2)] = 0 112 | dominantNormal_2 = averageNormals[np.argmax(bins)] 113 | 114 | # dot_2 = np.tensordot(normalBins, dominantNormal_2, axes=([1], [0])) 115 | # bins[np.logical_or(dot_2 < dotThreshold_1, dot_2 > dotThreshold_2)] = 0 116 | # dominantNormal_3 = averageNormals[np.argmax(bins)] 117 | dominantNormal_3 = np.cross(dominant_normal_1, dominant_normal_2) 118 | dominantNormal_2 = np.cross(dominant_normal_3, dominant_normal_1) 119 | rotation_matrix = np.stack([dominantNormal_1, dominantNormal_2, dominantNormal_3], axis=0) 120 | rotated_points = points.copy() 121 | rotated_points[:, :3] = np.matmul(points[:, :3], rotation_matrix.transpose()) 122 | if points.shape[-1] >= 9: 123 | rotated_points[:, 6:9] = np.matmul(points[:, 6:9], rotation_matrix.transpose()) 124 | pass 125 | ## Rectify the rotation matrix 126 | return rotated_points 127 | 128 | def scalePoints(self, points, rotation_matrix=None): 129 | """Function to scale the point cloud to range [0, 1].""" 130 | XYZ = points[:, :3] 131 | mins = XYZ.min(0, keepdims=True) 132 | maxs = XYZ.max(0, keepdims=True) 133 | maxRange = (maxs - mins)[:, :2].max() 134 | padding = maxRange * 0.05 135 | maxRange += padding * 2 136 | mins = (maxs + mins) / 2 - maxRange / 2 137 | 138 | XYZ = (XYZ - mins) / maxRange 139 | points[:, :3] = XYZ 140 | return points 141 | 142 | def computeCoordinates(self, points): 143 | """Compute the image coordinate for each point""" 144 | coordinates = np.minimum(np.maximum(np.round(points[:, :2] * imageSize).astype(np.int32), 0), imageSize - 1) 145 | coordinates = coordinates[:, 1] * WIDTH + coordinates[:, 0] 146 | return coordinates 147 | 148 | def load(self, name, filename): 149 | """Function to load info""" 150 | assert(False, "Not implemented") 151 | return 152 | 153 | def writeExample(self, filename_dict): 154 | """Write one data entry""" 155 | assert('point_cloud' in filename_dict) 156 | 157 | ## Implement a function to load the point cloud as a numpy array of size [NxC] where N is the number of points and C is the number of channels 158 | points = loadPoints(filename_dict['point_cloud']) 159 | 160 | points = points[:, :self.numInputChannels] 161 | if len(points) < self.numPoints: 162 | indices = np.arange(len(points)) 163 | points = np.concatenate([points, points[np.random.choice(indices, self.numPoints - len(points))]], axis=0) 164 | elif len(points) > self.numPoints: 165 | sampledInds = np.arange(len(points)) 166 | points = points[np.random.choice(indices, self.numPoints)] 167 | pass 168 | 169 | points = self.rotatedPoints(points) 170 | points = self.scalePoints(points) 171 | 172 | imageSize = np.array([WIDTH, HEIGHT]) 173 | indicesMap = self.computeCoordinates() 174 | 175 | info_dict = {} 176 | for info_name in ['corner_gt', 'icon_gt', 'room_gt', 'image_feature']: 177 | if info_name in filename_dict: 178 | info = self.load(info_name, filename_dict[info_name]) 179 | else: 180 | if info_name == 'corner_gt': 181 | info = np.zeros((1, 3)) 182 | elif info_name in ['icon_gt', 'room_gt']: 183 | info = np.zeros((HEIGHT, WIDTH), dtype=uint8) 184 | else: 185 | info = np.zeros(sum([size * size * numChannels for size, numChannels in zip(SIZES, NUM_CHANNELS)[1:]])) 186 | pass 187 | pass 188 | info_dict[info_name] = info 189 | continue 190 | 191 | example = tf.train.Example(features=tf.train.Features(feature={ 192 | 'image_path': _bytes_feature(filename_dict['point_cloud']), 193 | 'points': _float_feature(points.reshape(-1)), 194 | 'point_indices': _int64_feature(indicesMap.reshape(-1)), 195 | 'corner': _int64_feature(info['corner_gt'].reshape(-1)), 196 | 'num_corners': _int64_feature([len(info['corner_gt'])]), 197 | 'icon': _bytes_feature(info['icon_gt'].tostring()), 198 | 'room': _bytes_feature(info['room_gt'].tostring()), 199 | 'image': _float_feature(info['image_feature'].reshape(-1)), 200 | 'flags': _int64_feature(np.zeros((2, dtype=np.int64))), 201 | })) 202 | self.writer.write(example.SerializeToString()) 203 | return 204 | 205 | 206 | if __name__ == "__main__": 207 | RecordWriterCustom() 208 | exit(1) 209 | -------------------------------------------------------------------------------- /RecordWriterTango.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import numpy as np 3 | import cv2 4 | import random 5 | import math 6 | import os 7 | import time 8 | import zlib 9 | import socket 10 | import threading 11 | import Queue 12 | import sys 13 | import cPickle as pickle 14 | import PIL 15 | import json 16 | import glob 17 | import torchfile 18 | import math 19 | import os.path 20 | import zipfile 21 | import lupa 22 | #from drnseg import DRNSeg 23 | from utils import * 24 | from floorplan_utils import * 25 | import argparse 26 | from skimage import measure 27 | from augmentation import augment 28 | import tensorflow as tf 29 | 30 | #import show3d 31 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 32 | 33 | 34 | #BATCH_SIZE=16 35 | FETCH_BATCH_SIZE=16 36 | HEIGHT=256 37 | WIDTH=256 38 | NUM_POINTS = 50000 39 | ROOT_FOLDER = '/home/chenliu/Projects/Data/Tango/' 40 | #CORNER_POINTS_FOLDER = '../../../Data/Floorplan_syn/test_pntnetpp_new/data/oldway/real/combined_out/' 41 | 42 | TEST_INDICES = [153, 102, 115, 57, 104, 156, 154, 310, 70, 134, 98, 8, 51, 76, 87, 36, 18, 106, 10, 73] 43 | 44 | 45 | def _bytes_feature(value): 46 | return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) 47 | 48 | def _int64_feature(value): 49 | return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) 50 | 51 | def _float_feature(value): 52 | return tf.train.Feature(float_list=tf.train.FloatList(value=value)) 53 | 54 | def load_pnt(pnt_file): 55 | all_pnts = [] 56 | texture_file = pnt_file.replace('.obj', '.mtl') 57 | texture_exists = os.path.exists(texture_file) 58 | print(pnt_file, texture_exists) 59 | if texture_exists: 60 | with open(texture_file) as f: 61 | for line in f: 62 | values = line.split(' ') 63 | if values[0] == 'map_Kd': 64 | paths = texture_file.split('/') 65 | paths[-1] = values[1].strip() 66 | image_file = '/'.join(paths) 67 | #print(image_file) 68 | texture = cv2.imread(image_file) 69 | texture_width = texture.shape[1] 70 | texture_height = texture.shape[0] 71 | pass 72 | continue 73 | pass 74 | pass 75 | 76 | with open(pnt_file) as f: 77 | vertex_index = 0 78 | for line in f: 79 | eles = line.split() 80 | if len(eles) == 0: 81 | continue 82 | if eles[0] == "v": 83 | if texture_exists: 84 | pnt = [float(eles[1]), float(eles[2]), float(eles[3])] 85 | else: 86 | assert len(eles) >= 7, str(vertex_index) + line 87 | pnt = [float(eles[1]), float(eles[2]), float(eles[3]), float(eles[4]), float(eles[5]), float(eles[6])] 88 | pass 89 | all_pnts.append(pnt) 90 | elif eles[0] == 'vt' and texture_exists: 91 | u = float(eles[1]) 92 | v = float(eles[2]) 93 | if u < 0 or u > 1 or v < 0 or v > 1: 94 | color = [0, 0, 0] 95 | else: 96 | u = int(min(np.round(u * texture_width), texture_width - 1)) 97 | v = int(min(np.round((1 - v) * texture_height), texture_height - 1)) 98 | color = texture[v][u] 99 | color = color.astype(np.float32) / 255 100 | pass 101 | all_pnts[vertex_index] += [color[2], color[1], color[0]] 102 | vertex_index += 1 103 | pass 104 | continue 105 | pass 106 | #if texture_exists: 107 | #writePointCloud('test/pointcloud.ply', all_pnts) 108 | #exit(1) 109 | # print(all_pnts) 110 | return all_pnts 111 | 112 | class PntMapper: 113 | def __init__(self, start, size, newstart, newsize): 114 | self.start = start 115 | self.size = size 116 | self.newsize = newsize 117 | self.newstart=newstart 118 | def map_pnts(self, x, y): 119 | # print(self.size) 120 | return (x - self.start[0])/(self.size[0]-1)*(self.newsize[0]-1) + self.newstart[0], (y- self.start[1])/(self.size[1]-1)*(self.newsize[1]-1) + self.newstart[1] 121 | def pnts_to_heatmap_transformation_matrix(pnts, is_binary, label, transform=None, img=None, bd_box=None): 122 | if not isinstance(pnts, np.ndarray): 123 | pnts = np.array(pnts) 124 | if pnts.shape[0] == 0: 125 | return None 126 | if transform is not None: 127 | if not isinstance(transform, np.ndarray): 128 | transform = np.array(transform) 129 | aug_pnts = np.ones((pnts.shape[0], 6)) 130 | aug_pnts[:, :-1] = pnts 131 | prj_pnts = np.transpose(np.dot(transform, np.transpose(aug_pnts, (1, 0))), (1, 0)) 132 | prj_pnts /= prj_pnts[:, -1][:,np.newaxis] 133 | else: 134 | prj_pnts = pnts 135 | if bd_box==None: 136 | x_max = prj_pnts[:, 0].max() 137 | x_min = prj_pnts[:, 0].min() 138 | y_max = prj_pnts[:, 1].max() 139 | y_min = prj_pnts[:, 1].min() 140 | else: 141 | x_min = bd_box[0] 142 | y_min = bd_box[1] 143 | x_max = bd_box[2] 144 | y_max = bd_box[3] 145 | #print(x_min, y_min, x_max, y_max) 146 | x_size = math.ceil(x_max - x_min + 1) 147 | y_size = math.ceil(y_max - y_min + 1) 148 | 149 | #print(x_min, y_min, x_size, y_size) 150 | mx_size = max(x_size, y_size) 151 | 152 | pnt_mapper = PntMapper((x_min, y_min), (x_size, y_size), (0,0), (HEIGHT, WIDTH)) 153 | if img is None: 154 | img = np.zeros((HEIGHT, WIDTH), dtype=np.uint8) 155 | for pnt in prj_pnts: 156 | # print(pnt) 157 | m_pnt = pnt_mapper.map_pnts(pnt[0], pnt[1]) 158 | #if len(pnts) < 30: 159 | # print(pnt, m_pnt, (x_min, y_min, x_max, y_max)) 160 | try: 161 | if is_binary: 162 | img[int(round(m_pnt[1])), int(round(m_pnt[0]))] = label 163 | else: 164 | img[int(round(m_pnt[1])), int(round(m_pnt[0]))] += 1 165 | except Exception as e: 166 | print(pnt, m_pnt, e) 167 | raise e 168 | return img, (x_min, y_min, x_max, y_max) 169 | 170 | def getZRotationMatrix(deg): 171 | return np.array([[math.cos(deg), -math.sin(deg), 0], 172 | [-math.sin(deg), -math.cos(deg), 0], 173 | [0, 0, 1]]) 174 | 175 | class RecordWriterTango(): 176 | def __init__(self, options): 177 | #super(RecordWriterTango, self).__init__() 178 | self.dataPath = 'data/' 179 | self.stopped=False 180 | self.split = options.task 181 | self.imagePaths = self.loadImagePaths() 182 | 183 | self.numImages = len(self.imagePaths) 184 | self.numPoints = options.numPoints 185 | self.numInputChannels = options.numInputChannels 186 | 187 | self.labelMap = self.loadLabelMap() 188 | self.gap = 3 189 | 190 | self.roomLabelMap = getRoomLabelMap() 191 | 192 | # numTrainingImages = int(round(self.numImages * 0.9)) 193 | # if options.task == 'train': 194 | # self.imagePaths = self.imagePaths[:numTrainingImages] 195 | # else: 196 | # self.imagePaths = self.imagePaths[numTrainingImages:] 197 | # pass 198 | 199 | if options.task == 'train': 200 | self.imagePaths = [imagePath for imagePath in self.imagePaths if int(imagePath[5:]) not in TEST_INDICES] 201 | else: 202 | self.imagePaths = [imagePath for imagePath in self.imagePaths if int(imagePath[5:]) in TEST_INDICES] 203 | pass 204 | 205 | filename = os.path.basename(__file__).rstrip('.py')[12:] 206 | self.writer = tf.python_io.TFRecordWriter('data/' + filename + '_' + options.task + '.tfrecords') 207 | for imagePath in self.imagePaths: 208 | self.writeExample(imagePath) 209 | continue 210 | self.writer.close() 211 | return 212 | 213 | def getNumBatchesTesting(self): 214 | #print(self.numBatches) 215 | #print(self.numBatchesTraining) 216 | return self.numBatches - self.numBatchesTraining 217 | 218 | def loadImagePaths(self): 219 | scene_ids = [] 220 | for scene_id in os.listdir(ROOT_FOLDER): 221 | annotation_filename = ROOT_FOLDER + scene_id + '/annotation/floorplan.txt' 222 | exists1 = os.path.exists(annotation_filename) 223 | filename_glob = ROOT_FOLDER + scene_id + '/' + 'dataset/mesh/*.obj' 224 | filename = list(glob.glob(filename_glob)) 225 | exists2 = len(filename) > 0 226 | if exists1 and exists2: 227 | scene_ids.append(scene_id) 228 | else: 229 | print(scene_id) 230 | pass 231 | continue 232 | return scene_ids 233 | 234 | def loadLabelMap(self): 235 | roomMap = getRoomLabelMap() 236 | 237 | iconMap = getIconLabelMap() 238 | 239 | # iconMap['washing_basin'] = iconMap['washbasin'] 240 | # iconMap['sofa'] = 6 241 | # iconMap['chair'] = 5 242 | # iconMap['TV'] = iconMap['tv'] 243 | #for icon in ['cooking_counter', 'bathtub', 'toilet', 'washing_basin', 'sofa', 'cabinet', 'bed', 'table', 'desk', 'refrigerator', 'TV', 'entrance', 'chair']: 244 | 245 | labelMap = {} 246 | for icon in ['cooking_counter', 'bathtub', 'toilet', 'washing_basin', 'sofa', 'cabinet', 'bed', 'table', 'desk', 'refrigerator', 'TV', 'entrance', 'chair']: 247 | if icon not in iconMap: 248 | print(icon) 249 | exit(1) 250 | labelMap[icon] = ('icons', iconMap[icon]) 251 | continue 252 | for room in ['living_room', 'kitchen', 'bedroom', 'bathroom', 'office', 'closet', 'balcony', 'corridor', 'dining_room', 'stairs']: 253 | if room not in roomMap: 254 | print(room) 255 | exit(1) 256 | labelMap[room] = ('rooms', roomMap[room]) 257 | continue 258 | 259 | labelMap['door'] = 11 260 | labelMap['window'] = 12 261 | return labelMap 262 | 263 | def writeExample(self, scene_id): 264 | filename_glob = ROOT_FOLDER + scene_id + '/' + 'dataset/mesh/*.obj' 265 | filename = list(glob.glob(filename_glob))[0] 266 | points = load_pnt(filename) 267 | points = np.array(points) 268 | 269 | 270 | #segmentation = segmentation[sampledInds[:NUM_POINTS]] 271 | filename = ROOT_FOLDER + scene_id + '/annotation/metadata.t7' 272 | metadata = torchfile.load(filename) 273 | topDownViewTransformation = metadata["topDownTransformation"] 274 | 275 | #degree = metadata["topDownViewAngle"] 276 | rotMat = getZRotationMatrix(-metadata["topDownViewAngle"]) 277 | #print(rotMat, topDownViewTransformation) 278 | #exit(1) 279 | 280 | #XYZ_rotated = np.transpose(np.dot(rotMat, np.transpose(points[:, :3]))) 281 | #XYZ = np.tensordot(np.concatenate([points[:, :3], np.ones((points.shape[0], 1))], axis=1), topDownViewTransformation, axes=((1), (1))) 282 | #XYZ[:, 2] = points[:, 2] 283 | 284 | #ratio_1 = (XYZ[:, 0].max() - XYZ[:, 0].min()) / (XYZ_rotated[:, 0].max() - XYZ_rotated[:, 0].min()) 285 | #ratio_2 = (XYZ[:, 1].max() - XYZ[:, 1].min()) / (XYZ_rotated[:, 1].max() - XYZ_rotated[:, 1].min()) 286 | 287 | #XYZ[2, 2] *= np.sqrt(ratio_1 * ratio_2) 288 | 289 | #ratio = pow(np.abs(rotMat[0][0] / topDownViewTransformation[0][0] * rotMat[1][0] / topDownViewTransformation[1][0] * rotMat[0][1] / topDownViewTransformation[0][1] * rotMat[1][1] / topDownViewTransformation[1][1]), 0.25) 290 | ratio = 0 291 | for i in xrange(2): 292 | for j in xrange(2): 293 | if rotMat[i][j] != 0: 294 | ratio = max(ratio, np.abs(topDownViewTransformation[i][j] / rotMat[i][j])) 295 | pass 296 | continue 297 | continue 298 | 299 | globalTransformation = topDownViewTransformation 300 | globalTransformation[2, 2] = ratio 301 | globalTransformation[2, 3] = 0 302 | globalTransformation = np.concatenate([globalTransformation, np.zeros((1, 4))], axis=0) 303 | globalTransformation[3, 3] = 1 304 | 305 | 306 | XYZ = np.tensordot(np.concatenate([points[:, :3], np.ones((points.shape[0], 1))], axis=1), globalTransformation, axes=((1), (1))) 307 | XYZ = XYZ[:, :3] / XYZ[:, 3:] 308 | 309 | 310 | mins = XYZ.min(0, keepdims=True) 311 | maxs = XYZ.max(0, keepdims=True) 312 | maxRange = (maxs - mins)[:, :2].max() 313 | padding = maxRange * 0.05 314 | mins = (maxs + mins) / 2 - maxRange / 2 315 | mins -= padding 316 | maxRange += padding * 2 317 | minXY = mins[:, :2] 318 | 319 | #XYZ[:, :2] = (XYZ[:, :2] - minXY) / maxRange 320 | XYZ = (XYZ - mins) / maxRange 321 | points[:, :3] = XYZ 322 | 323 | originalWidth = 700. 324 | 325 | if points.shape[0] < NUM_POINTS: 326 | indices = np.arange(points.shape[0]) 327 | points = np.concatenate([points, points[np.random.choice(indices, NUM_POINTS - points.shape[0])]], axis=0) 328 | elif points.shape[0] > NUM_POINTS: 329 | sampledInds = np.arange(points.shape[0]) 330 | np.random.shuffle(sampledInds) 331 | points = points[sampledInds[:NUM_POINTS]] 332 | pass 333 | 334 | points[:, 3:] = points[:, 3:] / 255 - 0.5 335 | 336 | coordinates = np.clip(np.round(points[:, :2] * HEIGHT).astype(np.int32), 0, HEIGHT - 1) 337 | 338 | self.indicesMaps = np.zeros((NUM_POINTS), dtype=np.int64) 339 | self.projectIndices(np.concatenate([coordinates, np.arange(NUM_POINTS).reshape(-1, 1)], axis=1), 0, WIDTH, 0, HEIGHT) 340 | 341 | 342 | filename = ROOT_FOLDER + scene_id + '/annotation/floorplan.txt' 343 | walls = [] 344 | doors = [] 345 | windows = [] 346 | semantics = {} 347 | def transformPoint(v, c): 348 | return max(min(round((float(v) - minXY[0, c]) / maxRange * WIDTH), WIDTH - 1), 0) 349 | 350 | with open(filename) as info_file: 351 | line_index = 0 352 | for line in info_file: 353 | line = line.split('\t') 354 | if line[4] == 'wall': 355 | walls.append(((transformPoint(line[0], 0), transformPoint(line[1], 1)), (transformPoint(line[2], 0), transformPoint(line[3], 1)))) 356 | elif line[4] == 'door': 357 | doors.append(((transformPoint(line[0], 0), transformPoint(line[1], 1)), (transformPoint(line[2], 0), transformPoint(line[3], 1)))) 358 | elif line[4] == 'window': 359 | windows.append(((transformPoint(line[0], 0), transformPoint(line[1], 1)), (transformPoint(line[2], 0), transformPoint(line[3], 1)))) 360 | else: 361 | if line[4] not in semantics: 362 | semantics[line[4]] = [] 363 | pass 364 | semantics[line[4]].append(((transformPoint(line[0], 0), transformPoint(line[1], 1)), (transformPoint(line[2], 0), transformPoint(line[3], 1)))) 365 | pass 366 | continue 367 | pass 368 | 369 | 370 | roomSegmentation = np.zeros((HEIGHT, WIDTH), dtype=np.uint8) 371 | for line in walls: 372 | cv2.line(roomSegmentation, (int(round(line[0][0])), int(round(line[0][1]))), (int(round(line[1][0])), int(round(line[1][1]))), color = 15 + calcLineDirection(line), thickness=self.gap) 373 | #cv2.line(roomSegmentation, (int(round(line[0][0])), int(round(line[0][1]))), (int(round(line[1][0])), int(round(line[1][1]))), color = 15, thickness=self.gap) 374 | continue 375 | 376 | rooms = measure.label(roomSegmentation == 0, background=0) 377 | 378 | 379 | corners = lines2Corners(walls, gap=self.gap) 380 | #corner_gt = np.zeros((HEIGHT, WIDTH), dtype=np.uint8) 381 | corner_gt = [] 382 | for corner in corners: 383 | #corner_gt[int(round(corner[0][1])), int(round(corner[0][0]))] = corner[1] + 1 384 | corner_gt.append((int(round(corner[0][0])), int(round(corner[0][1])), corner[1] + 1)) 385 | continue 386 | 387 | openingCornerMap = [[3, 1], [0, 2]] 388 | for openingType, openings in enumerate([doors, windows]): 389 | for opening in openings: 390 | direction = calcLineDirection(opening) 391 | for cornerIndex, corner in enumerate(opening): 392 | #corner_gt[int(round(corner[1])), int(round(corner[0]))] = 14 + openingCornerMap[direction][cornerIndex] 393 | corner_gt.append((int(round(corner[0])), int(round(corner[1])), 14 + openingCornerMap[direction][cornerIndex])) 394 | continue 395 | continue 396 | continue 397 | 398 | 399 | wallIndex = rooms.min() 400 | for pixel in [(0, 0), (0, HEIGHT - 1), (WIDTH - 1, 0), (WIDTH - 1, HEIGHT - 1)]: 401 | backgroundIndex = rooms[pixel[1]][pixel[0]] 402 | if backgroundIndex != wallIndex: 403 | break 404 | continue 405 | 406 | iconSegmentation = np.zeros((HEIGHT, WIDTH), dtype=np.uint8) 407 | for line in doors: 408 | cv2.line(iconSegmentation, (int(round(line[0][0])), int(round(line[0][1]))), (int(round(line[1][0])), int(round(line[1][1]))), color = self.labelMap['door'], thickness=self.gap - 1) 409 | continue 410 | for line in windows: 411 | cv2.line(iconSegmentation, (int(round(line[0][0])), int(round(line[0][1]))), (int(round(line[1][0])), int(round(line[1][1]))), color = self.labelMap['window'], thickness=self.gap - 1) 412 | continue 413 | 414 | roomLabelMap = {} 415 | for semantic, items in semantics.iteritems(): 416 | group, label = self.labelMap[semantic] 417 | for corners in items: 418 | if group == 'icons': 419 | if label == 0: 420 | continue 421 | cv2.rectangle(iconSegmentation, (int(round(corners[0][0])), int(round(corners[0][1]))), (int(round(corners[1][0])), int(round(corners[1][1]))), color=label, thickness=-1) 422 | # corner_gt[int(round(corners[0][1])), int(round(corners[0][0]))] = 18 + 2 423 | # corner_gt[int(round(corners[1][1])), int(round(corners[0][0]))] = 18 + 1 424 | # corner_gt[int(round(corners[0][1])), int(round(corners[1][0]))] = 18 + 3 425 | # corner_gt[int(round(corners[1][1])), int(round(corners[1][0]))] = 18 + 0 426 | corner_gt.append((int(round(corners[0][0])), int(round(corners[0][1])), 18 + 2)) 427 | corner_gt.append((int(round(corners[0][0])), int(round(corners[1][1])), 18 + 1)) 428 | corner_gt.append((int(round(corners[1][0])), int(round(corners[0][1])), 18 + 3)) 429 | corner_gt.append((int(round(corners[1][0])), int(round(corners[1][1])), 18 + 0)) 430 | else: 431 | roomIndex = rooms[int(round((corners[0][1] + corners[1][1]) / 2))][int(round((corners[0][0] + corners[1][0]) / 2))] 432 | if roomIndex == wallIndex or roomIndex == backgroundIndex: 433 | print('label on background') 434 | exit(1) 435 | pass 436 | if roomIndex in roomLabelMap: 437 | print('room has more than one labels', label) 438 | exit(1) 439 | pass 440 | roomLabelMap[roomIndex] = label 441 | roomSegmentation[rooms == roomIndex] = label 442 | pass 443 | continue 444 | continue 445 | for roomIndex in xrange(rooms.min(), rooms.max() + 1): 446 | if roomIndex == wallIndex or roomIndex == backgroundIndex: 447 | continue 448 | if roomIndex not in roomLabelMap: 449 | print('room has no label') 450 | print(roomIndex, rooms.max()) 451 | pass 452 | continue 453 | flags = np.zeros(2, np.int64) 454 | flags[0] = 1 455 | 456 | corner_feature_file = ROOT_FOLDER + scene_id + '/corner_acc.npy' 457 | icon_feature_file = ROOT_FOLDER + scene_id + '/topdown_acc.npy' 458 | image_features = [[], []] 459 | if os.path.exists(corner_feature_file) and os.path.exists(icon_feature_file): 460 | flags[1] = 1 461 | corner_feature = np.load(corner_feature_file).reshape((HEIGHT, WIDTH, -1)) 462 | icon_feature = np.load(icon_feature_file).reshape((HEIGHT, WIDTH, -1)) 463 | for featureIndex, feature in enumerate([corner_feature, icon_feature]): 464 | for size, numChannels in zip(SIZES, NUM_CHANNELS)[1:]: 465 | feature = cv2.resize(feature, (size, size)) 466 | image_features[featureIndex].append(feature.reshape((size, size, numChannels, -1)).mean(-1).reshape(-1)) 467 | continue 468 | image_features[featureIndex] = np.concatenate(image_features[featureIndex], axis=0) 469 | continue 470 | image_features = image_features[0] + image_features[1] 471 | else: 472 | image_features = np.zeros(sum([size * size * numChannels for size, numChannels in zip(SIZES, NUM_CHANNELS)[1:]])) 473 | pass 474 | 475 | if False: 476 | cv2.imwrite('test/density.png', drawDensityImage(getDensity(points, HEIGHT, WIDTH))) 477 | cv2.imwrite('test/density_indices.png', drawDensityImage(getDensityFromIndices(self.indicesMaps, HEIGHT, WIDTH))) 478 | cv2.imwrite('test/icon_segmentation.png', drawSegmentationImage(iconSegmentation)) 479 | cv2.imwrite('test/room_segmentation.png', drawSegmentationImage(roomSegmentation)) 480 | cv2.imwrite('test/corner_segmentation.png', drawSegmentationImage(corner_gt, blackIndex=0)) 481 | if flags[1]: 482 | cv2.imwrite('test/topdown_corner.png', cv2.imread(ROOT_FOLDER + scene_id + '/corner_pred.png')) 483 | cv2.imwrite('test/topdown_icon.png', cv2.imread(ROOT_FOLDER + scene_id + '/topdown_pred_nonzero.png')) 484 | exit(1) 485 | pass 486 | pass 487 | 488 | corner_gt = np.array(corner_gt, dtype=np.int64) 489 | numCorners = len(corner_gt) 490 | print('num corners', numCorners) 491 | if numCorners > MAX_NUM_CORNERS: 492 | exit(1) 493 | elif numCorners < MAX_NUM_CORNERS: 494 | corner_gt = np.concatenate([corner_gt, np.zeros((MAX_NUM_CORNERS - numCorners, 3), dtype=np.int64)], axis=0) 495 | pass 496 | 497 | example = tf.train.Example(features=tf.train.Features(feature={ 498 | 'image_path': _bytes_feature(scene_id), 499 | 'points': _float_feature(points.reshape(-1)), 500 | 'point_indices': _int64_feature(self.indicesMaps.reshape(-1)), 501 | 'corner': _int64_feature(corner_gt.reshape(-1)), 502 | 'num_corners': _int64_feature([numCorners]), 503 | 'icon': _bytes_feature(iconSegmentation.tostring()), 504 | 'room': _bytes_feature(roomSegmentation.tostring()), 505 | 'image': _float_feature(image_features.reshape(-1)), 506 | 'flags': _int64_feature(flags), 507 | })) 508 | self.writer.write(example.SerializeToString()) 509 | return 510 | 511 | 512 | def projectSegmentation(self, pointSegmentation, min_x, max_x, min_y, max_y): 513 | if max_x - min_x == 1 and max_y - min_y == 1: 514 | segments, counts = np.unique(pointSegmentation[:, 2], return_counts=True) 515 | segmentList = zip(segments.tolist(), counts.tolist()) 516 | segmentList = [segment for segment in segmentList if segment[0] not in [0, 2]] 517 | label = 0 518 | if 2 in segments: 519 | label = 2 520 | pass 521 | if len(segmentList) > 0: 522 | segment = max(segmentList, key=lambda x: x[1]) 523 | if segment[1] > 0: 524 | label = segment[0] 525 | pass 526 | pass 527 | self.segmentation[min_y][min_x] = label 528 | elif max_x - min_x >= max_y - min_y: 529 | middle_x = int((max_x + min_x + 1) / 2) 530 | mask_1 = pointSegmentation[:, 1] < middle_x 531 | self.projectSegmentation(pointSegmentation[mask_1], min_x, middle_x, min_y, max_y) 532 | mask_2 = pointSegmentation[:, 1] >= middle_x 533 | self.projectSegmentation(pointSegmentation[mask_2], middle_x, max_x, min_y, max_y) 534 | else: 535 | middle_y = int((max_y + min_y + 1) / 2) 536 | mask_1 = pointSegmentation[:, 0] < middle_y 537 | self.projectSegmentation(pointSegmentation[mask_1], min_x, max_x, min_y, middle_y) 538 | mask_2 = pointSegmentation[:, 0] >= middle_y 539 | self.projectSegmentation(pointSegmentation[mask_2], min_x, max_x, middle_y, max_y) 540 | pass 541 | return 542 | 543 | 544 | def projectIndices(self, pointSegmentation, min_x, max_x, min_y, max_y): 545 | # for strideIndex in xrange(6): 546 | # stride = pow(2, strideIndex) 547 | # if max_x - min_x == stride and max_y - min_y == stride: 548 | # self.indicesMaps[strideIndex][pointSegmentation[:, 2]] = min_y / stride * WIDTH / stride + min_x / stride + self.batchIndex * (HEIGHT / stride * WIDTH / stride) 549 | # pass 550 | # continue 551 | 552 | if max_x - min_x == 1 and max_y - min_y == 1: 553 | self.indicesMaps[pointSegmentation[:, 2]] = min_y * WIDTH + min_x 554 | return 555 | elif max_x - min_x >= max_y - min_y: 556 | middle_x = int((max_x + min_x + 1) / 2) 557 | mask_1 = pointSegmentation[:, 0] < middle_x 558 | self.projectIndices(pointSegmentation[mask_1], min_x, middle_x, min_y, max_y) 559 | mask_2 = pointSegmentation[:, 0] >= middle_x 560 | self.projectIndices(pointSegmentation[mask_2], middle_x, max_x, min_y, max_y) 561 | else: 562 | middle_y = int((max_y + min_y + 1) / 2) 563 | mask_1 = pointSegmentation[:, 1] < middle_y 564 | self.projectIndices(pointSegmentation[mask_1], min_x, max_x, min_y, middle_y) 565 | mask_2 = pointSegmentation[:, 1] >= middle_y 566 | self.projectIndices(pointSegmentation[mask_2], min_x, max_x, middle_y, max_y) 567 | pass 568 | return 569 | 570 | if __name__ == "__main__": 571 | parser = argparse.ArgumentParser(description='Planenet') 572 | args = parser.parse_args() 573 | args.numPoints = 50000 574 | args.numInputChannels = 7 575 | args.augmentation = '' 576 | for split in ['val']: 577 | print(split) 578 | args.task = split 579 | RecordWriterTango(args) 580 | continue 581 | exit(1) 582 | 583 | # reader = BatchFetcher('train') 584 | # for i in xrange(12): 585 | # print(i) 586 | # reader.work(i) 587 | # continue 588 | # exit(1) 589 | -------------------------------------------------------------------------------- /augmentation_tf.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from utils import * 3 | import tensorflow as tf 4 | from floorplan_utils import * 5 | 6 | 7 | def warpIndices(xs, ys, gridStride, gridWidth, gridHeight, width, height, gridXsTarget, gridYsTarget): 8 | numPoints = xs.shape[0] 9 | 10 | minXs = xs / gridStride 11 | minYs = ys / gridStride 12 | maxXs = xs / gridStride + 1 13 | maxYs = ys / gridStride + 1 14 | 15 | topLeft = tf.expand_dims(minYs * gridWidth + minXs, -1) 16 | topRight = tf.expand_dims(minYs * gridWidth + maxXs, -1) 17 | bottomLeft = tf.expand_dims(maxYs * gridWidth + minXs, -1) 18 | bottomRight = tf.expand_dims(maxYs * gridWidth + maxXs, -1) 19 | 20 | topLeftXsTarget = tf.gather_nd(gridXsTarget, topLeft) 21 | topLeftYsTarget = tf.gather_nd(gridYsTarget, topLeft) 22 | topRightXsTarget = tf.gather_nd(gridXsTarget, topRight) 23 | topRightYsTarget = tf.gather_nd(gridYsTarget, topRight) 24 | bottomLeftXsTarget = tf.gather_nd(gridXsTarget, bottomLeft) 25 | bottomLeftYsTarget = tf.gather_nd(gridYsTarget, bottomLeft) 26 | bottomRightXsTarget = tf.gather_nd(gridXsTarget, bottomRight) 27 | bottomRightYsTarget = tf.gather_nd(gridYsTarget, bottomRight) 28 | 29 | ratioX = tf.cast(xs - minXs * gridStride, tf.float32) / float(gridStride) 30 | ratioY = tf.cast(ys - minYs * gridStride, tf.float32) / float(gridStride) 31 | topLeftRatio = (1 - ratioX) * (1 - ratioY) 32 | topRightRatio = ratioX * (1 - ratioY) 33 | bottomLeftRatio = (1 - ratioX) * ratioY 34 | bottomRightRatio = ratioX * ratioY 35 | 36 | xsTarget = topLeftXsTarget * topLeftRatio + topRightXsTarget * topRightRatio + bottomLeftXsTarget * bottomLeftRatio + bottomRightXsTarget * bottomRightRatio 37 | ysTarget = topLeftYsTarget * topLeftRatio + topRightYsTarget * topRightRatio + bottomLeftYsTarget * bottomLeftRatio + bottomRightYsTarget * bottomRightRatio 38 | 39 | xsTarget = tf.clip_by_value(tf.cast(tf.round(xsTarget), tf.int32), 0, width - 1) 40 | ysTarget = tf.clip_by_value(tf.cast(tf.round(ysTarget), tf.int32), 0, height - 1) 41 | return xsTarget, ysTarget 42 | 43 | def scaleIndices(xs, ys, min_x, min_y, max_x, max_y, width, height): 44 | xsTarget = (tf.cast(xs, tf.float32) - min_x) / (max_x - min_x + 1) * width 45 | ysTarget = (tf.cast(ys, tf.float32) - min_y) / (max_y - min_y + 1) * height 46 | xsTarget = tf.clip_by_value(tf.cast(tf.round(xsTarget), tf.int32), 0, width - 1) 47 | ysTarget = tf.clip_by_value(tf.cast(tf.round(ysTarget), tf.int32), 0, height - 1) 48 | return xsTarget, ysTarget 49 | 50 | #Get coarse indices maps from 256x256 indices map 51 | def getCoarseIndicesMaps(indicesMap, width=256, height=256, batchIndex=0): 52 | indicesMaps = [] 53 | for strideIndex in xrange(6): 54 | stride = pow(2, strideIndex) 55 | if strideIndex == 0: 56 | indicesMaps.append(indicesMap + batchIndex * width * height) 57 | else: 58 | indicesMaps.append(indicesMap / (width * stride) * (width / stride) + indicesMap % width / stride + batchIndex * width / stride * height / stride) 59 | pass 60 | #print(indicesMaps) 61 | continue 62 | indicesMaps = tf.stack(indicesMaps, axis=0) 63 | return indicesMaps 64 | 65 | #Get coarse indices maps from 256x256 indices map 66 | def getCoarseIndicesMapsBatch(indicesMap, width=256, height=256): 67 | indicesMaps = [] 68 | for strideIndex in xrange(6): 69 | stride = pow(2, strideIndex) 70 | if strideIndex == 0: 71 | indicesMaps.append(indicesMap) 72 | else: 73 | indicesMaps.append(indicesMap / (width * stride) * (width / stride) + indicesMap % width / stride) 74 | pass 75 | #print(indicesMaps) 76 | continue 77 | indicesMaps = tf.stack(indicesMaps, axis=0) 78 | return indicesMaps 79 | 80 | 81 | def augmentWarping(pointcloudIndices, corners, heatmaps, gridStride=16, randomScale=4): 82 | width = WIDTH 83 | height = HEIGHT 84 | gridWidth = int(width / gridStride + 1) 85 | gridHeight = int(height / gridStride + 1) 86 | 87 | gridXs = tf.reshape(tf.tile(tf.expand_dims(tf.range(gridWidth) * gridStride, 0), [gridHeight, 1]), [-1]) 88 | gridYs= tf.reshape(tf.tile(tf.expand_dims(tf.range(gridHeight) * gridStride, -1), [1, gridWidth]), [-1]) 89 | 90 | gridXsTarget = tf.cast(gridXs, tf.float32) + tf.random_normal(stddev=randomScale, shape=[gridHeight * gridWidth]) 91 | gridYsTarget = tf.cast(gridYs, tf.float32) + tf.random_normal(stddev=randomScale, shape=[gridHeight * gridWidth]) 92 | 93 | xsTarget, ysTarget = warpIndices(pointcloudIndices % width, pointcloudIndices / width, gridStride, gridWidth, gridHeight, width, height, gridXsTarget, gridYsTarget) 94 | 95 | newPointcloudIndices = tf.clip_by_value(ysTarget, 0, height - 1) * width + tf.clip_by_value(xsTarget, 0, width - 1) 96 | 97 | xsTarget, ysTarget = warpIndices(corners[:, 0], corners[:, 1], gridStride, gridWidth, gridHeight, width, height, gridXsTarget, gridYsTarget) 98 | newCorners = tf.stack([xsTarget, ysTarget, corners[:, 2]], axis=1) 99 | 100 | return newPointcloudIndices, newCorners, heatmaps 101 | 102 | def augmentScaling(pointcloud, pointcloudIndices, corners, heatmaps, imageFeatures): 103 | width = WIDTH 104 | height = HEIGHT 105 | xs = pointcloudIndices % width 106 | ys = pointcloudIndices / width 107 | 108 | imageSize = tf.constant((height, width), dtype=np.float32) 109 | #randomScale = pow(2.0, tf.random.uniform([1]) - 1) 110 | randomScale = tf.random_uniform(shape=[1], minval=0.5, maxval=1.5)[0] 111 | xsTarget = tf.clip_by_value(tf.cast(tf.round((tf.cast(xs, tf.float32) - width / 2) * randomScale + width / 2), tf.int32), 0, width - 1) 112 | ysTarget = tf.clip_by_value(tf.cast(tf.round((tf.cast(ys, tf.float32) - height / 2) * randomScale + height / 2), tf.int32), 0, height - 1) 113 | 114 | 115 | newPointcloudIndices = ysTarget * width + xsTarget 116 | 117 | pointcloud = (pointcloud - 0.5) * randomScale + 0.5 118 | 119 | newHeatmaps = tf.image.resize_nearest_neighbor(heatmaps, size = tf.cast(tf.round(imageSize * randomScale), tf.int32)) 120 | newHeatmaps = tf.image.resize_image_with_crop_or_pad(newHeatmaps, height, width) 121 | 122 | xsTarget = tf.cast(tf.round((tf.cast(corners[:, 0], tf.float32) - width / 2) * randomScale + width / 2), tf.int32) 123 | ysTarget = tf.cast(tf.round((tf.cast(corners[:, 1], tf.float32) - height / 2) * randomScale + height / 2), tf.int32) 124 | newCorners = tf.stack([xsTarget, ysTarget, corners[:, 2]], axis=1) 125 | validMask = tf.logical_and(tf.logical_and(tf.greater_equal(newCorners[:, 0], 0), tf.greater_equal(newCorners[:, 1], 0)), tf.logical_and(tf.less(newCorners[:, 0], WIDTH), tf.less(newCorners[:, 1], HEIGHT))) 126 | newCorners = tf.boolean_mask(newCorners, validMask) 127 | 128 | for index, (featureSize, numChannels) in enumerate(zip(SIZES, NUM_CHANNELS)[1:]): 129 | if index in imageFeatures: 130 | imageFeatures[index] = tf.image.resize_image_with_crop_or_pad(tf.image.resize_nearest_neighbor(tf.expand_dims(imageFeatures[index], 0), size = tf.cast(tf.round(tf.constant((featureSize, featureSize), dtype=tf.float32) * randomScale), tf.int32)), featureSize, featureSize)[0] 131 | pass 132 | continue 133 | 134 | return pointcloud, newPointcloudIndices, newCorners, newHeatmaps, imageFeatures 135 | 136 | 137 | def augmentFlipping(pointcloud, pointcloudIndices, corners, heatmaps, imageFeatures): 138 | width = WIDTH 139 | height = HEIGHT 140 | 141 | orientation = tf.cast(tf.random_uniform(shape=[1], maxval=4)[0], tf.int32) 142 | 143 | #if orientation == 0: 144 | #return pointcloud, pointcloudIndices, newCorners, heatmaps 145 | 146 | xsTarget = pointcloudIndices % width 147 | ysTarget = pointcloudIndices / width 148 | 149 | reverseChannelsY = tf.constant([-1, 2, 1, 0, 3, 7, 6, 5, 4, 10, 9, 8, 11, 12, 15, 14, 13, 16, 20, 19, 18, 17], dtype=tf.int32) + 1 150 | reverseChannelsX = tf.constant([-1, 0, 3, 2, 1, 5, 4, 7, 6, 8, 11, 10, 9, 12, 13, 16, 15, 14, 18, 17, 20, 19], dtype=tf.int32) + 1 151 | 152 | xsTarget = tf.cond(orientation >= 2, lambda: width - 1 - xsTarget, lambda: xsTarget) 153 | pointcloud = tf.cond(orientation >= 2, lambda: tf.concat([1 - pointcloud[:, 0:1], pointcloud[:, 1:]], axis=1), lambda: pointcloud) 154 | heatmaps = tf.cond(orientation >= 2, lambda: heatmaps[:, :, ::-1], lambda: heatmaps) 155 | corners = tf.cond(orientation >= 2, lambda: tf.stack([width - 1 - corners[:, 0], corners[:, 1], tf.gather_nd(reverseChannelsX, corners[:, 2:3])], axis=1), lambda: corners) 156 | 157 | ysTarget = tf.cond(tf.equal(orientation % 2, 1), lambda: height - 1 - ysTarget, lambda: ysTarget) 158 | pointcloud = tf.cond(tf.equal(orientation % 2, 1), lambda: tf.concat([pointcloud[:, :1], 1 - pointcloud[:, 1:2], pointcloud[:, 2:]], axis=1), lambda: pointcloud) 159 | heatmaps = tf.cond(tf.equal(orientation % 2, 1), lambda: heatmaps[:, ::-1], lambda: heatmaps) 160 | corners = tf.cond(tf.equal(orientation % 2, 1), lambda: tf.stack([corners[:, 0], height - 1 - corners[:, 1], tf.gather_nd(reverseChannelsY, corners[:, 2:3])], axis=1), lambda: corners) 161 | 162 | 163 | for index, (size, numChannels) in enumerate(zip(SIZES, NUM_CHANNELS)[1:]): 164 | if index in imageFeatures: 165 | imageFeatures[index] = tf.cond(orientation >= 2, lambda: imageFeatures[index][:, ::-1], lambda: imageFeatures[index]) 166 | imageFeatures[index] = tf.cond(tf.equal(orientation % 2, 1), lambda: imageFeatures[index][::-1], lambda: imageFeatures[index]) 167 | pass 168 | continue 169 | 170 | newPointcloudIndices = ysTarget * width + xsTarget 171 | 172 | return pointcloud, newPointcloudIndices, corners, heatmaps, imageFeatures 173 | 174 | 175 | def augmentDropping(pointcloud, pointcloud_indices, changeIndices): 176 | p = tf.random.random() * 0.5 + 0.5 177 | indices = tf.range(pointcloud.shape[0], dtype='int32') 178 | out_shape = int(pointcloud.shape[0] * p) 179 | chosen_indices = tf.random.choice(indices, (out_shape, ), replace=True) 180 | rest_mask = tf.ones(indices.shape, dtype=tf.bool) 181 | rest_mask[chosen_indices] = 0 182 | rest_indices = indices[rest_mask] 183 | #rest_indices = tf.array(list(set(indices) - set(chosen_indices))) 184 | #rest = pointcloud[rest_indices] 185 | 186 | #rest_chosen_indices = tf.random.choice(rest_indices, (pointcloud.shape[0] - out_shape, ), replace=True) 187 | #rest_chosen = rest[rest_chosen_indices] 188 | #aug_pointcloud = pointcloud[chosen_indices] = rest_chosen 189 | 190 | rest_indices = tf.random.choice(rest_indices, chosen_indices.shape[0], replace=True) 191 | pointcloud[chosen_indices] = pointcloud[rest_indices] 192 | if changeIndices: 193 | pointcloud_indices[chosen_indices] = pointcloud_indices[rest_indices] 194 | pass 195 | return pointcloud, pointcloud_indices 196 | 197 | 198 | 199 | def augment(pointcloud_inp, pointcloud_indices_0_inp, heatmapBatches, augmentation, numPoints=50000, numInputChannels=7): 200 | pointcloud_indices_inp = tf.zeros((FETCH_BATCH_SIZE, 6, NUM_POINTS),dtype='int32') 201 | newHeatmapBatches = [[] for heatmapIndex in xrange(len(heatmapBatches))] 202 | 203 | for imageIndex in xrange(pointcloud_inp.shape[0]): 204 | # pointcloud = pointcloud_inp[imageIndex] 205 | # pointcloud_indices_0 = pointcloud_indices_0_inp[imageIndex] 206 | # corner = corner_gt[imageIndex] 207 | # icon = icon_gt[imageIndex] 208 | # room = room_gt[imageIndex] 209 | # feature = feature_inp[imageIndex] 210 | # if 'w' in augmentation: 211 | # pointcloud_indices_0, [corner, icon, room, feature] = augmentWarping(pointcloud_indices_0, [corner, icon, room, feature], gridStride=32., randomScale=4) 212 | # pass 213 | # if 's' in augmentation: 214 | # pointcloud_indices_0, [corner, icon, room, feature] = augmentScaling(pointcloud_indices_0, [corner, icon, room, feature], randomScale=0) 215 | # pass 216 | # if 'f' in augmentation: 217 | # pointcloud_indices_0, [corner, icon, room, feature] = augmentFlipping(pointcloud_indices_0, [corner, icon, room, feature]) 218 | # pass 219 | # if 'd' in augmentation: 220 | # pointcloud, pointcloud_indices_0 = augmentDropping(pointcloud, pointcloud_indices_0, changeIndices=True) 221 | # pass 222 | # if 'p' in augmentation: 223 | # pointcloud, pointcloud_indices_0 = augmentDropping(pointcloud, pointcloud_indices_0, changeIndices=False) 224 | # pass 225 | 226 | # pointcloud_inp[imageIndex] = pointcloud 227 | # pointcloud_indices_inp[imageIndex] = getCoarseIndicesMaps(pointcloud_indices_0, WIDTH, HEIGHT, 0) 228 | # corner_gt[imageIndex] = corner 229 | # icon_gt[imageIndex] = icon 230 | # room_gt[imageIndex] = room 231 | # feature_inp[imageIndex] = feature 232 | 233 | 234 | newHeatmaps = [heatmapBatch[imageIndex] for heatmapBatch in heatmapBatches] 235 | if 'w' in augmentation: 236 | pointcloud_indices_0_inp[imageIndex], newHeatmaps = augmentWarping(pointcloud_indices_0_inp[imageIndex], newHeatmaps, gridStride=32, randomScale=4) 237 | pass 238 | if 's' in augmentation: 239 | pointcloud_inp[imageIndex], pointcloud_indices_0_inp[imageIndex], newHeatmaps = augmentScaling(pointcloud_inp[imageIndex], pointcloud_indices_0_inp[imageIndex], newHeatmaps) 240 | pass 241 | if 'f' in augmentation: 242 | pointcloud_inp[imageIndex], pointcloud_indices_0_inp[imageIndex], newHeatmaps = augmentFlipping(pointcloud_inp[imageIndex], pointcloud_indices_0_inp[imageIndex], newHeatmaps) 243 | pass 244 | if 'd' in augmentation: 245 | pointcloud_inp[imageIndex], pointcloud_indices_0_inp[imageIndex] = augmentDropping(pointcloud_inp[imageIndex], pointcloud_indices_0_inp[imageIndex], changeIndices=True) 246 | pass 247 | if 'p' in augmentation: 248 | pointcloud_inp[imageIndex], pointcloud_indices_0_inp[imageIndex] = augmentDropping(pointcloud_inp[imageIndex], pointcloud_indices_0_inp[imageIndex], changeIndices=False) 249 | pass 250 | 251 | #print(pointcloud_indices_0_inp[imageIndex].shape, pointcloud_indices_inp[imageIndex].shape) 252 | pointcloud_indices_inp[imageIndex] = getCoarseIndicesMaps(pointcloud_indices_0_inp[imageIndex], WIDTH, HEIGHT, 0) 253 | for heatmapIndex, newHeatmap in enumerate(newHeatmaps): 254 | newHeatmapBatches[heatmapIndex].append(newHeatmap) 255 | continue 256 | continue 257 | newHeatmapBatches = [tf.array(newHeatmapBatch) for newHeatmapBatch in newHeatmapBatches] 258 | pointcloud_inp = tf.concatenate([pointcloud_inp, tf.ones((FETCH_BATCH_SIZE, NUM_POINTS, 1))], axis=2) 259 | #print(pointcloud_itf.shape) 260 | #writePointCloud('test/pointcloud.ply', pointcloud_inp[0, :, :6]) 261 | #exit(1) 262 | 263 | if numPoints < pointcloud_itf.shape[1]: 264 | sampledInds = tf.range(pointcloud_itf.shape[1]) 265 | tf.random.shuffle(sampledInds) 266 | sampledInds = sampledInds[:numPoints] 267 | pointcloud_inp = pointcloud_inp[:, sampledInds] 268 | pointcloud_indices_inp = pointcloud_indices_inp[:, :, sampledInds] 269 | pass 270 | 271 | if numInputChannels == 4: 272 | pointcloud_inp = tf.concatenate([pointcloud_inp[:, :, :3], pointcloud_inp[:, :, 6:]], axis=2) 273 | pass 274 | 275 | 276 | return pointcloud_inp, pointcloud_indices_inp, newHeatmapBatches 277 | -------------------------------------------------------------------------------- /checkpoint/floornet_hybrid1_branch0123_wsf/checkpoint: -------------------------------------------------------------------------------- 1 | model_checkpoint_path: "checkpoint.ckpt" 2 | all_model_checkpoint_paths: "checkpoint.ckpt" 3 | -------------------------------------------------------------------------------- /checkpoint/floornet_hybrid1_branch0123_wsf/checkpoint.ckpt.data-00000-of-00001: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/FloorNet/93db5bdf12f8b60c9d788aa8e124fbe852e50b4a/checkpoint/floornet_hybrid1_branch0123_wsf/checkpoint.ckpt.data-00000-of-00001 -------------------------------------------------------------------------------- /checkpoint/floornet_hybrid1_branch0123_wsf/checkpoint.ckpt.index: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/FloorNet/93db5bdf12f8b60c9d788aa8e124fbe852e50b4a/checkpoint/floornet_hybrid1_branch0123_wsf/checkpoint.ckpt.index -------------------------------------------------------------------------------- /checkpoint/floornet_hybrid1_branch0123_wsf/checkpoint.ckpt.meta: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/FloorNet/93db5bdf12f8b60c9d788aa8e124fbe852e50b4a/checkpoint/floornet_hybrid1_branch0123_wsf/checkpoint.ckpt.meta -------------------------------------------------------------------------------- /evaluate.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | from QP import reconstructFloorplan, findMatches 4 | from RecordReader import * 5 | from train import * 6 | 7 | # Disable 8 | def blockPrint(): 9 | return 10 | sys.stdout = open(os.devnull, 'w') 11 | 12 | # Restore 13 | def enablePrint(): 14 | return 15 | sys.stdout = sys.__stdout__ 16 | 17 | def evaluate(options): 18 | if not os.path.exists(options.test_dir): 19 | os.system("mkdir -p %s"%options.test_dir) 20 | pass 21 | if os.path.exists(options.test_dir + '/dummy'): 22 | #os.rmdir(options.test_dir + '/dummy') 23 | pass 24 | else: 25 | os.mkdir(options.test_dir + '/dummy') 26 | pass 27 | 28 | if options.useCache == 2 and os.path.exists(options.test_dir + '/dummy/gt_dict.npy') and os.path.exists(options.test_dir + '/dummy/pred_dict.npy'): 29 | return 30 | 31 | if options.useCache == 1 and os.path.exists(options.test_dir + '/dummy/gt_dict.npy') and os.path.exists(options.test_dir + '/dummy/pred_dict.npy'): 32 | evaluateBatch(options) 33 | return 34 | 35 | if not os.path.exists(options.test_dir): 36 | os.system("mkdir -p %s"%options.test_dir) 37 | pass 38 | 39 | #print(options.checkpoint_dir) 40 | tf.reset_default_graph() 41 | 42 | filenames = [] 43 | if '0' in options.dataset: 44 | filenames.append('data/Syn_val.tfrecords') 45 | if '1' in options.dataset: 46 | filenames.append('data/Tango_val.tfrecords') 47 | pass 48 | if '2' in options.dataset: 49 | filenames.append('data/ScanNet_val.tfrecords') 50 | pass 51 | if '3' in options.dataset: 52 | filenames.append('data/Matterport_val.tfrecords') 53 | pass 54 | if '4' in options.dataset: 55 | filenames.append('data/SUNCG_val.tfrecords') 56 | pass 57 | 58 | dataset = getDatasetVal(filenames, '', '4' in options.branches, options.batchSize) 59 | 60 | 61 | iterator = dataset.make_one_shot_iterator() 62 | input_dict, gt_dict = iterator.get_next() 63 | 64 | 65 | pred_dict, debug_dict = build_graph(options, input_dict) 66 | dataset_flag = input_dict['flags'][0, 0] 67 | flags = input_dict['flags'][:, 1] 68 | loss, loss_list = build_loss(options, pred_dict, gt_dict, dataset_flag, debug_dict, input_dict['flags']) 69 | 70 | var_to_restore = [v for v in tf.global_variables()] 71 | 72 | config = tf.ConfigProto() 73 | config.gpu_options.allow_growth=True 74 | config.allow_soft_placement=True 75 | #config.log_device_placement=True 76 | 77 | statisticsSum = {k: [0.0, 0.0, 0.0] for k in ['wall', 'door', 'icon', 'room']} 78 | 79 | numbers = {} 80 | 81 | with tf.Session(config=config) as sess: 82 | sess.run(tf.global_variables_initializer()) 83 | tflearn.is_training(False) 84 | #var_to_restore = [v for v in var_to_restore if 'pred_room' not in v.name] 85 | var_to_restore = [v for v in var_to_restore if 'is_training' not in v.name] 86 | loader = tf.train.Saver(var_to_restore) 87 | if options.startIteration <= 0: 88 | loader.restore(sess,"%s/checkpoint.ckpt"%(options.checkpoint_dir)) 89 | else: 90 | loader.restore(sess,"%s/checkpoint_%d.ckpt"%(options.checkpoint_dir, options.startIteration)) 91 | pass 92 | 93 | #if tf.train.checkpoint_exists("%s/%s.ckpt"%(dumpdir,keyname)): 94 | #saver.restore(sess,"%s/%s.ckpt"%(dumpdir,keyname)) 95 | #pass 96 | 97 | MOVING_AVERAGE_DECAY = 1 98 | losses = [0., 0., 0.] 99 | acc = [1e-4, 1e-4, 1e-4] 100 | 101 | cornerCounters = {} 102 | for cornerType in CORNER_RANGES.keys(): 103 | cornerCounters[cornerType] = np.zeros(3) 104 | pass 105 | 106 | globalCornerCounter = np.zeros(3) 107 | iconCounter = np.zeros(2) 108 | roomCounter = np.zeros(2) 109 | 110 | numImages = 0 111 | try: 112 | for iteration in xrange(options.numTestingImages): 113 | total_loss, losses, dataset, image_flags, gt, pred, debug, inp = sess.run([loss, loss_list, dataset_flag, flags, gt_dict, pred_dict, debug_dict, input_dict]) 114 | 115 | for lossIndex, value in enumerate(losses): 116 | losses[lossIndex] = losses[lossIndex] * MOVING_AVERAGE_DECAY + value 117 | acc[lossIndex] = acc[lossIndex] * MOVING_AVERAGE_DECAY + 1 118 | continue 119 | print('testing', losses[0] / acc[0], losses[1] / acc[1], losses[2] / acc[2]) 120 | 121 | gt = {'corner': gt['corner'], 'corner_values': gt['corner_values'], 'icon': gt['icon'], 'room': gt['room'], 'density': debug['x0_topdown'][:, :, :, -1], 'image_path': inp['image_path'], 'num_corners': gt['num_corners'], 'image_flags': image_flags} 122 | if iteration == 0: 123 | gtAll = gt 124 | predAll = pred 125 | else: 126 | for k, v in gt.iteritems(): 127 | gtAll[k] = np.concatenate([gtAll[k], v], axis=0) 128 | continue 129 | for k, v in pred.iteritems(): 130 | predAll[k] = np.concatenate([predAll[k], v], axis=0) 131 | continue 132 | pass 133 | continue 134 | except tf.errors.OutOfRangeError: 135 | print('Finish testing') 136 | pass 137 | 138 | pass 139 | 140 | if options.useCache != -1: 141 | np.save(options.test_dir + '/dummy/gt_dict.npy', gtAll) 142 | np.save(options.test_dir + '/dummy/pred_dict.npy', predAll) 143 | pass 144 | if options.useCache == -2: 145 | return 146 | 147 | evaluateBatch(options, gtAll, predAll) 148 | return 149 | 150 | def evaluateBatch(options, gt_dict=None, pred_dict=None): 151 | datasetFlag = 1 152 | if options.useCache != -1: 153 | if options.loss != '5': 154 | gt_dict = np.load(options.test_dir + '/dummy/gt_dict.npy')[()] 155 | pred_dict = np.load(options.test_dir + '/dummy/pred_dict.npy')[()] 156 | else: 157 | gt_dict = np.load(options.test_dir.replace('loss5', 'loss0') + '/dummy/gt_dict.npy')[()] 158 | pred_wc = np.load(options.test_dir.replace('loss5', 'loss0') + '/dummy/pred_dict.npy')[()]['corner'][:, :, :, :NUM_WALL_CORNERS] 159 | pred_oc = np.load(options.test_dir.replace('loss5', 'loss1') + '/dummy/pred_dict.npy')[()]['corner'][:, :, :, NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4] 160 | pred_ic = np.load(options.test_dir.replace('loss5', 'loss2').replace('hybrid14', 'hybrid1') + '/dummy/pred_dict.npy')[()]['corner'][:, :, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8] 161 | pred_icon = np.load(options.test_dir.replace('loss5', 'loss3').replace('hybrid14', 'hybrid1') + '/dummy/pred_dict.npy')[()]['icon'] 162 | pred_room = np.load(options.test_dir.replace('loss5', 'loss4').replace('hybrid14', 'hybrid1') + '/dummy/pred_dict.npy')[()]['room'] 163 | pred_dict = {'corner': np.concatenate([pred_wc, pred_oc, pred_ic], axis=-1), 'icon': pred_icon, 'room': pred_room} 164 | pass 165 | 166 | if options.separateIconLoss: 167 | pred_icon_separate = softmax(np.load(options.test_dir.replace('wsf', 'wsf_loss3') + '/dummy/pred_dict.npy')[()]['icon']) 168 | pass 169 | #pred_dict['icon'] = np.load(options.test_dir.replace('wsf', 'wsf_loss3').replace('hybrid1', 'hybrid14').replace('dataset_1', '') + '/dummy/pred_dict.npy')[()]['icon'] 170 | #pred_dict['corner'][:, :, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8] = np.load(options.test_dir.replace('wsf', 'wsf_loss2') + '/dummy/pred_dict.npy')[()]['corner'][:, :, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8] 171 | #pass 172 | 173 | if options.cornerLossType != 'mse': 174 | threshold = np.ones((HEIGHT, WIDTH, 1)) * 0.5 175 | else: 176 | threshold = np.ones((HEIGHT, WIDTH, 1)) * 0.5# HEATMAP_SCALE / 2 177 | pass 178 | 179 | statisticsSum = {k: [0.0, 0.0, 0.0] for k in ['wall', 'door', 'icon', 'room', 'neighbor', 'neighbor_all']} 180 | #print(pred_dict['corner'].max()) 181 | pred_wc = pred_dict['corner'][:, :, :, :NUM_WALL_CORNERS] 182 | pred_oc = pred_dict['corner'][:, :, :, NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4] 183 | pred_ic = pred_dict['corner'][:, :, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8] 184 | 185 | if options.branches != '5': 186 | pred_wc = sigmoid(pred_wc) 187 | pred_oc = sigmoid(pred_oc) 188 | pred_ic = sigmoid(pred_ic) 189 | else: 190 | threshold = np.ones((HEIGHT, WIDTH, 1)) * 0.3 191 | pass 192 | 193 | gt_wc = gt_dict['corner'][:, :, :, :NUM_WALL_CORNERS] 194 | gt_oc = gt_dict['corner'][:, :, :, NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4] 195 | gt_ic = gt_dict['corner'][:, :, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8] 196 | 197 | names = [] 198 | 199 | for batchIndex in xrange(gt_dict['corner'].shape[0]): 200 | #if batchIndex == 0: 201 | #continue 202 | 203 | #if options.branches == '4' and gt_dict['image_flags'][batchIndex] == 0: 204 | if options.evaluateImage and gt_dict['image_flags'][batchIndex] == 0: 205 | continue 206 | 207 | density = np.minimum(gt_dict['density'][batchIndex] * 255, 255).astype(np.uint8) 208 | density = np.stack([density, density, density], axis=2) 209 | 210 | pred_icon = softmax(pred_dict['icon'][batchIndex]) 211 | pred_room = softmax(pred_dict['room'][batchIndex]) 212 | if options.separateIconLoss: 213 | pred_icon[:, :, :-2] = pred_icon_separate[batchIndex][:, :, :-2] 214 | #pred_icon = pred_icon_separate[batchIndex] 215 | pass 216 | 217 | if False: 218 | #print('batch index', batchIndex) 219 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_density.png', density) 220 | #print('heatmap max value', pred_wc[batchIndex].max()) 221 | 222 | if datasetFlag in [0, 1, 4]: 223 | cornerImage = drawSegmentationImage(np.concatenate([threshold, pred_wc[batchIndex]], axis=2), blackIndex=0) 224 | cornerImage[cornerImage == 0] = density[cornerImage == 0] 225 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_corner_pred.png', cornerImage) 226 | 227 | cornerImage = drawSegmentationImage(np.concatenate([threshold, gt_wc[batchIndex]], axis=2), blackIndex=0) 228 | cornerImage[cornerImage == 0] = density[cornerImage == 0] 229 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_corner_gt.png', cornerImage) 230 | pass 231 | 232 | 233 | if False: 234 | corner_heat = np.max(pred_wc[batchIndex], axis=-1) 235 | #print('corner_shape', corner_heat.shape) 236 | cmap = plt.get_cmap('jet') 237 | corner_rgba_img = cmap(corner_heat) 238 | corner_rgb_img = np.delete(corner_rgba_img, 3, 2) 239 | #print('rgb_out', corner_rgb_img.shape, corner_rgb_img.max(), corner_rgb_img.min()) 240 | corner_rgb_img = (corner_rgb_img * 255).round().astype('uint8') 241 | #print('rgb_out', corner_rgb_img.shape, corner_rgb_img.max()) 242 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_corner_heatmap.png', corner_rgb_img) 243 | pass 244 | 245 | if datasetFlag in [1, 4]: 246 | cornerImage = drawSegmentationImage(np.concatenate([threshold, pred_oc[batchIndex]], axis=2), blackIndex=0) 247 | cornerImage[cornerImage == 0] = density[cornerImage == 0] 248 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_opening_corner_pred.png', cornerImage) 249 | 250 | cornerImage = drawSegmentationImage(np.concatenate([threshold, pred_ic[batchIndex]], axis=2), blackIndex=0) 251 | cornerImage[cornerImage == 0] = density[cornerImage == 0] 252 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_icon_corner_pred.png', cornerImage) 253 | 254 | 255 | cornerImage = drawSegmentationImage(np.concatenate([threshold, gt_oc[batchIndex]], axis=2), blackIndex=0) 256 | cornerImage[cornerImage == 0] = density[cornerImage == 0] 257 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_opening_corner_gt.png', cornerImage) 258 | 259 | cornerImage = drawSegmentationImage(np.concatenate([threshold, gt_ic[batchIndex]], axis=2), blackIndex=0) 260 | cornerImage[cornerImage == 0] = density[cornerImage == 0] 261 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_icon_corner_gt.png', cornerImage) 262 | pass 263 | 264 | 265 | if datasetFlag in [1, 2, 3, 4]: 266 | icon_density = drawSegmentationImage(gt_dict['icon'][batchIndex], blackIndex=0) 267 | icon_density[icon_density == 0] = density[icon_density == 0] 268 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_icon_gt.png', icon_density) 269 | 270 | icon_density = drawSegmentationImage(pred_dict['icon'][batchIndex], blackIndex=0) 271 | icon_density[icon_density == 0] = density[icon_density == 0] 272 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_icon_pred.png', icon_density) 273 | pass 274 | 275 | if datasetFlag in [1, 3, 4]: 276 | room_density = drawSegmentationImage(gt_dict['room'][batchIndex], blackIndex=0) 277 | room_density[room_density == 0] = density[room_density == 0] 278 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_room_gt.png', room_density) 279 | 280 | room_density = drawSegmentationImage(pred_dict['room'][batchIndex], blackIndex=0) 281 | room_density[room_density == 0] = density[room_density == 0] 282 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_room_pred.png', room_density) 283 | pass 284 | 285 | 286 | if batchIndex == 0 and False: 287 | for c in xrange(22): 288 | cv2.imwrite(options.test_dir + '/mask_' + str(c) + '.png', cv2.dilate(drawMaskImage(corner_segmentation[batchIndex] == c), np.ones((3, 3)), 3)) 289 | continue 290 | continue 291 | 292 | 293 | if batchIndex < options.visualizeReconstruction or True: 294 | if options.debug >= 0 and batchIndex != options.debug: 295 | continue 296 | names.append((batchIndex, gt_dict['image_path'][batchIndex])) 297 | print(batchIndex, 'start reconstruction', gt_dict['image_path'][batchIndex]) 298 | if True: 299 | if options.debug == -1: 300 | blockPrint() 301 | pass 302 | 303 | # gtHeatmaps = gt_dict['corner'][batchIndex] 304 | #result_gt = reconstructFloorplan(gtHeatmaps[:, :, :NUM_WALL_CORNERS], gtHeatmaps[:, :, NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4], gtHeatmaps[:, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8], segmentation2Heatmaps(gt_dict['icon'][batchIndex], NUM_ICONS), segmentation2Heatmaps(gt_dict['room'][batchIndex], NUM_ROOMS), density[:, :, 0], gt=True) 305 | orientationCorners = getOrientationCorners(gt_dict['corner_values'][batchIndex][:gt_dict['num_corners'][batchIndex]]) 306 | result_gt = reconstructFloorplan(orientationCorners[:NUM_WALL_CORNERS], orientationCorners[NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4], orientationCorners[NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8], segmentation2Heatmaps(gt_dict['icon'][batchIndex], NUM_ICONS), segmentation2Heatmaps(gt_dict['room'][batchIndex], NUM_ROOMS), density[:, :, 0], gt=True) 307 | 308 | #if batchIndex == 1: 309 | #exit(1) 310 | 311 | 312 | #pred_debug_dir = options.test_dir + '/' + str(batchIndex) + '_debug' 313 | pred_debug_dir = options.test_dir 314 | try: 315 | os.mkdir(pred_debug_dir) 316 | pass 317 | except OSError as e: 318 | pass 319 | 320 | result_pred = reconstructFloorplan(pred_wc[batchIndex], pred_oc[batchIndex], pred_ic[batchIndex], pred_icon, pred_room, density[:, :, 0], gt_dict=result_gt, gt=False, debug_prefix=pred_debug_dir) 321 | 322 | if True: 323 | try: 324 | newWidth = newHeight = 1000 325 | resizeResult(result_gt, newWidth, newHeight, WIDTH, HEIGHT) 326 | resultImageGT = drawResultImageFinal(newWidth, newHeight, result_gt) 327 | #cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_result_gt.png', resultImageGT) 328 | cv2.imwrite(options.test_dir + '/' + gt_dict['image_path'][batchIndex] + '_gt.png', resultImageGT) 329 | 330 | resizeResult(result_pred, newWidth, newHeight, WIDTH, HEIGHT) 331 | resultImagePred = drawResultImageFinal(newWidth, newHeight, result_pred) 332 | cv2.imwrite(options.test_dir + '/' + gt_dict['image_path'][batchIndex] + '_pred.png', resultImagePred) 333 | except: 334 | continue 335 | continue 336 | 337 | if 'wall' not in result_pred or 'wall' not in result_gt: 338 | print('invalid result') 339 | continue 340 | 341 | statistics = findMatches(result_pred, result_gt, distanceThreshold=10) 342 | 343 | if options.drawFinal: 344 | newWidth = newHeight = 1000 345 | resizeResult(result_gt, newWidth, newHeight, WIDTH, HEIGHT) 346 | resultImageGT = drawResultImageFinal(newWidth, newHeight, result_gt) 347 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_result_gt.png', resultImageGT) 348 | 349 | resizeResult(result_pred, newWidth, newHeight, WIDTH, HEIGHT) 350 | resultImagePred = drawResultImageFinal(newWidth, newHeight, result_pred) 351 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_result_pred.png', resultImagePred) 352 | 353 | writeRepresentation('popup/data/floorplan_' + str(batchIndex) + '_gt.txt', newWidth, newHeight, result_gt) 354 | writeRepresentation('popup/data/floorplan_' + str(batchIndex) + '_pred.txt', newWidth, newHeight, result_pred) 355 | cv2.imwrite('popup/data/floorplan_' + str(batchIndex) + '_gt.png', resultImageGT) 356 | cv2.imwrite('popup/data/floorplan_' + str(batchIndex) + '_pred.png', resultImagePred) 357 | exit(1) 358 | else: 359 | resultImage, iconImage = drawResultImage(WIDTH, HEIGHT, result_gt) 360 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_reconstruction_wall_gt.png', resultImage) 361 | iconImage[iconImage == 0] = density[iconImage == 0] 362 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_reconstruction_icon_gt.png', iconImage) 363 | resultImage, iconImage = drawResultImage(WIDTH, HEIGHT, result_pred) 364 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_reconstruction_wall_pred.png', resultImage) 365 | iconImage[iconImage == 0] = density[iconImage == 0] 366 | cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_reconstruction_icon_pred.png', iconImage) 367 | pass 368 | 369 | 370 | if options.debug == -1: 371 | enablePrint() 372 | pass 373 | if len(result_pred) == 0: 374 | continue 375 | 376 | # print(result_pred) 377 | # print(result_pred['door']) 378 | # print('gt') 379 | # print(result_gt) 380 | # print(result_gt['door']) 381 | # exit(1) 382 | 383 | print('find predictions among ground-truths') 384 | #print(result_pred['wall'][2]) 385 | #statistics = findMatches(result_pred, result_gt, distanceThreshold=10) 386 | #statistics = findMatches(result_gt, result_pred, distanceThreshold=10) 387 | 388 | #print('find ground-truths among predictions') 389 | #statistics = findMatches(result_gt, result_pred, distanceThreshold=10) 390 | #print(statistics) 391 | print('statistics', [(k, float(v[0]) / max(v[1], 1), float(v[0]) / max(v[2], 1)) for k, v in statistics.iteritems()]) 392 | #print('topology statistics', [(k, float(v[0]) / max(v[1], 1), float(v[0]) / max(v[2], 1)) for k, v in topologyStatistics.iteritems()]) 393 | print('finish reconstruction', gt_dict['image_path'][batchIndex]) 394 | for k, v in statistics.iteritems(): 395 | if k in statisticsSum: 396 | for c in xrange(3): 397 | statisticsSum[k][c] += v[c] 398 | continue 399 | else: 400 | print(k, 'not in', statisticsSum) 401 | continue 402 | if options.debug >= 0: 403 | exit(1) 404 | pass 405 | pass 406 | 407 | # except Exception as e: 408 | # #traceback.print_tb(e) 409 | # print('exception-----------: ', e) 410 | # #raise e 411 | continue 412 | print(names) 413 | print('final statistics', [(k, float(v[0]) / max(v[1], 1), float(v[0]) / max(v[2], 1)) for k, v in statisticsSum.iteritems()]) 414 | np.save(options.test_dir + '/numbers.npy', statisticsSum) 415 | #print(statisticsSum) 416 | return 417 | -------------------------------------------------------------------------------- /floorplan_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from skimage import measure 3 | import cv2 4 | import copy 5 | from utils import * 6 | 7 | NUM_WALL_CORNERS = 13 8 | NUM_CORNERS = 21 9 | CORNER_RANGES = {'wall': (0, 13), 'opening': (13, 17), 'icon': (17, 21)} 10 | MAX_NUM_CORNERS = 300 11 | 12 | NUM_FINAL_ICONS = 10 13 | NUM_FINAL_ROOMS = 15 14 | NUM_ICONS = 13 15 | NUM_ROOMS = 16 16 | HEIGHT=256 17 | WIDTH=256 18 | NUM_POINTS = 50000 19 | 20 | NUM_INPUT_CHANNELS = 7 21 | NUM_CHANNELS = [7, 64, 64, 64, 128, 256] 22 | SIZES = [WIDTH, WIDTH // 2, WIDTH // 4, WIDTH // 8, WIDTH // 16, WIDTH // 32] 23 | 24 | 25 | 26 | POINT_ORIENTATIONS = [[(2, ), (3, ), (0, ), (1, )], [(0, 3), (0, 1), (1, 2), (2, 3)], [(1, 2, 3), (0, 2, 3), (0, 1, 3), (0, 1, 2)], [(0, 1, 2, 3)]] 27 | 28 | 29 | def getOrientationRanges(width, height): 30 | orientationRanges = [[width, 0, 0, 0], [width, height, width, 0], [width, height, 0, height], [0, height, 0, 0]] 31 | return orientationRanges 32 | 33 | def getIconNames(): 34 | iconNames = [] 35 | iconLabelMap = getIconLabelMap() 36 | for iconName, _ in iconLabelMap.iteritems(): 37 | iconNames.append(iconName) 38 | continue 39 | return iconNames 40 | 41 | def getRoomLabelMap(): 42 | labelMap = {} 43 | labelMap['living_room'] = 1 44 | labelMap['kitchen'] = 2 45 | labelMap['bedroom'] = 3 46 | labelMap['bathroom'] = 4 47 | labelMap['restroom'] = 4 48 | labelMap['office'] = 3 49 | labelMap['closet'] = 6 50 | labelMap['balcony'] = 7 51 | labelMap['corridor'] = 8 52 | labelMap['dining_room'] = 9 53 | labelMap['laundry_room'] = 10 54 | labelMap['garage'] = 11 55 | labelMap['recreation_room'] = 12 56 | labelMap['stairs'] = 13 57 | labelMap['other'] = 14 58 | labelMap['wall'] = 15 59 | return labelMap 60 | 61 | def getLabelRoomMap(): 62 | labelMap = {} 63 | labelMap[1] = 'living room' 64 | labelMap[2] = 'kitchen' 65 | labelMap[3] = 'bedroom' 66 | labelMap[4] = 'bathroom' 67 | labelMap[6] = 'closet' 68 | labelMap[7] = 'balcony' 69 | labelMap[8] = 'corridor' 70 | labelMap[9] = 'dining room' 71 | return labelMap 72 | 73 | def getIconLabelMap(): 74 | labelMap = {} 75 | labelMap['cooking_counter'] = 1 76 | labelMap['bathtub'] = 2 77 | labelMap['toilet'] = 3 78 | labelMap['washing_basin'] = 4 79 | labelMap['sofa'] = 5 80 | labelMap['cabinet'] = 6 81 | labelMap['bed'] = 7 82 | labelMap['table'] = 8 83 | labelMap['desk'] = 8 84 | labelMap['refrigerator'] = 9 85 | labelMap['TV'] = 0 86 | labelMap['entrance'] = 0 87 | labelMap['chair'] = 0 88 | labelMap['door'] = 11 89 | labelMap['window'] = 12 90 | return labelMap 91 | 92 | def getLabelIconMap(): 93 | labelMap = {} 94 | labelMap[1] = 'cooking_counter' 95 | labelMap[2] = 'bathtub' 96 | labelMap[3] = 'toilet' 97 | labelMap[4] = 'washing_basin' 98 | labelMap[5] = 'sofa' 99 | labelMap[6] = 'cabinet' 100 | labelMap[7] = 'bed' 101 | labelMap[8] = 'table' 102 | labelMap[9] = 'refrigerator' 103 | return labelMap 104 | 105 | def getLabelMapNYU40(): 106 | labelMap = {} 107 | labelMap[1] = 'wall' 108 | labelMap[2] = 'floor' 109 | labelMap[3] = 'cabinet' 110 | labelMap[4] = 'bed' 111 | labelMap[5] = 'chair' 112 | labelMap[6] = 'sofa' 113 | labelMap[7] = 'table' 114 | labelMap[8] = 'door' 115 | labelMap[9] = 'window' 116 | labelMap[10] = 'bookshelf' 117 | labelMap[11] = 'picture' 118 | labelMap[12] = 'cooking_counter' 119 | labelMap[13] = 'blinds' 120 | labelMap[14] = 'desk' 121 | labelMap[15] = 'shelf' 122 | labelMap[16] = 'curtain' 123 | labelMap[17] = 'dresser' 124 | labelMap[18] = 'pillow' 125 | labelMap[19] = 'mirror' 126 | labelMap[20] = 'entrance' #mat 127 | labelMap[21] = 'clothes' 128 | labelMap[22] = 'ceiling' 129 | labelMap[23] = 'book' 130 | labelMap[24] = 'refrigerator' 131 | labelMap[25] = 'TV' 132 | labelMap[26] = 'paper' 133 | labelMap[27] = 'towel' 134 | labelMap[28] = 'shower_curtain' 135 | labelMap[29] = 'box' 136 | labelMap[30] = 'whiteboard' 137 | labelMap[31] = 'person' 138 | labelMap[32] = 'nightstand' 139 | labelMap[33] = 'toilet' 140 | labelMap[34] = 'washing_basin' 141 | labelMap[35] = 'lamp' 142 | labelMap[36] = 'bathtub' 143 | labelMap[37] = 'bag' 144 | labelMap[38] = 'otherprop' 145 | labelMap[39] = 'otherstructure' 146 | labelMap[40] = 'unannotated' 147 | return labelMap 148 | 149 | def getNYUScanNetMap(): 150 | labelMap = np.zeros(41, dtype=np.int32) 151 | labelMap[1] = 1 152 | labelMap[2] = 2 153 | labelMap[3] = 19 154 | labelMap[4] = 6 155 | labelMap[5] = 3 156 | labelMap[6] = 8 157 | labelMap[7] = 4 158 | labelMap[8] = 14 159 | labelMap[9] = 15 160 | labelMap[10] = 7 161 | labelMap[11] = 18 162 | labelMap[12] = 13 163 | labelMap[13] = 12 #20 Blinds 164 | labelMap[14] = 5 165 | labelMap[15] = 7 166 | labelMap[16] = 12 167 | labelMap[17] = 19 168 | labelMap[18] = 20 169 | labelMap[19] = 20 170 | labelMap[20] = 20 171 | labelMap[21] = 20 172 | labelMap[22] = 1 173 | labelMap[23] = 20 174 | labelMap[24] = 17 175 | labelMap[25] = 20 176 | labelMap[26] = 20 177 | labelMap[27] = 20 178 | labelMap[28] = 16 179 | labelMap[29] = 20 180 | labelMap[30] = 20 181 | labelMap[31] = 20 182 | labelMap[32] = 20 183 | labelMap[33] = 11 184 | labelMap[34] = 9 185 | labelMap[35] = 20 186 | labelMap[36] = 10 187 | labelMap[37] = 20 188 | labelMap[38] = 20 189 | labelMap[39] = 20 190 | labelMap[40] = 0 191 | return labelMap 192 | 193 | def getMatterportClassMap(): 194 | classMap = { 195 | 'a': 1, 196 | 'b': 2, 197 | 'c': 3, 198 | 'd': 4, 199 | 'e': 5, 200 | 'f': 6, 201 | 'g': 7, 202 | 'h': 8, 203 | 'i': 9, 204 | 'j': 10, 205 | 'k': 11, 206 | 'l': 12, 207 | 'm': 13, 208 | 'n': 14, 209 | 'o': 15, 210 | 'p': 16, 211 | 'r': 17, 212 | 's': 18, 213 | 't': 19, 214 | 'u': 20, 215 | 'v': 21, 216 | 'w': 22, 217 | 'x': 23, 218 | 'y': 24, 219 | 'z': 25, 220 | 'B': 26, 221 | 'C': 27, 222 | 'D': 28, 223 | 'S': 29, 224 | 'Z': 30 225 | } 226 | return classMap 227 | 228 | def calcLineDim(points, line): 229 | point_1 = points[line[0]] 230 | point_2 = points[line[1]] 231 | if abs(point_2[0] - point_1[0]) > abs(point_2[1] - point_1[1]): 232 | lineDim = 0 233 | else: 234 | lineDim = 1 235 | return lineDim 236 | 237 | def calcLineDirection(line): 238 | return int(abs(line[0][0] - line[1][0]) < abs(line[0][1] - line[1][1])) 239 | 240 | def calcLineDirectionPoints(points, line): 241 | point_1 = points[line[0]] 242 | point_2 = points[line[1]] 243 | if isinstance(point_1[0], tuple): 244 | point_1 = point_1[0] 245 | pass 246 | if isinstance(point_2[0], tuple): 247 | point_2 = point_2[0] 248 | pass 249 | return calcLineDirection((point_1, point_2)) 250 | 251 | def pointDistance(point_1, point_2): 252 | #return np.sqrt(pow(point_1[0] - point_2[0], 2) + pow(point_1[1] - point_2[1], 2)) 253 | return max(abs(point_1[0] - point_2[0]), abs(point_1[1] - point_2[1])) 254 | 255 | def sortLines(lines): 256 | newLines = [] 257 | for line in lines: 258 | direction = calcLineDirection(line) 259 | if line[0][direction] < line[1][direction]: 260 | newLines.append((line[0], line[1])) 261 | else: 262 | newLines.append((line[1], line[0])) 263 | pass 264 | continue 265 | return newLines 266 | 267 | def lineRange(line): 268 | direction = calcLineDirection(line) 269 | fixedValue = (line[0][1 - direction] + line[1][1 - direction]) / 2 270 | minValue = min(line[0][direction], line[1][direction]) 271 | maxValue = max(line[0][direction], line[1][direction]) 272 | return direction, fixedValue, minValue, maxValue 273 | 274 | def findConnections(line_1, line_2, gap): 275 | connection_1 = -1 276 | connection_2 = -1 277 | pointConnected = False 278 | for c_1 in xrange(2): 279 | if pointConnected: 280 | break 281 | for c_2 in xrange(2): 282 | if pointDistance(line_1[c_1], line_2[c_2]) > gap: 283 | continue 284 | 285 | connection_1 = c_1 286 | connection_2 = c_2 287 | connectionPoint = ((line_1[c_1][0] + line_2[c_2][0]) / 2, (line_1[c_1][1] + line_2[c_2][1]) / 2) 288 | pointConnected = True 289 | break 290 | continue 291 | if pointConnected: 292 | return [connection_1, connection_2], connectionPoint 293 | direction_1, fixedValue_1, min_1, max_1 = lineRange(line_1) 294 | direction_2, fixedValue_2, min_2, max_2 = lineRange(line_2) 295 | if direction_1 == direction_2: 296 | return [-1, -1], (0, 0) 297 | 298 | #print(fixedValue_1, min_1, max_1, fixedValue_2, min_2, max_2) 299 | if min(fixedValue_1, max_2) < max(fixedValue_1, min_2) - gap or min(fixedValue_2, max_1) < max(fixedValue_2, min_1) - gap: 300 | return [-1, -1], (0, 0) 301 | 302 | if abs(min_1 - fixedValue_2) <= gap: 303 | return [0, 2], (fixedValue_2, fixedValue_1) 304 | if abs(max_1 - fixedValue_2) <= gap: 305 | return [1, 2], (fixedValue_2, fixedValue_1) 306 | if abs(min_2 - fixedValue_1) <= gap: 307 | return [2, 0], (fixedValue_2, fixedValue_1) 308 | if abs(max_2 - fixedValue_1) <= gap: 309 | return [2, 1], (fixedValue_2, fixedValue_1) 310 | return [2, 2], (fixedValue_2, fixedValue_1) 311 | 312 | def lines2Corners(lines, gap, getSingularCorners=False): 313 | corners = [] 314 | lineConnections = [] 315 | for _ in xrange(len(lines)): 316 | lineConnections.append({}) 317 | continue 318 | 319 | connectionCornerMap = {} 320 | connectionCornerMap[(1, 1)] = 4 321 | connectionCornerMap[(0, 1)] = 5 322 | connectionCornerMap[(0, 0)] = 6 323 | connectionCornerMap[(1, 0)] = 7 324 | connectionCornerMap[(2, 0)] = 8 325 | connectionCornerMap[(1, 2)] = 9 326 | connectionCornerMap[(2, 1)] = 10 327 | connectionCornerMap[(0, 2)] = 11 328 | connectionCornerMap[(2, 2)] = 12 329 | corners = [] 330 | for lineIndex_1, line_1 in enumerate(lines): 331 | for lineIndex_2, line_2 in enumerate(lines): 332 | if lineIndex_2 == lineIndex_1: 333 | continue 334 | connections, connectionPoint = findConnections(line_1, line_2, gap=gap) 335 | if connections[0] == -1 and connections[1] == -1: 336 | continue 337 | if calcLineDirection(line_1) == calcLineDirection(line_2): 338 | print('overlap', line_1, line_2, connections) 339 | exit(1) 340 | pass 341 | if calcLineDirection(line_1) == 1: 342 | continue 343 | 344 | indices = [lineIndex_1, lineIndex_2] 345 | #print(lineIndex_1, lineIndex_2, connections) 346 | for c in xrange(2): 347 | if connections[c] in [0, 1] and connections[c] in lineConnections[indices[c]]: 348 | print('duplicate corner', line_1, line_2, connections) 349 | exit(1) 350 | pass 351 | lineConnections[indices[c]][connections[c]] = True 352 | continue 353 | corners.append((connectionPoint, connectionCornerMap[tuple(connections)])) 354 | continue 355 | continue 356 | 357 | if getSingularCorners: 358 | singularCorners = [] 359 | for lineIndex, connections in enumerate(lineConnections): 360 | if 0 not in connections: 361 | print('single corner', lines[lineIndex], connections) 362 | singularCorners.append((lineIndex, 0)) 363 | pass 364 | if 1 not in connections: 365 | print('single corner', lines[lineIndex], connections) 366 | singularCorners.append((lineIndex, 1)) 367 | pass 368 | continue 369 | return corners, singularCorners 370 | 371 | return corners 372 | 373 | def drawWallMask(walls, width, height, thickness=3, indexed=False): 374 | if indexed: 375 | wallMask = np.full((height, width), -1, dtype=np.int32) 376 | for wallIndex, wall in enumerate(walls): 377 | cv2.line(wallMask, (int(wall[0][0]), int(wall[0][1])), (int(wall[1][0]), int(wall[1][1])), color=wallIndex, thickness=thickness) 378 | continue 379 | else: 380 | wallMask = np.zeros((height, width), dtype=np.int32) 381 | for wall in walls: 382 | cv2.line(wallMask, (int(wall[0][0]), int(wall[0][1])), (int(wall[1][0]), int(wall[1][1])), color=1, thickness=thickness) 383 | continue 384 | wallMask = wallMask.astype(np.bool) 385 | pass 386 | return wallMask 387 | 388 | def mergeLines(line_1, line_2): 389 | direction_1, fixedValue_1, min_1, max_1 = lineRange(line_1) 390 | direction_2, fixedValue_2, min_2, max_2 = lineRange(line_2) 391 | fixedValue = (fixedValue_1 + fixedValue_2) / 2 392 | if direction_1 == 0: 393 | return [[min(min_1, min_2), fixedValue], [max(max_1, max_2), fixedValue]] 394 | else: 395 | return [[fixedValue, min(min_1, min_2)], [fixedValue, max(max_1, max_2)]] 396 | return 397 | 398 | def findIntersection(line_1, line_2): 399 | direction_1, fixedValue_1, min_1, max_1 = lineRange(line_1) 400 | direction_2, fixedValue_2, min_2, max_2 = lineRange(line_2) 401 | if direction_1 == 0: 402 | return (fixedValue_2, fixedValue_1) 403 | else: 404 | return (fixedValue_1, fixedValue_2) 405 | return 406 | 407 | def extendLine(line, point): 408 | direction, fixedValue, min_value, max_value = lineRange(line) 409 | if direction == 0: 410 | return ((min(min_value, point[direction]), fixedValue), (max(max_value, point[direction]), fixedValue)) 411 | else: 412 | return ((fixedValue, min(min_value, point[direction])), (fixedValue, max(max_value, point[direction]))) 413 | return 414 | 415 | def divideWalls(walls): 416 | horizontalWalls = [] 417 | verticalWalls = [] 418 | for wall in walls: 419 | if calcLineDirection(wall) == 0: 420 | horizontalWalls.append(wall) 421 | else: 422 | verticalWalls.append(wall) 423 | pass 424 | continue 425 | return horizontalWalls, verticalWalls 426 | 427 | 428 | def connectWalls(walls, roomSegmentation, gap=3): 429 | width = roomSegmentation.shape[1] 430 | height = roomSegmentation.shape[0] 431 | roomBoundary = np.zeros(roomSegmentation.shape, dtype=np.bool) 432 | for direction in xrange(2): 433 | for shift in [-1, 1]: 434 | roomBoundary = np.logical_or(roomBoundary, roomSegmentation != np.roll(roomSegmentation, shift, axis=direction)) 435 | continue 436 | continue 437 | roomBoundary = roomBoundary.astype(np.uint8) 438 | roomBoundary[0] = roomBoundary[-1] = roomBoundary[:, 0] = roomBoundary[:, -1] = 0 439 | 440 | uncoveredBoundary = roomBoundary.copy() 441 | wallGroups = divideWalls(walls) 442 | wallMasks = [drawWallMask(walls, width, height, indexed=True, thickness=gap * 2) for walls in wallGroups] 443 | 444 | uncoveredBoundary[wallMasks[0] >= 0] = 0 445 | uncoveredBoundary[wallMasks[1] >= 0] = 0 446 | uncoveredBoundary = cv2.dilate(uncoveredBoundary, np.ones((3, 3)), iterations=gap) 447 | components = measure.label(uncoveredBoundary, background=0) 448 | 449 | connectedWalls = [] 450 | 451 | for walls, wallMask in zip(wallGroups, wallMasks): 452 | 453 | newWalls = copy.deepcopy(walls) 454 | invalidWallIndices = [] 455 | for label in xrange(components.min() + 1, components.max() + 1): 456 | mask = components == label 457 | wallIndices = np.unique(wallMask[mask]).tolist() 458 | if -1 in wallIndices: 459 | wallIndices.remove(-1) 460 | pass 461 | if len(wallIndices) != 2: 462 | continue 463 | 464 | wall_1 = newWalls[wallIndices[0]] 465 | wall_2 = newWalls[wallIndices[1]] 466 | 467 | direction_1, fixedValue_1, min_1, max_1 = lineRange(wall_1) 468 | direction_2, fixedValue_2, min_2, max_2 = lineRange(wall_2) 469 | 470 | if direction_1 == direction_2: 471 | if abs(fixedValue_1 - fixedValue_2) < gap: 472 | newWallIndex = len(newWalls) 473 | wallMask[wallMask == wallIndices[0]] = newWallIndex 474 | wallMask[wallMask == wallIndices[1]] = newWallIndex 475 | newWall = mergeLines(wall_1, wall_2) 476 | newWalls.append(newWall) 477 | invalidWallIndices.append(wallIndices[0]) 478 | invalidWallIndices.append(wallIndices[1]) 479 | pass 480 | pass 481 | # else: 482 | # print(wall_1, wall_2) 483 | # ys, xs = mask.nonzero() 484 | # newWall = [[xs.min(), ys.min()], [xs.max(), ys.max()]] 485 | # newWallDirection = calcLineDirection(newWall) 486 | # if newWallDirection != direction_1 and newWall[1][1 - newWallDirection] - newWall[0][1 - newWallDirection] < gap * 2 + 1: 487 | # fixedValue = (newWall[1][1 - newWallDirection] + newWall[0][1 - newWallDirection]) / 2 488 | # newWall[1][1 - newWallDirection] = newWall[0][1 - newWallDirection] = fixedValue 489 | # newWalls.append(newWall) 490 | # pass 491 | # pass 492 | # else: 493 | # assert(False) 494 | # intersectionPoint = findIntersection(wall_1, wall_2) 495 | # newWalls[wallIndices[0]] = extendLine(wall_1, intersectionPoint) 496 | # newWalls[wallIndices[1]] = extendLine(wall_2, intersectionPoint) 497 | # pass 498 | continue 499 | 500 | #print(invalidWallIndices) 501 | invalidWallIndices = sorted(invalidWallIndices, key=lambda x: -x) 502 | for index in invalidWallIndices: 503 | del newWalls[index] 504 | continue 505 | connectedWalls += newWalls 506 | continue 507 | 508 | newWalls = connectedWalls 509 | wallMask = drawWallMask(newWalls, width, height, indexed=True, thickness=gap * 2) 510 | uncoveredBoundary = roomBoundary.copy() 511 | uncoveredBoundary[wallMask >= 0] = 0 512 | uncoveredBoundary = cv2.dilate(uncoveredBoundary, np.ones((3, 3)), iterations=gap) 513 | components = measure.label(uncoveredBoundary, background=0) 514 | 515 | #cv2.imwrite('test/segmentation.png', drawSegmentationImage(components)) 516 | 517 | for label in xrange(components.min() + 1, components.max() + 1): 518 | mask = components == label 519 | #cv2.imwrite('test/mask_' + str(label) + '.png', drawMaskImage(mask)) 520 | wallIndices = np.unique(wallMask[mask]).tolist() 521 | if -1 in wallIndices: 522 | wallIndices.remove(-1) 523 | pass 524 | 525 | lines = [newWalls[index] for index in wallIndices] 526 | #cv2.imwrite('test/mask_' + str(label) + '_segment.png', drawMaskImage(mask)) 527 | #cv2.imwrite('test/mask_' + str(label) + '.png', drawMaskImage(drawWallMask(lines, width, height))) 528 | 529 | horizontalLines, verticalLines = divideWalls(lines) 530 | if len(horizontalLines) > 0 and len(verticalLines) > 0: 531 | continue 532 | #print(label, wallIndices, len(horizontalLines), len(verticalLines)) 533 | for direction, lines in enumerate([horizontalLines, verticalLines]): 534 | if len(lines) < 2: 535 | continue 536 | #wall_1 = lines[0] 537 | #wall_2 = lines[1] 538 | #print(wall_1, wall_2) 539 | #direction_1, fixedValue_1, min_1, max_1 = lineRange(wall_1) 540 | #direction_2, fixedValue_2, min_2, max_2 = lineRange(wall_2) 541 | #values = [line[direction] for line in lines] 542 | 543 | #print(wall_1, wall_2) 544 | ys, xs = mask.nonzero() 545 | newWall = [[xs.min(), ys.min()], [xs.max(), ys.max()]] 546 | newWallDirection = calcLineDirection(newWall) 547 | #print(label, wallIndices, newWallDirection, direction, newWall[1][1 - newWallDirection] - newWall[0][1 - newWallDirection]) 548 | if newWallDirection != direction and newWall[1][1 - newWallDirection] - newWall[0][1 - newWallDirection] <= (gap * 2 + 2) * 2: 549 | fixedValue = (newWall[1][1 - newWallDirection] + newWall[0][1 - newWallDirection]) / 2 550 | newWall[1][1 - newWallDirection] = newWall[0][1 - newWallDirection] = fixedValue 551 | values = [line[0][newWallDirection] for line in lines] + [line[1][newWallDirection] for line in lines] 552 | min_value = min(values) 553 | max_value = max(values) 554 | newWall[0][newWallDirection] = min_value 555 | newWall[1][newWallDirection] = max_value 556 | 557 | newWalls.append(newWall) 558 | #print('new orthogonal wall', newWall) 559 | pass 560 | 561 | continue 562 | continue 563 | 564 | wallMask = drawWallMask(newWalls, width, height, indexed=True, thickness=gap * 2) 565 | uncoveredBoundary = roomBoundary.copy() 566 | uncoveredBoundary[wallMask >= 0] = 0 567 | uncoveredBoundary = cv2.dilate(uncoveredBoundary, np.ones((3, 3)), iterations=gap) 568 | components = measure.label(uncoveredBoundary, background=0) 569 | 570 | for label in xrange(components.min() + 1, components.max() + 1): 571 | mask = components == label 572 | wallIndices = np.unique(wallMask[mask]).tolist() 573 | if -1 in wallIndices: 574 | wallIndices.remove(-1) 575 | pass 576 | if len(wallIndices) != 2: 577 | continue 578 | 579 | wall_1 = newWalls[wallIndices[0]] 580 | wall_2 = newWalls[wallIndices[1]] 581 | 582 | #print(wall_1, wall_2) 583 | 584 | direction_1 = calcLineDirection(wall_1) 585 | direction_2 = calcLineDirection(wall_2) 586 | 587 | if direction_1 != direction_2: 588 | intersectionPoint = findIntersection(wall_1, wall_2) 589 | newWalls[wallIndices[0]] = extendLine(wall_1, intersectionPoint) 590 | newWalls[wallIndices[1]] = extendLine(wall_2, intersectionPoint) 591 | pass 592 | continue 593 | 594 | # try: 595 | # _, singularCorners = lines2Corners(newWalls, gap=gap, getSingularCorners=True) 596 | # for _, singularCorner_1 in enumerate(singularCorners): 597 | # for singularCorner_2 in singularCorners[_ + 1:]: 598 | # wall_1 = newWalls[singularCorner_1[0]] 599 | # wall_2 = newWalls[singularCorner_2[0]] 600 | # corner_1 = wall_1[singularCorner_1[1]] 601 | # corner_2 = wall_2[singularCorner_2[1]] 602 | # if pointDistance(corner_1, corner_2) < (gap * 2 + 1) * 2: 603 | # intersectionPoint = findIntersection(wall_1, wall_2) 604 | # newWalls[singularCorner_1[0]] = extendLine(wall_1, intersectionPoint) 605 | # newWalls[singularCorner_2[0]] = extendLine(wall_2, intersectionPoint) 606 | # pass 607 | # continue 608 | # continue 609 | # except: 610 | # pass 611 | 612 | return newWalls 613 | 614 | def extractLines(lineMask, lengthThreshold=11, widthThreshold=5): 615 | lines = [] 616 | components = measure.label(lineMask, background=0) 617 | for label in xrange(components.min() + 1, components.max() + 1): 618 | mask = components == label 619 | ys, xs = mask.nonzero() 620 | line = [[xs.min(), ys.min()], [xs.max(), ys.max()]] 621 | direction = calcLineDirection(line) 622 | if abs(line[1][1 - direction] - line[0][1 - direction]) > widthThreshold or abs(line[1][direction] - line[0][direction]) < lengthThreshold: 623 | continue 624 | fixedValue = (line[1][1 - direction] + line[0][1 - direction]) / 2 625 | line[1][1 - direction] = line[0][1 - direction] = fixedValue 626 | lines.append(line) 627 | continue 628 | return lines 629 | 630 | 631 | def drawPoints(filename, width, height, points, backgroundImage=None, pointSize=5, pointColor=None): 632 | colorMap = ColorPalette(NUM_CORNERS).getColorMap() 633 | if np.all(np.equal(backgroundImage, None)): 634 | image = np.zeros((height, width, 3), np.uint8) 635 | else: 636 | if backgroundImage.ndim == 2: 637 | image = np.tile(np.expand_dims(backgroundImage, -1), [1, 1, 3]) 638 | else: 639 | image = backgroundImage 640 | pass 641 | pass 642 | no_point_color = pointColor is None 643 | for point in points: 644 | if no_point_color: 645 | pointColor = colorMap[point[2] * 4 + point[3]] 646 | pass 647 | #print('used', pointColor) 648 | #print('color', point[2] , point[3]) 649 | image[max(int(round(point[1])) - pointSize, 0):min(int(round(point[1])) + pointSize, height), max(int(round(point[0])) - pointSize, 0):min(int(round(point[0])) + pointSize, width)] = pointColor 650 | continue 651 | 652 | if filename != '': 653 | cv2.imwrite(filename, image) 654 | return 655 | else: 656 | return image 657 | 658 | def drawPointsSeparately(path, width, height, points, backgroundImage=None, pointSize=5): 659 | if np.all(np.equal(backgroundImage, None)): 660 | image = np.zeros((height, width, 13), np.uint8) 661 | else: 662 | image = np.tile(np.expand_dims(backgroundImage, -1), [1, 1, 13]) 663 | pass 664 | 665 | for point in points: 666 | image[max(int(round(point[1])) - pointSize, 0):min(int(round(point[1])) + pointSize, height), max(int(round(point[0])) - pointSize, 0):min(int(round(point[0])) + pointSize, width), int(point[2] * 4 + point[3])] = 255 667 | continue 668 | for channel in xrange(13): 669 | cv2.imwrite(path + '_' + str(channel) + '.png', image[:, :, channel]) 670 | continue 671 | return 672 | 673 | def drawLineMask(width, height, points, lines, lineWidth = 5, backgroundImage = None): 674 | lineMask = np.zeros((height, width)) 675 | 676 | for lineIndex, line in enumerate(lines): 677 | point_1 = points[line[0]] 678 | point_2 = points[line[1]] 679 | direction = calcLineDirectionPoints(points, line) 680 | 681 | fixedValue = int(round((point_1[1 - direction] + point_2[1 - direction]) / 2)) 682 | minValue = int(min(point_1[direction], point_2[direction])) 683 | maxValue = int(max(point_1[direction], point_2[direction])) 684 | if direction == 0: 685 | lineMask[max(fixedValue - lineWidth, 0):min(fixedValue + lineWidth + 1, height), minValue:maxValue + 1] = 1 686 | else: 687 | lineMask[minValue:maxValue + 1, max(fixedValue - lineWidth, 0):min(fixedValue + lineWidth + 1, width)] = 1 688 | pass 689 | continue 690 | return lineMask 691 | 692 | 693 | 694 | def drawLines(filename, width, height, points, lines, lineLabels = [], backgroundImage = None, lineWidth = 5, lineColor = None): 695 | colorMap = ColorPalette(len(lines)).getColorMap() 696 | if backgroundImage is None: 697 | image = np.ones((height, width, 3), np.uint8) * 0 698 | else: 699 | if backgroundImage.ndim == 2: 700 | image = np.stack([backgroundImage, backgroundImage, backgroundImage], axis=2) 701 | else: 702 | image = backgroundImage 703 | pass 704 | pass 705 | 706 | for lineIndex, line in enumerate(lines): 707 | point_1 = points[line[0]] 708 | point_2 = points[line[1]] 709 | direction = calcLineDirectionPoints(points, line) 710 | 711 | 712 | fixedValue = int(round((point_1[1 - direction] + point_2[1 - direction]) / 2)) 713 | minValue = int(round(min(point_1[direction], point_2[direction]))) 714 | maxValue = int(round(max(point_1[direction], point_2[direction]))) 715 | if len(lineLabels) == 0: 716 | if np.any(lineColor == None): 717 | lineColor = np.random.rand(3) * 255 718 | pass 719 | if direction == 0: 720 | image[max(fixedValue - lineWidth, 0):min(fixedValue + lineWidth + 1, height), minValue:maxValue + 1, :] = lineColor 721 | else: 722 | image[minValue:maxValue + 1, max(fixedValue - lineWidth, 0):min(fixedValue + lineWidth + 1, width), :] = lineColor 723 | else: 724 | labels = lineLabels[lineIndex] 725 | isExterior = False 726 | if direction == 0: 727 | for c in xrange(3): 728 | image[max(fixedValue - lineWidth, 0):min(fixedValue, height), minValue:maxValue, c] = colorMap[labels[0]][c] 729 | image[max(fixedValue, 0):min(fixedValue + lineWidth + 1, height), minValue:maxValue, c] = colorMap[labels[1]][c] 730 | continue 731 | else: 732 | for c in xrange(3): 733 | image[minValue:maxValue, max(fixedValue - lineWidth, 0):min(fixedValue, width), c] = colorMap[labels[1]][c] 734 | image[minValue:maxValue, max(fixedValue, 0):min(fixedValue + lineWidth + 1, width), c] = colorMap[labels[0]][c] 735 | continue 736 | pass 737 | pass 738 | continue 739 | 740 | if filename == '': 741 | return image 742 | else: 743 | cv2.imwrite(filename, image) 744 | 745 | 746 | def drawRectangles(filename, width, height, points, rectangles, labels, lineWidth = 2, backgroundImage = None, rectangleColor = None): 747 | colorMap = ColorPalette(NUM_ICONS).getColorMap() 748 | if backgroundImage is None: 749 | image = np.ones((height, width, 3), np.uint8) * 0 750 | else: 751 | image = backgroundImage 752 | pass 753 | 754 | for rectangleIndex, rectangle in enumerate(rectangles): 755 | point_1 = points[rectangle[0]] 756 | point_2 = points[rectangle[1]] 757 | point_3 = points[rectangle[2]] 758 | point_4 = points[rectangle[3]] 759 | 760 | 761 | if len(labels) == 0: 762 | if rectangleColor is None: 763 | color = np.random.rand(3) * 255 764 | else: 765 | color = rectangleColor 766 | else: 767 | color = colorMap[labels[rectangleIndex]] 768 | pass 769 | 770 | x_1 = int(round((point_1[0] + point_3[0]) / 2)) 771 | x_2 = int(round((point_2[0] + point_4[0]) / 2)) 772 | y_1 = int(round((point_1[1] + point_2[1]) / 2)) 773 | y_2 = int(round((point_3[1] + point_4[1]) / 2)) 774 | 775 | cv2.rectangle(image, (x_1, y_1), (x_2, y_2), color=tuple(color.tolist()), thickness = 2) 776 | 777 | # point_1 = (int(point_1[0]), int(point_1[1])) 778 | # point_2 = (int(point_2[0]), int(point_2[1])) 779 | # point_3 = (int(point_3[0]), int(point_3[1])) 780 | # point_4 = (int(point_4[0]), int(point_4[1])) 781 | 782 | # image[max(point_1[1] - lineWidth, 0):min(point_1[1] + lineWidth, height), point_1[0]:point_2[0] + 1, :] = color 783 | # image[max(point_3[1] - lineWidth, 0):min(point_3[1] + lineWidth, height), point_3[0]:point_4[0] + 1, :] = color 784 | # image[point_1[1]:point_3[1] + 1, max(point_1[0] - lineWidth, 0):min(point_1[0] + lineWidth, width), :] = color 785 | # image[point_2[1]:point_4[1] + 1, max(point_2[0] - lineWidth, 0):min(point_2[0] + lineWidth, width), :] = color 786 | 787 | continue 788 | 789 | if filename == '': 790 | return image 791 | else: 792 | cv2.imwrite(filename, image) 793 | pass 794 | 795 | 796 | def drawResultImage(width, height, result): 797 | resultImage = drawLines('', width, height, result['wall'][0], result['wall'][1], result['wall'][2], None, lineWidth=3) 798 | resultImage = drawLines('', width, height, result['door'][0], result['door'][1], [], resultImage, lineWidth=2, lineColor=0) 799 | iconImage = drawRectangles('', width, height, result['icon'][0], result['icon'][1], result['icon'][2], lineWidth=2) 800 | return resultImage, iconImage 801 | 802 | def resizeResult(result, width, height, oriWidth=256, oriHeight=256): 803 | result['wall'][0] = [[float(point[0]) / oriWidth * width, float(point[1]) / oriHeight * height, point[2], point[3]] for point in result['wall'][0]] 804 | result['door'][0] = [[float(point[0]) / oriWidth * width, float(point[1]) / oriHeight * height, point[2], point[3]] for point in result['door'][0]] 805 | result['icon'][0] = [[float(point[0]) / oriWidth * width, float(point[1]) / oriHeight * height, point[2], point[3]] for point in result['icon'][0]] 806 | #result['room'][0] = [(cv2.resize(mask, (width, height), interpolation=cv2.INTER_NEAREST) > 0).astype(np.uint8) for mask in result['room'][0]] 807 | return 808 | 809 | def drawResultImageFinal(width, height, result): 810 | colorMap = np.array([(224, 255, 192), (255, 160, 96), (255, 224, 128), (192, 255, 255), (192, 255, 255), (192, 255, 255), (192, 192, 224), (224, 255, 192), (255, 224, 224), (224, 224, 224)]) 811 | borderColorMap = np.array([(128, 192, 64), (192, 64, 64), (192, 128, 64), (0, 128, 192), (0, 128, 192), (0, 128, 192), (128, 64, 160), (128, 192, 64), (192, 64, 0), (255, 255, 255)]) 812 | colorMap = np.concatenate([np.full(shape=(1, 3), fill_value=0), colorMap, borderColorMap], axis=0).astype(np.uint8) 813 | colorMap = colorMap[:, ::-1] 814 | 815 | labelRoomMap = getLabelRoomMap() 816 | 817 | roomSegmentation = np.zeros((height, width), dtype=np.int32) 818 | 819 | roomsInfo = [] 820 | wall_dict = result['wall'] 821 | wallMask = drawWallMask([(wall_dict[0][line[0]], wall_dict[0][line[1]]) for line in wall_dict[1]], width, height, thickness=3) 822 | roomRegions = measure.label(1 - wallMask, background=0) 823 | #cv2.imwrite('test/' + str(dictIndex) + '_segmentation_regions.png', drawSegmentationImage(roomRegions)) 824 | backgroundIndex = roomRegions.min() 825 | wallPoints = wall_dict[0] 826 | roomLabels = {} 827 | sizes = np.array([width, height]) 828 | for wallIndex, wallLabels in enumerate(wall_dict[2]): 829 | wallLine = wall_dict[1][wallIndex] 830 | lineDim = calcLineDim(wallPoints, wallLine) 831 | #print('wall', wallIndex, wallPoints[wallLine[0]][:2], wallPoints[wallLine[1]][:2]) 832 | center = np.round((np.array(wallPoints[wallLine[0]][:2]) + np.array(wallPoints[wallLine[1]][:2])) / 2).astype(np.int32) 833 | 834 | for c in xrange(2): 835 | direction = c * 2 - 1 836 | if lineDim == 1: 837 | direction *= -1 838 | pass 839 | point = center 840 | for offset in xrange(10): 841 | point[1 - lineDim] += direction 842 | if point[lineDim] < 0 or point[lineDim] >= sizes[lineDim]: 843 | break 844 | roomIndex = roomRegions[point[1], point[0]] 845 | if roomIndex != backgroundIndex: 846 | #print(roomIndex, wallLabels[c], wallLabels, point.tolist()) 847 | #mask = roomRegions == roomIndex 848 | #mask = cv2.dilate(mask.astype(np.uint8), np.ones((3, 3)), iterations=1) 849 | #roomSegmentation[mask] = wallLabels[c] 850 | #rooms[wallLabels[c]].append(cv2.dilate(mask.astype(np.uint8), np.ones((3, 3)), iterations=wallLineWidth)) 851 | #roomRegions[mask] = backgroundIndex 852 | if roomIndex not in roomLabels: 853 | roomLabels[roomIndex] = {} 854 | pass 855 | roomLabels[roomIndex][wallLabels[c]] = True 856 | break 857 | continue 858 | continue 859 | continue 860 | 861 | 862 | rooms = [] 863 | indexMap = {} 864 | for roomIndex, labels in roomLabels.iteritems(): 865 | #print(roomIndex, labels) 866 | if roomIndex == roomRegions[0][0]: 867 | continue 868 | indexMap[roomIndex] = len(rooms) 869 | mask = roomRegions == roomIndex 870 | mask = cv2.dilate(mask.astype(np.uint8), np.ones((3, 3)), iterations=3) 871 | 872 | # if 7 in labels and 2 not in labels: 873 | # labels[2] = True 874 | # pass 875 | # if 5 in labels and 3 not in labels: 876 | # labels[3] = True 877 | # pass 878 | # if 9 in labels and 1 not in labels: 879 | # labels[1] = True 880 | # pass 881 | rooms.append((mask, labels)) 882 | continue 883 | 884 | wallLineWidth = 5 885 | # foregroundMask = roomSegmentation > 0 886 | # foregroundMask = cv2.dilate(foregroundMask, np.ones((3, 3)), iterations=wallLineWidth) 887 | # roomSegmentation[foregroundMask] = 888 | for mask, labels in rooms: 889 | label = min([label for label in labels]) 890 | if label < 0: 891 | continue 892 | kernel = np.zeros((3, 3)) 893 | kernel[1:, 1:] = 1 894 | #mask = cv2.erode(mask.astype(np.uint8), kernel.astype(np.uint8), iterations=1) 895 | erodedMask = cv2.erode(mask, np.ones((3, 3)), iterations=wallLineWidth) 896 | roomSegmentation[mask.astype(np.bool)] = label + 10 897 | roomSegmentation[erodedMask.astype(np.bool)] = label 898 | 899 | continue 900 | image = colorMap[roomSegmentation.reshape(-1)].reshape((height, width, 3)) 901 | 902 | pointColor = tuple((np.array([0.3, 0.3, 0.9]) * 255).astype(np.uint8).tolist()) 903 | for wallLine in result['wall'][1]: 904 | for pointIndex in wallLine: 905 | point = result['wall'][0][pointIndex] 906 | cv2.circle(image, (int(point[0]), int(point[1])), color=pointColor, radius=8, thickness=-1) 907 | cv2.circle(image, (int(point[0]), int(point[1])), color=(255, 255, 255), radius=4, thickness=-1) 908 | continue 909 | continue 910 | 911 | lineSegmentLength = 20.0 912 | for doorLine in result['door'][1]: 913 | point_1 = np.array(result['door'][0][doorLine[0]][:2]).astype(np.float32) 914 | point_2 = np.array(result['door'][0][doorLine[1]][:2]).astype(np.float32) 915 | lineDim = calcLineDim(result['door'][0], doorLine) 916 | for i in xrange(int(abs(point_1[lineDim] - point_2[lineDim]) / lineSegmentLength + 1)): 917 | ratio = i * lineSegmentLength / abs(point_1[lineDim] - point_2[lineDim]) 918 | if ratio >= 1: 919 | break 920 | startPoint = point_1 + ratio * (point_2 - point_1) 921 | ratio = (i + 0.5) * lineSegmentLength / abs(point_1[lineDim] - point_2[lineDim]) 922 | ratio = min(ratio, 1) 923 | endPoint = point_1 + ratio * (point_2 - point_1) 924 | cv2.line(image, (startPoint[0], startPoint[1]), (endPoint[0], endPoint[1]), color=(0, 0, 0), thickness=4) 925 | continue 926 | for point in [point_1, point_2]: 927 | startPoint = point.copy() 928 | startPoint[1 - lineDim] += lineSegmentLength / 2 929 | endPoint = point.copy() 930 | endPoint[1 - lineDim] -= lineSegmentLength / 2 931 | cv2.line(image, (startPoint[0], startPoint[1]), (endPoint[0], endPoint[1]), color=(0, 0, 0), thickness=2) 932 | continue 933 | continue 934 | 935 | 936 | 937 | labelIconMap = getLabelIconMap() 938 | iconPos = [] 939 | for iconIndex, (icon, label) in enumerate(zip(result['icon'][1], result['icon'][2])): 940 | name = labelIconMap[label + 1] 941 | iconImage = cv2.imread('icons/' + name + '.jpg') 942 | 943 | points = [result['icon'][0][pointIndex] for pointIndex in icon] 944 | x_1 = int(round((points[0][0] + points[2][0]) / 2)) 945 | x_2 = int(round((points[1][0] + points[3][0]) / 2)) 946 | y_1 = int(round((points[0][1] + points[1][1]) / 2)) 947 | y_2 = int(round((points[2][1] + points[3][1]) / 2)) 948 | 949 | iconSize = iconImage.shape #(y, x) 950 | #print('icon_size', iconSize) 951 | icon_is_landscape = iconSize[1] > iconSize[0] 952 | 953 | slot_size = (x_2 - x_1 + 1, y_2 - y_1 + 1) 954 | slot_center = np.array((x_1 + slot_size[0]/2, y_1 + slot_size[1] / 2)) 955 | slot_is_landscape = slot_size[0] > slot_size[1] 956 | 957 | min_dist = float('inf') 958 | line = None 959 | close_line_dim = 0 960 | for wallIndex, wallLabels in enumerate(wall_dict[2]): 961 | wallLine = wall_dict[1][wallIndex] 962 | lineDim = calcLineDim(wallPoints, wallLine) 963 | center = np.round((np.array(wallPoints[wallLine[0]][:2]) + np.array(wallPoints[wallLine[1]][:2])) / 2).astype(np.int32) 964 | point1=np.array(wallPoints[wallLine[0]][:2]) 965 | point2 = np.array(wallPoints[wallLine[1]][:2]) 966 | n = point2 - point1 967 | dist = np.dot((point1 - slot_center) - np.dot((point1 - slot_center), n), n) 968 | #print('indices', wallIndex, wallLabels, wallLine) 969 | # print('points', wallPoints[wallLine[0]], wallPoints[wallLine[1]]) 970 | # pass 971 | 972 | if dist < 5: 973 | min_dist = dist 974 | line = (point1, point2) 975 | close_line_dim = lineDim 976 | pass 977 | pass 978 | 979 | #sys.stderr.write("{}, {}, {}, {}, {}\n".format(y_1, y_2, x_1, x_2, iconImage.shape)) 980 | print('has line: ', line, name, close_line_dim) 981 | if name == "toilet": 982 | if line is not None: 983 | if close_line_dim == 0: #x 984 | y_pos = (line[0][1] + line[1][1]) / 2 985 | if y_pos > y_2: #toilet is below 986 | print('first case rot') 987 | iconImage = rotateImage(iconImage, 2) 988 | elif y_pos < y_1: # toilet is above 989 | pass # do nothing 990 | else: 991 | print("bad case", x_1, x_2, y_1, y_2, line) 992 | pass 993 | else: #y 994 | x_pos = (line[0][0] + line[1][0])/2 995 | print('here', x_pos, x_1, x_2) 996 | if x_pos > x_2: #toilet is to the left 997 | pass # do nothing 998 | elif x_pos < x_1: # toilet is to the right 999 | print(slot_is_landscape, icon_is_landscape) 1000 | if slot_is_landscape: 1001 | iconImage = rotateImage(iconImage, 2) 1002 | pass # do nothing 1003 | else: 1004 | print("bad case", x_1, x_2, y_1, y_2, line) 1005 | pass 1006 | pass 1007 | elif name == "washing_basin": 1008 | if line is not None: 1009 | if close_line_dim == 0: #x 1010 | 1011 | y_pos = (line[0][1] + line[1][1]) / 2 1012 | print(y_pos, y_1, y_2, 'y') 1013 | if y_pos > y_2: #toilet is below 1014 | iconImage = rotateImage(iconImage, 2) 1015 | pass 1016 | elif y_pos < y_1: # toilet is above 1017 | pass # do nothing 1018 | else: 1019 | print("bad case", x_1, x_2, y_1, y_2, line) 1020 | pass 1021 | else: #y 1022 | x_pos = (line[0][0] + line[1][0])/2 1023 | print(x_pos, x_1, x_2 , 'x') 1024 | if x_pos > x_2: #toilet is to the left 1025 | pass # do nothing 1026 | elif x_pos < x_1: # toilet is to the right 1027 | if not slot_is_landscape: 1028 | iconImage = rotateImage(iconImage, 2) 1029 | pass # do nothing 1030 | pass 1031 | else: 1032 | print("bad case", x_1, x_2, y_1, y_2, line) 1033 | pass 1034 | pass 1035 | pass 1036 | pass 1037 | pass 1038 | if slot_is_landscape != icon_is_landscape: 1039 | iconImage = rotateImage(iconImage, 1) 1040 | 1041 | 1042 | iconImage = cv2.resize(iconImage, slot_size) 1043 | 1044 | 1045 | image[y_1:y_2 + 1, x_1:x_2 + 1] = iconImage 1046 | if name == "washing_basin": 1047 | print('basin pose', [x_1, y_1, x_2, y_2]) 1048 | iconPos.append([x_1,y_1, x_2, y_2]) 1049 | continue 1050 | 1051 | fontSize = 0.7 1052 | for mask, labels in rooms: 1053 | label = min([label for label in labels]) 1054 | if label <= 0: 1055 | continue 1056 | ys, xs = mask.nonzero() 1057 | print(xs.mean(), ys.mean(), label) 1058 | 1059 | #label_half_size_x = int(fontSize * len(labelRoomMap[label]) / 2 * 20) 1060 | #label_half_size_y = int(fontSize / 2 * 20) 1061 | ret, baseline = cv2.getTextSize(labelRoomMap[label], fontFace=cv2.FONT_HERSHEY_TRIPLEX, fontScale=fontSize, thickness=1) 1062 | print(labelRoomMap[label]) 1063 | #print('ret', ret) 1064 | center = findBestTextLabelCenter(iconPos, xs, ys, ret[0]/2, ret[1]/2) 1065 | print('comp', [xs.mean(), ys.mean()], center) 1066 | #center = np.round([xs.mean(), ys.mean()]).astype(np.int32) 1067 | if center is not None: 1068 | cv2.putText(image, labelRoomMap[label], (center[0] - ret[0]/2, center[1] + ret[1]/2), fontFace=cv2.FONT_HERSHEY_TRIPLEX, fontScale=fontSize, color=(0, 0, 0), thickness=1) 1069 | else: 1070 | if label != 4: 1071 | import sys 1072 | sys.stderr.write("panic! I cannot find valid position to put label in room: {}, {}\n".format(label, labelRoomMap[label])) 1073 | continue 1074 | print('end draw') 1075 | #cv2.imwrite('test/result.png', image) 1076 | #exit(1) 1077 | #cv2.imwrite('test/region.png', drawSegmentationImage(roomRegions)) 1078 | # for regionIndex in xrange(roomRegions.max() + 1): 1079 | # cv2.imwrite('test/mask_' + str(regionIndex) + '.png', drawMaskImage(roomRegions == regionIndex)) 1080 | # continue 1081 | #resultImage = drawLines('', width, height, result['wall'][0], result['wall'][1], result['wall'][2], None, lineWidth=3) 1082 | #resultImage = drawLines('', width, height, result['door'][0], result['door'][1], [], resultImage, lineWidth=2, lineColor=0) 1083 | #iconImage = drawRectangles('', width, height, result['icon'][0], result['icon'][1], result['icon'][2], lineWidth=2) 1084 | return image 1085 | def findBestTextLabelCenter( icon_pos, xs, ys, label_half_size_x, label_half_size_y): 1086 | center = np.array([xs.mean(), ys.mean()]) 1087 | icon_pos = np.array(icon_pos) 1088 | room_points = np.array([xs, ys]).transpose() 1089 | min_point = room_points.min(axis=0, keepdims=True) 1090 | max_point = room_points.max(axis=0, keepdims=True) 1091 | size = np.array([label_half_size_x, label_half_size_y]) 1092 | print('size', size) 1093 | avail_min_point = min_point + size 1094 | avail_max_point = max_point - size 1095 | 1096 | avail_points = np.logical_and(room_points > avail_min_point, room_points < avail_max_point) 1097 | avail_points = np.all(avail_points, axis=1) 1098 | 1099 | room_points_aug = np.tile(room_points[:, :, np.newaxis], [1, 1, icon_pos.shape[0]]) 1100 | 1101 | room_points_gt_tl_x = room_points_aug[:, 0, :] > icon_pos[:, 0] - size[0] - 5 1102 | room_points_lt_br_x = room_points_aug[:, 0, :] < icon_pos[:, 2] + size[0] + 5 1103 | 1104 | room_points_gt_tl_y = room_points_aug[:, 1, :] > icon_pos[:, 1] - size[1] - 5 1105 | room_points_lt_br_y = room_points_aug[:, 1, :] < icon_pos[:, 3] + size[1] + 5 1106 | 1107 | room_points_in_square = np.logical_and(room_points_gt_tl_x, room_points_lt_br_x) 1108 | room_points_in_square = np.logical_and(room_points_in_square, room_points_gt_tl_y) 1109 | room_points_in_square = np.logical_and(room_points_in_square, room_points_lt_br_y) 1110 | 1111 | #room_points_in_square = np.all(room_points_in_square, axis=1) 1112 | room_points_in_square = np.any(room_points_in_square, axis=1) 1113 | 1114 | room_points_not_in_square = np.logical_not(room_points_in_square) 1115 | 1116 | good_points_mask = np.logical_and(avail_points, room_points_not_in_square) 1117 | good_points = room_points[good_points_mask] 1118 | good_points_center_dist = np.linalg.norm(good_points - center, axis=1) 1119 | #good_points_center_dist = np.sum(np.abs(good_points - center), axis=1) 1120 | #print('icon_pos') 1121 | #print(icon_pos) 1122 | #print('goodpoints') 1123 | #print(center) 1124 | #print(good_points) 1125 | #print(good_points_center_dist) 1126 | if len(good_points) == 0: 1127 | #print('give up') 1128 | return None 1129 | #return np.round(center).astype(np.int32) 1130 | best_point_idx = np.argmin(good_points_center_dist, axis=0) 1131 | #print('cost', good_points_center_dist[best_point_idx]) 1132 | #print('best points') 1133 | #print(good_points[best_point_idx]) 1134 | return good_points[best_point_idx] 1135 | 1136 | def rotateImage(image, orientation): 1137 | if orientation == 0: 1138 | return image 1139 | elif orientation == 1: 1140 | return np.flip(image.transpose((1, 0, 2)), axis=0) 1141 | elif orientation == 2: 1142 | return np.flip(np.flip(image, axis=1), axis=0) 1143 | else: 1144 | return np.flip(image.transpose(1, 0, 2), axis=1) 1145 | return 1146 | 1147 | def projectIndices(pointIndices, pointSegmentation, min_x, max_x, min_y, max_y): 1148 | if max_x - min_x == 1 and max_y - min_y == 1: 1149 | pointIndices[pointSegmentation[:, 2]] = min_y * WIDTH + min_x 1150 | return 1151 | elif max_x - min_x >= max_y - min_y: 1152 | middle_x = int((max_x + min_x + 1) / 2) 1153 | mask_1 = pointSegmentation[:, 0] < middle_x 1154 | projectIndices(pointIndices, pointSegmentation[mask_1], min_x, middle_x, min_y, max_y) 1155 | mask_2 = pointSegmentation[:, 0] >= middle_x 1156 | projectIndices(pointIndices, pointSegmentation[mask_2], middle_x, max_x, min_y, max_y) 1157 | else: 1158 | middle_y = int((max_y + min_y + 1) / 2) 1159 | mask_1 = pointSegmentation[:, 1] < middle_y 1160 | projectIndices(pointIndices, pointSegmentation[mask_1], min_x, max_x, min_y, middle_y) 1161 | mask_2 = pointSegmentation[:, 1] >= middle_y 1162 | projectIndices(pointIndices, pointSegmentation[mask_2], min_x, max_x, middle_y, max_y) 1163 | pass 1164 | return 1165 | 1166 | def drawCornerSegmentation(corners, radius=1, width=WIDTH, height=HEIGHT): 1167 | cornerSegmentation = np.zeros((height, width), dtype=np.int64) 1168 | for corner in corners: 1169 | cornerSegmentation[max(corner[1] - radius + 1, 0):min(corner[1] + radius, height - 1), max(corner[0] - radius + 1, 0):min(corner[0] + radius, width - 1)] = corner[2] 1170 | continue 1171 | return cornerSegmentation 1172 | 1173 | def getOrientationCorners(corners, cornerSize=3): 1174 | orientationCorners = [[] for _ in xrange(NUM_CORNERS)] 1175 | for corner in corners: 1176 | orientationCorners[corner[2] - 1].append(((corner[0], corner[1]), (corner[0] - cornerSize, corner[1] - cornerSize), (corner[0] + cornerSize, corner[1] + cornerSize))) 1177 | continue 1178 | return orientationCorners 1179 | 1180 | def getGTPrimitives(gt_dict): 1181 | result_dict = {'wall': [wallPoints, filteredWallLines, filteredWallLabels], 'door': [doorPoints, filteredDoorLines, []], 'icon': [iconPoints, filteredIcons, filteredIconTypes]} 1182 | return 1183 | 1184 | def writeRepresentation(filename, width, height, result_dict): 1185 | labelMap = [11, 1, 2, 3, 4, 3, 6, 7, 8, 2,] 1186 | labelIconMap = getLabelIconMap() 1187 | with open(filename, 'w') as f: 1188 | f.write(str(width) + '\t' + str(height) + '\n') 1189 | f.write(str(len(result_dict['wall'][1])) + '\n') 1190 | for wallLine, wallLabels in zip(result_dict['wall'][1], result_dict['wall'][2]): 1191 | point_1 = result_dict['wall'][0][wallLine[0]] 1192 | point_2 = result_dict['wall'][0][wallLine[1]] 1193 | lineDim = calcLineDim(result_dict['wall'][0], wallLine) 1194 | if point_1[lineDim] > point_2[lineDim]: 1195 | point_1[lineDim], point_2[lineDim] = point_2[lineDim], point_1[lineDim] 1196 | pass 1197 | f.write(str(int(point_1[0])) + '\t' + str(int(point_1[1])) + '\t' + str(int(point_2[0])) + '\t' + str(int(point_2[1])) + '\t' + str(labelMap[wallLabels[0]]) + '\t' + str(labelMap[wallLabels[1]]) + '\n') 1198 | continue 1199 | for doorLine in result_dict['door'][1]: 1200 | point_1 = result_dict['door'][0][doorLine[0]] 1201 | point_2 = result_dict['door'][0][doorLine[1]] 1202 | lineDim = calcLineDim(result_dict['door'][0], doorLine) 1203 | if point_1[lineDim] > point_2[lineDim]: 1204 | point_1[lineDim], point_2[lineDim] = point_2[lineDim], point_1[lineDim] 1205 | pass 1206 | f.write(str(int(point_1[0])) + '\t' + str(int(point_1[1])) + '\t' + str(int(point_2[0])) + '\t' + str(int(point_2[1])) + '\tdoor\t1\t1\n') 1207 | continue 1208 | #print(len(result_dict['icon'][1])) 1209 | for icon, iconLabel in zip(result_dict['icon'][1], result_dict['icon'][2]): 1210 | #print(iconLabel, labelIconMap[iconLabel + 1]) 1211 | points = np.array([result_dict['icon'][0][pointIndex][:2] for pointIndex in icon]).astype(np.int32) 1212 | mins = points.min(0) 1213 | maxs = points.max(0) 1214 | f.write(str(int(mins[0])) + '\t' + str(int(mins[1])) + '\t' + str(int(maxs[0])) + '\t' + str(int(maxs[1])) + '\t' + labelIconMap[iconLabel + 1] + '\t1\t1\n') 1215 | continue 1216 | f.close() 1217 | pass 1218 | return 1219 | -------------------------------------------------------------------------------- /icons/bathtub.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/FloorNet/93db5bdf12f8b60c9d788aa8e124fbe852e50b4a/icons/bathtub.jpg -------------------------------------------------------------------------------- /icons/bed.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/FloorNet/93db5bdf12f8b60c9d788aa8e124fbe852e50b4a/icons/bed.jpg -------------------------------------------------------------------------------- /icons/cabinet.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/FloorNet/93db5bdf12f8b60c9d788aa8e124fbe852e50b4a/icons/cabinet.jpg -------------------------------------------------------------------------------- /icons/cooking_counter.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/FloorNet/93db5bdf12f8b60c9d788aa8e124fbe852e50b4a/icons/cooking_counter.jpg -------------------------------------------------------------------------------- /icons/refrigerator.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/FloorNet/93db5bdf12f8b60c9d788aa8e124fbe852e50b4a/icons/refrigerator.jpg -------------------------------------------------------------------------------- /icons/sofa.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/FloorNet/93db5bdf12f8b60c9d788aa8e124fbe852e50b4a/icons/sofa.jpg -------------------------------------------------------------------------------- /icons/table.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/FloorNet/93db5bdf12f8b60c9d788aa8e124fbe852e50b4a/icons/table.jpg -------------------------------------------------------------------------------- /icons/toilet.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/FloorNet/93db5bdf12f8b60c9d788aa8e124fbe852e50b4a/icons/toilet.jpg -------------------------------------------------------------------------------- /icons/washing_basin.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/FloorNet/93db5bdf12f8b60c9d788aa8e124fbe852e50b4a/icons/washing_basin.jpg -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import glob 3 | import cv2 4 | if 'xrange' not in globals(): 5 | xrange = range 6 | 7 | #consistent color mapping 8 | class ColorPalette: 9 | def __init__(self, numColors): 10 | np.random.seed(1) 11 | #self.colorMap = np.random.randint(255, size = (numColors, 3), dtype=np.uint8) 12 | #self.colorMap[0] = 0 13 | #self.colorMap[0] = np.maximum(self.colorMap[0], 1) 14 | 15 | # [128, 128, 128], 16 | # [0, 0, 255], 17 | # [64, 128, 192], 18 | # [0, 128, 0], 19 | # [192, 0, 0], 20 | # [128, 0, 128], 21 | # [128, 128, 192], 22 | # [128, 192, 192], 23 | # [0, 128, 0], 24 | # [0, 0, 128], 25 | # [128, 128, 0], 26 | # [0, 128, 128] 27 | 28 | self.colorMap = np.array([[255, 0, 0], 29 | [50, 150, 0], 30 | [0, 0, 255], 31 | [80, 128, 255], 32 | [255, 230, 180], 33 | [255, 0, 255], 34 | [0, 255, 255], 35 | [255, 255, 0], 36 | [0, 255, 0], 37 | [200, 255, 255], 38 | [255, 200, 255], 39 | [100, 0, 0], 40 | [0, 100, 0], 41 | [128, 128, 80], 42 | [0, 50, 128], 43 | [0, 100, 100], 44 | [0, 255, 128], 45 | [0, 128, 255], 46 | [255, 0, 128], 47 | [128, 0, 255], 48 | [255, 128, 0], 49 | [128, 255, 0], 50 | ], dtype=np.uint8) 51 | 52 | self.colorMap = np.maximum(self.colorMap, 1) 53 | 54 | if numColors > self.colorMap.shape[0]: 55 | self.colorMap = np.concatenate([self.colorMap, np.random.randint(255, size = (numColors, 3), dtype=np.uint8)], axis=0) 56 | pass 57 | 58 | return 59 | 60 | def getColorMap(self): 61 | return self.colorMap 62 | 63 | def getColor(self, index): 64 | if index >= colorMap.shape[0]: 65 | return np.random.randint(255, size = (3), dtype=np.uint8) 66 | else: 67 | return self.colorMap[index] 68 | pass 69 | 70 | def sigmoid(values): 71 | return 1 / (1 + np.exp(-values)) 72 | 73 | def softmax(values): 74 | values_exp = np.exp(values) 75 | return values_exp / np.sum(values_exp, axis=-1, keepdims=True) 76 | 77 | 78 | #Draw density image 79 | def drawDensityImage(density, maxDensity=-1, nChannels=1): 80 | if maxDensity < 0: 81 | maxDensity = density.max() / 2 82 | pass 83 | densityImage = np.minimum(np.round(density / maxDensity * 255).astype(np.uint8), 255) 84 | if nChannels == 3: 85 | densityImage = np.stack([densityImage, densityImage, densityImage], axis=2) 86 | pass 87 | return densityImage 88 | 89 | #Draw segmentation image. The input could be either HxW or HxWxC 90 | def drawSegmentationImage(segmentations, numColors=42, blackIndex=-1): 91 | if segmentations.ndim == 2: 92 | numColors = max(numColors, segmentations.max() + 2) 93 | else: 94 | numColors = max(numColors, segmentations.shape[2] + 2) 95 | pass 96 | randomColor = ColorPalette(numColors).getColorMap() 97 | if blackIndex >= 0: 98 | randomColor[blackIndex] = 0 99 | pass 100 | width = segmentations.shape[1] 101 | height = segmentations.shape[0] 102 | if segmentations.ndim == 3: 103 | #segmentation = (np.argmax(segmentations, 2) + 1) * (np.max(segmentations, 2) > 0.5) 104 | segmentation = np.argmax(segmentations, 2) 105 | else: 106 | segmentation = segmentations 107 | pass 108 | segmentation = segmentation.astype(np.int) 109 | return randomColor[segmentation.reshape(-1)].reshape((height, width, 3)) 110 | 111 | def drawMaskImage(mask): 112 | return (np.clip(mask * 255, 0, 255)).astype(np.uint8) 113 | 114 | 115 | 116 | def projectIndices(indicesMap, pointSegmentation, min_x, max_x, min_y, max_y, width): 117 | if max_x - min_x == 1 and max_y - min_y == 1: 118 | indicesMap[pointSegmentation[:, 2]] = min_y * width + min_x 119 | return 120 | elif max_x - min_x >= max_y - min_y: 121 | middle_x = int((max_x + min_x + 1) / 2) 122 | mask_1 = pointSegmentation[:, 0] < middle_x 123 | projectIndices(indicesMap, pointSegmentation[mask_1], min_x, middle_x, min_y, max_y, width) 124 | mask_2 = pointSegmentation[:, 0] >= middle_x 125 | projectIndices(indicesMap, pointSegmentation[mask_2], middle_x, max_x, min_y, max_y, width) 126 | else: 127 | middle_y = int((max_y + min_y + 1) / 2) 128 | mask_1 = pointSegmentation[:, 1] < middle_y 129 | projectIndices(indicesMap, pointSegmentation[mask_1], min_x, max_x, min_y, middle_y, width) 130 | mask_2 = pointSegmentation[:, 1] >= middle_y 131 | projectIndices(indicesMap, pointSegmentation[mask_2], min_x, max_x, middle_y, max_y, width) 132 | pass 133 | return 134 | 135 | 136 | #Extract corners from heatmaps 137 | def extractCornersFromHeatmaps(heatmaps, heatmapThreshold=0.5, numPixelsThreshold=5, returnRanges=True): 138 | from skimage import measure 139 | heatmaps = (heatmaps > heatmapThreshold).astype(np.float32) 140 | orientationPoints = [] 141 | #kernel = np.ones((3, 3), np.float32) 142 | for heatmapIndex in xrange(0, heatmaps.shape[-1]): 143 | heatmap = heatmaps[:, :, heatmapIndex] 144 | #heatmap = cv2.dilate(cv2.erode(heatmap, kernel), kernel) 145 | components = measure.label(heatmap, background=0) 146 | points = [] 147 | for componentIndex in xrange(components.min() + 1, components.max() + 1): 148 | ys, xs = (components == componentIndex).nonzero() 149 | if ys.shape[0] <= numPixelsThreshold: 150 | continue 151 | #print(heatmapIndex, xs.shape, ys.shape, componentIndex) 152 | if returnRanges: 153 | points.append(((xs.mean(), ys.mean()), (xs.min(), ys.min()), (xs.max(), ys.max()))) 154 | else: 155 | points.append((xs.mean(), ys.mean())) 156 | pass 157 | continue 158 | orientationPoints.append(points) 159 | continue 160 | return orientationPoints 161 | 162 | #Extract corners from heatmaps 163 | def extractCornersFromSegmentation(segmentation, cornerTypeRange=[0, 13]): 164 | from skimage import measure 165 | orientationPoints = [] 166 | for heatmapIndex in xrange(cornerTypeRange[0], cornerTypeRange[1]): 167 | heatmap = segmentation == heatmapIndex 168 | #heatmap = cv2.dilate(cv2.erode(heatmap, kernel), kernel) 169 | components = measure.label(heatmap, background=0) 170 | points = [] 171 | for componentIndex in xrange(components.min()+1, components.max() + 1): 172 | ys, xs = (components == componentIndex).nonzero() 173 | points.append((xs.mean(), ys.mean())) 174 | continue 175 | orientationPoints.append(points) 176 | continue 177 | return orientationPoints 178 | 179 | #Extract corners from heatmaps 180 | def getSegmentationFromCorners(width, height, orientationCorners): 181 | segmentation = np.zeros((height, width)) 182 | for orientation, corners in enumerate(orientationCorners): 183 | for corner in corners: 184 | segmentation[int(round(corner[1]))][int(round(corner[0]))] = orientation + 1 185 | continue 186 | continue 187 | return segmentation 188 | 189 | #Evaluate corner predictions 190 | def evaluateCorners(cornersPred, cornersGT, distanceThreshold = 15): 191 | numGT = 0 192 | numPred = 0 193 | numMatches = 0 194 | for cornerType, gt_c in enumerate(cornersGT): 195 | pred_c = cornersPred[cornerType] 196 | gt_c = np.array(gt_c) 197 | pred_c = np.array(pred_c) 198 | numGT += gt_c.shape[0] 199 | numPred += pred_c.shape[0] 200 | if gt_c.shape[0] == 0 or pred_c.shape[0] == 0: 201 | continue 202 | diff = np.linalg.norm(np.expand_dims(gt_c, 1) - np.expand_dims(pred_c, 0), axis=2) 203 | numMatches += (diff.min(axis=1) < distanceThreshold).sum() 204 | continue 205 | return np.array([numMatches, numGT, numPred]) 206 | 207 | 208 | #Evaluate segmentation predictions 209 | def evaluateSegmentation(segmentationPred, segmentationGT, numSegments = 12): 210 | #print("hack in evaluate! remove this!") 211 | #print(segmentationGT.shape) 212 | #segmentationGT = segmentationGT.transpose(1, 0) 213 | height = segmentationPred.shape[0] 214 | width = segmentationPred.shape[1] 215 | nonemptyMask = segmentationGT > 0 216 | correctMask = segmentationPred == segmentationGT 217 | #accuracy = float(correctMask.sum()) / (width * height) 218 | accuracy = float(correctMask[nonemptyMask].sum()) / max(nonemptyMask.sum(), 1) 219 | #(width * height) 220 | 221 | sumIOU = 0. 222 | numIOU = 0 223 | for segmentIndex in xrange(numSegments): 224 | gt_s = segmentationGT == segmentIndex 225 | pred_s = segmentationPred == segmentIndex 226 | union = np.logical_or(pred_s, gt_s) 227 | unionSum = union.sum() 228 | if unionSum == 0: 229 | continue 230 | intersection = np.logical_and(pred_s, gt_s) 231 | IOU = float(intersection.sum()) / unionSum 232 | sumIOU += IOU 233 | numIOU += 1 234 | continue 235 | meanIOU = sumIOU / numIOU 236 | return np.array([accuracy, meanIOU]) 237 | 238 | def evaluateDetection(segmentationPred, segmentationGT, numSegments = 12, IOUThreshold = 0.5): 239 | from skimage import measure 240 | numGT = 0 241 | numPred = 0 242 | numMatches = 0 243 | 244 | for segmentIndex in xrange(numSegments): 245 | gt_s = segmentationGT == segmentIndex 246 | if gt_s.sum() > 0: 247 | numGT += 1 248 | pass 249 | pred_s = segmentationPred == segmentIndex 250 | if pred_s.sum() > 0: 251 | numPred += 1 252 | pass 253 | IOU = float(np.logical_and(pred_s, gt_s).sum()) / np.logical_or(pred_s, gt_s).sum() 254 | if IOU > IOUThreshold: 255 | numMatches += 1 256 | pass 257 | continue 258 | return (numMatches, numGT, numPred) 259 | 260 | def fitPlane(points): 261 | if points.shape[0] == points.shape[1]: 262 | return np.linalg.solve(points, np.ones(points.shape[0])) 263 | else: 264 | return np.linalg.lstsq(points, np.ones(points.shape[0]))[0] 265 | 266 | def rotatePoints(points, segmentation, numSampledPoints = 10000): 267 | sampledInds = np.arange(points.shape[0]) 268 | np.random.shuffle(sampledInds) 269 | sampledPoints = points[sampledInds[:numSampledPoints]] 270 | sampledSegmentation = segmentation[sampledInds[:numSampledPoints]] 271 | segments = np.unique(sampledSegmentation).tolist() 272 | binSize = 3 273 | numAngleBins = 90 // 3 + 1 274 | angleSums = np.zeros(numAngleBins) 275 | angleCounts = np.zeros(numAngleBins) 276 | for segmentIndex in segments: 277 | segmentPoints = sampledPoints[sampledSegmentation == segmentIndex] 278 | if segmentPoints.shape[0] < 3: 279 | continue 280 | try: 281 | plane = fitPlane(segmentPoints) 282 | except: 283 | continue 284 | if np.argmax(np.abs(plane)) == 2: 285 | continue 286 | angle = np.arctan2(plane[0], plane[1]) 287 | angle = np.rad2deg(angle) % 90 288 | numPoints = segmentPoints.shape[0] 289 | angleSums[int(np.round(angle / 3))] += angle * numPoints 290 | angleCounts[int(np.round(angle / 3))] += numPoints 291 | continue 292 | angles = angleSums / np.maximum(angleCounts, 1) 293 | angle = angles[np.argmax(angleCounts)] 294 | angle = np.deg2rad(angle) 295 | rotationMatrix = np.zeros((2, 2)) 296 | rotationMatrix[0][0] = np.cos(angle) 297 | rotationMatrix[0][1] = np.sin(angle) 298 | rotationMatrix[1][0] = -np.sin(angle) 299 | rotationMatrix[1][1] = np.cos(angle) 300 | 301 | points[:, :2] = np.matmul(points[:, :2], rotationMatrix) 302 | return points 303 | 304 | def rotatePointsWithMatrix(points, segmentation, numSampledPoints = 10000): 305 | sampledInds = np.arange(points.shape[0]) 306 | np.random.shuffle(sampledInds) 307 | sampledPoints = points[sampledInds[:numSampledPoints]] 308 | sampledSegmentation = segmentation[sampledInds[:numSampledPoints]] 309 | segments = np.unique(sampledSegmentation).tolist() 310 | binSize = 3 311 | numAngleBins = 90 // 3 + 1 312 | #print("rotate!", numAngleBins, flush=True) 313 | assert isinstance(numAngleBins, int), numAngleBins 314 | 315 | angleSums = np.zeros(numAngleBins) 316 | angleCounts = np.zeros(numAngleBins) 317 | for segmentIndex in segments: 318 | segmentPoints = sampledPoints[sampledSegmentation == segmentIndex] 319 | if segmentPoints.shape[0] < 3: 320 | continue 321 | try: 322 | plane = fitPlane(segmentPoints) 323 | except: 324 | continue 325 | if np.argmax(np.abs(plane)) == 2: 326 | continue 327 | angle = np.arctan2(plane[0], plane[1]) 328 | angle = np.rad2deg(angle) % 90 329 | numPoints = segmentPoints.shape[0] 330 | angleSums[int(np.round(angle / 3))] += angle * numPoints 331 | angleCounts[int(np.round(angle / 3))] += numPoints 332 | continue 333 | angles = angleSums / np.maximum(angleCounts, 1) 334 | angle = angles[np.argmax(angleCounts)] 335 | angle = np.deg2rad(angle) 336 | rotationMatrix = np.zeros((2, 2)) 337 | rotationMatrix[0][0] = np.cos(angle) 338 | rotationMatrix[0][1] = np.sin(angle) 339 | rotationMatrix[1][0] = -np.sin(angle) 340 | rotationMatrix[1][1] = np.cos(angle) 341 | 342 | points[:, :2] = np.matmul(points[:, :2], rotationMatrix) 343 | return points, rotationMatrix 344 | 345 | 346 | def drawTopDownView(points, width, height): 347 | coordinates = points[:, :2] 348 | mins = coordinates.min(0, keepdims=True) 349 | maxs = coordinates.max(0, keepdims=True) 350 | ranges = maxs - mins 351 | padding = ranges * 0.05 352 | mins -= padding 353 | ranges += padding * 2 354 | 355 | maxRange = ranges.max() 356 | mins = (maxs + mins) / 2 - maxRange / 2 357 | 358 | coordinates = ((coordinates - mins) / ranges * height).astype(np.int32) 359 | coordinates = np.minimum(coordinates, height - 1) 360 | image = np.zeros((height, width)) 361 | for coordinate in coordinates: 362 | image[coordinate[1]][coordinate[0]] += 1 363 | continue 364 | print(image.max()) 365 | image /= min(image.max(), 300) 366 | image = (np.minimum(image * 255, 255)).astype(np.uint8) 367 | return image 368 | 369 | def writePointCloud(filename, pointCloud): 370 | with open(filename, 'w') as f: 371 | header = """ply 372 | format ascii 1.0 373 | element vertex """ 374 | header += str(len(pointCloud)) 375 | header += """ 376 | property float x 377 | property float y 378 | property float z 379 | property uchar red { start of vertex color } 380 | property uchar green 381 | property uchar blue 382 | end_header 383 | """ 384 | f.write(header) 385 | for point in pointCloud: 386 | for valueIndex, value in enumerate(point): 387 | if valueIndex < 3: 388 | f.write(str(value) + ' ') 389 | else: 390 | f.write(str(int(value * 255)) + ' ') 391 | pass 392 | continue 393 | f.write('\n') 394 | continue 395 | f.close() 396 | pass 397 | return 398 | 399 | 400 | # def convertPointCloudIndices(folder, width=256, height=256): 401 | # filenames = glob.glob(folder + '/pointcloud_indices_*.npy') 402 | # for filename in filenames: 403 | # pointcloud_indices = np.load(filename) 404 | # new_pointcloud_indices = [] 405 | # for imageIndex in xrange(pointcloud_indices.shape[0]): 406 | # pointcloud_indices_0 = pointcloud_indices[imageIndex][0] - imageIndex * width * height 407 | # new_pointcloud_indices.append(pointcloud_indices_0) 408 | # continue 409 | # np.save(filename, np.stack(new_pointcloud_indices, 0)) 410 | # continue 411 | # return 412 | 413 | 414 | def getDensity(points, width=256, height=256): 415 | imageSizes = np.array([width, height]).reshape((-1, 2)) 416 | # mins = points.min(0, keepdims=True) 417 | # maxs = points.max(0, keepdims=True) 418 | # maxRange = (maxs - mins)[:, :2].max() 419 | # padding = maxRange * 0.05 420 | # mins = (maxs + mins) / 2 - maxRange / 2 421 | # mins -= padding 422 | # maxRange += padding * 2 423 | # coordinates = np.round((points - mins) / maxRange * imageSizes).astype(np.int32) 424 | coordinates = np.round(points[:, :2] * imageSizes).astype(np.int32) 425 | coordinates = np.minimum(np.maximum(coordinates, 0), imageSizes - 1) 426 | density = np.zeros((height, width)) 427 | for uv in coordinates: 428 | density[uv[1], uv[0]] += 1 429 | continue 430 | return density 431 | 432 | def getDensityFromIndices(indices, width=256, height=256): 433 | density = np.zeros((height, width)) 434 | for index in indices: 435 | #print(index, index / width, index % width) 436 | density[index / width, index % width] += 1 437 | continue 438 | return density 439 | 440 | def drawCornerImages(segmentations, numColors=42, blackIndex=0): 441 | if segmentations.ndim == 2: 442 | numColors = max(numColors, segmentations.max() + 2) 443 | else: 444 | numColors = max(numColors, segmentations.shape[2] + 2) 445 | pass 446 | randomColor = ColorPalette(numColors).getColorMap() 447 | if blackIndex >= 0: 448 | randomColor[blackIndex] = 0 449 | pass 450 | width = segmentations.shape[1] 451 | height = segmentations.shape[0] 452 | if segmentations.ndim == 3: 453 | #segmentation = (np.argmax(segmentations, 2) + 1) * (np.max(segmentations, 2) > 0.5) 454 | segmentation = np.argmax(segmentations, 2) 455 | else: 456 | segmentation = segmentations 457 | pass 458 | segmentation = segmentation.astype(np.int32).reshape(-1) 459 | images = [] 460 | print(np.unique(segmentation)) 461 | for segment in [(1, 14), (14, 18), (18, 22)]: 462 | colorMap = randomColor.copy() 463 | colorMap[:segment[0]] = 0 464 | colorMap[segment[1]:] = 0 465 | image = colorMap[segmentation].reshape((height, width, 3)) 466 | image = cv2.dilate(image, np.ones((3, 3), dtype=np.uint8), iterations=3) 467 | images.append(image) 468 | continue 469 | return images 470 | 471 | def segmentation2Heatmaps(segmentation, numLabels): 472 | width = segmentation.shape[1] 473 | height = segmentation.shape[0] 474 | labels = np.arange(numLabels, dtype=np.int32).reshape((1, 1, -1)) 475 | heatmaps = (np.expand_dims(segmentation, -1) == labels).astype(np.float32) 476 | return heatmaps 477 | 478 | def heatmaps2Segmentation(heatmaps): 479 | return np.argmax(heatmaps, axis=2) 480 | 481 | def calcIOU(rectangle_1, rectangle_2): 482 | # mins_1 = rectangle_1.min(0) 483 | # maxs_1 = rectangle_1.max(0) 484 | # area_1 = (maxs_1[0] - mins_1[0] + 1) * (maxs_1[1] - mins_1[1] + 1) 485 | # mins_2 = rectangle_2.min(0) 486 | # maxs_2 = rectangle_2.max(0) 487 | # area_2 = (maxs_2[0] - mins_2[0] + 1) * (maxs_2[1] - mins_2[1] + 1) 488 | # intersection = (min(maxs_1[0], maxs_2[0]) - max(mins_1[0], mins_2[0]) + 1) * (min(maxs_1[1], maxs_2[1]) - max(mins_1[1], mins_2[1]) + 1) 489 | 490 | rectangles = [rectangle_1, rectangle_2] 491 | 492 | x_1 = max([int(round((rectangle[0][0] + rectangle[2][0]) / 2)) for rectangle in rectangles]) 493 | x_2 = min([int(round((rectangle[1][0] + rectangle[3][0]) / 2)) for rectangle in rectangles]) 494 | y_1 = max([int(round((rectangle[0][1] + rectangle[1][1]) / 2)) for rectangle in rectangles]) 495 | y_2 = min([int(round((rectangle[2][1] + rectangle[3][1]) / 2)) for rectangle in rectangles]) 496 | if x_1 >= x_2 or y_1 >= y_2: 497 | return 0 498 | intersection = (x_2 - x_1 + 1) * (y_2 - y_1 + 1) 499 | 500 | area_1, area_2 = ((int(round((rectangle[1][0] + rectangle[3][0]) / 2)) - int(round((rectangle[0][0] + rectangle[2][0]) / 2)) + 1) * (int(round((rectangle[2][1] + rectangle[3][1]) / 2)) - int(round((rectangle[0][1] + rectangle[1][1]) / 2)) + 1) for rectangle in rectangles) 501 | 502 | union = area_1 + area_2 - intersection 503 | return float(intersection) / union 504 | 505 | def calcIOUMask(mask_1, mask_2): 506 | intersection = (mask_1 * mask_2).sum() 507 | union = mask_1.sum() + mask_2.sum() - intersection 508 | return float(intersection) / max(union, 1) 509 | 510 | 511 | def gaussian(k=5, sig=0): 512 | """ 513 | creates gaussian kernel with side length l and a sigma of sig 514 | v """ 515 | if sig == 0: 516 | sig = 0.3 * ((k - 1) * 0.5 - 1) + 0.8 517 | pass 518 | 519 | ax = np.arange(-k // 2 + 1., k // 2 + 1.) 520 | xx, yy = np.meshgrid(ax, ax) 521 | 522 | kernel = np.exp(-(xx**2 + yy**2) / (2. * sig**2)) 523 | 524 | return kernel / np.sum(kernel) 525 | 526 | def disk(k): 527 | """ 528 | creates gaussian kernel with side length l and a sigma of sig 529 | """ 530 | ax = np.arange(-k // 2 + 1., k // 2 + 1.) 531 | xx, yy = np.meshgrid(ax, ax) 532 | 533 | kernel = (np.sqrt(pow(xx, 2) + pow(yy, 2)) <= (k - 1) / 2).astype(np.float32) 534 | 535 | return kernel 536 | 537 | if __name__ == '__main__': 538 | rect_1 = np.array([[54.56637168141593, 91.65781710914455], [85.51592356687898, 90.92993630573248], [64.36440677966101, 123.17372881355932], [84.18115942028986, 109.3840579710145]]) 539 | rect_2 = np.array([[64.0, 92.0, 1, 2], [86.0, 92.0, 1, 3], [64.0, 111.0, 1, 1], [86.0, 111.0, 1, 0]]) 540 | print('IOU', calcIOU(rect_1, rect_2)) 541 | pass 542 | --------------------------------------------------------------------------------