├── LICENSE ├── README.md ├── datasets ├── scannet_dataset.py ├── semantic_counts_pixelwise.npy ├── split_test.txt ├── split_train.txt └── split_val.txt ├── inference.py ├── models └── instance.py ├── options.py ├── requirements.txt ├── scripts ├── evaluate_semantic_instance.py ├── evaluate_semantic_label.py ├── export_script.py ├── export_train_mesh_for_evaluation.py ├── prepare_data.py ├── prepare_data_ori.py ├── util.py └── util_3d.py ├── train.py ├── train_confidence.py └── utils.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 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 | # MASC: Multi-scale Affinity with Sparse Convolution for 3D Instance Segmentation (Technical Report) 2 | ## Introduction 3 | This is the PyTorch implementation for our [technical report](https://arxiv.org/abs/1902.04478) which achieves the state-of-the-art performance on the 3D instance segmentation task of the ScanNet benchmark. 4 | 5 | ## Installation 6 | ``` 7 | pip install -r requirements.txt 8 | ``` 9 | We are using Python 3.5.2. And as pointed out by [Issue #3](https://github.com/art-programmer/MASC/issues/3), please consider using Python 3.6 and refer to [SparseConvNet](https://github.com/facebookresearch/SparseConvNet) for related issues. 10 | 11 | ## Data preparation 12 | To prepare training data from ScanNet mesh models, please run: 13 | ``` 14 | python train.py --task=prepare --dataFolder=[SCANNET_PATH] --labelFile=[SCANNET_LABEL_FILE_PATH (i.e., scannetv2-labels.combined.tsv)] 15 | ``` 16 | 17 | ## Training 18 | To train the main model which predict semantics and affinities, please run: 19 | ``` 20 | python train.py --restore=0 --dataFolder=[SCANNET_PATH] 21 | ``` 22 | 23 | ## Validation 24 | To validate the trained model, please run: 25 | ``` 26 | python train.py --restore=1 --dataFolder=[SCANNET_PATH] --task=test 27 | ``` 28 | 29 | ## Inference 30 | To run the inference using the trained model, please run: 31 | 32 | ``` 33 | python inference.py --dataFolder=[SCANNET_PATH] --task=predict_cluster split=val 34 | ``` 35 | 36 | The task option indicates: 37 | - "predict": predict semantics and affinities 38 | - "cluster": run the clustering algorithm based on the predicted affinities 39 | - "write": write instance segmentation results 40 | 41 | The "task" option can contain any combinations of these three tasks, but the earlier task must be run before later tasks. And a task only needs to be run once. The "split" option specifies the data split to run the inference. 42 | 43 | ## Write results for the final evaluation 44 | To train the instance confidence model, please first generate the instance segmentation results: 45 | ``` 46 | python inference.py --dataFolder=[SCANNET_PATH] --task=predict_cluster --split=val 47 | python inference.py --dataFolder=[SCANNET_PATH] --task=predict_cluster --split=train 48 | ``` 49 | 50 | Then train the confidence model: 51 | ``` 52 | python train_confidence.py --restore=0 --dataFolder=[SCANNET_PATH] 53 | ``` 54 | 55 | Predict instance confidence, add additional instances for certain semantic labels, and write instance segmentation results: 56 | ``` 57 | python inference.py --dataFolder=[SCANNET_PATH] --task=predict_cluster_write split=test 58 | ``` 59 | -------------------------------------------------------------------------------- /datasets/scannet_dataset.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data import Dataset 2 | import torch 3 | import numpy as np 4 | import time 5 | import os 6 | import scipy.ndimage 7 | from utils import remapper 8 | 9 | blur0=np.ones((3,1,1)).astype('float32')/3 10 | blur1=np.ones((1,3,1)).astype('float32')/3 11 | blur2=np.ones((1,1,3)).astype('float32')/3 12 | 13 | ## Elastic augmentation from https://github.com/facebookresearch/SparseConvNet/blob/master/examples/ScanNet/data.py 14 | def elastic(x,gran,mag): 15 | bb=np.abs(x).max(0).astype(np.int32)//gran+3 16 | noise=[np.random.randn(bb[0],bb[1],bb[2]).astype('float32') for _ in range(3)] 17 | noise=[scipy.ndimage.filters.convolve(n,blur0,mode='constant',cval=0) for n in noise] 18 | noise=[scipy.ndimage.filters.convolve(n,blur1,mode='constant',cval=0) for n in noise] 19 | noise=[scipy.ndimage.filters.convolve(n,blur2,mode='constant',cval=0) for n in noise] 20 | noise=[scipy.ndimage.filters.convolve(n,blur0,mode='constant',cval=0) for n in noise] 21 | noise=[scipy.ndimage.filters.convolve(n,blur1,mode='constant',cval=0) for n in noise] 22 | noise=[scipy.ndimage.filters.convolve(n,blur2,mode='constant',cval=0) for n in noise] 23 | ax=[np.linspace(-(b-1)*gran,(b-1)*gran,b) for b in bb] 24 | 25 | interp=[scipy.interpolate.RegularGridInterpolator(ax,n,bounds_error=0,fill_value=0) for n in noise] 26 | def g(x_): 27 | return np.hstack([i(x_)[:,None] for i in interp]) 28 | noise = g(x) 29 | return x+g(x)*mag 30 | 31 | class ScanNetDataset(Dataset): 32 | """ ScanNet data loader """ 33 | def __init__(self, options, split, load_confidence=False, random=True): 34 | self.options = options 35 | self.split = split 36 | self.random = random 37 | self.imagePaths = [] 38 | self.dataFolder = options.dataFolder 39 | self.load_confidence = load_confidence 40 | 41 | with open('datasets/split_' + split + '.txt', 'r') as f: 42 | for line in f: 43 | scene_id = line.strip() 44 | if len(scene_id) < 5 or scene_id[:5] != 'scene': 45 | continue 46 | if options.scene_id != '' and options.scene_id not in scene_id: 47 | continue 48 | if load_confidence: 49 | confidence_filename = options.test_dir + '/inference/' + split + '/cache/' + scene_id + '.pth' 50 | if not os.path.exists(confidence_filename): 51 | continue 52 | pass 53 | filename = self.dataFolder + '/' + scene_id + '/' + scene_id + '_vh_clean_2.pth' 54 | if os.path.exists(filename): 55 | info = torch.load(filename) 56 | if len(info) == 5: 57 | self.imagePaths.append(filename) 58 | #np.savetxt('semantic_val/' + scene_id + '.txt', info[2], fmt='%d') 59 | pass 60 | pass 61 | if split != 'train' and len(self.imagePaths) >= options.numTestingImages: 62 | break 63 | continue 64 | pass 65 | 66 | #self.imagePaths = [filename for filename in self.imagePaths if 'scene0217_00' in filename] 67 | print('the number of images', split, len(self.imagePaths)) 68 | 69 | if options.numTrainingImages > 0 and split == 'train': 70 | self.numImages = options.numTrainingImages 71 | else: 72 | self.numImages = len(self.imagePaths) 73 | pass 74 | return 75 | 76 | def __len__(self): 77 | return self.numImages 78 | 79 | def __getitem__(self, index): 80 | if self.random: 81 | t = int(time.time() * 1000000) 82 | np.random.seed(((t & 0xff000000) >> 24) + 83 | ((t & 0x00ff0000) >> 8) + 84 | ((t & 0x0000ff00) << 8) + 85 | ((t & 0x000000ff) << 24)) 86 | index = np.random.randint(len(self.imagePaths)) 87 | else: 88 | index = index % len(self.imagePaths) 89 | pass 90 | 91 | debug = -1 92 | if debug >= 0: 93 | index = debug 94 | print(index, self.imagePaths[index]) 95 | pass 96 | 97 | coords, colors, labels, instances, faces = torch.load(self.imagePaths[index]) 98 | #invalid_instances, = torch.load(self.imagePaths[index].replace('.pth', '_invalid.pth')) 99 | 100 | labels = remapper[labels] 101 | 102 | if self.split == 'train': 103 | m = np.eye(3) + np.random.randn(3,3) * 0.1 104 | m[0][0] *= np.random.randint(2) * 2 - 1 105 | theta = np.random.rand() * 2 * np.pi 106 | else: 107 | m = np.eye(3) 108 | theta = 0 109 | pass 110 | 111 | scale = self.options.scanScale 112 | full_scale = self.options.inputScale 113 | m *= scale 114 | m = np.matmul(m, [[np.cos(theta), np.sin(theta),0], [-np.sin(theta), np.cos(theta),0], [0,0,1]]) 115 | coords = np.matmul(coords, m) 116 | if self.split == 'train': 117 | coords = elastic(coords, 6 * scale // 50,40 * scale / 50) 118 | #coords = elastic(coords, 20 * scale // 50, 160 * scale / 50) 119 | pass 120 | 121 | ## Load normals as input 122 | if 'normal' in self.options.suffix: 123 | points_1 = coords[faces[:, 0]] 124 | points_2 = coords[faces[:, 1]] 125 | points_3 = coords[faces[:, 2]] 126 | face_normals = np.cross(points_2 - points_1, points_3 - points_1) 127 | face_normals /= np.maximum(np.linalg.norm(face_normals, axis=-1, keepdims=True), 1e-4) 128 | normals = np.zeros((len(coords), 3)) 129 | for c in range(3): 130 | np.add.at(normals, faces[:, c], face_normals) 131 | continue 132 | normals /= np.maximum(np.linalg.norm(normals, axis=-1, keepdims=True), 1e-4) 133 | colors = np.concatenate([colors, normals], axis=-1) 134 | pass 135 | 136 | if self.split == 'train': 137 | colors[:, :3] = colors[:, :3] + np.random.randn(3) * 0.1 138 | pass 139 | 140 | ## Load instance segmentation results to train the confidence prediction network 141 | if self.load_confidence: 142 | scene_id = self.imagePaths[index].split('/')[-1].split('_vh_clean_2')[0] 143 | info = torch.load(self.options.test_dir + '/inference/' + self.split + '/cache/' + scene_id + '.pth') 144 | if len(info) == 2: 145 | semantic_pred, instance_pred = info 146 | else: 147 | semantic_pred, instance_pred = info[3], info[6] 148 | semantic_pred = semantic_pred[:len(coords)] 149 | instance_pred = instance_pred[:len(coords)] 150 | pass 151 | instance_pred += 1 152 | unique_instances, indices, counts = np.unique(instances, return_index=True, return_counts=True) 153 | 154 | instance_counts = np.zeros(unique_instances.max() + 1) 155 | instance_counts[unique_instances] = counts 156 | instance_semantics = np.zeros(unique_instances.max() + 1) 157 | instance_semantics[unique_instances] = labels[indices] 158 | confidence_gt = [] 159 | semantic_gt = [] 160 | instance_masks = [] 161 | new_coords = np.zeros(coords.shape, dtype=coords.dtype) 162 | for instance in range(instance_pred.max() + 1): 163 | instance_mask = instance_pred == instance 164 | if instance_mask.sum() == 0: 165 | print('sum = 0', instance, instance_pred.max() + 1, instance_mask.sum()) 166 | exit(1) 167 | info = np.unique(semantic_pred[instance_mask > 0.5], return_counts=True) 168 | label_pred = info[0][info[1].argmax()] 169 | info = np.unique(instances[instance_mask > 0.5], return_counts=True) 170 | instance_gt = info[0][info[1].argmax()] 171 | 172 | instance_coords = coords[instance_mask] 173 | mins = instance_coords.min(0) 174 | maxs = instance_coords.max(0) 175 | max_range = (maxs - mins).max() 176 | if self.split == 'train': 177 | padding = (maxs - mins) * np.random.random(3) * 0.1 178 | else: 179 | padding = max_range * 0.05 180 | pass 181 | max_range += padding * 2 182 | mins = (mins + maxs) / 2 - max_range / 2 183 | instance_coords = np.clip(np.round((instance_coords - mins) / max_range * full_scale), 0, full_scale - 1) 184 | new_coords[instance_mask] = instance_coords 185 | 186 | if instance > 0: 187 | confidence_gt.append(int(label_pred == instance_semantics[instance_gt] and info[1].max() > 0.5 * instance_counts[instance_gt])) 188 | semantic_gt.append(label_pred) 189 | instance_masks.append(instance_mask) 190 | pass 191 | continue 192 | coords = np.concatenate([new_coords, np.expand_dims(instance_pred, -1)], axis=-1) 193 | sample = [coords.astype(np.int64), colors.astype(np.float32), faces.astype(np.int64), np.stack(semantic_gt).astype(np.int64), np.stack(confidence_gt).astype(np.int64), self.imagePaths[index], np.stack(instance_masks).astype(np.int32)] 194 | return sample 195 | 196 | 197 | mins = coords.min(0) 198 | maxs = coords.max(0) 199 | #ranges = maxs - mins 200 | if self.split == 'train': 201 | offset = -mins + np.clip(full_scale - maxs + mins - 0.001, 0, None) * np.random.rand(3) + np.clip(full_scale - maxs + mins + 0.001, None, 0) * np.random.rand(3) 202 | coords += offset 203 | else: 204 | coords -= (mins + maxs) // 2 - full_scale // 2 205 | #coords -= mins 206 | pass 207 | 208 | coords = np.round(coords) 209 | coords = np.clip(coords, 0, full_scale - 1) 210 | 211 | coords = np.concatenate([coords, np.full((coords.shape[0], 1), fill_value=index)], axis=-1) 212 | #coords = np.concatenate([coords, np.expand_dims(instances, -1)], axis=-1) 213 | sample = [coords.astype(np.int64), colors.astype(np.float32), faces.astype(np.int64), labels.astype(np.int64), instances.astype(np.int64), self.imagePaths[index]] 214 | return sample 215 | -------------------------------------------------------------------------------- /datasets/semantic_counts_pixelwise.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/art-programmer/MASC/05204159c48941d21bb4d6b06c282bf4a9312f23/datasets/semantic_counts_pixelwise.npy -------------------------------------------------------------------------------- /datasets/split_test.txt: -------------------------------------------------------------------------------- 1 | scene0707_00 2 | scene0708_00 3 | scene0709_00 4 | scene0710_00 5 | scene0711_00 6 | scene0712_00 7 | scene0713_00 8 | scene0714_00 9 | scene0715_00 10 | scene0716_00 11 | scene0717_00 12 | scene0718_00 13 | scene0719_00 14 | scene0720_00 15 | scene0721_00 16 | scene0722_00 17 | scene0723_00 18 | scene0724_00 19 | scene0725_00 20 | scene0726_00 21 | scene0727_00 22 | scene0728_00 23 | scene0729_00 24 | scene0730_00 25 | scene0731_00 26 | scene0732_00 27 | scene0733_00 28 | scene0734_00 29 | scene0735_00 30 | scene0736_00 31 | scene0737_00 32 | scene0738_00 33 | scene0739_00 34 | scene0740_00 35 | scene0741_00 36 | scene0742_00 37 | scene0743_00 38 | scene0744_00 39 | scene0745_00 40 | scene0746_00 41 | scene0747_00 42 | scene0748_00 43 | scene0749_00 44 | scene0750_00 45 | scene0751_00 46 | scene0752_00 47 | scene0753_00 48 | scene0754_00 49 | scene0755_00 50 | scene0756_00 51 | scene0757_00 52 | scene0758_00 53 | scene0759_00 54 | scene0760_00 55 | scene0761_00 56 | scene0762_00 57 | scene0763_00 58 | scene0764_00 59 | scene0765_00 60 | scene0766_00 61 | scene0767_00 62 | scene0768_00 63 | scene0769_00 64 | scene0770_00 65 | scene0771_00 66 | scene0772_00 67 | scene0773_00 68 | scene0774_00 69 | scene0775_00 70 | scene0776_00 71 | scene0777_00 72 | scene0778_00 73 | scene0779_00 74 | scene0780_00 75 | scene0781_00 76 | scene0782_00 77 | scene0783_00 78 | scene0784_00 79 | scene0785_00 80 | scene0786_00 81 | scene0787_00 82 | scene0788_00 83 | scene0789_00 84 | scene0790_00 85 | scene0791_00 86 | scene0792_00 87 | scene0793_00 88 | scene0794_00 89 | scene0795_00 90 | scene0796_00 91 | scene0797_00 92 | scene0798_00 93 | scene0799_00 94 | scene0800_00 95 | scene0801_00 96 | scene0802_00 97 | scene0803_00 98 | scene0804_00 99 | scene0805_00 100 | scene0806_00 101 | 102 | -------------------------------------------------------------------------------- /datasets/split_train.txt: -------------------------------------------------------------------------------- 1 | scene0191_00 2 | scene0191_01 3 | scene0191_02 4 | scene0119_00 5 | scene0230_00 6 | scene0528_00 7 | scene0528_01 8 | scene0705_00 9 | scene0705_01 10 | scene0705_02 11 | scene0415_00 12 | scene0415_01 13 | scene0415_02 14 | scene0007_00 15 | scene0141_00 16 | scene0141_01 17 | scene0141_02 18 | scene0515_00 19 | scene0515_01 20 | scene0515_02 21 | scene0447_00 22 | scene0447_01 23 | scene0447_02 24 | scene0531_00 25 | scene0503_00 26 | scene0285_00 27 | scene0069_00 28 | scene0584_00 29 | scene0584_01 30 | scene0584_02 31 | scene0581_00 32 | scene0581_01 33 | scene0581_02 34 | scene0620_00 35 | scene0620_01 36 | scene0263_00 37 | scene0263_01 38 | scene0481_00 39 | scene0481_01 40 | scene0020_00 41 | scene0020_01 42 | scene0291_00 43 | scene0291_01 44 | scene0291_02 45 | scene0469_00 46 | scene0469_01 47 | scene0469_02 48 | scene0659_00 49 | scene0659_01 50 | scene0024_00 51 | scene0024_01 52 | scene0024_02 53 | scene0564_00 54 | scene0117_00 55 | scene0027_00 56 | scene0027_01 57 | scene0027_02 58 | scene0028_00 59 | scene0330_00 60 | scene0418_00 61 | scene0418_01 62 | scene0418_02 63 | scene0233_00 64 | scene0233_01 65 | scene0673_00 66 | scene0673_01 67 | scene0673_02 68 | scene0673_03 69 | scene0673_04 70 | scene0673_05 71 | scene0585_00 72 | scene0585_01 73 | scene0362_00 74 | scene0362_01 75 | scene0362_02 76 | scene0362_03 77 | scene0035_00 78 | scene0035_01 79 | scene0358_00 80 | scene0358_01 81 | scene0358_02 82 | scene0037_00 83 | scene0194_00 84 | scene0321_00 85 | scene0293_00 86 | scene0293_01 87 | scene0623_00 88 | scene0623_01 89 | scene0592_00 90 | scene0592_01 91 | scene0569_00 92 | scene0569_01 93 | scene0413_00 94 | scene0313_00 95 | scene0313_01 96 | scene0313_02 97 | scene0480_00 98 | scene0480_01 99 | scene0401_00 100 | scene0517_00 101 | scene0517_01 102 | scene0517_02 103 | scene0032_00 104 | scene0032_01 105 | scene0613_00 106 | scene0613_01 107 | scene0613_02 108 | scene0306_00 109 | scene0306_01 110 | scene0052_00 111 | scene0052_01 112 | scene0052_02 113 | scene0053_00 114 | scene0444_00 115 | scene0444_01 116 | scene0055_00 117 | scene0055_01 118 | scene0055_02 119 | scene0560_00 120 | scene0589_00 121 | scene0589_01 122 | scene0589_02 123 | scene0610_00 124 | scene0610_01 125 | scene0610_02 126 | scene0364_00 127 | scene0364_01 128 | scene0383_00 129 | scene0383_01 130 | scene0383_02 131 | scene0006_00 132 | scene0006_01 133 | scene0006_02 134 | scene0275_00 135 | scene0451_00 136 | scene0451_01 137 | scene0451_02 138 | scene0451_03 139 | scene0451_04 140 | scene0451_05 141 | scene0135_00 142 | scene0065_00 143 | scene0065_01 144 | scene0065_02 145 | scene0104_00 146 | scene0674_00 147 | scene0674_01 148 | scene0448_00 149 | scene0448_01 150 | scene0448_02 151 | scene0502_00 152 | scene0502_01 153 | scene0502_02 154 | scene0440_00 155 | scene0440_01 156 | scene0440_02 157 | scene0071_00 158 | scene0072_00 159 | scene0072_01 160 | scene0072_02 161 | scene0509_00 162 | scene0509_01 163 | scene0509_02 164 | scene0649_00 165 | scene0649_01 166 | scene0602_00 167 | scene0694_00 168 | scene0694_01 169 | scene0101_00 170 | scene0101_01 171 | scene0101_02 172 | scene0101_03 173 | scene0101_04 174 | scene0101_05 175 | scene0218_00 176 | scene0218_01 177 | scene0579_00 178 | scene0579_01 179 | scene0579_02 180 | scene0039_00 181 | scene0039_01 182 | scene0493_00 183 | scene0493_01 184 | scene0242_00 185 | scene0242_01 186 | scene0242_02 187 | scene0083_00 188 | scene0083_01 189 | scene0127_00 190 | scene0127_01 191 | scene0662_00 192 | scene0662_01 193 | scene0662_02 194 | scene0018_00 195 | scene0087_00 196 | scene0087_01 197 | scene0087_02 198 | scene0332_00 199 | scene0332_01 200 | scene0332_02 201 | scene0628_00 202 | scene0628_01 203 | scene0628_02 204 | scene0134_00 205 | scene0134_01 206 | scene0134_02 207 | scene0238_00 208 | scene0238_01 209 | scene0092_00 210 | scene0092_01 211 | scene0092_02 212 | scene0092_03 213 | scene0092_04 214 | scene0022_00 215 | scene0022_01 216 | scene0467_00 217 | scene0392_00 218 | scene0392_01 219 | scene0392_02 220 | scene0424_00 221 | scene0424_01 222 | scene0424_02 223 | scene0646_00 224 | scene0646_01 225 | scene0646_02 226 | scene0098_00 227 | scene0098_01 228 | scene0044_00 229 | scene0044_01 230 | scene0044_02 231 | scene0510_00 232 | scene0510_01 233 | scene0510_02 234 | scene0571_00 235 | scene0571_01 236 | scene0166_00 237 | scene0166_01 238 | scene0166_02 239 | scene0563_00 240 | scene0172_00 241 | scene0172_01 242 | scene0388_00 243 | scene0388_01 244 | scene0215_00 245 | scene0215_01 246 | scene0252_00 247 | scene0287_00 248 | scene0668_00 249 | scene0572_00 250 | scene0572_01 251 | scene0572_02 252 | scene0026_00 253 | scene0224_00 254 | scene0113_00 255 | scene0113_01 256 | scene0551_00 257 | scene0381_00 258 | scene0381_01 259 | scene0381_02 260 | scene0371_00 261 | scene0371_01 262 | scene0460_00 263 | scene0118_00 264 | scene0118_01 265 | scene0118_02 266 | scene0417_00 267 | scene0008_00 268 | scene0634_00 269 | scene0521_00 270 | scene0123_00 271 | scene0123_01 272 | scene0123_02 273 | scene0045_00 274 | scene0045_01 275 | scene0511_00 276 | scene0511_01 277 | scene0114_00 278 | scene0114_01 279 | scene0114_02 280 | scene0070_00 281 | scene0029_00 282 | scene0029_01 283 | scene0029_02 284 | scene0129_00 285 | scene0103_00 286 | scene0103_01 287 | scene0002_00 288 | scene0002_01 289 | scene0132_00 290 | scene0132_01 291 | scene0132_02 292 | scene0124_00 293 | scene0124_01 294 | scene0143_00 295 | scene0143_01 296 | scene0143_02 297 | scene0604_00 298 | scene0604_01 299 | scene0604_02 300 | scene0507_00 301 | scene0105_00 302 | scene0105_01 303 | scene0105_02 304 | scene0428_00 305 | scene0428_01 306 | scene0311_00 307 | scene0140_00 308 | scene0140_01 309 | scene0182_00 310 | scene0182_01 311 | scene0182_02 312 | scene0142_00 313 | scene0142_01 314 | scene0399_00 315 | scene0399_01 316 | scene0012_00 317 | scene0012_01 318 | scene0012_02 319 | scene0060_00 320 | scene0060_01 321 | scene0370_00 322 | scene0370_01 323 | scene0370_02 324 | scene0310_00 325 | scene0310_01 326 | scene0310_02 327 | scene0661_00 328 | scene0650_00 329 | scene0152_00 330 | scene0152_01 331 | scene0152_02 332 | scene0158_00 333 | scene0158_01 334 | scene0158_02 335 | scene0482_00 336 | scene0482_01 337 | scene0600_00 338 | scene0600_01 339 | scene0600_02 340 | scene0393_00 341 | scene0393_01 342 | scene0393_02 343 | scene0562_00 344 | scene0174_00 345 | scene0174_01 346 | scene0157_00 347 | scene0157_01 348 | scene0161_00 349 | scene0161_01 350 | scene0161_02 351 | scene0159_00 352 | scene0254_00 353 | scene0254_01 354 | scene0115_00 355 | scene0115_01 356 | scene0115_02 357 | scene0162_00 358 | scene0163_00 359 | scene0163_01 360 | scene0523_00 361 | scene0523_01 362 | scene0523_02 363 | scene0459_00 364 | scene0459_01 365 | scene0175_00 366 | scene0085_00 367 | scene0085_01 368 | scene0279_00 369 | scene0279_01 370 | scene0279_02 371 | scene0201_00 372 | scene0201_01 373 | scene0201_02 374 | scene0283_00 375 | scene0456_00 376 | scene0456_01 377 | scene0429_00 378 | scene0043_00 379 | scene0043_01 380 | scene0419_00 381 | scene0419_01 382 | scene0419_02 383 | scene0368_00 384 | scene0368_01 385 | scene0348_00 386 | scene0348_01 387 | scene0348_02 388 | scene0442_00 389 | scene0178_00 390 | scene0380_00 391 | scene0380_01 392 | scene0380_02 393 | scene0165_00 394 | scene0165_01 395 | scene0165_02 396 | scene0181_00 397 | scene0181_01 398 | scene0181_02 399 | scene0181_03 400 | scene0333_00 401 | scene0614_00 402 | scene0614_01 403 | scene0614_02 404 | scene0404_00 405 | scene0404_01 406 | scene0404_02 407 | scene0185_00 408 | scene0126_00 409 | scene0126_01 410 | scene0126_02 411 | scene0519_00 412 | scene0236_00 413 | scene0236_01 414 | scene0189_00 415 | scene0075_00 416 | scene0267_00 417 | scene0192_00 418 | scene0192_01 419 | scene0192_02 420 | scene0281_00 421 | scene0420_00 422 | scene0420_01 423 | scene0420_02 424 | scene0195_00 425 | scene0195_01 426 | scene0195_02 427 | scene0597_00 428 | scene0597_01 429 | scene0597_02 430 | scene0041_00 431 | scene0041_01 432 | scene0111_00 433 | scene0111_01 434 | scene0111_02 435 | scene0666_00 436 | scene0666_01 437 | scene0666_02 438 | scene0200_00 439 | scene0200_01 440 | scene0200_02 441 | scene0536_00 442 | scene0536_01 443 | scene0536_02 444 | scene0390_00 445 | scene0280_00 446 | scene0280_01 447 | scene0280_02 448 | scene0344_00 449 | scene0344_01 450 | scene0205_00 451 | scene0205_01 452 | scene0205_02 453 | scene0484_00 454 | scene0484_01 455 | scene0009_00 456 | scene0009_01 457 | scene0009_02 458 | scene0302_00 459 | scene0302_01 460 | scene0209_00 461 | scene0209_01 462 | scene0209_02 463 | scene0210_00 464 | scene0210_01 465 | scene0395_00 466 | scene0395_01 467 | scene0395_02 468 | scene0683_00 469 | scene0601_00 470 | scene0601_01 471 | scene0214_00 472 | scene0214_01 473 | scene0214_02 474 | scene0477_00 475 | scene0477_01 476 | scene0439_00 477 | scene0439_01 478 | scene0468_00 479 | scene0468_01 480 | scene0468_02 481 | scene0546_00 482 | scene0466_00 483 | scene0466_01 484 | scene0220_00 485 | scene0220_01 486 | scene0220_02 487 | scene0122_00 488 | scene0122_01 489 | scene0130_00 490 | scene0110_00 491 | scene0110_01 492 | scene0110_02 493 | scene0327_00 494 | scene0156_00 495 | scene0266_00 496 | scene0266_01 497 | scene0001_00 498 | scene0001_01 499 | scene0228_00 500 | scene0199_00 501 | scene0219_00 502 | scene0464_00 503 | scene0232_00 504 | scene0232_01 505 | scene0232_02 506 | scene0299_00 507 | scene0299_01 508 | scene0530_00 509 | scene0363_00 510 | scene0453_00 511 | scene0453_01 512 | scene0570_00 513 | scene0570_01 514 | scene0570_02 515 | scene0183_00 516 | scene0239_00 517 | scene0239_01 518 | scene0239_02 519 | scene0373_00 520 | scene0373_01 521 | scene0241_00 522 | scene0241_01 523 | scene0241_02 524 | scene0188_00 525 | scene0622_00 526 | scene0622_01 527 | scene0244_00 528 | scene0244_01 529 | scene0691_00 530 | scene0691_01 531 | scene0206_00 532 | scene0206_01 533 | scene0206_02 534 | scene0247_00 535 | scene0247_01 536 | scene0061_00 537 | scene0061_01 538 | scene0082_00 539 | scene0250_00 540 | scene0250_01 541 | scene0250_02 542 | scene0501_00 543 | scene0501_01 544 | scene0501_02 545 | scene0320_00 546 | scene0320_01 547 | scene0320_02 548 | scene0320_03 549 | scene0631_00 550 | scene0631_01 551 | scene0631_02 552 | scene0255_00 553 | scene0255_01 554 | scene0255_02 555 | scene0047_00 556 | scene0265_00 557 | scene0265_01 558 | scene0265_02 559 | scene0004_00 560 | scene0336_00 561 | scene0336_01 562 | scene0058_00 563 | scene0058_01 564 | scene0260_00 565 | scene0260_01 566 | scene0260_02 567 | scene0243_00 568 | scene0603_00 569 | scene0603_01 570 | scene0093_00 571 | scene0093_01 572 | scene0093_02 573 | scene0109_00 574 | scene0109_01 575 | scene0434_00 576 | scene0434_01 577 | scene0434_02 578 | scene0290_00 579 | scene0627_00 580 | scene0627_01 581 | scene0470_00 582 | scene0470_01 583 | scene0137_00 584 | scene0137_01 585 | scene0137_02 586 | scene0270_00 587 | scene0270_01 588 | scene0270_02 589 | scene0271_00 590 | scene0271_01 591 | scene0504_00 592 | scene0274_00 593 | scene0274_01 594 | scene0274_02 595 | scene0036_00 596 | scene0036_01 597 | scene0276_00 598 | scene0276_01 599 | scene0272_00 600 | scene0272_01 601 | scene0499_00 602 | scene0698_00 603 | scene0698_01 604 | scene0051_00 605 | scene0051_01 606 | scene0051_02 607 | scene0051_03 608 | scene0108_00 609 | scene0245_00 610 | scene0369_00 611 | scene0369_01 612 | scene0369_02 613 | scene0284_00 614 | scene0289_00 615 | scene0289_01 616 | scene0286_00 617 | scene0286_01 618 | scene0286_02 619 | scene0286_03 620 | scene0031_00 621 | scene0031_01 622 | scene0031_02 623 | scene0545_00 624 | scene0545_01 625 | scene0545_02 626 | scene0557_00 627 | scene0557_01 628 | scene0557_02 629 | scene0533_00 630 | scene0533_01 631 | scene0116_00 632 | scene0116_01 633 | scene0116_02 634 | scene0611_00 635 | scene0611_01 636 | scene0688_00 637 | scene0294_00 638 | scene0294_01 639 | scene0294_02 640 | scene0295_00 641 | scene0295_01 642 | scene0296_00 643 | scene0296_01 644 | scene0596_00 645 | scene0596_01 646 | scene0596_02 647 | scene0532_00 648 | scene0532_01 649 | scene0637_00 650 | scene0638_00 651 | scene0121_00 652 | scene0121_01 653 | scene0121_02 654 | scene0040_00 655 | scene0040_01 656 | scene0197_00 657 | scene0197_01 658 | scene0197_02 659 | scene0410_00 660 | scene0410_01 661 | scene0305_00 662 | scene0305_01 663 | scene0615_00 664 | scene0615_01 665 | scene0703_00 666 | scene0703_01 667 | scene0555_00 668 | scene0297_00 669 | scene0297_01 670 | scene0297_02 671 | scene0582_00 672 | scene0582_01 673 | scene0582_02 674 | scene0023_00 675 | scene0094_00 676 | scene0013_00 677 | scene0013_01 678 | scene0013_02 679 | scene0136_00 680 | scene0136_01 681 | scene0136_02 682 | scene0407_00 683 | scene0407_01 684 | scene0062_00 685 | scene0062_01 686 | scene0062_02 687 | scene0386_00 688 | scene0318_00 689 | scene0554_00 690 | scene0554_01 691 | scene0497_00 692 | scene0213_00 693 | scene0258_00 694 | scene0323_00 695 | scene0323_01 696 | scene0324_00 697 | scene0324_01 698 | scene0016_00 699 | scene0016_01 700 | scene0016_02 701 | scene0681_00 702 | scene0398_00 703 | scene0398_01 704 | scene0227_00 705 | scene0090_00 706 | scene0066_00 707 | scene0262_00 708 | scene0262_01 709 | scene0155_00 710 | scene0155_01 711 | scene0155_02 712 | scene0352_00 713 | scene0352_01 714 | scene0352_02 715 | scene0038_00 716 | scene0038_01 717 | scene0038_02 718 | scene0335_00 719 | scene0335_01 720 | scene0335_02 721 | scene0261_00 722 | scene0261_01 723 | scene0261_02 724 | scene0261_03 725 | scene0640_00 726 | scene0640_01 727 | scene0640_02 728 | scene0080_00 729 | scene0080_01 730 | scene0080_02 731 | scene0403_00 732 | scene0403_01 733 | scene0282_00 734 | scene0282_01 735 | scene0282_02 736 | scene0682_00 737 | scene0173_00 738 | scene0173_01 739 | scene0173_02 740 | scene0522_00 741 | scene0687_00 742 | scene0345_00 743 | scene0345_01 744 | scene0612_00 745 | scene0612_01 746 | scene0411_00 747 | scene0411_01 748 | scene0411_02 749 | scene0625_00 750 | scene0625_01 751 | scene0211_00 752 | scene0211_01 753 | scene0211_02 754 | scene0211_03 755 | scene0676_00 756 | scene0676_01 757 | scene0179_00 758 | scene0498_00 759 | scene0498_01 760 | scene0498_02 761 | scene0547_00 762 | scene0547_01 763 | scene0547_02 764 | scene0269_00 765 | scene0269_01 766 | scene0269_02 767 | scene0366_00 768 | scene0680_00 769 | scene0680_01 770 | scene0588_00 771 | scene0588_01 772 | scene0588_02 773 | scene0588_03 774 | scene0346_00 775 | scene0346_01 776 | scene0359_00 777 | scene0359_01 778 | scene0014_00 779 | scene0120_00 780 | scene0120_01 781 | scene0212_00 782 | scene0212_01 783 | scene0212_02 784 | scene0176_00 785 | scene0049_00 786 | scene0259_00 787 | scene0259_01 788 | scene0586_00 789 | scene0586_01 790 | scene0586_02 791 | scene0309_00 792 | scene0309_01 793 | scene0125_00 794 | scene0455_00 795 | scene0177_00 796 | scene0177_01 797 | scene0177_02 798 | scene0326_00 799 | scene0372_00 800 | scene0171_00 801 | scene0171_01 802 | scene0374_00 803 | scene0654_00 804 | scene0654_01 805 | scene0445_00 806 | scene0445_01 807 | scene0475_00 808 | scene0475_01 809 | scene0475_02 810 | scene0349_00 811 | scene0349_01 812 | scene0234_00 813 | scene0669_00 814 | scene0669_01 815 | scene0375_00 816 | scene0375_01 817 | scene0375_02 818 | scene0387_00 819 | scene0387_01 820 | scene0387_02 821 | scene0312_00 822 | scene0312_01 823 | scene0312_02 824 | scene0384_00 825 | scene0385_00 826 | scene0385_01 827 | scene0385_02 828 | scene0000_00 829 | scene0000_01 830 | scene0000_02 831 | scene0376_00 832 | scene0376_01 833 | scene0376_02 834 | scene0301_00 835 | scene0301_01 836 | scene0301_02 837 | scene0322_00 838 | scene0542_00 839 | scene0079_00 840 | scene0079_01 841 | scene0099_00 842 | scene0099_01 843 | scene0476_00 844 | scene0476_01 845 | scene0476_02 846 | scene0394_00 847 | scene0394_01 848 | scene0147_00 849 | scene0147_01 850 | scene0067_00 851 | scene0067_01 852 | scene0067_02 853 | scene0397_00 854 | scene0397_01 855 | scene0337_00 856 | scene0337_01 857 | scene0337_02 858 | scene0431_00 859 | scene0223_00 860 | scene0223_01 861 | scene0223_02 862 | scene0010_00 863 | scene0010_01 864 | scene0402_00 865 | scene0268_00 866 | scene0268_01 867 | scene0268_02 868 | scene0679_00 869 | scene0679_01 870 | scene0405_00 871 | scene0128_00 872 | scene0408_00 873 | scene0408_01 874 | scene0190_00 875 | scene0107_00 876 | scene0076_00 877 | scene0167_00 878 | scene0361_00 879 | scene0361_01 880 | scene0361_02 881 | scene0216_00 882 | scene0202_00 883 | scene0303_00 884 | scene0303_01 885 | scene0303_02 886 | scene0446_00 887 | scene0446_01 888 | scene0089_00 889 | scene0089_01 890 | scene0089_02 891 | scene0360_00 892 | scene0150_00 893 | scene0150_01 894 | scene0150_02 895 | scene0421_00 896 | scene0421_01 897 | scene0421_02 898 | scene0454_00 899 | scene0626_00 900 | scene0626_01 901 | scene0626_02 902 | scene0186_00 903 | scene0186_01 904 | scene0538_00 905 | scene0479_00 906 | scene0479_01 907 | scene0479_02 908 | scene0656_00 909 | scene0656_01 910 | scene0656_02 911 | scene0656_03 912 | scene0525_00 913 | scene0525_01 914 | scene0525_02 915 | scene0308_00 916 | scene0396_00 917 | scene0396_01 918 | scene0396_02 919 | scene0624_00 920 | scene0292_00 921 | scene0292_01 922 | scene0632_00 923 | scene0253_00 924 | scene0021_00 925 | scene0325_00 926 | scene0325_01 927 | scene0437_00 928 | scene0437_01 929 | scene0438_00 930 | scene0590_00 931 | scene0590_01 932 | scene0400_00 933 | scene0400_01 934 | scene0541_00 935 | scene0541_01 936 | scene0541_02 937 | scene0677_00 938 | scene0677_01 939 | scene0677_02 940 | scene0443_00 941 | scene0315_00 942 | scene0288_00 943 | scene0288_01 944 | scene0288_02 945 | scene0422_00 946 | scene0672_00 947 | scene0672_01 948 | scene0184_00 949 | scene0449_00 950 | scene0449_01 951 | scene0449_02 952 | scene0048_00 953 | scene0048_01 954 | scene0138_00 955 | scene0452_00 956 | scene0452_01 957 | scene0452_02 958 | scene0667_00 959 | scene0667_01 960 | scene0667_02 961 | scene0463_00 962 | scene0463_01 963 | scene0078_00 964 | scene0078_01 965 | scene0078_02 966 | scene0636_00 967 | scene0457_00 968 | scene0457_01 969 | scene0457_02 970 | scene0465_00 971 | scene0465_01 972 | scene0577_00 973 | scene0151_00 974 | scene0151_01 975 | scene0339_00 976 | scene0573_00 977 | scene0573_01 978 | scene0154_00 979 | scene0096_00 980 | scene0096_01 981 | scene0096_02 982 | scene0235_00 983 | scene0168_00 984 | scene0168_01 985 | scene0168_02 986 | scene0594_00 987 | scene0587_00 988 | scene0587_01 989 | scene0587_02 990 | scene0587_03 991 | scene0229_00 992 | scene0229_01 993 | scene0229_02 994 | scene0512_00 995 | scene0106_00 996 | scene0106_01 997 | scene0106_02 998 | scene0472_00 999 | scene0472_01 1000 | scene0472_02 1001 | scene0489_00 1002 | scene0489_01 1003 | scene0489_02 1004 | scene0425_00 1005 | scene0425_01 1006 | scene0641_00 1007 | scene0526_00 1008 | scene0526_01 1009 | scene0317_00 1010 | scene0317_01 1011 | scene0544_00 1012 | scene0017_00 1013 | scene0017_01 1014 | scene0017_02 1015 | scene0042_00 1016 | scene0042_01 1017 | scene0042_02 1018 | scene0576_00 1019 | scene0576_01 1020 | scene0576_02 1021 | scene0347_00 1022 | scene0347_01 1023 | scene0347_02 1024 | scene0436_00 1025 | scene0226_00 1026 | scene0226_01 1027 | scene0485_00 1028 | scene0486_00 1029 | scene0487_00 1030 | scene0487_01 1031 | scene0619_00 1032 | scene0097_00 1033 | scene0367_00 1034 | scene0367_01 1035 | scene0491_00 1036 | scene0492_00 1037 | scene0492_01 1038 | scene0005_00 1039 | scene0005_01 1040 | scene0543_00 1041 | scene0543_01 1042 | scene0543_02 1043 | scene0657_00 1044 | scene0341_00 1045 | scene0341_01 1046 | scene0534_00 1047 | scene0534_01 1048 | scene0319_00 1049 | scene0273_00 1050 | scene0273_01 1051 | scene0225_00 1052 | scene0198_00 1053 | scene0003_00 1054 | scene0003_01 1055 | scene0003_02 1056 | scene0409_00 1057 | scene0409_01 1058 | scene0331_00 1059 | scene0331_01 1060 | scene0505_00 1061 | scene0505_01 1062 | scene0505_02 1063 | scene0505_03 1064 | scene0505_04 1065 | scene0506_00 1066 | scene0057_00 1067 | scene0057_01 1068 | scene0074_00 1069 | scene0074_01 1070 | scene0074_02 1071 | scene0091_00 1072 | scene0112_00 1073 | scene0112_01 1074 | scene0112_02 1075 | scene0240_00 1076 | scene0102_00 1077 | scene0102_01 1078 | scene0513_00 1079 | scene0514_00 1080 | scene0514_01 1081 | scene0537_00 1082 | scene0516_00 1083 | scene0516_01 1084 | scene0495_00 1085 | scene0617_00 1086 | scene0133_00 1087 | scene0520_00 1088 | scene0520_01 1089 | scene0635_00 1090 | scene0635_01 1091 | scene0054_00 1092 | scene0473_00 1093 | scene0473_01 1094 | scene0524_00 1095 | scene0524_01 1096 | scene0379_00 1097 | scene0471_00 1098 | scene0471_01 1099 | scene0471_02 1100 | scene0566_00 1101 | scene0248_00 1102 | scene0248_01 1103 | scene0248_02 1104 | scene0529_00 1105 | scene0529_01 1106 | scene0529_02 1107 | scene0391_00 1108 | scene0264_00 1109 | scene0264_01 1110 | scene0264_02 1111 | scene0675_00 1112 | scene0675_01 1113 | scene0350_00 1114 | scene0350_01 1115 | scene0350_02 1116 | scene0450_00 1117 | scene0068_00 1118 | scene0068_01 1119 | scene0237_00 1120 | scene0237_01 1121 | scene0365_00 1122 | scene0365_01 1123 | scene0365_02 1124 | scene0605_00 1125 | scene0605_01 1126 | scene0539_00 1127 | scene0539_01 1128 | scene0539_02 1129 | scene0540_00 1130 | scene0540_01 1131 | scene0540_02 1132 | scene0170_00 1133 | scene0170_01 1134 | scene0170_02 1135 | scene0433_00 1136 | scene0340_00 1137 | scene0340_01 1138 | scene0340_02 1139 | scene0160_00 1140 | scene0160_01 1141 | scene0160_02 1142 | scene0160_03 1143 | scene0160_04 1144 | scene0059_00 1145 | scene0059_01 1146 | scene0059_02 1147 | scene0056_00 1148 | scene0056_01 1149 | scene0478_00 1150 | scene0478_01 1151 | scene0548_00 1152 | scene0548_01 1153 | scene0548_02 1154 | scene0204_00 1155 | scene0204_01 1156 | scene0204_02 1157 | scene0033_00 1158 | scene0145_00 1159 | scene0483_00 1160 | scene0508_00 1161 | scene0508_01 1162 | scene0508_02 1163 | scene0180_00 1164 | scene0148_00 1165 | scene0556_00 1166 | scene0556_01 1167 | scene0416_00 1168 | scene0416_01 1169 | scene0416_02 1170 | scene0416_03 1171 | scene0416_04 1172 | scene0073_00 1173 | scene0073_01 1174 | scene0073_02 1175 | scene0073_03 1176 | scene0034_00 1177 | scene0034_01 1178 | scene0034_02 1179 | scene0639_00 1180 | scene0561_00 1181 | scene0561_01 1182 | scene0298_00 1183 | scene0692_00 1184 | scene0692_01 1185 | scene0692_02 1186 | scene0692_03 1187 | scene0692_04 1188 | scene0642_00 1189 | scene0642_01 1190 | scene0642_02 1191 | scene0642_03 1192 | scene0630_00 1193 | scene0630_01 1194 | scene0630_02 1195 | scene0630_03 1196 | scene0630_04 1197 | scene0630_05 1198 | scene0630_06 1199 | scene0706_00 1200 | scene0567_00 1201 | scene0567_01 1202 | 1203 | -------------------------------------------------------------------------------- /datasets/split_val.txt: -------------------------------------------------------------------------------- 1 | scene0568_00 2 | scene0568_01 3 | scene0568_02 4 | scene0304_00 5 | scene0488_00 6 | scene0488_01 7 | scene0412_00 8 | scene0412_01 9 | scene0217_00 10 | scene0019_00 11 | scene0019_01 12 | scene0414_00 13 | scene0575_00 14 | scene0575_01 15 | scene0575_02 16 | scene0426_00 17 | scene0426_01 18 | scene0426_02 19 | scene0426_03 20 | scene0549_00 21 | scene0549_01 22 | scene0578_00 23 | scene0578_01 24 | scene0578_02 25 | scene0665_00 26 | scene0665_01 27 | scene0050_00 28 | scene0050_01 29 | scene0050_02 30 | scene0257_00 31 | scene0025_00 32 | scene0025_01 33 | scene0025_02 34 | scene0583_00 35 | scene0583_01 36 | scene0583_02 37 | scene0701_00 38 | scene0701_01 39 | scene0701_02 40 | scene0580_00 41 | scene0580_01 42 | scene0565_00 43 | scene0169_00 44 | scene0169_01 45 | scene0655_00 46 | scene0655_01 47 | scene0655_02 48 | scene0063_00 49 | scene0221_00 50 | scene0221_01 51 | scene0591_00 52 | scene0591_01 53 | scene0591_02 54 | scene0678_00 55 | scene0678_01 56 | scene0678_02 57 | scene0462_00 58 | scene0427_00 59 | scene0595_00 60 | scene0193_00 61 | scene0193_01 62 | scene0164_00 63 | scene0164_01 64 | scene0164_02 65 | scene0164_03 66 | scene0598_00 67 | scene0598_01 68 | scene0598_02 69 | scene0599_00 70 | scene0599_01 71 | scene0599_02 72 | scene0328_00 73 | scene0300_00 74 | scene0300_01 75 | scene0354_00 76 | scene0458_00 77 | scene0458_01 78 | scene0423_00 79 | scene0423_01 80 | scene0423_02 81 | scene0307_00 82 | scene0307_01 83 | scene0307_02 84 | scene0606_00 85 | scene0606_01 86 | scene0606_02 87 | scene0432_00 88 | scene0432_01 89 | scene0608_00 90 | scene0608_01 91 | scene0608_02 92 | scene0651_00 93 | scene0651_01 94 | scene0651_02 95 | scene0430_00 96 | scene0430_01 97 | scene0689_00 98 | scene0357_00 99 | scene0357_01 100 | scene0574_00 101 | scene0574_01 102 | scene0574_02 103 | scene0329_00 104 | scene0329_01 105 | scene0329_02 106 | scene0153_00 107 | scene0153_01 108 | scene0616_00 109 | scene0616_01 110 | scene0671_00 111 | scene0671_01 112 | scene0618_00 113 | scene0382_00 114 | scene0382_01 115 | scene0490_00 116 | scene0621_00 117 | scene0607_00 118 | scene0607_01 119 | scene0149_00 120 | scene0695_00 121 | scene0695_01 122 | scene0695_02 123 | scene0695_03 124 | scene0389_00 125 | scene0377_00 126 | scene0377_01 127 | scene0377_02 128 | scene0342_00 129 | scene0139_00 130 | scene0629_00 131 | scene0629_01 132 | scene0629_02 133 | scene0496_00 134 | scene0633_00 135 | scene0633_01 136 | scene0518_00 137 | scene0652_00 138 | scene0406_00 139 | scene0406_01 140 | scene0406_02 141 | scene0144_00 142 | scene0144_01 143 | scene0494_00 144 | scene0278_00 145 | scene0278_01 146 | scene0316_00 147 | scene0609_00 148 | scene0609_01 149 | scene0609_02 150 | scene0609_03 151 | scene0084_00 152 | scene0084_01 153 | scene0084_02 154 | scene0696_00 155 | scene0696_01 156 | scene0696_02 157 | scene0351_00 158 | scene0351_01 159 | scene0643_00 160 | scene0644_00 161 | scene0645_00 162 | scene0645_01 163 | scene0645_02 164 | scene0081_00 165 | scene0081_01 166 | scene0081_02 167 | scene0647_00 168 | scene0647_01 169 | scene0535_00 170 | scene0353_00 171 | scene0353_01 172 | scene0353_02 173 | scene0559_00 174 | scene0559_01 175 | scene0559_02 176 | scene0593_00 177 | scene0593_01 178 | scene0246_00 179 | scene0653_00 180 | scene0653_01 181 | scene0064_00 182 | scene0064_01 183 | scene0356_00 184 | scene0356_01 185 | scene0356_02 186 | scene0030_00 187 | scene0030_01 188 | scene0030_02 189 | scene0222_00 190 | scene0222_01 191 | scene0338_00 192 | scene0338_01 193 | scene0338_02 194 | scene0378_00 195 | scene0378_01 196 | scene0378_02 197 | scene0660_00 198 | scene0553_00 199 | scene0553_01 200 | scene0553_02 201 | scene0527_00 202 | scene0663_00 203 | scene0663_01 204 | scene0663_02 205 | scene0664_00 206 | scene0664_01 207 | scene0664_02 208 | scene0334_00 209 | scene0334_01 210 | scene0334_02 211 | scene0046_00 212 | scene0046_01 213 | scene0046_02 214 | scene0203_00 215 | scene0203_01 216 | scene0203_02 217 | scene0088_00 218 | scene0088_01 219 | scene0088_02 220 | scene0088_03 221 | scene0086_00 222 | scene0086_01 223 | scene0086_02 224 | scene0670_00 225 | scene0670_01 226 | scene0256_00 227 | scene0256_01 228 | scene0256_02 229 | scene0249_00 230 | scene0441_00 231 | scene0658_00 232 | scene0704_00 233 | scene0704_01 234 | scene0187_00 235 | scene0187_01 236 | scene0131_00 237 | scene0131_01 238 | scene0131_02 239 | scene0207_00 240 | scene0207_01 241 | scene0207_02 242 | scene0461_00 243 | scene0011_00 244 | scene0011_01 245 | scene0343_00 246 | scene0251_00 247 | scene0077_00 248 | scene0077_01 249 | scene0684_00 250 | scene0684_01 251 | scene0550_00 252 | scene0686_00 253 | scene0686_01 254 | scene0686_02 255 | scene0208_00 256 | scene0500_00 257 | scene0500_01 258 | scene0552_00 259 | scene0552_01 260 | scene0648_00 261 | scene0648_01 262 | scene0435_00 263 | scene0435_01 264 | scene0435_02 265 | scene0435_03 266 | scene0690_00 267 | scene0690_01 268 | scene0693_00 269 | scene0693_01 270 | scene0693_02 271 | scene0700_00 272 | scene0700_01 273 | scene0700_02 274 | scene0699_00 275 | scene0231_00 276 | scene0231_01 277 | scene0231_02 278 | scene0697_00 279 | scene0697_01 280 | scene0697_02 281 | scene0697_03 282 | scene0474_00 283 | scene0474_01 284 | scene0474_02 285 | scene0474_03 286 | scene0474_04 287 | scene0474_05 288 | scene0355_00 289 | scene0355_01 290 | scene0146_00 291 | scene0146_01 292 | scene0146_02 293 | scene0196_00 294 | scene0702_00 295 | scene0702_01 296 | scene0702_02 297 | scene0314_00 298 | scene0277_00 299 | scene0277_01 300 | scene0277_02 301 | scene0095_00 302 | scene0095_01 303 | scene0015_00 304 | scene0100_00 305 | scene0100_01 306 | scene0100_02 307 | scene0558_00 308 | scene0558_01 309 | scene0558_02 310 | scene0685_00 311 | scene0685_01 312 | scene0685_02 313 | 314 | -------------------------------------------------------------------------------- /inference.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | import multiprocessing as mp 3 | import sys 4 | import torch 5 | import numpy as np 6 | import glob 7 | 8 | from options import parse_args 9 | from utils import * 10 | from models.instance import Model, CoordAugmentation, NeighborGT 11 | from train_confidence import InstanceValidator 12 | 13 | 14 | def inference(filename, model, neighbor_model, augmentation_model, test_dir, augment=True, scan_scale=50, full_scale=4096, evaluate_loss=False, use_normal=False): 15 | """ Predict semantics and affinities """ 16 | 17 | coords, colors, labels, instances, faces = torch.load(filename) 18 | m = np.eye(3) * scan_scale 19 | theta = 0 20 | 21 | m = np.matmul(m, [[np.cos(theta), np.sin(theta),0], [-np.sin(theta), np.cos(theta),0], [0,0,1]]) 22 | coords = np.matmul(coords, m) 23 | if use_normal: 24 | points_1 = coords[faces[:, 0]] 25 | points_2 = coords[faces[:, 1]] 26 | points_3 = coords[faces[:, 2]] 27 | face_normals = np.cross(points_2 - points_1, points_3 - points_1) 28 | face_normals /= np.maximum(np.linalg.norm(face_normals, axis=-1, keepdims=True), 1e-4) 29 | normals = np.zeros((len(coords), 3)) 30 | for c in range(3): 31 | np.add.at(normals, faces[:, c], face_normals) 32 | continue 33 | normals /= np.maximum(np.linalg.norm(normals, axis=-1, keepdims=True), 1e-4) 34 | colors = np.concatenate([colors, normals], axis=-1) 35 | pass 36 | 37 | mins = coords.min(0) 38 | maxs = coords.max(0) 39 | coords -= (mins + maxs) // 2 - full_scale // 2 40 | 41 | coords = np.round(coords) 42 | coords = np.clip(coords, 0, full_scale - 1) 43 | coords = np.concatenate([coords, np.full((len(coords), 1), fill_value=0)], axis=-1) 44 | labels = remapper[labels] 45 | 46 | coords, colors, faces, semantic_gt, instance_gt = torch.from_numpy(coords.astype(np.int64)).cuda(), torch.from_numpy(colors.astype(np.float32)).cuda(), torch.from_numpy(faces.astype(np.int64)).cuda(), torch.from_numpy(labels.astype(np.int64)).cuda(), torch.from_numpy(instances.astype(np.int64)).cuda() 47 | 48 | edges = torch.cat([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]], dim=0) 49 | num_ori_coords = len(coords) 50 | if augment: 51 | augmented_coords, augmented_colors, augmented_instances, augmented_edges = augmentation_model(coords, faces, colors, instance_gt.unsqueeze(-1)) 52 | coords = torch.cat([coords, augmented_coords], dim=0) 53 | colors = torch.cat([colors, augmented_colors], dim=0) 54 | instance_gt = torch.cat([instance_gt, augmented_instances.squeeze(-1)], dim=0) 55 | edges = torch.cat([edges, augmented_edges], dim=0) 56 | pass 57 | 58 | semantic_pred, neighbor_pred = model(coords.reshape((-1, 4)), colors.reshape((-1, colors.shape[-1]))) 59 | #semantic_pred, neighbor_pred = semantic_pred[0], neighbor_pred[0] 60 | 61 | if evaluate_loss: 62 | semantic_loss = torch.nn.functional.cross_entropy(semantic_pred[:num_ori_coords].view((-1, int(semantic_pred.shape[-1]))), semantic_gt.view(-1)) 63 | pass 64 | semantic_pred = semantic_pred.max(-1)[1] 65 | 66 | for neighbor in neighbor_pred: 67 | neighbor.features = torch.sigmoid(neighbor.features) 68 | continue 69 | 70 | if evaluate_loss: 71 | neighbor_gt = neighbor_model(coords.reshape((-1, 4)), instance_gt.reshape((-1, ))) 72 | 73 | neighbor_losses = [] 74 | for scale in range(len(neighbor_gt)): 75 | pred = neighbor_pred[scale].features 76 | gt = neighbor_gt[scale].features[:, :pred.shape[-1]] 77 | mask = neighbor_gt[scale].features[:, pred.shape[-1]:] 78 | neighbor_losses.append(torch.sum(torch.nn.functional.binary_cross_entropy(pred, (gt > 0.5).float(), weight=(1 - gt) * 9 + 1, reduce=False) * mask) / torch.clamp(mask.sum(), min=1)) 79 | continue 80 | pass 81 | 82 | 83 | if evaluate_loss: 84 | losses = [semantic_loss] + neighbor_losses 85 | loss = sum(losses) 86 | 87 | loss_values = [l.data.item() for l in losses] 88 | status = 'val loss: ' 89 | for l in loss_values: 90 | status += '%0.5f '%l 91 | continue 92 | #sys.stdout.write('\r' + status) 93 | print(status) 94 | # print(coords.min(0)[0], coords.max(0)[0], colors.min(), colors.max()) 95 | # print(np.unique(semantic_pred.detach().cpu().numpy()), (neighbor_pred[0].features < 0.5).sum(), neighbor_pred[0].features.shape, neighbor_pred[0].features.min(), neighbor_pred[0].features.max()) 96 | # exit(1) 97 | pass 98 | 99 | scene_id = filename.split('/')[-1].split('_vh_clean')[0] 100 | 101 | semantic_pred = semantic_pred.detach().cpu().numpy() 102 | semantic_gt = semantic_gt.detach().cpu().numpy() 103 | semantic_gt[semantic_gt == -100] = 20 104 | instance_gt = instance_gt.detach().cpu().numpy() 105 | neighbors = neighbor_model.toDense(neighbor_pred) 106 | neighbors = [neighbor.detach().cpu().numpy() for neighbor in neighbors] 107 | coords = coords.detach().cpu().numpy() 108 | coords = coords[:, :3] 109 | colors = colors.detach().cpu().numpy() 110 | colors = np.clip((colors + 1) * 127.5, 0, 255).astype(np.uint8) 111 | edges = edges.detach().cpu().numpy() 112 | if True: 113 | writeSemantics(test_dir + '/sem_pred/' + scene_id + '.txt', semantic_pred) 114 | #writeSemantics(test_dir + '/sem_gt/' + str(sample_index * options.batchSize + batch_index) + '.txt', semantic_gt[batch_index]) 115 | pass 116 | 117 | torch.save((coords, colors, edges, semantic_pred, neighbors, num_ori_coords), test_dir + '/cache/' + scene_id + '.pth') 118 | return 119 | 120 | 121 | def group(filename, test_dir, num_scales, augment=True, num_cross_scales=0, scan_scale=50, full_scale=4096): 122 | """ Run the clustering algorithm """ 123 | 124 | scene_id = filename.split('/')[-1].split('_vh_clean')[0] 125 | coords, colors, edges, semantic_pred, neighbors, num_ori_coords = torch.load(test_dir + '/cache/' + scene_id + '.pth') 126 | #print(coords.min(0), coords.max(0), np.unique(semantic_pred), neighbors[0].min(), neighbors[0].max()) 127 | instance_pred, _ = findInstances(coords, edges, semantic_pred, neighbors, num_scales, num_cross_scales, print_info=True) 128 | 129 | if augment: 130 | instance_labels, counts = np.unique(instance_pred[:num_ori_coords], return_counts=True) 131 | valid_labels = instance_labels[np.logical_and(counts > 100, instance_labels >= 0)] 132 | #print('num valid instances', len(valid_labels)) 133 | label_map = np.full(instance_pred.max() + 1, fill_value=-1, dtype=np.int32) 134 | for index, label in enumerate(valid_labels): 135 | label_map[label] = index 136 | continue 137 | instance_pred = label_map[instance_pred] 138 | pass 139 | 140 | if 'train' in test_dir: 141 | torch.save((semantic_pred[:num_ori_coords], instance_pred[:num_ori_coords]), test_dir + '/cache/' + scene_id + '.pth') 142 | else: 143 | torch.save((coords, colors, edges, semantic_pred, neighbors, num_ori_coords, instance_pred), test_dir + '/cache/' + scene_id + '.pth') 144 | pass 145 | return 146 | 147 | def write(filename, model, validator, test_dir, num_scales, augment=True, num_cross_scales=0, scan_scale=50, full_scale=4096): 148 | """ Predict instance confidence, add additional instances, and write instances """ 149 | scene_id = filename.split('/')[-1].split('_vh_clean')[0] 150 | coords, colors, edges, semantic_pred, neighbors, num_ori_coords, instance_pred = torch.load(test_dir + '/cache/' + scene_id + '.pth') 151 | 152 | coords_inp = torch.from_numpy(np.concatenate([coords, np.zeros((len(coords), 1))], axis=-1)).cuda() 153 | colors_inp = torch.from_numpy(colors.astype(np.float32) / 127.5 - 1).cuda() 154 | semantic_pred, neighbor_pred = model(coords_inp, colors_inp) 155 | semantic_pred = semantic_pred.max(-1)[1].detach().cpu().numpy() 156 | 157 | instance_info = validator.validate(coords, colors[:, :3], instance_pred, semantic_pred) 158 | 159 | if True: 160 | ## Find connected components of certain labels to add additional instances which are likely to be coplanar with larger objects 161 | semantic_instances, num_semantic_instances = findInstancesSemanticsLabels(edges, semantic_pred, labels=[10, 13, 15, 17, 18]) 162 | if num_semantic_instances > 0: 163 | instance_info += validator.validate(coords, colors[:, :3], semantic_instances, semantic_pred) 164 | pass 165 | pass 166 | instance_info = [(mask[:num_ori_coords], label, confidence) for mask, label, confidence in instance_info] 167 | 168 | instance_pred = instance_pred[:num_ori_coords] 169 | semantic_pred = semantic_pred[:num_ori_coords] 170 | writeInstances(test_dir, scene_id, instance_pred, semantic_pred, instance_info=instance_info) 171 | return 172 | 173 | if __name__ == '__main__': 174 | options = parse_args() 175 | options.keyname = 'instance' 176 | #args.keyname += '_' + args.dataset 177 | 178 | if options.suffix != '': 179 | options.keyname += '_' + options.suffix 180 | pass 181 | if options.numScales != 1: 182 | options.keyname += '_' + str(options.numScales) 183 | pass 184 | 185 | options.checkpoint_dir = 'checkpoint/' + options.keyname 186 | 187 | filenames = [] 188 | split = options.split 189 | 190 | data_folder = options.dataFolder 191 | 192 | test_dir = 'test/' + options.keyname + '/inference/' + split + '/' 193 | if not os.path.exists(test_dir): 194 | os.system("mkdir -p %s"%test_dir) 195 | pass 196 | if not os.path.exists(test_dir + '/pred_mask'): 197 | os.system("mkdir -p %s"%(test_dir + '/pred_mask')) 198 | pass 199 | if not os.path.exists(test_dir + '/sem_pred'): 200 | os.system("mkdir -p %s"%(test_dir + '/sem_pred')) 201 | pass 202 | if not os.path.exists(test_dir + '/cache'): 203 | os.system("mkdir -p %s"%(test_dir + '/cache')) 204 | pass 205 | 206 | label_counts = np.zeros(42) 207 | with open('datasets/split_' + split + '.txt', 'r') as f: 208 | for line in f: 209 | scene_id = line.strip() 210 | if len(scene_id) < 5 or scene_id[:5] != 'scene': 211 | continue 212 | if options.scene_id != '' and options.scene_id not in scene_id: 213 | continue 214 | filename = data_folder + '/' + scene_id + '/' + scene_id + '_vh_clean_2.pth' 215 | info = torch.load(filename) 216 | 217 | if os.path.exists(filename) and len(info) == 5: 218 | filenames.append(filename) 219 | if split == 'val' and len(filenames) >= options.numTestingImages: 220 | break 221 | pass 222 | continue 223 | pass 224 | 225 | model = Model(options) 226 | model.cuda() 227 | model.eval() 228 | 229 | if options.startEpoch >= 0: 230 | model.load_state_dict(torch.load(options.checkpoint_dir + '/checkpoint_' + str(options.startEpoch) + '.pth')) 231 | else: 232 | model.load_state_dict(torch.load(options.checkpoint_dir + '/checkpoint.pth')) 233 | pass 234 | 235 | neighbor_model = NeighborGT(options) 236 | neighbor_model.cuda() 237 | 238 | augmentation_model = CoordAugmentation(options) 239 | augmentation_model.cuda() 240 | augmentation_model.eval() 241 | 242 | 243 | num_scales = options.numScales 244 | num_cross_scales = options.numCrossScales 245 | scan_scale = options.scanScale 246 | full_scale = options.inputScale 247 | 248 | if options.useCache > 0: 249 | filenames = [filename for filename in filenames if not os.path.exists(test_dir + '/' + filename + '.txt')] 250 | pass 251 | print('num images', len(filenames)) 252 | 253 | assert('predict' in options.task or 'cluster' in options.task or 'write' in options.task) 254 | 255 | if 'predict' in options.task: 256 | for filename in filenames: 257 | inference(filename, model, neighbor_model, augmentation_model, test_dir, scan_scale=scan_scale, full_scale=full_scale, evaluate_loss=split == 'val', use_normal='normal' in options.suffix) 258 | continue 259 | pass 260 | 261 | if 'cluster' in options.task: 262 | #num_processes = mp.cpu_count() 263 | num_processes = 8 264 | p = mp.Pool(processes=num_processes) 265 | p.map(partial(group, test_dir=test_dir, num_scales=num_scales, num_cross_scales=num_cross_scales, scan_scale=scan_scale), filenames) 266 | p.close() 267 | p.join() 268 | pass 269 | 270 | 271 | if 'write' in options.task: 272 | validator = InstanceValidator(options.checkpoint_dir) 273 | for filename in filenames: 274 | write(filename, model, validator, test_dir=test_dir, num_scales=num_scales, num_cross_scales=num_cross_scales, scan_scale=scan_scale) 275 | continue 276 | if split == 'val': 277 | os.system('python scripts/evaluate_semantic_instance.py --pred_path=' + test_dir + ' --num_testing_images=' + str(options.numTestingImages) + ' --scene_id=' + options.scene_id) 278 | pass 279 | pass 280 | -------------------------------------------------------------------------------- /models/instance.py: -------------------------------------------------------------------------------- 1 | import sparseconvnet as scn 2 | import torch.nn as nn 3 | import torch 4 | import numpy as np 5 | import time 6 | 7 | class NeighborGT(nn.Module): 8 | """ Generate affinity supervision on the fly based on gt instance labels """ 9 | def __init__(self, options): 10 | nn.Module.__init__(self) 11 | self.options = options 12 | dimension = 3 13 | self.input_layer = scn.InputLayer(dimension, options.inputScale, mode=4) 14 | self.conv = scn.SubmanifoldConvolution(dimension, 1, options.numNeighbors, 3, bias=False) 15 | 16 | self.pool_1 = scn.AveragePooling(dimension, 2, 2) 17 | self.pool_2 = scn.AveragePooling(dimension, 4, 4) 18 | self.pool_3 = scn.AveragePooling(dimension, 8, 8) 19 | self.pool_4 = scn.AveragePooling(dimension, 16, 16) 20 | self.pool_5 = scn.AveragePooling(dimension, 32, 32) 21 | self.unpool_1 = scn.UnPooling(dimension, 2, 2) 22 | self.unpool_2 = scn.UnPooling(dimension, 4, 4) 23 | self.unpool_3 = scn.UnPooling(dimension, 8, 8) 24 | self.unpool_4 = scn.UnPooling(dimension, 16, 16) 25 | self.unpool_5 = scn.UnPooling(dimension, 32, 32) 26 | 27 | with torch.no_grad(): 28 | weight = torch.zeros(27, 1, options.numNeighbors).cuda() 29 | if options.numNeighbors == 6: 30 | offsets = [4, 22, 10, 16, 12, 14] 31 | pass 32 | for index, offset in enumerate(offsets): 33 | weight[offset, 0, index] = 1 34 | continue 35 | self.conv.weight = nn.Parameter(weight) 36 | pass 37 | self.output_layer = scn.OutputLayer(dimension) 38 | return 39 | 40 | def toDense(self, neighbors): 41 | """ Convert prediction of sparse locations to dense representation """ 42 | if len(neighbors) >= 2: 43 | neighbors[1] = self.unpool_1(neighbors[1]) 44 | pass 45 | if len(neighbors) >= 3: 46 | neighbors[2] = self.unpool_2(neighbors[2]) 47 | pass 48 | if len(neighbors) >= 4: 49 | neighbors[3] = self.unpool_3(neighbors[3]) 50 | pass 51 | if len(neighbors) >= 5: 52 | neighbors[4] = self.unpool_4(neighbors[4]) 53 | pass 54 | if len(neighbors) >= 6: 55 | neighbors[5] = self.unpool_5(neighbors[5]) 56 | pass 57 | 58 | for neighbor in neighbors: 59 | neighbor.features = neighbor.features[:, :6] 60 | continue 61 | 62 | #print(neighbor.features[torch.all(neighbor.get_spatial_locations()[:, :3] == (torch.Tensor([7, 8, 24]).long() // pow(2, scale)), dim=-1).cpu()]) 63 | #print(neighbor.features[torch.all(neighbor.get_spatial_locations()[:, :3] == (torch.Tensor([8, 8, 24]).long() // pow(2, scale)), dim=-1).cpu()]) 64 | neighbors = [self.output_layer(neighbor) for neighbor in neighbors] 65 | return neighbors 66 | 67 | def forward(self, coords, instance_gt, use_gpu=True): 68 | #instance = self.input_layer((coords, torch.unsqueeze(instance_gt + 1, -1).float())) 69 | instance = self.input_layer((coords, torch.stack([instance_gt.float() + 1, torch.ones(len(instance_gt)).cuda()], dim=-1))) 70 | instances = [] 71 | if self.options.numScales >= 1: 72 | instances.append(instance) 73 | pass 74 | if self.options.numScales >= 2: 75 | instances.append(self.pool_1(instance)) 76 | pass 77 | if self.options.numScales >= 3: 78 | instances.append(self.pool_2(instance)) 79 | pass 80 | if self.options.numScales >= 4: 81 | instances.append(self.pool_3(instance)) 82 | pass 83 | if self.options.numScales >= 5: 84 | instances.append(self.pool_4(instance)) 85 | pass 86 | if self.options.numScales >= 6: 87 | instances.append(self.pool_5(instance)) 88 | pass 89 | 90 | # for instance in instances: 91 | # instance.features = instance.features / torch.clamp(torch.sum(instance.features, dim=-1, keepdim=True), min=1e-4) 92 | # continue 93 | 94 | # for scale, instance in enumerate(instances): 95 | # print(instance.features[torch.all(instance.get_spatial_locations()[:, :3] == (torch.Tensor([7, 8, 24]).long() // pow(2, scale)), dim=-1).cpu()]) 96 | # print(instance.features[torch.all(instance.get_spatial_locations()[:, :3] == (torch.Tensor([8, 8, 24]).long() // pow(2, scale)), dim=-1).cpu()]) 97 | # continue 98 | 99 | instance_counts = [instance.features[:, 1:2] for instance in instances] 100 | for instance, count in zip(instances, instance_counts): 101 | instance.features = instance.features[:, :1] / count 102 | continue 103 | scale_count_thresholds = pow(4.0, np.arange(self.options.numScales)) / pow(8, np.arange(self.options.numScales)) 104 | 105 | for scale, instance in enumerate(instances): 106 | neighbors = [] 107 | 108 | #feature_counts = torch.sum(instance.features, dim=-1, keepdim=True) 109 | #instance.features = instance.features / torch.clamp(feature_counts, min=1e-4) 110 | shifted_instances = self.conv(instance) 111 | features = ((shifted_instances.features - instance.features).abs() < 0.01).float() 112 | masks = shifted_instances.features > 0 113 | # if scale >= 1: 114 | # masks = masks & (feature_counts >= scale_count_thresholds[scale]) 115 | # pass 116 | instance.features = torch.cat([features.detach(), masks.float().detach()], dim=-1) 117 | continue 118 | 119 | # for scale, instance in enumerate(instances): 120 | # print(instance.features[torch.all(instance.get_spatial_locations()[:, :3] == (torch.Tensor([7, 8, 24]).long() // pow(2, scale)), dim=-1).cpu()]) 121 | # print(instance.features[torch.all(instance.get_spatial_locations()[:, :3] == (torch.Tensor([8, 8, 24]).long() // pow(2, scale)), dim=-1).cpu()]) 122 | # continue 123 | return instances 124 | 125 | 126 | class CoordAugmentation(nn.Module): 127 | """ Augment input mesh by randomly sampling points inside large faces """ 128 | def __init__(self, options): 129 | nn.Module.__init__(self) 130 | self.options = options 131 | dimension = 3 132 | self.input_layer = scn.InputLayer(dimension, options.inputScale, mode=4) 133 | 134 | self.coord_input_layer = scn.InputLayer(dimension, options.inputScale, mode=2) 135 | self.num_instances = 1 136 | self.conv_1 = scn.SubmanifoldConvolution(dimension, self.num_instances, self.num_instances, 3, bias=False) 137 | self.conv_2 = scn.SubmanifoldConvolution(dimension, self.num_instances, self.num_instances, 3, bias=False) 138 | self.conv_3 = scn.SubmanifoldConvolution(dimension, self.num_instances, self.num_instances, 3, bias=False) 139 | self.conv_4 = scn.SubmanifoldConvolution(dimension, self.num_instances, self.num_instances, 3, bias=False) 140 | self.conv_5 = scn.SubmanifoldConvolution(dimension, self.num_instances, self.num_instances, 3, bias=False) 141 | self.conv_6 = scn.SubmanifoldConvolution(dimension, self.num_instances, self.num_instances, 3, bias=False) 142 | 143 | with torch.no_grad(): 144 | weights = [torch.zeros(27, self.num_instances, self.num_instances).cuda() for _ in range(6)] 145 | for index, offset in enumerate([4, 22, 10, 16, 12, 14]): 146 | for instance_index in range(self.num_instances): 147 | weights[index][offset, instance_index, instance_index] = 1 148 | continue 149 | continue 150 | self.conv_1.weight = nn.Parameter(weights[0]) 151 | self.conv_2.weight = nn.Parameter(weights[1]) 152 | self.conv_3.weight = nn.Parameter(weights[2]) 153 | self.conv_4.weight = nn.Parameter(weights[3]) 154 | self.conv_5.weight = nn.Parameter(weights[4]) 155 | self.conv_6.weight = nn.Parameter(weights[5]) 156 | pass 157 | self.coord_output_layer = scn.OutputLayer(dimension) 158 | return 159 | 160 | def eval(self): 161 | torch.manual_seed(0) 162 | return 163 | 164 | def train(self): 165 | t = int(time.time() * 1000000) 166 | seed = ((t & 0xff000000) >> 24) + ((t & 0x00ff0000) >> 8) + ((t & 0x0000ff00) << 8) + ((t & 0x000000ff) << 24) 167 | torch.manual_seed(seed) 168 | return 169 | 170 | def forward(self, coords, faces, colors, instances): 171 | edges_1 = coords[faces[:, 0]] - coords[faces[:, 1]] 172 | edges_2 = coords[faces[:, 1]] - coords[faces[:, 2]] 173 | edges_3 = coords[faces[:, 2]] - coords[faces[:, 0]] 174 | edge_lengths = torch.max(torch.max(edges_1.abs().max(dim=-1)[0], edges_2.abs().max(dim=-1)[0]), edges_3.abs().max(dim=-1)[0]) 175 | large_faces = faces[edge_lengths >= 2] 176 | if len(large_faces) == 0: 177 | return coords[0:1] 178 | #cross_product = torch.stack([edges_1[:, 1] * edges_2[:, 2] - edges_1[:, 2] * edges_2[:, 1], edges_1[:, 2] * edges_2[:, 0] - edges_1[:, 0] * edges_2[:, 2], edges_1[:, 0] * edges_2[:, 1] - edges_1[:, 1] * edges_2[:, 0]], axis=-1) 179 | num_sampled_points = 5 180 | #torch.backends.cudnn.deterministic=True 181 | alphas = torch.rand(num_sampled_points, len(large_faces), 3).cuda() 182 | alphas = alphas / alphas.sum(-1, keepdim=True) 183 | vertices = torch.stack([coords[large_faces[:, 0]], coords[large_faces[:, 1]], coords[large_faces[:, 2]]], dim=1).float() 184 | sampled_coords = (alphas.unsqueeze(-1) * vertices).sum(2).view((-1, 4)).round().long() 185 | all_coords = torch.cat([coords, sampled_coords], dim=0) 186 | 187 | vertex_colors = torch.stack([colors[large_faces[:, 0]], colors[large_faces[:, 1]], colors[large_faces[:, 2]]], dim=1).float() 188 | sampled_colors = (alphas.unsqueeze(-1) * vertex_colors).sum(2).view((-1, colors.shape[-1])) 189 | all_colors = torch.cat([colors, sampled_colors], dim=0) 190 | 191 | vertex_instances = torch.stack([instances[large_faces[:, 0]], instances[large_faces[:, 1]], instances[large_faces[:, 2]]], dim=1).float() 192 | sampled_instances = vertex_instances[torch.arange(len(vertex_instances)).cuda().repeat(num_sampled_points), alphas.max(-1)[1].view(-1)] 193 | all_instances = torch.cat([instances.float(), sampled_instances], dim=0) 194 | 195 | all_flags = torch.cat([-100 * torch.ones((len(coords), 1)).cuda(), torch.ones((len(sampled_coords), 1)).cuda()], dim=0) 196 | 197 | all_values = torch.cat([all_colors, all_instances, all_flags], dim=-1) 198 | 199 | voxel = self.input_layer((all_coords, all_values)) 200 | 201 | valid_mask = voxel.features[:, -1] > 0 202 | valid_values = voxel.features[valid_mask] 203 | 204 | augmented_coords = voxel.get_spatial_locations()[valid_mask].cuda() 205 | 206 | all_coords = torch.cat([coords, augmented_coords], dim=0) 207 | #augmented_coord_indices = torch.arange(len(coords), len(coords) + len(augmented_coords)).cuda() 208 | all_coord_indices = torch.arange(1, 1 + len(all_coords)).cuda() 209 | coord_indices = self.coord_input_layer((all_coords, all_coord_indices.unsqueeze(-1).float())) 210 | neighbors = [] 211 | shifted_coord_indices = [self.conv_1(coord_indices), self.conv_2(coord_indices), self.conv_3(coord_indices), self.conv_4(coord_indices), self.conv_5(coord_indices), self.conv_6(coord_indices)] 212 | shifted_coord_indices = [self.coord_output_layer(indices) for indices in shifted_coord_indices] 213 | shifted_coord_indices = torch.cat(shifted_coord_indices, dim=0).long().view(-1) 214 | all_edges = torch.stack([all_coord_indices.repeat(6), shifted_coord_indices], dim=-1) 215 | augmented_edges = all_edges[(shifted_coord_indices > 0) & (shifted_coord_indices.max(-1)[0] > len(coords))] - 1 216 | return augmented_coords, valid_values[:, :colors.shape[-1]], valid_values[:, colors.shape[-1]:-1].round().long(), augmented_edges 217 | 218 | class Model(nn.Module): 219 | """ Base model based on sparse convolutions """ 220 | def __init__(self, options): 221 | nn.Module.__init__(self) 222 | 223 | self.options = options 224 | 225 | dimension = 3 226 | m = 32 # 16 or 32 227 | residual_blocks = True #True or False 228 | block_reps = 2 #Conv block repetition factor: 1 or 2 229 | 230 | self.outputs = [] 231 | def hook(module, input, output): 232 | self.outputs.append(output) 233 | return 234 | 235 | self.sparseModel = scn.Sequential().add( 236 | scn.InputLayer(dimension, options.inputScale, mode=4)).add( 237 | scn.SubmanifoldConvolution(dimension, 3 + 3 * int('normal' in self.options.suffix), m, 3, False)).add( 238 | scn.UNet(dimension, block_reps, [m, 2 * m, 3 * m, 4 * m, 5 * m, 6 * m, 7 * m], residual_blocks)).add( 239 | scn.BatchNormReLU(m)) 240 | #print(self.sparseModel[2]) 241 | #exit(1) 242 | 243 | if options.numScales >= 2: 244 | #list(self.sparseModel[2][1].children())[1][2][3].register_forward_hook(hook) 245 | list(self.sparseModel[2][4].children())[1][3].register_forward_hook(hook) 246 | pass 247 | if options.numScales >= 3: 248 | list(list(self.sparseModel[2][4].children())[1][2][4].children())[1][3].register_forward_hook(hook) 249 | pass 250 | if options.numScales >= 4: 251 | list(list(list(self.sparseModel[2][4].children())[1][2][4].children())[1][2][4].children())[1][3].register_forward_hook(hook) 252 | pass 253 | if options.numScales >= 5: 254 | list(list(list(list(self.sparseModel[2][4].children())[1][2][4].children())[1][2][4].children())[1][2][4].children())[1][3].register_forward_hook(hook) 255 | pass 256 | if options.numScales >= 6: 257 | list(list(list(list(list(self.sparseModel[2][4].children())[1][2][4].children())[1][2][4].children())[1][2][4].children())[1][2][4].children())[1][3].register_forward_hook(hook) 258 | pass 259 | 260 | # list(list(list(list(list(self.sparseModel[2][1].children())[1][2][1].children())[1][2][1].children())[1][2][1].children())[1][2][1].children())[1][2][3].register_forward_hook(hook) 261 | # list(list(list(list(self.sparseModel[2][1].children())[1][2][1].children())[1][2][1].children())[1][2][1].children())[1][2][3].register_forward_hook(hook) 262 | # list(list(list(self.sparseModel[2][1].children())[1][2][1].children())[1][2][1].children())[1][2][3].register_forward_hook(hook) 263 | # list(list(self.sparseModel[2][1].children())[1][2][1].children())[1][2][3].register_forward_hook(hook) 264 | # list(self.sparseModel[2][1].children())[1][2][3].register_forward_hook(hook) 265 | 266 | self.sparsify = scn.Sparsify(dimension) 267 | self.output_layer = scn.OutputLayer(dimension) 268 | 269 | self.linear = nn.Linear(m, 20) 270 | 271 | self.neighbor_linear_0 = nn.Linear(m, 6 + 7 * min(self.options.numCrossScales, max(self.options.numScales - 1, 0))) 272 | self.neighbor_linear_1 = nn.Linear(m * 2, 6 + 7 * min(self.options.numCrossScales, max(self.options.numScales - 2, 0))) 273 | self.neighbor_linear_2 = nn.Linear(m * 3, 6 + 7 * min(self.options.numCrossScales, max(self.options.numScales - 3, 0))) 274 | self.neighbor_linear_3 = nn.Linear(m * 4, 6 + 7 * min(self.options.numCrossScales, max(self.options.numScales - 4, 0))) 275 | self.neighbor_linear_4 = nn.Linear(m * 5, 6 + 7 * min(self.options.numCrossScales, max(self.options.numScales - 5, 0))) 276 | self.neighbor_linear_5 = nn.Linear(m * 6, 6 + 7 * min(self.options.numCrossScales, max(self.options.numScales - 6, 0))) 277 | 278 | # self.neighbor_linear_0 = nn.Linear(m, 6 + 7 * 5) 279 | # self.neighbor_linear_1 = nn.Linear(m * 2, 6 + 7 * 4) 280 | # self.neighbor_linear_2 = nn.Linear(m * 3, 6 + 7 * 3) 281 | # self.neighbor_linear_3 = nn.Linear(m * 4, 6 + 7 * 2) 282 | # self.neighbor_linear_4 = nn.Linear(m * 5, 6 + 7 * 1) 283 | # self.neighbor_linear_5 = nn.Linear(m * 6, 6 + 7 * 0) 284 | 285 | self.unpool_1 = scn.UnPooling(dimension, 2, 2) 286 | self.unpool_2 = scn.UnPooling(dimension, 4, 4) 287 | self.unpool_3 = scn.UnPooling(dimension, 8, 8) 288 | self.unpool_4 = scn.UnPooling(dimension, 16, 16) 289 | self.unpool_5 = scn.UnPooling(dimension, 32, 32) 290 | return 291 | 292 | def forward(self, coords, colors): 293 | # for scale_index, neighbor in enumerate(neighbors): 294 | # for index, _ in enumerate(neighbor): 295 | # print(scale_index, index, _.shape, _.min().item(), _.max().item()) 296 | # continue 297 | # continue 298 | # exit(1) 299 | x = self.sparseModel((coords, colors)) 300 | # print(x.shape, len(self.outputs)) 301 | # exit(1) 302 | if 'maxpool' not in self.options.suffix: 303 | semantic_pred = self.linear(self.output_layer(x)) 304 | else: 305 | semantic_pred = self.output_layer(x) 306 | pass 307 | 308 | if False: 309 | self.outputs = self.outputs[-5:] 310 | outputs = [x] + self.outputs[::-1] 311 | outputs[0].features, outputs[1].features, outputs[2].features, outputs[3].features, outputs[4].features, outputs[5].features = self.neighbor_linear_0(outputs[0].features), self.neighbor_linear_1(outputs[1].features), self.neighbor_linear_2(outputs[2].features), self.neighbor_linear_3(outputs[3].features), self.neighbor_linear_4(outputs[4].features), self.neighbor_linear_5(outputs[5].features) 312 | neighbor_pred = outputs 313 | else: 314 | if self.options.numScales > 1: 315 | self.outputs = self.outputs[-(self.options.numScales - 1):] 316 | outputs = [x] + self.outputs[::-1] 317 | elif self.options.numScales == 1: 318 | outputs = [x] 319 | else: 320 | outputs = [] 321 | pass 322 | 323 | if self.options.numScales >= 1: 324 | outputs[0].features = self.neighbor_linear_0(outputs[0].features)[:, :6 + min(self.options.numCrossScales, self.options.numScales - 1) * 7] 325 | pass 326 | if self.options.numScales >= 2: 327 | outputs[1].features = self.neighbor_linear_1(outputs[1].features)[:, :6 + min(self.options.numCrossScales, self.options.numScales - 2) * 7] 328 | pass 329 | if self.options.numScales >= 3: 330 | outputs[2].features = self.neighbor_linear_2(outputs[2].features)[:, :6 + min(self.options.numCrossScales, self.options.numScales - 3) * 7] 331 | pass 332 | if self.options.numScales >= 4: 333 | outputs[3].features = self.neighbor_linear_3(outputs[3].features)[:, :6 + min(self.options.numCrossScales, self.options.numScales - 4) * 7] 334 | pass 335 | if self.options.numScales >= 5: 336 | outputs[4].features = self.neighbor_linear_4(outputs[4].features)[:, :6 + min(self.options.numCrossScales, self.options.numScales - 5) * 7] 337 | pass 338 | if self.options.numScales >= 6: 339 | outputs[5].features = self.neighbor_linear_5(outputs[5].features)[:, :6 + min(self.options.numCrossScales, self.options.numScales - 6) * 7] 340 | pass 341 | neighbor_pred = outputs 342 | pass 343 | 344 | #neighbor_pred = self.neighbor_linear_0(outputs[0].features), self.neighbor_linear_1(outputs[1].features), self.neighbor_linear_2(outputs[2].features), self.neighbor_linear_3(outputs[3].features), self.neighbor_linear_4(outputs[4].features), self.neighbor_linear_5(outputs[5].features) 345 | 346 | # outputs = [self.unpool_1(outputs[0]), self.unpool_2(outputs[1]), self.unpool_3(outputs[2]), self.unpool_4(outputs[3]), self.unpool_5(outputs[4])] 347 | # outputs = [self.output_layer(output) for output in outputs] 348 | # outputs = [x] + outputs 349 | # neighbor_pred = [self.neighbor_linear_0(outputs[0]), self.neighbor_linear_1(outputs[1]), self.neighbor_linear_2(outputs[2]), self.neighbor_linear_3(outputs[3]), self.neighbor_linear_4(outputs[4]), self.neighbor_linear_5(outputs[5])] 350 | 351 | #neighbor_pred_0 = self.neighbor_linear_0(x) 352 | return semantic_pred, neighbor_pred 353 | 354 | 355 | class Validator(nn.Module): 356 | def __init__(self, full_scale=127, use_normal=False): 357 | nn.Module.__init__(self) 358 | 359 | dimension = 3 360 | m = 32 # 16 or 32 361 | residual_blocks = True #True or False 362 | block_reps = 2 #Conv block repetition factor: 1 or 2 363 | 364 | blocks = [['b', m * k, 2, 2] for k in [1, 2, 3, 4, 5]] 365 | self.num_final_channels = m * len(blocks) 366 | 367 | self.sparseModel = scn.Sequential().add( 368 | scn.InputLayer(dimension, full_scale, mode=4)).add( 369 | scn.SubmanifoldConvolution(dimension, 3 + 3 * int(use_normal), m, 3, False)).add( 370 | scn.MaxPooling(dimension, 3, 2)).add( 371 | scn.SparseResNet(dimension, m, blocks)).add( 372 | scn.BatchNormReLU(self.num_final_channels)).add( 373 | scn.SparseToDense(dimension, self.num_final_channels)) 374 | 375 | self.num_labels = 20 376 | self.label_encoder = nn.Sequential(nn.Linear(self.num_labels, 64), nn.ReLU()) 377 | #self.pred = nn.Linear(m * len(blocks), 1) 378 | self.pred = nn.Sequential(nn.Linear(self.num_final_channels + 64, 64), nn.ReLU(), nn.Linear(64, 1)) 379 | return 380 | 381 | def forward(self, coords, colors, labels): 382 | x = self.sparseModel((coords, colors)) 383 | labels = (labels.unsqueeze(-1) == torch.arange(self.num_labels).cuda()).float() 384 | label_x = self.label_encoder(labels) 385 | x = torch.cat([x.view((-1, self.num_final_channels))[1:], label_x], dim=1) 386 | pred = self.pred(x) 387 | return pred.view(-1) 388 | -------------------------------------------------------------------------------- /options.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | def parse_args(): 4 | """ 5 | Parse input arguments 6 | """ 7 | parser = argparse.ArgumentParser(description='PlaneFlow') 8 | 9 | parser.add_argument('--task', dest='task', 10 | help='task type: [train, test, predict]', 11 | default='train', type=str) 12 | parser.add_argument('--restore', dest='restore', 13 | help='how to restore the model', 14 | default=1, type=int) 15 | parser.add_argument('--batchSize', dest='batchSize', 16 | help='batch size', 17 | default=1, type=int) 18 | parser.add_argument('--numTrainingImages', dest='numTrainingImages', 19 | help='the number of images to train', 20 | default=0, type=int) 21 | parser.add_argument('--numTestingImages', dest='numTestingImages', 22 | help='the number of images to test/predict', 23 | default=20, type=int) 24 | parser.add_argument('--scene_id', dest='scene_id', 25 | help='scene id to test', 26 | default='', type=str) 27 | parser.add_argument('--LR', dest='LR', 28 | help='learning rate', 29 | default=2.5e-4, type=float) 30 | parser.add_argument('--numEpochs', dest='numEpochs', 31 | help='the number of epochs', 32 | default=50, type=int) 33 | parser.add_argument('--startEpoch', dest='startEpoch', 34 | help='starting epoch index', 35 | default=-1, type=int) 36 | parser.add_argument('--inputScale', dest='inputScale', 37 | help='input scale', 38 | default=4096, type=int) 39 | parser.add_argument('--scanScale', dest='scanScale', 40 | help='scan scale', 41 | default=50, type=int) 42 | parser.add_argument('--numScales', dest='numScales', 43 | help='the number of scales', 44 | default=2, type=int) 45 | parser.add_argument('--numCrossScales', dest='numCrossScales', 46 | help='the number of cross scales', 47 | default=0, type=int) 48 | parser.add_argument('--numNeighbors', dest='numNeighbors', 49 | help='the number of neighbors', 50 | default=6, type=int) 51 | parser.add_argument('--negativeWeights', dest='negativeWeights', 52 | help='negative weights', 53 | default='531111', type=str) 54 | parser.add_argument('--visualizeMode', dest='visualizeMode', 55 | help='visualization mode', 56 | default='', type=str) 57 | parser.add_argument('--suffix', dest='suffix', 58 | help='suffix to distinguish experiments', 59 | default='normal_augment', type=str) 60 | parser.add_argument('--useCache', dest='useCache', 61 | help='use cache instead of re-computing existing examples', 62 | default=0, type=int) 63 | parser.add_argument('--dataFolder', dest='dataFolder', 64 | help='data folder', 65 | default='/gruvi/Data/chenliu/ScanNet/scans/', type=str) 66 | parser.add_argument('--labelFile', dest='labelFile', 67 | help='path to scannetv2-labels.combined.tsv', 68 | default='/gruvi/Data/chenliu/ScanNet/tasks/scannetv2-labels.combined.tsv', type=str) 69 | parser.add_argument('--split', dest='split', 70 | help='data split: [train, test, val]', 71 | default='val', type=str) 72 | 73 | args = parser.parse_args() 74 | return args 75 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | imageio==2.5.0 2 | numpy==1.16.1 3 | opencv-python==4.0.0.21 4 | Pillow==5.3.0 5 | plyfile==0.7 6 | scipy==1.2.1 7 | -e git+https://github.com/facebookresearch/SparseConvNet.git@bd9f2c46e7105de6ac7a4e6dd9247ea0f204cd42#egg=sparseconvnet 8 | torch==1.0.0 9 | torchvision==0.2.1 10 | tqdm==4.31.1 11 | -------------------------------------------------------------------------------- /scripts/evaluate_semantic_instance.py: -------------------------------------------------------------------------------- 1 | # Evaluates semantic instance task 2 | # Adapted from the CityScapes evaluation: https://github.com/mcordts/cityscapesScripts/tree/master/cityscapesscripts/evaluation 3 | # Input: 4 | # - path to .txt prediction files 5 | # - path to .txt ground truth files 6 | # - output file to write results to 7 | # Each .txt prediction file look like: 8 | # [(pred0) rel. path to pred. mask over verts as .txt] [(pred0) label id] [(pred0) confidence] 9 | # [(pred1) rel. path to pred. mask over verts as .txt] [(pred1) label id] [(pred1) confidence] 10 | # [(pred2) rel. path to pred. mask over verts as .txt] [(pred2) label id] [(pred2) confidence] 11 | # ... 12 | # 13 | # NOTE: The prediction files must live in the root of the given prediction path. 14 | # Predicted mask .txt files must live in a subfolder. 15 | # Additionally, filenames must not contain spaces. 16 | # The relative paths to predicted masks must contain one integer per line, 17 | # where each line corresponds to vertices in the *_vh_clean_2.ply (in that order). 18 | # Non-zero integers indicate part of the predicted instance. 19 | # The label ids specify the class of the corresponding mask. 20 | # Confidence is a float confidence score of the mask. 21 | # 22 | # Note that only the valid classes are used for evaluation, 23 | # i.e., any ground truth label not in the valid label set 24 | # is ignored in the evaluation. 25 | # 26 | # example usage: evaluate_semantic_instance.py --scan_path [path to scan data] --output_file [output file] 27 | 28 | # python imports 29 | import math 30 | import os, sys, argparse 31 | import inspect 32 | from copy import deepcopy 33 | 34 | try: 35 | import numpy as np 36 | except: 37 | print("Failed to import numpy package.") 38 | sys.exit(-1) 39 | 40 | currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) 41 | parentdir = os.path.dirname(currentdir) 42 | sys.path.insert(0,parentdir) 43 | import util 44 | import util_3d 45 | 46 | parser = argparse.ArgumentParser() 47 | parser.add_argument('--pred_path', required=True, help='path to directory of predicted .txt files') 48 | parser.add_argument('--gt_path', default='instance_val', help='path to directory of gt .txt files') 49 | parser.add_argument('--output_file', default='', help='output file [default: pred_path/semantic_instance_evaluation.txt]') 50 | parser.add_argument('--scene_id', default='', type=str, help='scene id to test') 51 | parser.add_argument('--num_testing_images', default=100, type=int, help='num images to test') 52 | 53 | opt = parser.parse_args() 54 | 55 | if opt.output_file == '': 56 | opt.output_file = os.path.join(opt.pred_path, 'semantic_instance_evaluation.txt') 57 | 58 | 59 | # ---------- Label info ---------- # 60 | CLASS_LABELS = ['cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window', 'bookshelf', 'picture', 'counter', 'desk', 'curtain', 'refrigerator', 'shower curtain', 'toilet', 'sink', 'bathtub', 'otherfurniture'] 61 | VALID_CLASS_IDS = np.array([3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39]) 62 | ID_TO_LABEL = {} 63 | LABEL_TO_ID = {} 64 | for i in range(len(VALID_CLASS_IDS)): 65 | LABEL_TO_ID[CLASS_LABELS[i]] = VALID_CLASS_IDS[i] 66 | ID_TO_LABEL[VALID_CLASS_IDS[i]] = CLASS_LABELS[i] 67 | # ---------- Evaluation params ---------- # 68 | # overlaps for evaluation 69 | opt.overlaps = np.append(np.arange(0.5,0.95,0.05), 0.25) 70 | # minimum region size for evaluation [verts] 71 | opt.min_region_sizes = np.array( [ 100 ] ) 72 | # distance thresholds [m] 73 | opt.distance_threshes = np.array( [ float('inf') ] ) 74 | # distance confidences 75 | opt.distance_confs = np.array( [ -float('inf') ] ) 76 | 77 | 78 | def evaluate_matches(matches): 79 | overlaps = opt.overlaps 80 | min_region_sizes = [ opt.min_region_sizes[0] ] 81 | dist_threshes = [ opt.distance_threshes[0] ] 82 | dist_confs = [ opt.distance_confs[0] ] 83 | 84 | # results: class x overlap 85 | ap = np.zeros( (len(dist_threshes) , len(CLASS_LABELS) , len(overlaps)) , np.float ) 86 | for di, (min_region_size, distance_thresh, distance_conf) in enumerate(zip(min_region_sizes, dist_threshes, dist_confs)): 87 | for oi, overlap_th in enumerate(overlaps): 88 | pred_visited = {} 89 | for m in matches: 90 | for p in matches[m]['pred']: 91 | for label_name in CLASS_LABELS: 92 | for p in matches[m]['pred'][label_name]: 93 | if 'filename' in p: 94 | pred_visited[p['filename']] = False 95 | for li, label_name in enumerate(CLASS_LABELS): 96 | y_true = np.empty(0) 97 | y_score = np.empty(0) 98 | hard_false_negatives = 0 99 | has_gt = False 100 | has_pred = False 101 | for m in matches: 102 | pred_instances = matches[m]['pred'][label_name] 103 | gt_instances = matches[m]['gt'][label_name] 104 | #print(label_name, 'num', len(gt_instances), len(pred_instances)) 105 | # filter groups in ground truth 106 | gt_instances = [ gt for gt in gt_instances if gt['instance_id']>=1000 and gt['vert_count']>=min_region_size and gt['med_dist']<=distance_thresh and gt['dist_conf']>=distance_conf ] 107 | if gt_instances: 108 | has_gt = True 109 | if pred_instances: 110 | has_pred = True 111 | 112 | cur_true = np.ones ( len(gt_instances) ) 113 | cur_score = np.ones ( len(gt_instances) ) * (-float("inf")) 114 | cur_match = np.zeros( len(gt_instances) , dtype=np.bool ) 115 | # collect matches 116 | for (gti,gt) in enumerate(gt_instances): 117 | found_match = False 118 | num_pred = len(gt['matched_pred']) 119 | #print(gti, num_pred, len(gt['matched_pred'])) 120 | for pred in gt['matched_pred']: 121 | # greedy assignments 122 | if pred_visited[pred['filename']]: 123 | continue 124 | overlap = float(pred['intersection']) / (gt['vert_count']+pred['vert_count']-pred['intersection']) 125 | if overlap > overlap_th: 126 | confidence = pred['confidence'] 127 | # if already have a prediction for this gt, 128 | # the prediction with the lower score is automatically a false positive 129 | if cur_match[gti]: 130 | max_score = max( cur_score[gti] , confidence ) 131 | min_score = min( cur_score[gti] , confidence ) 132 | cur_score[gti] = max_score 133 | # append false positive 134 | cur_true = np.append(cur_true,0) 135 | cur_score = np.append(cur_score,min_score) 136 | cur_match = np.append(cur_match,True) 137 | # otherwise set score 138 | else: 139 | found_match = True 140 | cur_match[gti] = True 141 | cur_score[gti] = confidence 142 | pred_visited[pred['filename']] = True 143 | if not found_match: 144 | hard_false_negatives += 1 145 | if len(gt_instances) > 0 and False: 146 | exit(1) 147 | # remove non-matched ground truth instances 148 | cur_true = cur_true [ cur_match==True ] 149 | cur_score = cur_score[ cur_match==True ] 150 | 151 | # collect non-matched predictions as false positive 152 | for pred_index, pred in enumerate(pred_instances): 153 | found_gt = False 154 | #print(pred_index, pred) 155 | for gt in pred['matched_gt']: 156 | overlap = float(gt['intersection']) / (gt['vert_count']+pred['vert_count']-gt['intersection']) 157 | #print(overlap) 158 | if overlap > overlap_th: 159 | found_gt = True 160 | break 161 | if not found_gt: 162 | num_ignore = pred['void_intersection'] 163 | for gt in pred['matched_gt']: 164 | # group? 165 | if gt['instance_id'] < 1000: 166 | num_ignore += gt['intersection'] 167 | # small ground truth instances 168 | if gt['vert_count'] < min_region_size or gt['med_dist']>distance_thresh or gt['dist_conf'] 0: 183 | # compute precision recall curve first 184 | 185 | # sorting and cumsum 186 | score_arg_sort = np.argsort(y_score) 187 | y_score_sorted = y_score[score_arg_sort] 188 | y_true_sorted = y_true[score_arg_sort] 189 | y_true_sorted_cumsum = np.cumsum(y_true_sorted) 190 | 191 | # unique thresholds 192 | (thresholds,unique_indices) = np.unique( y_score_sorted , return_index=True ) 193 | num_prec_recall = len(unique_indices) + 1 194 | 195 | # prepare precision recall 196 | num_examples = len(y_score_sorted) 197 | num_true_examples = y_true_sorted_cumsum[-1] 198 | precision = np.zeros(num_prec_recall) 199 | recall = np.zeros(num_prec_recall) 200 | 201 | # deal with the first point 202 | y_true_sorted_cumsum = np.append( y_true_sorted_cumsum , 0 ) 203 | # deal with remaining 204 | for idx_res,idx_scores in enumerate(unique_indices): 205 | cumsum = y_true_sorted_cumsum[idx_scores-1] 206 | tp = num_true_examples - cumsum 207 | fp = num_examples - idx_scores - tp 208 | fn = cumsum + hard_false_negatives 209 | p = float(tp)/(tp+fp) 210 | r = float(tp)/(tp+fn) 211 | precision[idx_res] = p 212 | recall [idx_res] = r 213 | 214 | # first point in curve is artificial 215 | precision[-1] = 1. 216 | recall [-1] = 0. 217 | 218 | # compute average of precision-recall curve 219 | recall_for_conv = np.copy(recall) 220 | recall_for_conv = np.append(recall_for_conv[0], recall_for_conv) 221 | recall_for_conv = np.append(recall_for_conv, 0.) 222 | 223 | stepWidths = np.convolve(recall_for_conv,[-0.5,0,0.5],'valid') 224 | # integrate is now simply a dot product 225 | ap_current = np.dot(precision, stepWidths) 226 | 227 | elif has_gt: 228 | ap_current = 0.0 229 | else: 230 | ap_current = float('nan') 231 | ap[di,li,oi] = ap_current 232 | #exit(1) 233 | return ap 234 | 235 | def compute_averages(aps): 236 | d_inf = 0 237 | o50 = np.where(np.isclose(opt.overlaps,0.5)) 238 | o25 = np.where(np.isclose(opt.overlaps,0.25)) 239 | oAllBut25 = np.where(np.logical_not(np.isclose(opt.overlaps,0.25))) 240 | avg_dict = {} 241 | #avg_dict['all_ap'] = np.nanmean(aps[ d_inf,:,: ]) 242 | avg_dict['all_ap'] = np.nanmean(aps[ d_inf,:,oAllBut25]) 243 | avg_dict['all_ap_50%'] = np.nanmean(aps[ d_inf,:,o50]) 244 | avg_dict['all_ap_25%'] = np.nanmean(aps[ d_inf,:,o25]) 245 | avg_dict["classes"] = {} 246 | for (li,label_name) in enumerate(CLASS_LABELS): 247 | avg_dict["classes"][label_name] = {} 248 | #avg_dict["classes"][label_name]["ap"] = np.average(aps[ d_inf,li, :]) 249 | avg_dict["classes"][label_name]["ap"] = np.average(aps[ d_inf,li,oAllBut25]) 250 | avg_dict["classes"][label_name]["ap50%"] = np.average(aps[ d_inf,li,o50]) 251 | avg_dict["classes"][label_name]["ap25%"] = np.average(aps[ d_inf,li,o25]) 252 | return avg_dict 253 | 254 | 255 | def assign_instances_for_scan(pred_file, gt_file, pred_path): 256 | pred_info = util_3d.read_instance_prediction_file(pred_file, pred_path) 257 | try: 258 | pred_info = util_3d.read_instance_prediction_file(pred_file, pred_path) 259 | except: 260 | util.print_error('unable to load ' + pred_file) 261 | try: 262 | gt_ids = util_3d.load_ids(gt_file) 263 | except: 264 | util.print_error('unable to load ' + gt_file) 265 | # get gt instances 266 | gt_instances = util_3d.get_instances(gt_ids, VALID_CLASS_IDS, CLASS_LABELS, ID_TO_LABEL) 267 | # associate 268 | gt2pred = deepcopy(gt_instances) 269 | for label in gt2pred: 270 | for gt in gt2pred[label]: 271 | gt['matched_pred'] = [] 272 | pred2gt = {} 273 | for label in CLASS_LABELS: 274 | pred2gt[label] = [] 275 | num_pred_instances = 0 276 | # mask of void labels in the groundtruth 277 | bool_void = np.logical_not(np.in1d(gt_ids//1000, VALID_CLASS_IDS)) 278 | # go thru all prediction masks 279 | for pred_index, pred_mask_file in enumerate(pred_info): 280 | label_id = int(pred_info[pred_mask_file]['label_id']) 281 | conf = pred_info[pred_mask_file]['conf'] 282 | if not label_id in ID_TO_LABEL: 283 | continue 284 | label_name = ID_TO_LABEL[label_id] 285 | # read the mask 286 | pred_mask = util_3d.load_ids(pred_mask_file) 287 | if len(pred_mask) != len(gt_ids): 288 | util.print_error('wrong number of lines in ' + pred_mask_file + '(%d) vs #mesh vertices (%d), please double check and/or re-download the mesh' % (len(pred_mask), len(gt_ids))) 289 | # convert to binary 290 | pred_mask = np.not_equal(pred_mask, 0) 291 | num = np.count_nonzero(pred_mask) 292 | if num < opt.min_region_sizes[0]: 293 | continue # skip if empty 294 | 295 | pred_instance = {} 296 | pred_instance['filename'] = pred_mask_file 297 | pred_instance['pred_id'] = num_pred_instances 298 | pred_instance['label_id'] = label_id 299 | pred_instance['vert_count'] = num 300 | pred_instance['confidence'] = conf 301 | pred_instance['void_intersection'] = np.count_nonzero(np.logical_and(bool_void, pred_mask)) 302 | 303 | # matched gt instances 304 | matched_gt = [] 305 | # go thru all gt instances with matching label 306 | #print(gt2pred[label_name]) 307 | for (gt_num, gt_inst) in enumerate(gt2pred[label_name]): 308 | intersection = np.count_nonzero(np.logical_and(gt_ids == gt_inst['instance_id'], pred_mask)) 309 | #print((gt_ids == gt_inst['instance_id']).sum(), pred_mask.sum(), intersection) 310 | if intersection > 0: 311 | gt_copy = gt_inst.copy() 312 | pred_copy = pred_instance.copy() 313 | gt_copy['intersection'] = intersection 314 | pred_copy['intersection'] = intersection 315 | matched_gt.append(gt_copy) 316 | gt2pred[label_name][gt_num]['matched_pred'].append(pred_copy) 317 | #print(label_name, pred_index, num_pred_instances, num, conf, pred_instance['void_intersection'], len(matched_gt)) 318 | pred_instance['matched_gt'] = matched_gt 319 | num_pred_instances += 1 320 | pred2gt[label_name].append(pred_instance) 321 | return gt2pred, pred2gt 322 | 323 | 324 | def print_results(avgs): 325 | sep = "" 326 | col1 = ":" 327 | lineLen = 64 328 | 329 | print("") 330 | print("#"*lineLen) 331 | line = "" 332 | line += "{:<15}".format("what" ) + sep + col1 333 | line += "{:>15}".format("AP" ) + sep 334 | line += "{:>15}".format("AP_50%" ) + sep 335 | line += "{:>15}".format("AP_25%" ) + sep 336 | print(line) 337 | print("#"*lineLen) 338 | 339 | for (li,label_name) in enumerate(CLASS_LABELS): 340 | ap_avg = avgs["classes"][label_name]["ap"] 341 | ap_50o = avgs["classes"][label_name]["ap50%"] 342 | ap_25o = avgs["classes"][label_name]["ap25%"] 343 | line = "{:<15}".format(label_name + ' ' + str(VALID_CLASS_IDS[li])) + sep + col1 344 | line += sep + "{:>15.3f}".format(ap_avg ) + sep 345 | line += sep + "{:>15.3f}".format(ap_50o ) + sep 346 | line += sep + "{:>15.3f}".format(ap_25o ) + sep 347 | print(line) 348 | 349 | all_ap_avg = avgs["all_ap"] 350 | all_ap_50o = avgs["all_ap_50%"] 351 | all_ap_25o = avgs["all_ap_25%"] 352 | 353 | print("-"*lineLen) 354 | line = "{:<15}".format("average") + sep + col1 355 | line += "{:>15.3f}".format(all_ap_avg) + sep 356 | line += "{:>15.3f}".format(all_ap_50o) + sep 357 | line += "{:>15.3f}".format(all_ap_25o) + sep 358 | print(line) 359 | print("") 360 | 361 | 362 | def write_result_file(avgs, filename): 363 | _SPLITTER = ',' 364 | with open(filename, 'w') as f: 365 | f.write(_SPLITTER.join(['class', 'class id', 'ap', 'ap50', 'ap25']) + '\n') 366 | for i in range(len(VALID_CLASS_IDS)): 367 | class_name = CLASS_LABELS[i] 368 | class_id = VALID_CLASS_IDS[i] 369 | ap = avgs["classes"][class_name]["ap"] 370 | ap50 = avgs["classes"][class_name]["ap50%"] 371 | ap25 = avgs["classes"][class_name]["ap25%"] 372 | f.write(_SPLITTER.join([str(x) for x in [class_name, class_id, ap, ap50, ap25]]) + '\n') 373 | 374 | 375 | def evaluate(pred_files, gt_files, pred_path, output_file): 376 | print('evaluating', len(pred_files), 'scans...') 377 | matches = {} 378 | for i in range(len(pred_files)): 379 | matches_key = os.path.abspath(gt_files[i]) 380 | # assign gt to predictions 381 | gt2pred, pred2gt = assign_instances_for_scan(pred_files[i], gt_files[i], pred_path) 382 | matches[matches_key] = {} 383 | matches[matches_key]['gt'] = gt2pred 384 | matches[matches_key]['pred'] = pred2gt 385 | sys.stdout.write("\rscans processed: {}".format(i+1)) 386 | sys.stdout.flush() 387 | print('') 388 | ap_scores = evaluate_matches(matches) 389 | avgs = compute_averages(ap_scores) 390 | 391 | # print 392 | print_results(avgs) 393 | write_result_file(avgs, output_file) 394 | 395 | 396 | def main(): 397 | pred_files = [f for f in os.listdir(opt.pred_path) if f.endswith('.txt') and f != 'semantic_instance_evaluation.txt'] 398 | #print(pred_files) 399 | #['scene0217_00.txt', 'scene0488_01.txt', 'scene0412_00.txt', 'scene0412_01.txt', 'scene0488_00.txt', 'scene0304_00.txt', 'scene0019_00.txt', 'scene0568_00.txt', 'scene0414_00.txt', 'scene0019_01.txt', 'scene0426_00.txt', 'scene0568_02.txt', 'scene0426_02.txt', 'scene0426_01.txt', 'scene0426_03.txt', 'scene0568_01.txt', 'scene0575_02.txt', 'scene0575_00.txt', 'scene0575_01.txt', 'scene0549_00.txt'] 400 | if opt.scene_id != '': 401 | pred_files = [filename for filename in pred_files if opt.scene_id in filename] 402 | pass 403 | if opt.num_testing_images > 0: 404 | pred_files = pred_files[:opt.num_testing_images] 405 | print(pred_files) 406 | pass 407 | 408 | gt_files = [] 409 | if len(pred_files) == 0: 410 | util.print_error('No result files found.', user_fault=True) 411 | for i in range(len(pred_files)): 412 | gt_file = os.path.join(opt.gt_path, pred_files[i]) 413 | if not os.path.isfile(gt_file): 414 | util.print_error('Result file {} does not match any gt file'.format(pred_files[i]), user_fault=True) 415 | gt_files.append(gt_file) 416 | pred_files[i] = os.path.join(opt.pred_path, pred_files[i]) 417 | 418 | # evaluate 419 | evaluate(pred_files, gt_files, opt.pred_path, opt.output_file) 420 | 421 | 422 | if __name__ == '__main__': 423 | main() 424 | -------------------------------------------------------------------------------- /scripts/evaluate_semantic_label.py: -------------------------------------------------------------------------------- 1 | # Evaluates semantic label task 2 | # Input: 3 | # - path to .txt prediction files 4 | # - path to .txt ground truth files 5 | # - output file to write results to 6 | # Note that only the valid classes are used for evaluation, 7 | # i.e., any ground truth label not in the valid label set 8 | # is ignored in the evaluation. 9 | # 10 | # example usage: evaluate_semantic_label.py --scan_path [path to scan data] --output_file [output file] 11 | 12 | # python imports 13 | import math 14 | import os, sys, argparse 15 | import inspect 16 | 17 | try: 18 | import numpy as np 19 | except: 20 | print("Failed to import numpy package.") 21 | sys.exit(-1) 22 | try: 23 | from itertools import izip 24 | except ImportError: 25 | izip = zip 26 | 27 | currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) 28 | parentdir = os.path.dirname(currentdir) 29 | sys.path.insert(0,parentdir) 30 | import util 31 | import util_3d 32 | 33 | parser = argparse.ArgumentParser() 34 | parser.add_argument('--pred_path', required=True, help='path to directory of predicted .txt files') 35 | parser.add_argument('--gt_path', default='semantic_val', help='path to gt files') 36 | parser.add_argument('--output_file', default='', help='output file [default: pred_path/semantic_label_evaluation.txt]') 37 | opt = parser.parse_args() 38 | 39 | if opt.output_file == '': 40 | opt.output_file = os.path.join(opt.pred_path, 'semantic_label_evaluation.txt') 41 | 42 | 43 | CLASS_LABELS = ['wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window', 'bookshelf', 'picture', 'counter', 'desk', 'curtain', 'refrigerator', 'shower curtain', 'toilet', 'sink', 'bathtub', 'otherfurniture'] 44 | VALID_CLASS_IDS = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39]) 45 | UNKNOWN_ID = np.max(VALID_CLASS_IDS) + 1 46 | 47 | 48 | def evaluate_scan(pred_file, gt_file, confusion): 49 | try: 50 | pred_ids = util_3d.load_ids(pred_file) 51 | except: 52 | util.print_error('unable to load ' + pred_file + ': ' + str(e)) 53 | try: 54 | gt_ids = util_3d.load_ids(gt_file) 55 | except: 56 | util.print_error('unable to load ' + gt_file + ': ' + str(e)) 57 | # sanity checks 58 | if not pred_ids.shape == gt_ids.shape: 59 | util.print_error('%s: number of predicted values does not match number of vertices' % pred_file, user_fault=True) 60 | for (gt_val,pred_val) in izip(gt_ids.flatten(),pred_ids.flatten()): 61 | if gt_val not in VALID_CLASS_IDS: 62 | continue 63 | if pred_val not in VALID_CLASS_IDS: 64 | pred_val = UNKNOWN_ID 65 | confusion[gt_val][pred_val] += 1 66 | 67 | 68 | def get_iou(label_id, confusion): 69 | if not label_id in VALID_CLASS_IDS: 70 | return float('nan') 71 | # #true positives 72 | tp = np.longlong(confusion[label_id, label_id]) 73 | # #false negatives 74 | fn = np.longlong(confusion[label_id, :].sum()) - tp 75 | # #false positives 76 | not_ignored = [l for l in VALID_CLASS_IDS if not l == label_id] 77 | fp = np.longlong(confusion[not_ignored, label_id].sum()) 78 | 79 | denom = (tp + fp + fn) 80 | if denom == 0: 81 | #return float('nan') 82 | return (-1, 0, 0) 83 | return (float(tp) / denom, tp, denom) 84 | 85 | 86 | def write_result_file(confusion, ious, filename): 87 | with open(filename, 'w') as f: 88 | f.write('iou scores\n') 89 | for i in range(len(VALID_CLASS_IDS)): 90 | label_id = VALID_CLASS_IDS[i] 91 | label_name = CLASS_LABELS[i] 92 | #print(label_name, ious[label_name]) 93 | iou = ious[label_name][0] 94 | f.write('{0:<14s}({1:<2d}): {2:>5.3f}\n'.format(label_name, label_id, iou)) 95 | f.write('\nconfusion matrix\n') 96 | f.write('\t\t\t') 97 | for i in range(len(VALID_CLASS_IDS)): 98 | #f.write('\t{0:<14s}({1:<2d})'.format(CLASS_LABELS[i], VALID_CLASS_IDS[i])) 99 | f.write('{0:<8d}'.format(VALID_CLASS_IDS[i])) 100 | f.write('\n') 101 | for r in range(len(VALID_CLASS_IDS)): 102 | f.write('{0:<14s}({1:<2d})'.format(CLASS_LABELS[r], VALID_CLASS_IDS[r])) 103 | for c in range(len(VALID_CLASS_IDS)): 104 | f.write('\t{0:>5.3f}'.format(confusion[VALID_CLASS_IDS[r],VALID_CLASS_IDS[c]])) 105 | f.write('\n') 106 | print('wrote results to', filename) 107 | 108 | 109 | def evaluate(pred_files, gt_files, output_file): 110 | max_id = UNKNOWN_ID 111 | confusion = np.zeros((max_id+1, max_id+1), dtype=np.ulonglong) 112 | 113 | print('evaluating', len(pred_files), 'scans...') 114 | for i in range(len(pred_files)): 115 | evaluate_scan(pred_files[i], gt_files[i], confusion) 116 | sys.stdout.write("\rscans processed: {}".format(i+1)) 117 | sys.stdout.flush() 118 | print('') 119 | 120 | class_ious = {} 121 | for i in range(len(VALID_CLASS_IDS)): 122 | label_name = CLASS_LABELS[i] 123 | label_id = VALID_CLASS_IDS[i] 124 | class_ious[label_name] = get_iou(label_id, confusion) 125 | # print 126 | print('classes IoU') 127 | print('----------------------------') 128 | for i in range(len(VALID_CLASS_IDS)): 129 | label_name = CLASS_LABELS[i] 130 | print(label_name, class_ious[label_name]) 131 | continue 132 | print('mean IOU', np.mean(list(class_ious.values()), axis=0)) 133 | # for i in range(len(VALID_CLASS_IDS)): 134 | # label_name = CLASS_LABELS[i] 135 | # #print('{{0:<14s}: 1:>5.3f}'.format(label_name, class_ious[label_name][0])) 136 | # #print('{0:<14s}: {1:>5.3f} ({2:>6d}/{3:<6d})'.format(label_name, class_ious[label_name][0], class_ious[label_name][1], class_ious[label_name][2])) 137 | # print(label_name, class_ious) 138 | # print(label_name, class_ious[label_name]) 139 | write_result_file(confusion, class_ious, output_file) 140 | 141 | 142 | def main(): 143 | pred_files = [f for f in os.listdir(opt.pred_path) if f.endswith('.txt') and f != 'semantic_label_evaluation.txt'] 144 | gt_files = [] 145 | if len(pred_files) == 0: 146 | util.print_error('No result files found.', user_fault=True) 147 | for i in range(len(pred_files)): 148 | gt_file = os.path.join(opt.gt_path, pred_files[i]) 149 | if not os.path.isfile(gt_file): 150 | util.print_error('Result file {} does not match any gt file'.format(pred_files[i]), user_fault=True) 151 | gt_files.append(gt_file) 152 | pred_files[i] = os.path.join(opt.pred_path, pred_files[i]) 153 | 154 | # evaluate 155 | evaluate(pred_files, gt_files, opt.output_file) 156 | 157 | 158 | if __name__ == '__main__': 159 | main() 160 | -------------------------------------------------------------------------------- /scripts/export_script.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | dataFolder = '/gruvi/Data/chenliu/ScanNet/scans/' 4 | split = 'val' 5 | with open('split_' + split + '.txt', 'r') as f: 6 | for line in f: 7 | scene_id = line.strip() 8 | if len(scene_id) != 12: 9 | continue 10 | filename = dataFolder + '/' + scene_id 11 | output_filename = 'instance_val/' + scene_id + '.txt' 12 | print('python scripts/export_train_mesh_for_evaluation.py --scan_path=' + filename + ' --output_file=' + output_filename) 13 | os.system('python scripts/export_train_mesh_for_evaluation.py --scan_path=' + filename + ' --output_file=' + output_filename) 14 | continue 15 | pass 16 | -------------------------------------------------------------------------------- /scripts/export_train_mesh_for_evaluation.py: -------------------------------------------------------------------------------- 1 | # Exports a train scan in the evaluation format using: 2 | # - the *_vh_clean_2.ply mesh 3 | # - the labels defined by the *.aggregation.json and *_vh_clean_2.0.010000.segs.json files 4 | # 5 | # example usage: export_train_mesh_for_evaluation.py --scan_path [path to scan data] --output_file [output file] --type label 6 | # Note: technically does not need to load in the ply file, since the ScanNet annotations are defined against the mesh vertices, but we load it in here as an example. 7 | 8 | # python imports 9 | import math 10 | import os, sys, argparse 11 | import inspect 12 | import json 13 | import glob 14 | 15 | try: 16 | import numpy as np 17 | except: 18 | print("Failed to import numpy package.") 19 | sys.exit(-1) 20 | 21 | currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) 22 | parentdir = os.path.dirname(currentdir) 23 | sys.path.insert(0,parentdir) 24 | import util 25 | import util_3d 26 | import torch 27 | import multiprocessing as mp 28 | 29 | TASK_TYPES = {'label', 'instance'} 30 | 31 | parser = argparse.ArgumentParser() 32 | parser.add_argument('--scan_path', required=True, help='path to scannet scene (e.g., data/ScanNet/v2/scene0000_00') 33 | parser.add_argument('--output_file', required=True, help='output file') 34 | parser.add_argument('--label_map_file', default='/gruvi/Data/chenliu/ScanNet/tasks/scannetv2-labels.combined.tsv', help='path to scannetv2-labels.combined.tsv') 35 | parser.add_argument('--type', default='instance', help='task type [label or instance]') 36 | opt = parser.parse_args() 37 | assert opt.type in TASK_TYPES 38 | 39 | label_map = util.read_label_mapping(opt.label_map_file, label_from='raw_category', label_to='nyu40id') 40 | # remapper=np.ones(150)*(-100) 41 | # for i,x in enumerate([1,2,3,4,5,6,7,8,9,10,11,12,14,16,24,28,33,34,36,39]): 42 | # remapper[x]=i 43 | 44 | 45 | def read_aggregation(filename): 46 | assert os.path.isfile(filename) 47 | object_id_to_segs = {} 48 | label_to_segs = {} 49 | with open(filename) as f: 50 | data = json.load(f) 51 | num_objects = len(data['segGroups']) 52 | for i in range(num_objects): 53 | object_id = data['segGroups'][i]['objectId'] + 1 # instance ids should be 1-indexed 54 | label = data['segGroups'][i]['label'] 55 | segs = data['segGroups'][i]['segments'] 56 | object_id_to_segs[object_id] = segs 57 | if label in label_to_segs: 58 | label_to_segs[label].extend(segs) 59 | else: 60 | label_to_segs[label] = segs 61 | return object_id_to_segs, label_to_segs 62 | 63 | 64 | def read_segmentation(filename): 65 | assert os.path.isfile(filename) 66 | seg_to_verts = {} 67 | with open(filename) as f: 68 | data = json.load(f) 69 | num_verts = len(data['segIndices']) 70 | for i in range(num_verts): 71 | seg_id = data['segIndices'][i] 72 | if seg_id in seg_to_verts: 73 | seg_to_verts[seg_id].append(i) 74 | else: 75 | seg_to_verts[seg_id] = [i] 76 | return seg_to_verts, num_verts 77 | 78 | def export(mesh_file, agg_file, seg_file, label_map_file, type, output_file): 79 | label_map = util.read_label_mapping(opt.label_map_file, label_from='raw_category', label_to='nyu40id') 80 | mesh_vertices = util_3d.read_mesh_vertices(mesh_file) 81 | object_id_to_segs, label_to_segs = read_aggregation(agg_file) 82 | seg_to_verts, num_verts = read_segmentation(seg_file) 83 | label_ids = np.zeros(shape=(num_verts), dtype=np.uint32) # 0: unannotated 84 | for label, segs in label_to_segs.items(): 85 | label_id = label_map[label] 86 | for seg in segs: 87 | verts = seg_to_verts[seg] 88 | label_ids[verts] = label_id 89 | if type == 'label': 90 | util_3d.export_ids(output_file, label_ids) 91 | elif type == 'instance': 92 | instance_ids = np.zeros(shape=(num_verts), dtype=np.uint32) # 0: unannotated 93 | for object_id, segs in object_id_to_segs.items(): 94 | for seg in segs: 95 | verts = seg_to_verts[seg] 96 | instance_ids[verts] = object_id 97 | util_3d.export_ids(output_file, label_ids * 1000 + instance_ids) 98 | #util_3d.export_instance_ids_for_eval(output_file, label_ids, instance_ids) 99 | else: 100 | raise 101 | 102 | 103 | def main(): 104 | scan_name = os.path.split(opt.scan_path)[-1] 105 | mesh_file = os.path.join(opt.scan_path, scan_name + '_vh_clean_2.ply') 106 | agg_file = os.path.join(opt.scan_path, scan_name + '.aggregation.json') 107 | seg_file = os.path.join(opt.scan_path, scan_name + '_vh_clean_2.0.010000.segs.json') 108 | export(mesh_file, agg_file, seg_file, opt.label_map_file, opt.type, opt.output_file) 109 | 110 | 111 | if __name__ == '__main__': 112 | main() 113 | -------------------------------------------------------------------------------- /scripts/prepare_data.py: -------------------------------------------------------------------------------- 1 | # Exports a train scan in the evaluation format using: 2 | # - the *_vh_clean_2.ply mesh 3 | # - the labels defined by the *.aggregation.json and *_vh_clean_2.0.010000.segs.json files 4 | # 5 | # example usage: export_train_mesh_for_evaluation.py --scan_path [path to scan data] --output_file [output file] --type label 6 | # Note: technically does not need to load in the ply file, since the ScanNet annotations are defined against the mesh vertices, but we load it in here as an example. 7 | 8 | # python imports 9 | import math 10 | import os, sys 11 | import inspect 12 | import json 13 | import glob 14 | 15 | try: 16 | import numpy as np 17 | except: 18 | print("Failed to import numpy package.") 19 | sys.exit(-1) 20 | 21 | # currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) 22 | # parentdir = os.path.dirname(currentdir) 23 | # sys.path.insert(0,parentdir) 24 | from scripts.util import read_label_mapping 25 | from scripts.util_3d import read_mesh_vertices 26 | import torch 27 | import multiprocessing as mp 28 | import functools 29 | #from utils import write_ply_label 30 | 31 | def read_aggregation(filename): 32 | assert os.path.isfile(filename) 33 | object_id_to_segs = {} 34 | label_to_segs = {} 35 | with open(filename) as f: 36 | data = json.load(f) 37 | num_objects = len(data['segGroups']) 38 | for i in range(num_objects): 39 | object_id = data['segGroups'][i]['objectId'] + 1 # instance ids should be 1-indexed 40 | label = data['segGroups'][i]['label'] 41 | segs = data['segGroups'][i]['segments'] 42 | object_id_to_segs[object_id] = segs 43 | if label in label_to_segs: 44 | label_to_segs[label].extend(segs) 45 | else: 46 | label_to_segs[label] = segs 47 | return object_id_to_segs, label_to_segs 48 | 49 | 50 | def read_segmentation(filename): 51 | assert os.path.isfile(filename) 52 | seg_to_verts = {} 53 | with open(filename) as f: 54 | data = json.load(f) 55 | num_verts = len(data['segIndices']) 56 | for i in range(num_verts): 57 | seg_id = data['segIndices'][i] 58 | if seg_id in seg_to_verts: 59 | seg_to_verts[seg_id].append(i) 60 | else: 61 | seg_to_verts[seg_id] = [i] 62 | return seg_to_verts, num_verts 63 | 64 | 65 | def export(filename, label_map): 66 | scan_name = filename.split('_vh')[0] 67 | mesh_file = os.path.join(scan_name + '_vh_clean_2.ply') 68 | agg_file = os.path.join(scan_name + '.aggregation.json') 69 | seg_file = os.path.join(scan_name + '_vh_clean_2.0.010000.segs.json') 70 | 71 | print(filename) 72 | if os.path.exists(mesh_file[:-4] + '.pth'): 73 | return 74 | 75 | mesh_vertices, mesh_colors, faces = read_mesh_vertices(mesh_file) 76 | if os.path.exists(agg_file): 77 | object_id_to_segs, label_to_segs = read_aggregation(agg_file) 78 | seg_to_verts, num_verts = read_segmentation(seg_file) 79 | label_ids = np.zeros(shape=(num_verts), dtype=np.uint32) # 0: unannotated 80 | instance_ids = np.zeros(shape=(num_verts), dtype=np.uint32) # 0: unannotated 81 | 82 | # write_ply_label('test/mesh.ply', mesh_vertices, faces, label_ids) 83 | # exit(1) 84 | 85 | for label, segs in label_to_segs.items(): 86 | label_id = label_map[label] 87 | for seg in segs: 88 | verts = seg_to_verts[seg] 89 | label_ids[verts] = label_id 90 | 91 | for object_id, segs in object_id_to_segs.items(): 92 | for seg in segs: 93 | verts = seg_to_verts[seg] 94 | instance_ids[verts] = object_id 95 | else: 96 | num_verts = len(mesh_vertices) 97 | label_ids = np.zeros(shape=(num_verts), dtype=np.uint32) # 0: unannotated 98 | instance_ids = np.zeros(shape=(num_verts), dtype=np.uint32) # 0: unannotated 99 | pass 100 | mesh_vertices = np.ascontiguousarray(mesh_vertices - mesh_vertices.mean(0)) 101 | mesh_colors = np.ascontiguousarray(mesh_colors) / 127.5 - 1 102 | torch.save((mesh_vertices, mesh_colors, label_ids, instance_ids, faces), mesh_file[:-4] + '.pth') 103 | return 104 | 105 | def prepare_data(options): 106 | ROOT_FOLDER = options.dataFolder 107 | files = sorted(glob.glob(options.dataFolder + '*/*_vh_clean_2.ply')) 108 | p = mp.Pool(processes=mp.cpu_count()) 109 | 110 | label_map = read_label_mapping(options.labelFile, label_from='raw_category', label_to='nyu40id') 111 | 112 | p.map(functools.partial(export, label_map=label_map), files) 113 | p.close() 114 | p.join() 115 | return 116 | -------------------------------------------------------------------------------- /scripts/prepare_data_ori.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016-present, Facebook, Inc. 2 | # All rights reserved. 3 | # 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import glob, plyfile, numpy as np, multiprocessing as mp, torch 8 | import os 9 | 10 | # Map relevant classes to {0,1,...,19}, and ignored classes to -100 11 | remapper=np.ones(150)*(-100) 12 | for i,x in enumerate([1,2,3,4,5,6,7,8,9,10,11,12,14,16,24,28,33,34,36,39]): 13 | remapper[x]=i 14 | 15 | ROOT_FOLDER = '/gruvi/Data/chenliu/ScanNet/scans/' 16 | files=sorted(glob.glob(ROOT_FOLDER + '*/*_vh_clean_2.ply')) 17 | #files2=sorted(glob.glob(ROOT_FOLDER + '*/*_vh_clean_2.labels.ply')) 18 | #assert len(files) == len(files2) 19 | 20 | def f(fn): 21 | print(fn[:-4] + '.pth') 22 | if os.path.exists(fn[:-4] + '.pth'): 23 | return 24 | #fn2 = fn[:-3]+'labels.ply' 25 | a=plyfile.PlyData().read(fn) 26 | v=np.array([list(x) for x in a.elements[0]]) 27 | coords=np.ascontiguousarray(v[:,:3]-v[:,:3].mean(0)) 28 | colors=np.ascontiguousarray(v[:,3:6])/127.5-1 29 | #a=plyfile.PlyData().read(fn2) 30 | #w=remapper[np.array(a.elements[0]['label'])] 31 | 32 | #filename = ROOT_FOLDER + scene_id + '/' + scene_id + '.aggregation.json' 33 | data = json.load(open(fn.replace('vh_clean_2.ply', 'aggregation.json'), 'r')) 34 | aggregation = np.array(data['segGroups']) 35 | 36 | data = json.load(open(fn.replace('vh_clean_2.ply', 'vh_clean_2.0.010000.segs.json'), 'r')) 37 | segmentation = np.array(data['segIndices']) 38 | 39 | groupSegments = [] 40 | groupLabels = [] 41 | for segmentIndex in xrange(len(aggregation)): 42 | groupSegments.append(aggregation[segmentIndex]['segments']) 43 | groupLabels.append(aggregation[segmentIndex]['label']) 44 | continue 45 | 46 | segmentation = segmentation.astype(np.int32) 47 | torch.save((coords,colors,w, instance),fn[:-4]+'.pth') 48 | 49 | p = mp.Pool(processes=mp.cpu_count()) 50 | p.map(f,files) 51 | p.close() 52 | p.join() 53 | -------------------------------------------------------------------------------- /scripts/util.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | import csv 3 | try: 4 | import numpy as np 5 | except: 6 | print("Failed to import numpy package.") 7 | sys.exit(-1) 8 | try: 9 | import imageio 10 | except: 11 | print("Please install the module 'imageio' for image processing, e.g.") 12 | print("pip install imageio") 13 | sys.exit(-1) 14 | 15 | # print an error message and quit 16 | def print_error(message, user_fault=False): 17 | sys.stderr.write('ERROR: ' + str(message) + '\n') 18 | if user_fault: 19 | sys.exit(2) 20 | sys.exit(-1) 21 | 22 | 23 | # if string s represents an int 24 | def represents_int(s): 25 | try: 26 | int(s) 27 | return True 28 | except ValueError: 29 | return False 30 | 31 | 32 | def read_label_mapping(filename, label_from='raw_category', label_to='nyu40id'): 33 | assert os.path.isfile(filename) 34 | mapping = dict() 35 | with open(filename) as csvfile: 36 | reader = csv.DictReader(csvfile, delimiter='\t') 37 | for row in reader: 38 | mapping[row[label_from]] = int(row[label_to]) 39 | # if ints convert 40 | if represents_int(list(mapping.keys())[0]): 41 | mapping = {int(k):v for k,v in mapping.items()} 42 | return mapping 43 | 44 | 45 | # input: scene_types.txt or scene_types_all.txt 46 | def read_scene_types_mapping(filename, remove_spaces=True): 47 | assert os.path.isfile(filename) 48 | mapping = dict() 49 | lines = open(filename).read().splitlines() 50 | lines = [line.split('\t') for line in lines] 51 | if remove_spaces: 52 | mapping = { x[1].strip():int(x[0]) for x in lines } 53 | else: 54 | mapping = { x[1]:int(x[0]) for x in lines } 55 | return mapping 56 | 57 | 58 | # color by label 59 | def visualize_label_image(filename, image): 60 | height = image.shape[0] 61 | width = image.shape[1] 62 | vis_image = np.zeros([height, width, 3], dtype=np.uint8) 63 | color_palette = create_color_palette() 64 | for idx, color in enumerate(color_palette): 65 | vis_image[image==idx] = color 66 | imageio.imwrite(filename, vis_image) 67 | 68 | 69 | # color by different instances (mod length of color palette) 70 | def visualize_instance_image(filename, image): 71 | height = image.shape[0] 72 | width = image.shape[1] 73 | vis_image = np.zeros([height, width, 3], dtype=np.uint8) 74 | color_palette = create_color_palette() 75 | instances = np.unique(image) 76 | for idx, inst in enumerate(instances): 77 | vis_image[image==inst] = color_palette[inst%len(color_palette)] 78 | imageio.imwrite(filename, vis_image) 79 | 80 | 81 | # color palette for nyu40 labels 82 | def create_color_palette(): 83 | return [ 84 | (0, 0, 0), 85 | (174, 199, 232), # wall 86 | (152, 223, 138), # floor 87 | (31, 119, 180), # cabinet 88 | (255, 187, 120), # bed 89 | (188, 189, 34), # chair 90 | (140, 86, 75), # sofa 91 | (255, 152, 150), # table 92 | (214, 39, 40), # door 93 | (197, 176, 213), # window 94 | (148, 103, 189), # bookshelf 95 | (196, 156, 148), # picture 96 | (23, 190, 207), # counter 97 | (178, 76, 76), 98 | (247, 182, 210), # desk 99 | (66, 188, 102), 100 | (219, 219, 141), # curtain 101 | (140, 57, 197), 102 | (202, 185, 52), 103 | (51, 176, 203), 104 | (200, 54, 131), 105 | (92, 193, 61), 106 | (78, 71, 183), 107 | (172, 114, 82), 108 | (255, 127, 14), # refrigerator 109 | (91, 163, 138), 110 | (153, 98, 156), 111 | (140, 153, 101), 112 | (158, 218, 229), # shower curtain 113 | (100, 125, 154), 114 | (178, 127, 135), 115 | (120, 185, 128), 116 | (146, 111, 194), 117 | (44, 160, 44), # toilet 118 | (112, 128, 144), # sink 119 | (96, 207, 209), 120 | (227, 119, 194), # bathtub 121 | (213, 92, 176), 122 | (94, 106, 211), 123 | (82, 84, 163), # otherfurn 124 | (100, 85, 144) 125 | ] 126 | -------------------------------------------------------------------------------- /scripts/util_3d.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | import json 3 | 4 | try: 5 | import numpy as np 6 | except: 7 | print("Failed to import numpy package.") 8 | sys.exit(-1) 9 | 10 | try: 11 | from plyfile import PlyData, PlyElement 12 | except: 13 | print("Please install the module 'plyfile' for PLY i/o, e.g.") 14 | print("pip install plyfile") 15 | sys.exit(-1) 16 | 17 | import scripts.util 18 | 19 | 20 | # matrix: 4x4 np array 21 | # points Nx3 np array 22 | def transform_points(matrix, points): 23 | assert len(points.shape) == 2 and points.shape[1] == 3 24 | num_points = points.shape[0] 25 | p = np.concatenate([points, np.ones((num_points, 1))], axis=1) 26 | p = np.matmul(matrix, np.transpose(p)) 27 | p = np.transpose(p) 28 | p[:,:3] /= p[:,3,None] 29 | return p[:,:3] 30 | 31 | 32 | def export_ids(filename, ids): 33 | with open(filename, 'w') as f: 34 | for id in ids: 35 | f.write('%d\n' % id) 36 | 37 | 38 | def load_ids(filename): 39 | ids = open(filename).read().splitlines() 40 | ids = np.array(ids, dtype=np.int64) 41 | return ids 42 | 43 | 44 | def read_mesh_vertices(filename): 45 | assert os.path.isfile(filename) 46 | with open(filename, 'rb') as f: 47 | plydata = PlyData.read(f) 48 | num_verts = plydata['vertex'].count 49 | vertices = np.zeros(shape=[num_verts, 3], dtype=np.float32) 50 | vertices[:,0] = plydata['vertex'].data['x'] 51 | vertices[:,1] = plydata['vertex'].data['y'] 52 | vertices[:,2] = plydata['vertex'].data['z'] 53 | colors = np.zeros(shape=[num_verts, 3], dtype=np.float32) 54 | colors[:,0] = plydata['vertex'].data['red'] 55 | colors[:,1] = plydata['vertex'].data['green'] 56 | colors[:,2] = plydata['vertex'].data['blue'] 57 | faces = np.stack(plydata['face'].data['vertex_indices'].tolist(), axis=0) 58 | return vertices, colors, faces 59 | 60 | 61 | # export 3d instance labels for instance evaluation 62 | def export_instance_ids_for_eval(filename, label_ids, instance_ids): 63 | assert label_ids.shape[0] == instance_ids.shape[0] 64 | output_mask_path_relative = 'pred_mask' 65 | name = os.path.splitext(os.path.basename(filename))[0] 66 | output_mask_path = os.path.join(os.path.dirname(filename), output_mask_path_relative) 67 | if not os.path.isdir(output_mask_path): 68 | os.mkdir(output_mask_path) 69 | insts = np.unique(instance_ids) 70 | zero_mask = np.zeros(shape=(instance_ids.shape[0]), dtype=np.int32) 71 | with open(filename, 'w') as f: 72 | for idx, inst_id in enumerate(insts): 73 | if inst_id == 0: # 0 -> no instance for this vertex 74 | continue 75 | #output_mask_file = os.path.join(output_mask_path_relative, name + '_' + str(idx) + '.txt') 76 | output_mask_file = os.path.join(output_mask_path, name + '_' + str(idx) + '.txt') 77 | loc = np.where(instance_ids == inst_id) 78 | label_id = label_ids[loc[0][0]] 79 | f.write('%s %d %f\n' % (output_mask_file, label_id, 1.0)) 80 | # write mask 81 | mask = np.copy(zero_mask) 82 | mask[loc[0]] = 1 83 | export_ids(output_mask_file, mask) 84 | 85 | 86 | # ------------ Instance Utils ------------ # 87 | 88 | class Instance(object): 89 | instance_id = 0 90 | label_id = 0 91 | vert_count = 0 92 | med_dist = -1 93 | dist_conf = 0.0 94 | 95 | def __init__(self, mesh_vert_instances, instance_id): 96 | if (instance_id == -1): 97 | return 98 | self.instance_id = int(instance_id) 99 | self.label_id = int(self.get_label_id(instance_id)) 100 | self.vert_count = int(self.get_instance_verts(mesh_vert_instances, instance_id)) 101 | 102 | def get_label_id(self, instance_id): 103 | return int(instance_id // 1000) 104 | 105 | def get_instance_verts(self, mesh_vert_instances, instance_id): 106 | return (mesh_vert_instances == instance_id).sum() 107 | 108 | def to_json(self): 109 | return json.dumps(self, default=lambda o: o.__dict__, sort_keys=True, indent=4) 110 | 111 | def to_dict(self): 112 | dict = {} 113 | dict["instance_id"] = self.instance_id 114 | dict["label_id"] = self.label_id 115 | dict["vert_count"] = self.vert_count 116 | dict["med_dist"] = self.med_dist 117 | dict["dist_conf"] = self.dist_conf 118 | return dict 119 | 120 | def from_json(self, data): 121 | self.instance_id = int(data["instance_id"]) 122 | self.label_id = int(data["label_id"]) 123 | self.vert_count = int(data["vert_count"]) 124 | if ("med_dist" in data): 125 | self.med_dist = float(data["med_dist"]) 126 | self.dist_conf = float(data["dist_conf"]) 127 | 128 | def __str__(self): 129 | return "("+str(self.instance_id)+")" 130 | 131 | def read_instance_prediction_file(filename, pred_path): 132 | lines = open(filename).read().splitlines() 133 | instance_info = {} 134 | abs_pred_path = os.path.abspath(pred_path) 135 | for line in lines: 136 | parts = line.split(' ') 137 | if len(parts) != 3: 138 | util.print_error('invalid instance prediction file. Expected (per line): [rel path prediction] [label id prediction] [confidence prediction]') 139 | if os.path.isabs(parts[0]): 140 | util.print_error('invalid instance prediction file. First entry in line must be a relative path') 141 | mask_file = os.path.join(os.path.dirname(filename), parts[0]) 142 | mask_file = os.path.abspath(mask_file) 143 | # check that mask_file lives inside prediction path 144 | if os.path.commonprefix([mask_file, abs_pred_path]) != abs_pred_path: 145 | util.print_error('predicted mask {} in prediction text file {} points outside of prediction path.'.format(mask_file,filename)) 146 | 147 | #print(len(instance_info), int(float(parts[1])), np.loadtxt(mask_file).sum()) 148 | info = {} 149 | info["label_id"] = int(float(parts[1])) 150 | info["conf"] = float(parts[2]) 151 | instance_info[mask_file] = info 152 | return instance_info 153 | 154 | 155 | def get_instances(ids, class_ids, class_labels, id2label): 156 | instances = {} 157 | for label in class_labels: 158 | instances[label] = [] 159 | instance_ids = np.unique(ids) 160 | for id in instance_ids: 161 | if id == 0: 162 | continue 163 | inst = Instance(ids, id) 164 | if inst.label_id in class_ids: 165 | instances[id2label[inst.label_id]].append(inst.to_dict()) 166 | return instances 167 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DataLoader 3 | 4 | from tqdm import tqdm 5 | import numpy as np 6 | import os 7 | import cv2 8 | 9 | from utils import * 10 | from options import parse_args 11 | 12 | from models.instance import Model, CoordAugmentation, NeighborGT 13 | 14 | from datasets.scannet_dataset import ScanNetDataset 15 | from scripts.prepare_data import prepare_data 16 | 17 | torch.set_printoptions(threshold=5000) 18 | 19 | 20 | class_weights = np.zeros(20, dtype=np.float32) 21 | semantic_counts = np.load('datasets/semantic_counts_pixelwise.npy') 22 | for i, x in enumerate(label_subset): 23 | class_weights[i] = semantic_counts[x] 24 | continue 25 | class_weights = np.log(class_weights.sum() / class_weights) 26 | class_weights = class_weights / class_weights.sum() 27 | class_weights = torch.from_numpy(class_weights).cuda() 28 | 29 | 30 | def main(options): 31 | if not os.path.exists(options.checkpoint_dir): 32 | os.system("mkdir -p %s"%options.checkpoint_dir) 33 | pass 34 | if not os.path.exists(options.test_dir): 35 | os.system("mkdir -p %s"%options.test_dir) 36 | pass 37 | 38 | 39 | model = Model(options) 40 | model.cuda() 41 | model.train() 42 | 43 | neighbor_model = NeighborGT(options) 44 | neighbor_model.cuda() 45 | 46 | augmentation_model = CoordAugmentation(options) 47 | augmentation_model.cuda() 48 | augmentation_model.train() 49 | 50 | 51 | if options.restore == 1: 52 | print('restore') 53 | if options.startEpoch >= 0: 54 | model.load_state_dict(torch.load(options.checkpoint_dir + '/checkpoint_' + str(options.startEpoch) + '.pth')) 55 | else: 56 | model.load_state_dict(torch.load(options.checkpoint_dir + '/checkpoint.pth')) 57 | pass 58 | pass 59 | 60 | dataset_val = ScanNetDataset(options, split='val', random=False) 61 | if options.task == 'test': 62 | testOneEpoch(options, model, neighbor_model, augmentation_model, dataset_val, validation=False) 63 | exit(1) 64 | pass 65 | 66 | optimizer = torch.optim.Adam(model.parameters(), lr = options.LR) 67 | 68 | if options.restore == 1 and os.path.exists(options.checkpoint_dir + '/optim.pth'): 69 | optimizer.load_state_dict(torch.load(options.checkpoint_dir + '/optim.pth')) 70 | pass 71 | 72 | dataset = ScanNetDataset(options, split='train', random=True) 73 | print('the number of images', len(dataset)) 74 | dataloader = DataLoader(dataset, batch_size=options.batchSize, shuffle=True, num_workers=16) 75 | 76 | for epoch in range(options.numEpochs): 77 | epoch_losses = [] 78 | data_iterator = tqdm(dataloader, total=int(np.ceil(float(len(dataset)) / options.batchSize))) 79 | for sample_index, sample in enumerate(data_iterator): 80 | optimizer.zero_grad() 81 | 82 | coords, colors, faces, semantic_gt, instance_gt = sample[0].cuda(), sample[1].cuda(), sample[2].cuda(), sample[3].cuda(), sample[4].cuda() 83 | 84 | if 'augment' in options.suffix: 85 | num_coords = [len(c) for c in coords] 86 | new_coords = [] 87 | new_colors = [] 88 | new_instances = [] 89 | instances = instance_gt.unsqueeze(-1) 90 | for batch_index in range(len(coords)): 91 | augmented_coords, augmented_colors, augmented_instances, _ = augmentation_model(coords[batch_index], faces[batch_index], colors[batch_index], instances[batch_index]) 92 | new_coords.append(torch.cat([coords[batch_index], augmented_coords], dim=0)) 93 | new_colors.append(torch.cat([colors[batch_index], augmented_colors], dim=0)) 94 | new_instances.append(torch.cat([instances[batch_index], augmented_instances], dim=0)) 95 | continue 96 | coords = torch.stack(new_coords, 0) 97 | colors = torch.stack(new_colors, 0) 98 | new_instances = torch.stack(new_instances, 0) 99 | instance_gt = new_instances[:, :, 0] 100 | pass 101 | 102 | semantic_pred, neighbor_pred = model(coords.reshape((-1, 4)), colors.reshape((-1, colors.shape[-1]))) 103 | if 'augment' in options.suffix: 104 | semantic_pred = semantic_pred[:num_coords[0]] 105 | pass 106 | 107 | semantic_loss = torch.nn.functional.cross_entropy(semantic_pred.view((-1, int(semantic_pred.shape[-1]))), semantic_gt.view(-1), weight=class_weights) 108 | semantic_pred = semantic_pred.max(-1)[1].unsqueeze(0) 109 | 110 | if options.numScales > 0: 111 | neighbor_gt = neighbor_model(coords.reshape((-1, 4)), instance_gt.reshape((-1, ))) 112 | else: 113 | neighbor_gt = [] 114 | pass 115 | 116 | if 'mse' not in options.suffix: 117 | for neighbor in neighbor_pred: 118 | neighbor.features = torch.sigmoid(neighbor.features) 119 | continue 120 | pass 121 | neighbor_losses = [] 122 | for scale in range(len(neighbor_gt)): 123 | pred = neighbor_pred[scale].features 124 | gt = neighbor_gt[scale].features[:, :pred.shape[-1]] 125 | mask = neighbor_gt[scale].features[:, pred.shape[-1]:] 126 | if 'mse' in options.suffix: 127 | neighbor_losses.append(torch.sum(torch.nn.functional.mse_loss(pred, gt, reduce=False) * mask) / torch.clamp(mask.sum(), min=1) / options.numScales) 128 | else: 129 | neighbor_losses.append(torch.sum(torch.nn.functional.binary_cross_entropy(pred, (gt > 0.5).float(), weight=(1 - gt) * (int(options.negativeWeights[scale]) - 1) + 1, reduce=False) * mask) / torch.clamp(mask.sum(), min=1)) 130 | pass 131 | continue 132 | losses = [semantic_loss] + neighbor_losses 133 | loss = sum(losses) 134 | 135 | loss_values = [l.data.item() for l in losses] 136 | epoch_losses.append(loss_values) 137 | status = str(epoch + 1) + ' loss: ' 138 | for l in loss_values: 139 | status += '%0.5f '%l 140 | continue 141 | data_iterator.set_description(status) 142 | loss.backward() 143 | optimizer.step() 144 | 145 | continue 146 | print('loss', np.array(epoch_losses).mean(0)) 147 | 148 | if epoch % 10 == 0: 149 | torch.save(model.state_dict(), options.checkpoint_dir + '/checkpoint_' + str(epoch // 10) + '.pth') 150 | #torch.save(optimizer.state_dict(), options.checkpoint_dir + '/optim_' + str(epoch // 10) + '.pth') 151 | pass 152 | torch.save(model.state_dict(), options.checkpoint_dir + '/checkpoint.pth') 153 | torch.save(optimizer.state_dict(), options.checkpoint_dir + '/optim.pth') 154 | 155 | testOneEpoch(options, model, neighbor_model, augmentation_model, dataset_val, validation=True) 156 | continue 157 | return 158 | 159 | def testOneEpoch(options, model, neighbor_model, augmentation_model, dataset, validation=True): 160 | for split in ['pred', 'gt']: 161 | if not os.path.exists(options.test_dir + '/' + split): 162 | os.system("mkdir -p %s"%options.test_dir + '/' + split) 163 | pass 164 | if not os.path.exists(options.test_dir + '/' + split + '/pred_mask'): 165 | os.system("mkdir -p %s"%options.test_dir + '/' + split + '/pred_mask') 166 | pass 167 | continue 168 | 169 | model.eval() 170 | augmentation_model.eval() 171 | 172 | dataloader = DataLoader(dataset, batch_size=options.batchSize, shuffle=False, num_workers=1) 173 | 174 | epoch_losses = [] 175 | data_iterator = tqdm(dataloader, total=int(np.ceil(float(len(dataset)) / options.batchSize))) 176 | 177 | for sample_index, sample in enumerate(data_iterator): 178 | if sample_index == options.numTestingImages: 179 | break 180 | 181 | coords, colors, faces, semantic_gt, instance_gt, filenames = sample[0].cuda(), sample[1].cuda(), sample[2].cuda(), sample[3].cuda(), sample[4].cuda(), sample[5] 182 | 183 | edges = torch.cat([faces[:, :, [0, 1]], faces[:, :, [1, 2]], faces[:, :, [2, 0]]], dim=1) 184 | if 'augment' in options.suffix: 185 | num_coords = [len(c) for c in coords] 186 | new_coords = [] 187 | new_colors = [] 188 | new_edges = [] 189 | new_instances = [] 190 | instances = instance_gt.unsqueeze(-1) 191 | for batch_index in range(len(coords)): 192 | augmented_coords, augmented_colors, augmented_instances, augmented_edges = augmentation_model(coords[batch_index], faces[batch_index], colors[batch_index], instances[batch_index]) 193 | new_coords.append(torch.cat([coords[batch_index], augmented_coords], dim=0)) 194 | new_colors.append(torch.cat([colors[batch_index], augmented_colors], dim=0)) 195 | new_instances.append(torch.cat([instances[batch_index], augmented_instances], dim=0)) 196 | new_edges.append(torch.cat([edges[batch_index], augmented_edges], dim=0)) 197 | continue 198 | coords = torch.stack(new_coords, 0) 199 | colors = torch.stack(new_colors, 0) 200 | new_instances = torch.stack(new_instances, 0) 201 | instance_gt = new_instances[:, :, 0] 202 | edges = torch.stack(new_edges, 0) 203 | pass 204 | 205 | semantic_pred, neighbor_pred = model(coords.reshape((-1, 4)), colors.reshape((-1, colors.shape[-1]))) 206 | 207 | if 'augment' in options.suffix: 208 | semantic_loss = torch.nn.functional.cross_entropy(semantic_pred[:num_coords[0]].view((-1, int(semantic_pred.shape[-1]))), semantic_gt.view(-1), weight=class_weights) 209 | else: 210 | semantic_loss = torch.nn.functional.cross_entropy(semantic_pred.view((-1, int(semantic_pred.shape[-1]))), semantic_gt.view(-1), weight=class_weights) 211 | pass 212 | semantic_pred = semantic_pred.max(-1)[1].unsqueeze(0) 213 | 214 | if options.numScales > 0: 215 | neighbor_gt = neighbor_model(coords.reshape((-1, 4)), instance_gt.reshape((-1, ))) 216 | else: 217 | neighbor_gt = [] 218 | pass 219 | 220 | if 'mse' not in options.suffix: 221 | for neighbor in neighbor_pred: 222 | neighbor.features = torch.sigmoid(neighbor.features) 223 | continue 224 | pass 225 | neighbor_losses = [] 226 | for scale in range(len(neighbor_gt)): 227 | pred = neighbor_pred[scale].features 228 | gt = neighbor_gt[scale].features[:, :pred.shape[-1]] 229 | mask = neighbor_gt[scale].features[:, pred.shape[-1]:] 230 | if 'mse' in options.suffix: 231 | neighbor_losses.append(torch.sum(torch.nn.functional.mse_loss(pred, gt, reduce=False) * mask) / torch.clamp(mask.sum(), min=1) / options.numScales) 232 | else: 233 | neighbor_losses.append(torch.sum(torch.nn.functional.binary_cross_entropy(pred, (gt > 0.5).float(), weight=(1 - gt) * (int(options.negativeWeights[scale]) - 1) + 1, reduce=False) * mask) / torch.clamp(mask.sum(), min=1)) 234 | pass 235 | continue 236 | 237 | if not validation: 238 | for c in range(len(neighbor_pred)): 239 | mask_pred = neighbor_pred[c].features > 0.5 240 | mask_gt = neighbor_gt[c].features[:, :6] > 0.5 241 | neighbor_mask = neighbor_gt[c].features[:, 6:] > 0.5 242 | print(c, (mask_pred * mask_gt * neighbor_mask).sum(), ((1 - mask_pred) * mask_gt * neighbor_mask).sum(), (mask_pred * (1 - mask_gt) * neighbor_mask).sum(), ((1 - mask_pred) * (1 - mask_gt) * neighbor_mask).sum()) 243 | continue 244 | pass 245 | 246 | losses = [semantic_loss] + neighbor_losses 247 | loss = sum(losses) 248 | 249 | loss_values = [l.data.item() for l in losses] 250 | epoch_losses.append(loss_values) 251 | status = 'val loss: ' 252 | for l in loss_values: 253 | status += '%0.5f '%l 254 | continue 255 | data_iterator.set_description(status) 256 | 257 | #semantic_statistics.append(evaluateSemantics(semantic_pred.detach().cpu().numpy(), semantic_gt.detach().cpu().numpy())) 258 | if not validation: 259 | coords = coords.detach().cpu().numpy()[:, :, :3] 260 | colors = np.clip((colors.detach().cpu().numpy() + 1) * 127.5, 0, 255).astype(np.uint8) 261 | semantic_pred = semantic_pred.detach().cpu().numpy() 262 | semantic_gt = semantic_gt.detach().cpu().numpy() 263 | semantic_gt[semantic_gt == -100] = -1 264 | instance_gt = instance_gt.detach().cpu().numpy() 265 | faces = faces.detach().cpu().numpy() 266 | edges = edges.detach().cpu().numpy() 267 | 268 | neighbors = neighbor_model.toDense(neighbor_pred) 269 | neighbors = [neighbor.detach().cpu().numpy() for neighbor in neighbors] 270 | instance_pred = [] 271 | for batch_index in range(len(filenames)): 272 | scene_id = filenames[batch_index].split('/')[-1].split('_vh_clean')[0] 273 | instances, intermediate_instances = findInstances(coords[batch_index], edges[batch_index], np.zeros(len(coords[batch_index])).astype(np.int32), neighbors, options.numScales, options.numCrossScales, cache_filename=options.test_dir + '/pred/' + scene_id + '.txt' if options.useCache else '', scene_id=scene_id) 274 | instance_pred.append(instances) 275 | continue 276 | 277 | instance_info_array = [] 278 | if options.useCache <= 1: 279 | for batch_index in range(len(coords)): 280 | #writeSemantics(options.test_dir + '/sem_pred/' + str(sample_index * options.batchSize + batch_index) + '.txt', semantic_pred[batch_index]) 281 | #writeSemantics(options.test_dir + '/sem_gt/' + str(sample_index * options.batchSize + batch_index) + '.txt', semantic_gt[batch_index]) 282 | scene_id = filenames[batch_index].split('/')[-1].split('_vh_clean')[0] 283 | 284 | instances = instance_pred[batch_index] 285 | semantics = semantic_pred[batch_index] 286 | num_ori_coords = num_coords[batch_index] 287 | if 'augment' in options.suffix: 288 | instance_labels, counts = np.unique(instances[:num_ori_coords], return_counts=True) 289 | valid_labels = instance_labels[counts > 100] 290 | valid_labels = valid_labels[valid_labels >= 0] 291 | #print('num valid instances', len(valid_labels)) 292 | label_map = np.full(instances.max() + 1, fill_value=-1, dtype=np.int32) 293 | for index, label in enumerate(valid_labels): 294 | label_map[label] = index 295 | continue 296 | instances = label_map[instances] 297 | pass 298 | 299 | instance_info = [] 300 | semantic_instances, num_semantic_instances = findInstancesSemanticsLabels(edges[batch_index], semantics) 301 | if num_semantic_instances > 0: 302 | instances[semantic_instances >= 0] = semantic_instances[semantic_instances >= 0] + instances.max() + 1 303 | pass 304 | 305 | instances = instances[:num_ori_coords] 306 | semantics = semantics[:num_ori_coords] 307 | 308 | print('num instances', len(np.unique(instance_gt[batch_index])), instances.max() + 1) 309 | #writeInstances(options.test_dir + '/gt/', scene_id, instance_pred[batch_index], semantic_gt[batch_index]) 310 | 311 | instance_info = writeInstances(options.test_dir + '/pred/', scene_id, instances, semantics, instance_info) 312 | instance_labels = np.zeros(num_ori_coords, dtype=np.int32) 313 | for mask, label, confidence in instance_info: 314 | print(label, confidence) 315 | instance_labels[mask] = label 316 | continue 317 | 318 | unique_instances, first_indices, new_instance_gt = np.unique(instance_gt[batch_index], return_index=True, return_inverse=True) 319 | instance_semantics_gt = mapper[semantic_gt[batch_index][first_indices]] 320 | #print('num', (instance_semantics_gt == 8).sum()) 321 | instance_labels_gt = instance_semantics_gt[new_instance_gt] 322 | visualizeExample(options, coords[batch_index], faces[batch_index], colors[batch_index], num_ori_coords, [('pred', {'semantic': semantics, 'instance': instances, 'instance_label': instance_labels}), ('gt', {'semantic': semantic_gt[batch_index], 'instance': new_instance_gt, 'instance_label': instance_labels_gt})], index_offset=sample_index) 323 | continue 324 | pass 325 | if options.visualizeMode == 'debug': 326 | exit(1) 327 | pass 328 | pass 329 | continue 330 | print('validation loss', np.array(epoch_losses).mean(0)) 331 | #semantic_statistics = np.array(semantic_statistics).sum(0) 332 | #print(semantic_statistics[:len(semantic_statistics) // 2].astype(np.float32) / np.maximum(semantic_statistics[len(semantic_statistics) // 2:], 1)) 333 | model.train() 334 | augmentation_model.train() 335 | return 336 | 337 | 338 | def visualizeExample(options, coords, faces, colors, num_coords, dicts, index_offset=0, prefix=''): 339 | """ Visualize results for one example """ 340 | write_ply_color(options.test_dir + '/' + str(index_offset) + '_input_color.ply', coords, faces, colors[:, :3]) 341 | write_ply_color(options.test_dir + '/' + str(index_offset) + '_input_normal.ply', coords, faces, colors[:, 3:6]) 342 | for name, result_dict in dicts: 343 | semantics = result_dict['semantic'] 344 | 345 | filename = options.test_dir + '/' + str(index_offset) + '_' + name + '_semantic.ply' 346 | write_ply_label(filename, coords[:len(semantics)], faces, semantics) 347 | 348 | if 'instance' in result_dict: 349 | instances = result_dict['instance'] 350 | 351 | filename = options.test_dir + '/' + str(index_offset) + '_' + name + '_instance.ply' 352 | #print(name, len(instances), np.unique(instances, return_counts=True)) 353 | write_ply_label(filename, coords[:len(instances)], faces, instances, debug_index=-1) 354 | 355 | if False: 356 | filename = options.test_dir + '/' + str(index_offset) + '_' + name + '_edge.ply' 357 | write_ply_edge(filename, coords, faces, instances) 358 | pass 359 | pass 360 | print(result_dict.keys()) 361 | if 'instance_label' in result_dict: 362 | filename = options.test_dir + '/' + str(index_offset) + '_' + name + '_instance_semantic.ply' 363 | write_ply_label(filename, coords[:num_coords], faces, result_dict['instance_label'][:num_coords], debug_index=-1) 364 | pass 365 | continue 366 | return 367 | 368 | if __name__ == '__main__': 369 | args = parse_args() 370 | 371 | args.keyname = 'instance' 372 | 373 | if args.suffix != '': 374 | args.keyname += '_' + args.suffix 375 | pass 376 | if args.numScales != 1: 377 | args.keyname += '_' + str(args.numScales) 378 | pass 379 | 380 | args.checkpoint_dir = 'checkpoint/' + args.keyname 381 | args.test_dir = 'test/' + args.keyname 382 | 383 | print('keyname=%s task=%s started'%(args.keyname, args.task)) 384 | 385 | ## Prepare ScanNet data 386 | if args.task == 'prepare': 387 | prepare_data(args) 388 | exit(1) 389 | pass 390 | main(args) 391 | -------------------------------------------------------------------------------- /train_confidence.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DataLoader 3 | 4 | from tqdm import tqdm 5 | import numpy as np 6 | import os 7 | import cv2 8 | 9 | from utils import * 10 | from options import parse_args 11 | 12 | from models.instance import Model, CoordAugmentation, NeighborGT, Validator 13 | 14 | from datasets.scannet_dataset import ScanNetDataset 15 | 16 | torch.set_printoptions(threshold=5000) 17 | 18 | 19 | def main(options): 20 | if not os.path.exists(options.checkpoint_dir): 21 | os.system("mkdir -p %s"%options.checkpoint_dir) 22 | pass 23 | if not os.path.exists(options.test_dir): 24 | os.system("mkdir -p %s"%options.test_dir) 25 | pass 26 | 27 | confidence_model = Validator(options.inputScale, 'normal' in options.suffix) 28 | confidence_model.cuda() 29 | 30 | if options.restore == 1: 31 | print('restore') 32 | if options.startEpoch >= 0: 33 | confidence_model.load_state_dict(torch.load(options.checkpoint_dir + '/checkpoint_confidence_' + str(options.startEpoch) + '.pth')) 34 | else: 35 | confidence_model.load_state_dict(torch.load(options.checkpoint_dir + '/checkpoint_confidence.pth')) 36 | pass 37 | pass 38 | 39 | 40 | dataset_val = ScanNetDataset(options, split='val', load_confidence=True, random=False) 41 | if options.task == 'test': 42 | testOneEpoch(options, confidence_model, dataset_val, validation=False) 43 | exit(1) 44 | 45 | optimizer = torch.optim.Adam(confidence_model.parameters(), lr = options.LR) 46 | if options.restore == 1 and os.path.exists(options.checkpoint_dir + '/optim_confidence.pth'): 47 | optimizer.load_state_dict(torch.load(options.checkpoint_dir + '/optim_confidence.pth')) 48 | pass 49 | 50 | dataset = ScanNetDataset(options, split='train', load_confidence=True, random=True) 51 | dataloader = DataLoader(dataset, batch_size=options.batchSize, shuffle=True, num_workers=16) 52 | 53 | for epoch in range(options.numEpochs): 54 | epoch_losses = [] 55 | data_iterator = tqdm(dataloader, total=int(np.ceil(float(len(dataset)) / options.batchSize))) 56 | for sample_index, sample in enumerate(data_iterator): 57 | optimizer.zero_grad() 58 | 59 | coords, colors, faces, semantic_gt, confidence_gt = sample[0].cuda(), sample[1].cuda(), sample[2].cuda(), sample[3].cuda(), sample[4].cuda() 60 | 61 | confidence_pred = confidence_model(coords.reshape((-1, 4)), colors.reshape((-1, colors.shape[-1])), semantic_gt.view(-1)) 62 | #semantic_pred = semantic_pred.reshape((len(coords), -1)) 63 | confidence_pred = torch.sigmoid(confidence_pred) 64 | confidence_loss = torch.nn.functional.binary_cross_entropy(confidence_pred.view(-1), confidence_gt.view(-1).float()) 65 | 66 | losses = [confidence_loss] 67 | loss = sum(losses) 68 | 69 | loss_values = [l.data.item() for l in losses] 70 | epoch_losses.append(loss_values) 71 | status = str(epoch + 1) + ' loss: ' 72 | for l in loss_values: 73 | status += '%0.5f '%l 74 | continue 75 | data_iterator.set_description(status) 76 | loss.backward() 77 | optimizer.step() 78 | continue 79 | print('loss', np.array(epoch_losses).mean(0)) 80 | if True: 81 | if epoch % 10 == 0: 82 | torch.save(confidence_model.state_dict(), options.checkpoint_dir + '/checkpoint_confidence_' + str(epoch // 10) + '.pth') 83 | torch.save(optimizer.state_dict(), options.checkpoint_dir + '/optim_confidence_' + str(epoch // 10) + '.pth') 84 | pass 85 | torch.save(confidence_model.state_dict(), options.checkpoint_dir + '/checkpoint_confidence.pth') 86 | torch.save(optimizer.state_dict(), options.checkpoint_dir + '/optim_confidence.pth') 87 | pass 88 | testOneEpoch(options, confidence_model, dataset_val, validation=True) 89 | continue 90 | return 91 | 92 | def testOneEpoch(options, confidence_model, dataset, validation=True): 93 | for split in ['pred', 'gt']: 94 | if not os.path.exists(options.test_dir + '/' + split): 95 | os.system("mkdir -p %s"%options.test_dir + '/' + split) 96 | pass 97 | if not os.path.exists(options.test_dir + '/' + split + '/pred_mask'): 98 | os.system("mkdir -p %s"%options.test_dir + '/' + split + '/pred_mask') 99 | pass 100 | continue 101 | 102 | confidence_model.eval() 103 | 104 | # bn = list(list(model.children())[0].children())[3] 105 | # print(bn.running_var, bn.running_mean) 106 | # exit(1) 107 | dataloader = DataLoader(dataset, batch_size=options.batchSize, shuffle=False, num_workers=1) 108 | 109 | epoch_losses = [] 110 | data_iterator = tqdm(dataloader, total=int(np.ceil(float(len(dataset)) / options.batchSize))) 111 | 112 | for sample_index, sample in enumerate(data_iterator): 113 | if sample_index == options.numTestingImages: 114 | break 115 | 116 | coords, colors, faces, semantic_gt, confidence_gt, instance_masks = sample[0].cuda(), sample[1].cuda(), sample[2].cuda(), sample[3].cuda(), sample[4].cuda(), sample[6] 117 | 118 | confidence_pred = confidence_model(coords.reshape((-1, 4)), colors.reshape((-1, colors.shape[-1])), semantic_gt.view(-1)) 119 | #semantic_pred = semantic_pred.reshape((len(coords), -1)) 120 | confidence_pred = torch.sigmoid(confidence_pred) 121 | confidence_loss = torch.nn.functional.binary_cross_entropy(confidence_pred.view(-1), confidence_gt.view(-1).float()) 122 | 123 | losses = [confidence_loss] 124 | loss = sum(losses) 125 | 126 | loss_values = [l.data.item() for l in losses] 127 | epoch_losses.append(loss_values) 128 | status = 'val loss: ' 129 | for l in loss_values: 130 | status += '%0.5f '%l 131 | continue 132 | data_iterator.set_description(status) 133 | continue 134 | print('validation loss', np.array(epoch_losses).mean(0)) 135 | confidence_model.train() 136 | return 137 | 138 | 139 | class InstanceValidator(): 140 | """ Load trained model to predict confidence for instances """ 141 | def __init__(self, checkpoint_dir, full_scale=127, use_normal=False): 142 | self.full_scale = full_scale 143 | confidence_model = Validator(full_scale, use_normal) 144 | confidence_model.cuda() 145 | confidence_model.load_state_dict(torch.load(checkpoint_dir + '/checkpoint_confidence.pth')) 146 | confidence_model.eval() 147 | self.confidence_model = confidence_model 148 | return 149 | 150 | def validate(self, coords, colors, instances, semantics): 151 | instances += 1 152 | semantic_inp = [] 153 | instance_masks = [] 154 | new_coords = np.zeros(coords.shape, dtype=coords.dtype) 155 | confidence_by_counts = [] 156 | for instance in range(instances.max() + 1): 157 | instance_mask = instances == instance 158 | if instance_mask.sum() == 0: 159 | print('sum = 0', instance, instances.max() + 1, instance_mask.sum()) 160 | exit(1) 161 | info = np.unique(semantics[instance_mask > 0.5], return_counts=True) 162 | label_pred = info[0][info[1].argmax()] 163 | instance_coords = coords[instance_mask] 164 | mins = instance_coords.min(0) 165 | maxs = instance_coords.max(0) 166 | max_range = (maxs - mins).max() 167 | padding = max_range * 0.05 168 | max_range += padding * 2 169 | mins = (mins + maxs) / 2 - max_range / 2 170 | instance_coords = np.clip(np.round((instance_coords - mins) / max_range * self.full_scale), 0, self.full_scale - 1) 171 | new_coords[instance_mask] = instance_coords 172 | 173 | if instance > 0: 174 | semantic_inp.append(label_pred) 175 | instance_masks.append(instance_mask) 176 | confidence_by_counts.append(float(info[1].max()) / info[1].sum()) 177 | pass 178 | continue 179 | coords = np.concatenate([new_coords, np.expand_dims(instances, -1)], axis=-1) 180 | coords = torch.from_numpy(coords.astype(np.int64)).cuda() 181 | 182 | colors = colors.astype(np.float32) / 127.5 - 1 183 | colors = torch.from_numpy(colors.astype(np.float32)).cuda() 184 | 185 | semantic_inp = np.stack(semantic_inp).astype(np.int64) 186 | 187 | confidence_pred = self.confidence_model(coords.reshape((-1, 4)), colors.reshape((-1, colors.shape[-1])), torch.from_numpy(semantic_inp).view(-1).cuda()) 188 | confidence_pred = torch.sigmoid(confidence_pred) 189 | confidence_pred = confidence_pred.detach().cpu().numpy() 190 | 191 | instance_info = list(zip(instance_masks, semantic_inp, confidence_pred)) 192 | 193 | return instance_info 194 | 195 | if __name__ == '__main__': 196 | args = parse_args() 197 | 198 | args.keyname = 'instance' 199 | #args.keyname += '_' + args.dataset 200 | 201 | if args.suffix != '': 202 | args.keyname += '_' + args.suffix 203 | pass 204 | if args.numScales != 1: 205 | args.keyname += '_' + str(args.numScales) 206 | pass 207 | 208 | args.checkpoint_dir = 'checkpoint/' + args.keyname 209 | args.test_dir = 'test/' + args.keyname 210 | 211 | args.suffix = '' 212 | args.numScales = 0 213 | args.inputScale = 127 214 | 215 | print('keyname=%s task=%s started'%(args.keyname, args.task)) 216 | 217 | main(args) 218 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | from plyfile import PlyData, PlyElement 4 | import heapq 5 | import sys 6 | 7 | ## Mapping between NYU40 labels and ScanNet evaluation label subset 8 | remapper = np.full(150, fill_value=-100, dtype=np.int32) 9 | mapper = np.zeros(21, dtype=np.int32) 10 | label_subset = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39] 11 | for i, x in enumerate(label_subset): 12 | remapper[x] = i 13 | mapper[i] = x 14 | continue 15 | 16 | 17 | def write_ply_color(filename, coords, faces, colors): 18 | """ Write colored mesh model """ 19 | header = """ply 20 | format ascii 1.0 21 | element vertex """ 22 | header += str(len(coords)) 23 | header += """ 24 | property float x 25 | property float y 26 | property float z 27 | property uchar red 28 | property uchar green 29 | property uchar blue 30 | element face """ 31 | header += str(len(faces)) 32 | header += """ 33 | property list uchar int vertex_index 34 | end_header 35 | """ 36 | 37 | with open(filename, 'w') as f: 38 | f.write(header) 39 | for coord, color in zip(coords, colors): 40 | for value in coord: 41 | f.write(str(value) + ' ') 42 | continue 43 | for value in color[:3]: 44 | f.write(str(value) + ' ') 45 | continue 46 | f.write('\n') 47 | continue 48 | for face in faces: 49 | f.write('3 ' + str(face[0]) + ' ' + str(face[1]) + ' ' + str(face[2]) + '\n') 50 | continue 51 | pass 52 | return 53 | 54 | def write_ply_label(filename, coords, faces, labels, visualize_boundary=True, debug_index=-1): 55 | """ Write mesh model colored by segmentation labels """ 56 | if visualize_boundary: 57 | valid_indices = np.logical_and(labels[faces[:, 0]] == labels[faces[:, 1]], labels[faces[:, 0]] == labels[faces[:, 2]]) 58 | #print(coords.min(0), coords.max(0)) 59 | #valid_indices = np.logical_or(np.logical_or(np.all(coords[faces[:, 0]] == coords[faces[:, 1]], axis=-1), np.all(coords[faces[:, 0]] == coords[faces[:, 2]], axis=-1)), np.all(coords[faces[:, 1]] == coords[faces[:, 2]], axis=-1)) 60 | faces = faces[valid_indices] 61 | pass 62 | if debug_index != -1: 63 | valid_indices = np.logical_and(np.logical_and(labels[faces[:, 0]] == debug_index, labels[faces[:, 1]] == debug_index), labels[faces[:, 2]] == debug_index) 64 | faces = faces[valid_indices] 65 | #coords = coords[labels == debug_index] 66 | # print(len(faces), (labels == debug_index).sum(), [(labels[faces[:, c]] == debug_index).sum() for c in range(3)]) 67 | # print(len(np.unique(coords[:, 0] * 4096 * 4096 + coords[:, 1] * 4096 + coords[:, 2]))) 68 | # exit(1) 69 | if len(faces) == 0: 70 | return 71 | coords = np.concatenate([coords[face] for face in faces], axis=0) 72 | faces = np.arange(len(faces) * 3).reshape((-1, 3)) 73 | labels = np.full(len(coords), fill_value=debug_index) 74 | pass 75 | 76 | header = """ply 77 | format ascii 1.0 78 | element vertex """ 79 | header += str(len(coords)) 80 | header += """ 81 | property float x 82 | property float y 83 | property float z 84 | property uchar red 85 | property uchar green 86 | property uchar blue 87 | element face """ 88 | header += str(len(faces)) 89 | header += """ 90 | property list uchar int vertex_index 91 | end_header 92 | """ 93 | 94 | np.random.seed(1) 95 | color_map = np.random.randint(256, size=(labels.max() + 2, 3), dtype=np.uint8) 96 | #labels[labels == -100] = -1 97 | color_map[-1] = 255 98 | label_colors = color_map[labels] 99 | with open(filename, 'w') as f: 100 | f.write(header) 101 | for coord, color in zip(coords, label_colors): 102 | for value in coord: 103 | f.write(str(value) + ' ') 104 | continue 105 | for value in color: 106 | f.write(str(value) + ' ') 107 | continue 108 | f.write('\n') 109 | continue 110 | for face in faces: 111 | f.write('3 ' + str(face[0]) + ' ' + str(face[1]) + ' ' + str(face[2]) + '\n') 112 | continue 113 | pass 114 | return 115 | 116 | def write_ply_edge(filename, coords, edges, labels, augmented_edges=[]): 117 | """ Write mesh model with edges """ 118 | # edges = np.concatenate([faces[:, [0, 1]], faces[:, [0, 2]], faces[:, [1, 2]]], axis=0) 119 | # if len(augmented_edges) > 0: 120 | # edges = np.concatenate([edges, augmented_edges], axis=0) 121 | # pass 122 | if True: 123 | valid_indices = np.abs(coords[edges[:, 0]] - coords[edges[:, 1]]).sum(-1) == 1 124 | #valid_indices = np.logical_or(valid_indices, np.abs(coords[edges[:, 0]] // 2 - coords[edges[:, 1]] // 2).sum(-1) == 1) 125 | edges = edges[valid_indices] 126 | pass 127 | 128 | header = """ply 129 | format ascii 1.0 130 | element vertex """ 131 | header += str(len(coords)) 132 | header += """ 133 | property float x 134 | property float y 135 | property float z 136 | property uchar red 137 | property uchar green 138 | property uchar blue 139 | element edge """ 140 | header += str(len(edges)) 141 | header += """ 142 | property int vertex1 143 | property int vertex2 144 | property uchar red 145 | property uchar green 146 | property uchar blue 147 | end_header 148 | """ 149 | 150 | np.random.seed(1) 151 | color_map = np.random.randint(256, size=(labels.max() + 2, 3), dtype=np.uint8) 152 | #labels[labels == -100] = -1 153 | color_map[-1] = 255 154 | label_colors = color_map[labels] 155 | 156 | edge_colors = np.zeros((len(edges), 3), dtype=np.uint8) 157 | distances = np.abs(coords[edges[:, 0]] - coords[edges[:, 1]]).sum(-1) 158 | edge_colors[distances == 1] = 255 159 | edge_colors[distances > 1] = np.array([255, 0, 0]) 160 | with open(filename, 'w') as f: 161 | f.write(header) 162 | for coord, color in zip(coords, label_colors): 163 | for value in coord: 164 | f.write(str(value) + ' ') 165 | continue 166 | for value in color: 167 | f.write(str(value) + ' ') 168 | continue 169 | f.write('\n') 170 | continue 171 | for edge, color in zip(edges, edge_colors): 172 | f.write(str(edge[0]) + ' ' + str(edge[1]) + ' ' + str(color[0]) + ' ' + str(color[1]) + ' ' + str(color[2]) + '\n') 173 | continue 174 | pass 175 | return 176 | 177 | def write_ply_neighbor(filename, coords, neighbors, masks, size=4096): 178 | """ Write mesh model to visualize neighbors """ 179 | valid_mask = masks.sum(-1) > 0.5 180 | coords = coords[valid_mask] 181 | neighbors = (neighbors[valid_mask] > 0.95).astype(np.int32) 182 | masks = masks[valid_mask] 183 | 184 | index_map = {} 185 | for index, coord in enumerate(coords): 186 | index_map[toIndex(coord, size)] = index 187 | continue 188 | 189 | valid_edges = [] 190 | coord_offsets = np.array([[-1, 0, 0], [1, 0, 0], [0, -1, 0], [0, 1, 0], [0, 0, -1], [0, 0, 1]]) 191 | for index, coord in enumerate(coords): 192 | neighbor = neighbors[index] 193 | mask = masks[index] > 0.5 194 | for offset, label in zip(coord_offsets[mask], neighbor[mask]): 195 | neighbor_coord = coord + offset 196 | neighbor_index = toIndex(neighbor_coord, size) 197 | if neighbor_index not in index_map: 198 | continue 199 | neighbor_index = index_map[neighbor_index] 200 | if label == 0: 201 | continue 202 | valid_edges.append((index, neighbor_index, label)) 203 | continue 204 | continue 205 | 206 | edges = np.array(valid_edges) 207 | color_map = np.array([[255, 0, 0], [255, 255, 255]]) 208 | edge_colors = color_map[edges[:, 2]] 209 | edges = edges[:, :2] 210 | #edge_colors = np.full((len(edges), 3), fill_value=255, dtype=np.uint8) 211 | 212 | coords *= 4096 // size 213 | 214 | header = """ply 215 | format ascii 1.0 216 | element vertex """ 217 | header += str(len(coords)) 218 | header += """ 219 | property float x 220 | property float y 221 | property float z 222 | property uchar red 223 | property uchar green 224 | property uchar blue 225 | element edge """ 226 | header += str(len(edges)) 227 | header += """ 228 | property int vertex1 229 | property int vertex2 230 | property uchar red 231 | property uchar green 232 | property uchar blue 233 | end_header 234 | """ 235 | 236 | np.random.seed(1) 237 | #color_map = np.random.randint(256, size=(labels.max() + 2, 3), dtype=np.uint8) 238 | #labels[labels == -100] = -1 239 | #color_map[-1] = 255 240 | label_colors = np.full(coords.shape, fill_value=255, dtype=np.uint8) 241 | 242 | with open(filename, 'w') as f: 243 | f.write(header) 244 | for coord, color in zip(coords, label_colors): 245 | for value in coord: 246 | f.write(str(value) + ' ') 247 | continue 248 | for value in color: 249 | f.write(str(value) + ' ') 250 | continue 251 | f.write('\n') 252 | continue 253 | for edge, color in zip(edges, edge_colors): 254 | f.write(str(edge[0]) + ' ' + str(edge[1]) + ' ' + str(color[0]) + ' ' + str(color[1]) + ' ' + str(color[2]) + '\n') 255 | continue 256 | pass 257 | return 258 | 259 | 260 | def writeSemantics(filename, semantics): 261 | """ Write semantics """ 262 | semantics = mapper[semantics] 263 | np.savetxt(filename, semantics, fmt='%d') 264 | return 265 | 266 | def writeInstances(path, scene_id, instances, semantics, instance_info): 267 | """ Write instances """ 268 | semantics = mapper[semantics] 269 | print(scene_id, 'num instances', instances.max() + 1) 270 | #for instance_index in np.unique(instances): 271 | if len(instance_info) == 0: 272 | instance_info = [] 273 | for instance_index in range(instances.max() + 1): 274 | instance_mask = instances == instance_index 275 | semantic_labels = semantics[instance_mask] 276 | semantic_labels, counts = np.unique(semantic_labels, return_counts=True) 277 | instance_info.append((instance_mask, semantic_labels[counts.argmax()], float(counts.max()) / counts.sum())) 278 | continue 279 | else: 280 | valid_instance_info = [] 281 | for mask, label, confidence in instance_info: 282 | if label == 20: 283 | continue 284 | valid_instance_info.append((mask, mapper[label], confidence)) 285 | continue 286 | instance_info = valid_instance_info 287 | pass 288 | 289 | with open(path + '/' + scene_id + '.txt', 'w') as f: 290 | for instance_index, (mask, label, confidence) in enumerate(instance_info): 291 | f.write('pred_mask/' + scene_id + '_' + str(instance_index) + '.txt ' + str(label) + ' ' + str(confidence) + '\n') 292 | np.savetxt(path + '/pred_mask/' + scene_id + '_' + str(instance_index) + '.txt', mask, fmt='%d') 293 | continue 294 | pass 295 | return instance_info 296 | 297 | def findInstancesSemanticsLabels(edges, semantics, labels=[10, 13, 15, 17, 18], instance_info=[]): 298 | """ Find connected components based on semantic predictions """ 299 | ranges = np.arange(len(semantics), dtype=np.int32) 300 | valid_mask = semantics[edges[:, 0]] == semantics[edges[:, 1]] 301 | edges = edges[valid_mask] 302 | 303 | instance_masks = [] 304 | for label in labels: 305 | valid_mask = semantics[edges[:, 0]] == label 306 | semantic_edges = edges[valid_mask] 307 | existing_masks = [info[0] for info in instance_info if info[1] == label] 308 | existing_mask_areas = [mask.sum() for mask in existing_masks] 309 | if len(existing_masks) > 0: 310 | print(existing_masks[0].dtype) 311 | exit(1) 312 | while len(semantic_edges) >= 100: 313 | instances = np.zeros(len(semantics), dtype=np.bool) 314 | instances[semantic_edges[0]] = 1 315 | while True: 316 | valid_mask = np.logical_and(instances[semantic_edges[:, 0]], np.logical_not(instances[semantic_edges[:, 1]])) 317 | neighbors = semantic_edges[:, 1][valid_mask] 318 | if len(neighbors) == 0: 319 | semantic_edges = semantic_edges[np.logical_not(instances[semantic_edges[:, 0]])] 320 | break 321 | instances[neighbors] = 1 322 | continue 323 | instance_area = instances.sum() 324 | if instance_area < 100: 325 | continue 326 | if len(existing_masks) > 0: 327 | intersections = [np.logical_and(mask, instances).sum() for mask in existing_masks] 328 | IOUs = [float(intersection) / max(existing_mask_area + instance_area - intersection, 1) for intersection, existing_mask_area in zip(intersections, existing_mask_areas)] 329 | if max(IOUs) > 0.5: 330 | continue 331 | pass 332 | instance_masks.append(instances) 333 | continue 334 | continue 335 | 336 | instances = np.full(len(semantics), fill_value=-1, dtype=np.int32) 337 | for index, mask in enumerate(instance_masks): 338 | instances[mask] = index 339 | continue 340 | return instances, len(instance_masks) 341 | 342 | def toIndex(coord, size): 343 | """ Coordinate to index """ 344 | return coord[0] * size * size + coord[1] * size + coord[2] 345 | 346 | def toCoord(index, size): 347 | """ Index to coordinate """ 348 | return np.array([index // (size * size), index // size % size, index % size]) 349 | 350 | def loadInstance(filename): 351 | """ Load saved instance for later visualization """ 352 | with open(filename, 'r') as f: 353 | print('load instances', filename) 354 | for line_index, line in enumerate(f): 355 | mask = np.loadtxt('/'.join(filename.split('/')[:-1]) + '/' + line.split(' ')[0]).astype(np.int32) 356 | if line_index == 0: 357 | instances = mask - 1 358 | else: 359 | instances[mask > 0.5] = line_index 360 | pass 361 | continue 362 | pass 363 | return instances 364 | 365 | def findInstances(coords, edges, semantics, neighbors, num_scales, num_cross_scales, full_scale=4096, num_neighbors=6, num_dimensions=3, print_info=True, cache_filename='', scene_id=''): 366 | """ Clustering algorithm """ 367 | debug = False 368 | if cache_filename != '': 369 | instances = loadInstance(cache_filename) 370 | return instances, [] 371 | 372 | num_scales = num_scales 373 | 374 | for scale in range(num_scales): 375 | offsets = [0, num_neighbors] 376 | for _ in range(scale + 1, min(num_scales, scale + num_cross_scales)): 377 | offsets.append(offsets[-1] + num_neighbors + 1) 378 | continue 379 | neighbors[scale] = [np.array([])] * scale + [neighbors[scale][:, offsets[index]:offsets[index + 1]] for index in range(len(offsets) - 1)] 380 | #print(scale, [neighbor.shape for neighbor in neighbors[scale]]) 381 | continue 382 | coord_offsets = np.array([[-1, 0, 0], [1, 0, 0], [0, -1, 0], [0, 1, 0], [0, 0, -1], [0, 0, 1]]) 383 | 384 | 385 | scale_count_thresholds = pow(4, np.arange(num_scales)) 386 | connection_ratio_threshold = 0.2 387 | 388 | coord_node_map = [] 389 | for scale in range(num_scales): 390 | coord_node_map.append({toIndex(coord, full_scale // pow(2, scale)): node_index for node_index, coord in enumerate(coords // pow(2, scale))}) 391 | continue 392 | 393 | #print(faces.shape, faces.min(), faces.max(), semantics.shape) 394 | ori_node_mapping = np.arange(len(coords), dtype=np.int64) 395 | node_info = [({0: (np.expand_dims(coord, 0), np.array([node_index]))}, (np.arange(41) == semantics[node_index]).astype(np.float32)) for node_index, coord in enumerate(coords)] 396 | direction_multiplier = pow(2, np.arange(3)) 397 | direction_mapping_dict = {-1: 0, 1: 1, -2: 2, 2: 3, -4: 4, 4:5, 0: 6} 398 | direction_mapping = np.zeros(9, dtype=np.int32) 399 | direction_index_mapping = np.array([-1, 1, -2, 2, -4, 4, 0]) 400 | for direction, mapping in direction_mapping_dict.items(): 401 | direction_mapping[direction + 4] = mapping 402 | continue 403 | 404 | if debug: 405 | ori_node_mapping = loadInstance(cache_filename) 406 | node_instances = [[] for _ in range(ori_node_mapping.max() + 1)] 407 | for ori_node, node in enumerate(ori_node_mapping): 408 | node_instances[node].append(ori_node) 409 | continue 410 | 411 | node_instances = [np.array(instance) for instance in node_instances] 412 | node_info = [] 413 | for instance in node_instances: 414 | info = {} 415 | node_coords = coords[instance] 416 | labels = semantics[instance] 417 | unique_labels, counts = np.unique(labels, return_counts=True) 418 | label_counts = np.zeros(41) 419 | label_counts[unique_labels] = counts 420 | #print(instance, np.array([coord_node_map[0][toIndex(coord, full_scale)] for coord in node_coords])) 421 | #print('valid', np.all(neighbors[0][0][np.array([coord_node_map[0][toIndex(coord, full_scale)] for coord in node_coords])] == neighbors[0][0][instance])) 422 | #exit(1) 423 | info[0] = (node_coords, instance) 424 | for scale in range(1, num_scales): 425 | size = full_scale // pow(2, scale) 426 | node_coords = node_coords // 2 427 | indices = node_coords[:, 0] * size * size + node_coords[:, 1] * size + node_coords[:, 2] 428 | indices, mapping, counts = np.unique(indices, return_index=True, return_counts=True) 429 | valid_coords = [] 430 | for index, coord_index, count in zip(indices, mapping, counts): 431 | #if count >= count_dicts[scale][index] * 0.5: 432 | if count >= scale_count_thresholds[scale]: 433 | valid_coords.append(node_coords[coord_index]) 434 | pass 435 | continue 436 | if len(valid_coords) > 0: 437 | info[scale] = (valid_coords, np.array([coord_node_map[scale][toIndex(coord, full_scale // pow(2, scale))] for coord in valid_coords])) 438 | pass 439 | pass 440 | node_info.append((info, label_counts)) 441 | continue 442 | debug = True 443 | print('num instances', len(node_info)) 444 | for node_index, info in enumerate(node_info): 445 | print(node_index, info[1].argmax()) 446 | continue 447 | pass 448 | 449 | ori_edges = edges 450 | 451 | intermediate_instances = [] 452 | iteration = 0 453 | while True: 454 | node_scores = {} 455 | instance_coord_maps = {} 456 | 457 | edges = ori_node_mapping[ori_edges] 458 | edges.sort(-1) 459 | edges = np.unique(edges[:, 0] * len(coords) + edges[:, 1]) 460 | edges = np.stack([edges // len(coords), edges % len(coords)], axis=-1) 461 | edges = edges[np.logical_and(edges[:, 0] != edges[:, 1], np.all(edges >= 0, axis=-1))] 462 | # print(edges.max(), len(coords), len(node_info)) 463 | for edge_index, edge in enumerate(edges): 464 | node_pair = edge 465 | 466 | semantic_similarity = ((node_info[node_pair[0]][1] / max(node_info[node_pair[0]][1].sum(), 1)) * (node_info[node_pair[1]][1] / max(node_info[node_pair[1]][1].sum(), 1))).sum() 467 | 468 | scores = [] 469 | score_info = [] 470 | #largest_scale = min([max(node_info[node_pair[0]][0].keys()) for c in range(2)]) 471 | for scale_1, (coord_1, ori_node_indices_1) in node_info[node_pair[0]][0].items(): 472 | for scale_2, (coord_2, ori_node_indices_2) in node_info[node_pair[1]][0].items(): 473 | if print_info: 474 | sys.stdout.write('\r' + str(edge_index) + ' ' + str(len(edges)) + ' ' + str(scale_1) + ' ' + str(scale_2) + ' ' + str(len(coord_1)) + ' ' + str(len(coord_2)) + ' ') 475 | pass 476 | if scale_1 == scale_2: 477 | if scale_1 > 2: 478 | continue 479 | if len(coord_1) <= 1000 and len(coord_2) <= 1000: 480 | #if len(coord_1) * len(coord_2) <= 10000: 481 | directions = (np.expand_dims(coord_2, 0) - np.expand_dims(coord_1, 1)).reshape((-1, num_dimensions)) 482 | distances = np.abs(directions).sum(-1) 483 | direction_indices = np.dot(directions, direction_multiplier) 484 | valid_mask = distances == 1 485 | direction_indices = direction_indices[valid_mask] 486 | indices_1 = np.expand_dims(ori_node_indices_1, 1).repeat(len(ori_node_indices_2), axis=1).reshape(-1)[valid_mask] 487 | indices_2 = np.expand_dims(ori_node_indices_2, 0).repeat(len(ori_node_indices_1), axis=0).reshape(-1)[valid_mask] 488 | else: 489 | if len(coord_1) < len(coord_2): 490 | coord_1, coord_2 = coord_2, coord_1 491 | ori_node_indices_1, ori_node_indices_2 = ori_node_indices_2, ori_node_indices_1 492 | cache_pair_index = 0 493 | else: 494 | cache_pair_index = 1 495 | pass 496 | if scale_2 in instance_coord_maps and node_pair[cache_pair_index] in instance_coord_maps[scale_2]: 497 | instance_coord_map = instance_coord_maps[scale_2][node_pair[cache_pair_index]] 498 | else: 499 | if scale_2 not in instance_coord_maps: 500 | instance_coord_maps[scale_2] = {} 501 | pass 502 | instance_coord_map = {toIndex(coord, full_scale // pow(2, scale_2)): ori_node_index for coord, ori_node_index in zip(coord_2, ori_node_indices_2)} 503 | instance_coord_maps[scale_2][node_pair[cache_pair_index]] = instance_coord_map 504 | pass 505 | coord_neighbors = np.expand_dims(coord_1, 1) + coord_offsets 506 | size = full_scale // pow(2, scale_2) 507 | coord_neighbor_indices = coord_neighbors[:, :, 0] * size * size + coord_neighbors[:, :, 1] * size + coord_neighbors[:, :, 2] 508 | indices_1, indices_2, direction_indices = [], [], [] 509 | for neighbor_index, (coord_index, ori_node_index) in enumerate(zip(coord_neighbor_indices.reshape(-1), np.expand_dims(ori_node_indices_1, axis=-1).repeat(num_neighbors, axis=-1).reshape(-1))): 510 | if coord_index in instance_coord_map: 511 | indices_1.append(ori_node_index) 512 | indices_2.append(instance_coord_map[coord_index]) 513 | direction_indices.append(direction_index_mapping[neighbor_index % num_neighbors]) 514 | pass 515 | continue 516 | indices_1 = np.array(indices_1) 517 | indices_2 = np.array(indices_2) 518 | direction_indices = np.array(direction_indices) 519 | pass 520 | 521 | # if node_pair == (106827, 109911): 522 | # print(indices_1, direction_indices, neighbors[scale_1][scale_2][indices_1]) 523 | # print(neighbors[scale_1][scale_2][indices_1, direction_mapping[direction_indices + 4]]) 524 | # exit(1) 525 | #distances = distances[distances == 1] 526 | 527 | instance_size = np.sqrt(float(min(len(coord_1), len(coord_2)))) 528 | if len(direction_indices) > round(instance_size * connection_ratio_threshold): 529 | #scores.append((distances == 1) * (neighbors[scale_1][scale_2][edge[0]][direction_mapping[direction_indices]] + neighbors[scale_2][scale_1][edge[1]][direction_mapping[-direction_indices]]) / 2) 530 | #score_info.append(np.full((len(direction_indices), 2), fill_value=scale_1)) 531 | scores.append(((neighbors[scale_1][scale_2][indices_1, direction_mapping[direction_indices + 4]] + neighbors[scale_2][scale_1][indices_2, direction_mapping[-direction_indices + 4]]) / 2).mean()) 532 | score_info.append(np.array([scale_1, scale_1])) 533 | pass 534 | elif abs(scale_1 - scale_2) <= num_cross_scales: 535 | if scale_1 > scale_2: 536 | scale_1, scale_2 = scale_2, scale_1 537 | coord_1, coord_2 = coord_2, coord_1 538 | ori_node_indices_1, ori_node_indices_2 = ori_node_indices_2, ori_node_indices_1 539 | pass 540 | coord_1 = coord_1 // pow(2, scale_2 - scale_1) 541 | if len(coord_1) <= 1000 and len(coord_2) <= 1000: 542 | #if len(coord_1) * len(coord_2) <= 10000: 543 | directions = (np.expand_dims(coord_2, 0) - np.expand_dims(coord_1, 1)).reshape((-1, num_dimensions)) 544 | distances = np.abs(directions).sum(-1) 545 | direction_indices = np.dot(directions, direction_multiplier) 546 | valid_mask = distances <= 1 547 | direction_indices = direction_indices[valid_mask] 548 | indices_1 = np.expand_dims(ori_node_indices_1, 1).repeat(len(ori_node_indices_2), axis=1).reshape(-1)[valid_mask] 549 | else: 550 | if scale_2 in instance_coord_maps and node_pair[cache_pair_index] in instance_coord_maps[scale_2]: 551 | instance_coord_map = instance_coord_maps[scale_2][node_pair[cache_pair_index]] 552 | else: 553 | if scale_2 not in instance_coord_maps: 554 | instance_coord_maps[scale_2] = {} 555 | pass 556 | instance_coord_map = {toIndex(coord, full_scale // pow(2, scale_2)): ori_node_index for coord, ori_node_index in zip(coord_2, ori_node_indices_2)} 557 | instance_coord_maps[scale_2][node_pair[cache_pair_index]] = instance_coord_map 558 | pass 559 | coord_neighbors = np.expand_dims(coord_1, 1) + coord_offsets 560 | coord_neighbors = np.concatenate([coord_neighbors, np.expand_dims(coord_1, 1)], axis=1) 561 | size = full_scale // pow(2, scale_2) 562 | coord_neighbor_indices = coord_neighbors[:, :, 0] * size * size + coord_neighbors[:, :, 1] * size + coord_neighbors[:, :, 2] 563 | indices_1, indices_2, direction_indices = [], [], [] 564 | for neighbor_index, (coord_index, ori_node_index) in enumerate(zip(coord_neighbor_indices.reshape(-1), np.expand_dims(ori_node_indices_1, axis=-1).repeat(num_neighbors + 1, axis=-1).reshape(-1))): 565 | if coord_index in instance_coord_map: 566 | indices_1.append(ori_node_index) 567 | direction_indices.append(direction_index_mapping[neighbor_index % (num_neighbors + 1)]) 568 | pass 569 | continue 570 | indices_1 = np.array(indices_1) 571 | direction_indices = np.array(direction_indices) 572 | pass 573 | 574 | if len(direction_indices) > 0: 575 | scores.append(neighbors[scale_1][scale_2][indices_1, direction_mapping[direction_indices + 4]].mean()) 576 | score_info.append(np.array([scale_1, scale_2])) 577 | pass 578 | pass 579 | continue 580 | continue 581 | if len(scores) > 0: 582 | #scores = np.concatenate(scores, axis=0) 583 | scores = np.array(scores) 584 | 585 | #score = scores.mean() 586 | score = scores[np.array(score_info).sum(-1).argmax()] 587 | 588 | if debug: 589 | if node_info[node_pair[0]][1].argmax() == 4 and node_info[node_pair[1]][1].argmax() == 4: 590 | print('bed') 591 | print(scores) 592 | exit(1) 593 | pass 594 | 595 | for c in range(2): 596 | if node_pair[c] not in node_scores: 597 | node_scores[node_pair[c]] = {} 598 | pass 599 | #node_scores[node_pair[c]][node_pair[1 - c]] = score * semantic_similarity 600 | node_scores[node_pair[c]][node_pair[1 - c]] = score 601 | continue 602 | pass 603 | continue 604 | if print_info: 605 | print('') 606 | pass 607 | has_change = False 608 | node_mapping = np.arange(len(node_info), dtype=np.int64) 609 | for node, neighbor_scores in node_scores.items(): 610 | max_score_neighbor = (0.5, -1) 611 | for neighbor, score in neighbor_scores.items(): 612 | if score > max_score_neighbor[0]: 613 | max_score_neighbor = [score, neighbor] 614 | pass 615 | continue 616 | if max_score_neighbor[1] >= 0: 617 | node_mapping[node] = max_score_neighbor[1] 618 | has_change = True 619 | pass 620 | continue 621 | if not has_change: 622 | break 623 | 624 | new_node_mapping = np.full(node_mapping.shape, fill_value=-1, dtype=node_mapping.dtype) 625 | new_node_index = 0 626 | for ori_node, node in enumerate(node_mapping): 627 | if new_node_mapping[ori_node] != -1: 628 | continue 629 | instance = {ori_node: True} 630 | instance_index = -1 631 | while node not in instance: 632 | instance_index = new_node_mapping[node] 633 | if instance_index != -1: 634 | break 635 | instance[node] = True 636 | node = node_mapping[node] 637 | continue 638 | instance = list(instance.keys()) 639 | if instance_index != -1: 640 | for node_index in instance: 641 | new_node_mapping[node_index] = instance_index 642 | continue 643 | else: 644 | for node_index in instance: 645 | new_node_mapping[node_index] = new_node_index 646 | continue 647 | new_node_index += 1 648 | pass 649 | #print(node_instances) 650 | continue 651 | #node_mapping = new_node_mapping 652 | ori_node_mapping = new_node_mapping[ori_node_mapping] 653 | 654 | intermediate_instances.append(ori_node_mapping) 655 | 656 | node_instances = [[] for _ in range(new_node_index)] 657 | for ori_node, node in enumerate(ori_node_mapping): 658 | node_instances[node].append(ori_node) 659 | continue 660 | 661 | iteration += 1 662 | 663 | if print_info: 664 | print('num nodes', len(node_instances)) 665 | pass 666 | node_instances = [np.array(instance) for instance in node_instances] 667 | node_info = [] 668 | for instance in node_instances: 669 | info = {} 670 | node_coords = coords[instance] 671 | labels = semantics[instance] 672 | unique_labels, counts = np.unique(labels, return_counts=True) 673 | label_counts = np.zeros(41) 674 | label_counts[unique_labels] = counts 675 | #print(instance, np.array([coord_node_map[0][toIndex(coord, full_scale)] for coord in node_coords])) 676 | #print('valid', np.all(neighbors[0][0][np.array([coord_node_map[0][toIndex(coord, full_scale)] for coord in node_coords])] == neighbors[0][0][instance])) 677 | #exit(1) 678 | info[0] = (node_coords, instance) 679 | for scale in range(1, num_scales): 680 | size = full_scale // pow(2, scale) 681 | node_coords = node_coords // 2 682 | indices = node_coords[:, 0] * size * size + node_coords[:, 1] * size + node_coords[:, 2] 683 | indices, mapping, counts = np.unique(indices, return_index=True, return_counts=True) 684 | valid_coords = [] 685 | for index, coord_index, count in zip(indices, mapping, counts): 686 | #if count >= count_dicts[scale][index] * 0.5: 687 | if count >= scale_count_thresholds[scale]: 688 | valid_coords.append(node_coords[coord_index]) 689 | pass 690 | continue 691 | if len(valid_coords) > 0: 692 | info[scale] = (valid_coords, np.array([coord_node_map[scale][toIndex(coord, full_scale // pow(2, scale))] for coord in valid_coords])) 693 | pass 694 | pass 695 | node_info.append((info, label_counts)) 696 | continue 697 | continue 698 | instances = ori_node_mapping 699 | 700 | return instances, intermediate_instances 701 | --------------------------------------------------------------------------------