├── LICENSE ├── README.md └── scene-synth ├── __init__.py ├── cat.py ├── data ├── __init__.py ├── dataset.py ├── house.py ├── object.py ├── object_data.py ├── projection.py ├── rendered.py └── top_down.py ├── dims.py ├── fast_synth.py ├── filters ├── bedroom.py ├── collision.py ├── floor_node.py ├── global_category_filter.py ├── good_house.py ├── livingroom.py ├── office.py ├── renderable.py ├── room_type.py └── toilet.py ├── latent_dataset.py ├── loc.py ├── loc_dataset.py ├── math_utils ├── OBB.py ├── Simulator.py └── __init__.py ├── model_prior.py ├── models ├── __init__.py ├── resnet.py └── utils.py ├── obj_dims.py ├── orient.py ├── priors ├── arrangement.py ├── observations.py └── pairwise.py ├── requirements.txt ├── scene_filter.py ├── support_prior.py └── utils.py /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2018, Brown University, Providence, RI. 2 | 3 | All Rights Reserved 4 | 5 | Permission to use, copy, modify, and distribute this software and 6 | its documentation for any purpose other than its incorporation into a 7 | commercial product or service is hereby granted without fee, provided 8 | that the above copyright notice appear in all copies and that both 9 | that copyright notice and this permission notice appear in supporting 10 | documentation, and that the name of Brown University not be used in 11 | advertising or publicity pertaining to distribution of the software 12 | without specific, written prior permission. 13 | 14 | BROWN UNIVERSITY DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, 15 | INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR ANY 16 | PARTICULAR PURPOSE. IN NO EVENT SHALL BROWN UNIVERSITY BE LIABLE FOR 17 | ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES 18 | WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN 19 | ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF 20 | OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. 21 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Fast and Flexible Indoor Scene Synthesis via Deep Convolutional Generative Models 2 | PyTorch code for our CVPR Paper [Fast and Flexible Indoor Scene Synthesis via Deep Convolutional Generative Models] 3 | 4 | Requires PyTorch 0.4 to run, additional python library requirements could be found at [/scene-synth/requirements.txt](/scene-synth/requirements.txt). Run 5 | ```bash 6 | pip install -r requirements.txt 7 | ``` 8 | to install. 9 | 10 | We trained our models on the SUNCG dataset. Due to an ongoing legal dispute, SUNCG is currently unavailable to download. As such, we feel that we should not provide anything derived from the dataset, including the pre-trained models, as well as several metadata files that our code rely on. 11 | We will update this page in the event that this situation changes in the future. 12 | 13 | `cat.py`, `loc.py`, `orient.py`, `dims.py` contains the four components of our model respectively. `fast_synth.py` contains the main scene synthesis code. We will update with instructions on them as well. 14 | 15 | We have a follow-up of this paper set to appear at SIGGRAPH 2019. In that paper, we introduce a high level "planning" module, formulated as a graph convolutional network, and uses the modules introduced here to fill out/"instantiate" low level details. Please refer to https://github.com/brownvc/planit for more details. 16 | -------------------------------------------------------------------------------- /scene-synth/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/brownvc/fast-synth/863213bedee078cf531d1518708515ca3928bf63/scene-synth/__init__.py -------------------------------------------------------------------------------- /scene-synth/cat.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import torch 3 | import torch.optim as optim 4 | from latent_dataset import LatentDataset 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | from models import * 8 | import math 9 | import utils 10 | 11 | """ 12 | Module that predicts the category of the next object 13 | """ 14 | 15 | parser = argparse.ArgumentParser(description='cat') 16 | parser.add_argument('--data-folder', type=str, default="bedroom_6x6", metavar='S') 17 | parser.add_argument('--num-workers', type=int, default=6, metavar='N') 18 | parser.add_argument('--last-epoch', type=int, default=-1, metavar='N') 19 | parser.add_argument('--train-size', type=int, default=5000, metavar='N') 20 | parser.add_argument('--save-dir', type=str, default="cat_test", metavar='S') 21 | parser.add_argument('--save-every-n-epochs', type=int, default=5, metavar='N') 22 | parser.add_argument('--lr', type=float, default=0.0005, metavar='N') 23 | args = parser.parse_args() 24 | 25 | save_dir = args.save_dir 26 | save_every = args.save_every_n_epochs 27 | utils.ensuredir(save_dir) 28 | 29 | batch_size = 16 30 | latent_dim = 200 31 | epoch_size = 10000 32 | 33 | 34 | # --------------------------------------------------------------------------------------- 35 | class NextCategory(nn.Module): 36 | 37 | def __init__(self, n_input_channels, n_categories, bottleneck_size): 38 | super(NextCategory, self).__init__() 39 | 40 | activation = nn.LeakyReLU() 41 | 42 | self.cat_prior_img = nn.Sequential( 43 | resnet18(num_input_channels=n_input_channels, num_classes=bottleneck_size), 44 | nn.BatchNorm1d(bottleneck_size), 45 | activation 46 | ) 47 | self.cat_prior_counts = nn.Sequential( 48 | nn.Linear(n_categories, bottleneck_size), 49 | nn.BatchNorm1d(bottleneck_size), 50 | activation, 51 | nn.Linear(bottleneck_size, bottleneck_size), 52 | nn.BatchNorm1d(bottleneck_size), 53 | activation 54 | ) 55 | self.cat_prior_final = nn.Sequential( 56 | nn.Linear(2*bottleneck_size, bottleneck_size), 57 | nn.BatchNorm1d(bottleneck_size), 58 | activation, 59 | # +1 -> the 'stop' category 60 | nn.Linear(bottleneck_size, n_categories+1) 61 | ) 62 | 63 | def forward(self, input_scene, catcount): 64 | cat_img = self.cat_prior_img(input_scene) 65 | cat_count = self.cat_prior_counts(catcount) 66 | catlogits = self.cat_prior_final(torch.cat([cat_img, cat_count], dim=-1)) 67 | return catlogits 68 | 69 | def sample(self, input_scene, catcount, temp=1, return_logits=False): 70 | logits = self.forward(input_scene, catcount) 71 | if temp != 1.0: 72 | logits = logits * temp 73 | #logits[:,-1] += 1 74 | cat = torch.distributions.Categorical(logits=logits).sample() 75 | if return_logits: 76 | return cat, logits 77 | else: 78 | return cat 79 | # --------------------------------------------------------------------------------------- 80 | 81 | if __name__ == '__main__': 82 | 83 | data_root_dir = utils.get_data_root_dir() 84 | with open(f"{data_root_dir}/{args.data_folder}/final_categories_frequency", "r") as f: 85 | lines = f.readlines() 86 | num_categories = len(lines)-2 # -2 for 'window' and 'door' 87 | 88 | num_input_channels = num_categories+8 89 | 90 | logfile = open(f"{save_dir}/log.txt", 'w') 91 | def LOG(msg): 92 | print(msg) 93 | logfile.write(msg + '\n') 94 | logfile.flush() 95 | 96 | 97 | LOG('Building model...') 98 | model = NextCategory(num_input_channels, num_categories, latent_dim).cuda() 99 | 100 | 101 | LOG('Building datasets...') 102 | train_dataset = LatentDataset( 103 | data_folder = args.data_folder, 104 | scene_indices = (0, args.train_size), 105 | cat_only = True, 106 | importance_order = True 107 | ) 108 | validation_dataset = LatentDataset( 109 | data_folder = args.data_folder, 110 | scene_indices = (args.train_size, args.train_size+160), 111 | # seed = 42, 112 | cat_only = True, 113 | importance_order = True 114 | ) 115 | # train_dataset.build_cat2scene() 116 | 117 | 118 | LOG('Building data loader...') 119 | train_loader = torch.utils.data.DataLoader( 120 | train_dataset, 121 | batch_size = batch_size, 122 | num_workers = args.num_workers, 123 | shuffle = True 124 | ) 125 | validation_loader = torch.utils.data.DataLoader( 126 | validation_dataset, 127 | batch_size = batch_size, 128 | num_workers = 0, 129 | shuffle = True 130 | ) 131 | 132 | optimizer = optim.Adam(list(model.parameters()), 133 | lr = args.lr 134 | ) 135 | 136 | if args.last_epoch < 0: 137 | load = False 138 | starting_epoch = 0 139 | else: 140 | load = True 141 | last_epoch = args.last_epoch 142 | 143 | if load: 144 | LOG('Loading saved models...') 145 | model.load_state_dict(torch.load(f"{save_dir}/nextcat_{last_epoch}.pt")) 146 | optimizer.load_state_dict(torch.load(f"{save_dir}/nextcat_optim_backup.pt")) 147 | starting_epoch = last_epoch + 1 148 | 149 | current_epoch = starting_epoch 150 | num_seen = 0 151 | 152 | model.train() 153 | LOG(f'=========================== Epoch {current_epoch} ===========================') 154 | 155 | 156 | loss_running_avg = 0 157 | 158 | def train(): 159 | global num_seen, current_epoch, loss_running_avg 160 | 161 | for batch_idx, (input_img, t_cat, catcount) in enumerate(train_loader): 162 | 163 | # Get rid of singleton dimesion in t_cat (NLLLoss complains about this) 164 | t_cat = torch.squeeze(t_cat) 165 | 166 | input_img, t_cat, catcount = input_img.cuda(), t_cat.cuda(), catcount.cuda() 167 | 168 | optimizer.zero_grad() 169 | logits = model(input_img, catcount) 170 | 171 | loss = F.cross_entropy(logits, t_cat) 172 | loss.backward() 173 | optimizer.step() 174 | loss_precision = 4 175 | LOG(f'Loss: {loss.cpu().data.numpy():{loss_precision}.{loss_precision}}') 176 | 177 | num_seen += batch_size 178 | 179 | if num_seen % 800 == 0: 180 | LOG(f'Examples {num_seen}/{epoch_size}') 181 | if num_seen > 0 and num_seen % epoch_size == 0: 182 | validate() 183 | num_seen = 0 184 | if current_epoch % save_every == 0: 185 | torch.save(model.state_dict(), f"{save_dir}/nextcat_{current_epoch}.pt") 186 | torch.save(optimizer.state_dict(), f"{save_dir}/nextcat_optim_backup.pt") 187 | current_epoch += 1 188 | LOG(f'=========================== Epoch {current_epoch} ===========================') 189 | 190 | def validate(): 191 | LOG('Validating') 192 | model.eval() 193 | total_loss = 0 194 | num_correct_top1 = 0 195 | num_correct_top5 = 0 196 | num_batches = 0 197 | num_items = 0 198 | 199 | for batch_idx, (input_img, t_cat, catcount) in enumerate(validation_loader): 200 | 201 | # Get rid of singleton dimesion in t_cat (NLLLoss complains about this) 202 | t_cat = torch.squeeze(t_cat) 203 | 204 | input_img, t_cat, catcount = input_img.cuda(), t_cat.cuda(), catcount.cuda() 205 | 206 | with torch.no_grad(): 207 | logits = model(input_img, catcount) 208 | loss = F.cross_entropy(logits, t_cat) 209 | 210 | total_loss += loss.cpu().data.numpy() 211 | 212 | # TODO: Why is top1 always giving me zero...? 213 | _, argmax = logits.max(dim=-1) 214 | num_correct_top1 += (argmax == t_cat).sum() 215 | 216 | lsorted, lsorted_idx = torch.sort(logits, dim=-1, descending=True) 217 | lsorted_idx_top5 = lsorted_idx[:, 0:5] 218 | num_correct = torch.zeros(input_img.shape[0]).cuda() 219 | for i in range(0, 5): 220 | correct = lsorted_idx_top5[0:, i] == t_cat 221 | num_correct = torch.max(num_correct, correct.float()) 222 | num_correct_top5 += num_correct.sum() 223 | 224 | num_batches += 1 225 | num_items += input_img.shape[0] 226 | 227 | LOG(f'Average Loss: {total_loss / num_batches}') 228 | LOG(f'Top 1 Accuracy: {num_correct_top1 / num_items}') 229 | LOG(f'Top 5 Accuracy: {num_correct_top5 / num_items}') 230 | model.train() 231 | 232 | # Train forever (until we stop it) 233 | while True: 234 | train() 235 | -------------------------------------------------------------------------------- /scene-synth/data/__init__.py: -------------------------------------------------------------------------------- 1 | from data.object_data import * 2 | from data.projection import * 3 | from data.object import * 4 | from data.house import * 5 | from data.top_down import * 6 | from data.dataset import * 7 | from data.rendered import * 8 | -------------------------------------------------------------------------------- /scene-synth/data/house.py: -------------------------------------------------------------------------------- 1 | """ 2 | Three level House-Level-Node/Room representation of SUNCG 3 | """ 4 | import os 5 | import json 6 | import numpy as np 7 | from data import ObjectData 8 | import utils 9 | 10 | class House(): 11 | """ 12 | Represents a House 13 | See SUNCG Toolbox for a detailed description of the structure of the json file 14 | describing a house 15 | """ 16 | object_data = ObjectData() 17 | def __init__(self, index=0, id_=None, house_json=None, file_dir=None, 18 | include_support_information=True, include_arch_information=False): 19 | """ 20 | Get a set of rooms from the house which satisfies a certain criteria 21 | Parameters 22 | ---------- 23 | index (int): The index of the house among all houses sorted in alphabetical order 24 | default way of loading a house 25 | id_ (string, optional): If set, then the house with the specified directory name is chosen 26 | house_json(json, optional): If set, then the specified json object 27 | is used directly to initiate the house 28 | file_dir (string, optional): If set, then the json pointed to by file_dir will be loaded 29 | include_support_information(bool): If true, then support information is loaded from suncg_data/house_relations 30 | might not be available, so defaults to False 31 | include_arch_information (bool): If true, then arch information is loaded from suncg_data/wall 32 | might not be available, so defaults to False 33 | """ 34 | data_dir = utils.get_data_root_dir() 35 | if house_json is None: 36 | if file_dir is None: 37 | house_dir = data_dir + "/suncg_data/house/" 38 | 39 | if id_ is None: 40 | houses = dict(enumerate(os.listdir(house_dir))) 41 | self.__dict__ = json.loads(open(house_dir+houses[index]+"/house.json", 'r').read()) 42 | else: 43 | self.__dict__ = json.loads(open(house_dir+id_+"/house.json", 'r').read()) 44 | else: 45 | self.__dict__ = json.loads(open(file_dir, 'r').read()) 46 | else: 47 | self.__dict__ = house_json 48 | 49 | self.filters = [] 50 | self.levels = [Level(l,self) for l in self.levels] 51 | self.rooms = [r for l in self.levels for r in l.rooms] 52 | self.nodes = [n for l in self.levels for n in l.nodes] 53 | self.node_dict = {id_: n for l in self.levels for id_,n in l.node_dict.items()} 54 | if include_support_information: 55 | house_stats_dir = data_dir + "/suncg_data/house_relations/" 56 | stats = json.loads(open(house_stats_dir+self.id+"/"+self.id+".stats.json", 'r').read()) 57 | supports = [(s["parent"],s["child"]) for s in stats["relations"]["support"]] 58 | for parent, child in supports: 59 | if child not in self.node_dict: 60 | print(f'Warning: support relation {supports} involves not present {child} node') 61 | continue 62 | if "f" in parent: 63 | self.get_node(child).parent = "Floor" 64 | elif "c" in parent: 65 | self.get_node(child).parent = "Ceiling" 66 | elif len(parent.split("_")) > 2: 67 | self.get_node(child).parent = "Wall" 68 | else: 69 | if parent not in self.node_dict: 70 | print(f'Warning: support relation {supports} involves not present {parent} node') 71 | continue 72 | self.get_node(parent).child.append(self.get_node(child)) 73 | self.get_node(child).parent = self.get_node(parent) 74 | if include_arch_information: 75 | house_arch_dir = data_dir + '/suncg_data/wall/' 76 | arch = json.loads(open(house_arch_dir+self.id+'/'+self.id+'.arch.json', 'r').read()) 77 | self.walls = [w for w in arch['elements'] if w['type'] == 'Wall'] 78 | 79 | def get_node(self, id_): 80 | return self.node_dict[id_] 81 | 82 | def get_rooms(self, filters=None): 83 | """ 84 | Get a set of rooms from the house which satisfies a certain criteria 85 | Parameters 86 | ---------- 87 | filters (list[room_filter]): room_filter is tuple[Room,House] which returns 88 | if the Room should be included 89 | 90 | Returns 91 | ------- 92 | list[Room] 93 | """ 94 | if filters is None: filters = self.filters 95 | if not isinstance(filters, list): filters = [filters] 96 | rooms = self.rooms 97 | for filter_ in filters: 98 | rooms = [room for room in rooms if filter_(room, self)] 99 | return rooms 100 | 101 | def filter_rooms(self, filters): 102 | """ 103 | Similar to get_rooms, but overwrites self.node instead of returning a list 104 | """ 105 | self.rooms = self.get_rooms(filters) 106 | 107 | def trim(self): 108 | """ 109 | Get rid of some intermediate attributes 110 | """ 111 | nodes = list(self.node_dict.values()) 112 | if hasattr(self, 'rooms'): 113 | nodes.extend(self.rooms) 114 | for n in nodes: 115 | for attr in ['xform', 'obb', 'frame', 'model2world']: 116 | if hasattr(n, attr): 117 | delattr(n, attr) 118 | self.nodes = None 119 | self.walls = None 120 | for room in self.rooms: 121 | room.filters = None 122 | self.levels = None 123 | self.node_dict = None 124 | self.filters = None 125 | 126 | 127 | class Level(): 128 | """ 129 | Represents a floor level in the house 130 | Currently mostly just used to parse the list of nodes and rooms 131 | Might change in the future 132 | """ 133 | def __init__(self, dict_, house): 134 | self.__dict__ = dict_ 135 | self.house = house 136 | invalid_nodes = [n["id"] for n in self.nodes if (not n["valid"]) and "id" in n] 137 | self.nodes = [Node(n,self) for n in self.nodes if n["valid"]] 138 | self.node_dict = {n.id: n for n in self.nodes} 139 | self.nodes = list(self.node_dict.values()) # deduplicate nodes with same id 140 | self.rooms = [Room(n, ([self.node_dict[i] for i in [f"{self.id}_{j}" \ 141 | for j in list(set(n.nodeIndices))] if i not in invalid_nodes]), self) \ 142 | for n in self.nodes if n.isRoom() and hasattr(n, 'nodeIndices')] 143 | 144 | class Room(): 145 | """ 146 | Represents a room in the house 147 | """ 148 | def __init__(self, room, nodes, level): 149 | self.__dict__ = room.__dict__ 150 | #self.room = room 151 | self.nodes = nodes 152 | #self.level = level 153 | self.filters = [] 154 | self.house_id = level.house.id 155 | 156 | def get_nodes(self, filters=None): 157 | """ 158 | Get a set of nodes from the room which satisfies a certain criteria 159 | Parameters 160 | ---------- 161 | filters (list[node_filter]): node_filter is tuple[Node,Room] which returns 162 | if the Node should be included 163 | 164 | Returns 165 | ------- 166 | list[Node] 167 | """ 168 | if filters is None: filters = self.filters 169 | if not isinstance(filters, list): filters = [filters] 170 | nodes = self.nodes 171 | for filter_ in filters: 172 | nodes = [node for node in nodes if filter_(node, self)] 173 | return nodes 174 | 175 | def filter_nodes(self, filters): 176 | """ 177 | Similar to get_nodes, but overwrites self.node instead of returning a list 178 | """ 179 | self.nodes = self.get_nodes(filters) 180 | 181 | class Node(): 182 | """ 183 | Basic unit of representation of SUNCG 184 | Usually a room or an object 185 | Refer to SUNCG toolbox for the possible set of attributes 186 | """ 187 | warning = True 188 | def __init__(self, dict_, level): 189 | self.__dict__ = dict_ 190 | #self.level = level 191 | self.parent = None 192 | self.child = [] 193 | if hasattr(self, 'bbox'): 194 | (self.xmin, self.zmin, self.ymin) = self.bbox["min"] 195 | (self.xmax, self.zmax, self.ymax) = self.bbox["max"] 196 | (self.width, self.length) = sorted([self.xmax - self.xmin, self.ymax - self.ymin]) 197 | self.height = self.zmax - self.zmin 198 | else: # warn and populate with default bbox values 199 | if self.warning: 200 | print(f'Warning: node id={self.id} is valid but has no bbox, setting default values') 201 | (self.xmin, self.zmin, self.ymin) = (0, 0, 0) 202 | (self.xmax, self.zmax, self.ymax) = (0, 0, 0) 203 | if hasattr(self, 'transform') and hasattr(self, 'modelId'): 204 | t = np.asarray(self.transform).reshape(4,4) 205 | #Special cases of models with were not aligned in the way we want 206 | alignment_matrix = House.object_data.get_alignment_matrix(self.modelId) 207 | if alignment_matrix is not None: 208 | t = np.dot(np.linalg.inv(alignment_matrix), t) 209 | self.transform = list(t.flatten()) 210 | 211 | #If a reflection is present, switch to the mirrored model 212 | #And adjust transform accordingly 213 | if np.linalg.det(t) < 0: 214 | t_reflec = np.asarray([[-1, 0, 0, 0], \ 215 | [0, 1, 0, 0], \ 216 | [0, 0, 1, 0], \ 217 | [0, 0, 0, 1]]) 218 | t = np.dot(t_reflec, t) 219 | self.modelId += "_mirror" 220 | self.transform = list(t.flatten()) 221 | 222 | def isRoom(self): 223 | return self.type == "Room" 224 | 225 | if __name__ == "__main__": 226 | a = House(id_ = "f53c0878c4db848bfa43163473b74245") 227 | for room in a.rooms: 228 | if room.id == "0_7": 229 | print(room.__dict__) 230 | -------------------------------------------------------------------------------- /scene-synth/data/object.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | import os 3 | import numpy as np 4 | from data import ObjectData 5 | import utils 6 | 7 | """ 8 | Taking care of wavefront obj files 9 | Convert to pickle for faster loading 10 | Currently just geometric information. 11 | Call this file once to create a pickled version of the objects 12 | For faster loading in the future 13 | """ 14 | 15 | class Obj(): 16 | """ 17 | Standard vertex-face representation, triangulated 18 | Order: x, z, y 19 | """ 20 | object_data = ObjectData() 21 | def __init__(self, modelId, houseId=None, from_source=False, is_room=False, mirror=False): 22 | """ 23 | Parameters 24 | ---------- 25 | modelId (string): name of the object to be loaded 26 | houseId (string, optional): If loading a room, specify which house does the room belong to 27 | from_source (bool, optional): If false, loads the pickled version of the object 28 | need to call object.py once to create the pickled version. 29 | does not apply for rooms 30 | mirror (bool, optional): If true, loads the mirroed version 31 | """ 32 | if is_room: from_source = True #Don't want to save rooms... 33 | data_dir = utils.get_data_root_dir() 34 | self.vertices = [] 35 | self.faces = [] 36 | if from_source: 37 | if is_room: 38 | path = f"{data_dir}/suncg_data/room/{houseId}/{modelId}.obj" 39 | else: 40 | path = f"{data_dir}/suncg_data/object/{modelId}/{modelId}.obj" 41 | with open(path,"r") as f: 42 | for line in f: 43 | data = line.split() 44 | if len(data) > 0: 45 | if data[0] == "v": 46 | v = np.asarray([float(i) for i in data[1:4]]+[1]) 47 | self.vertices.append(v) 48 | if data[0] == "f": 49 | face = [int(i.split("/")[0])-1 for i in data[1:]] 50 | if len(face) == 4: 51 | self.faces.append([face[0],face[1],face[2]]) 52 | self.faces.append([face[0],face[2],face[3]]) 53 | elif len(face) == 3: 54 | self.faces.append([face[0],face[1],face[2]]) 55 | else: 56 | print(f"Found a face with {len(face)} edges!!!") 57 | 58 | self.vertices = np.asarray(self.vertices) 59 | if not is_room and Obj.object_data.get_alignment_matrix(modelId) is not None: 60 | self.transform(data.get_alignment_matrix(modelId)) 61 | else: 62 | with open(f"{data_dir}/object/{modelId}/vertices.pkl", "rb") as f: 63 | self.vertices = pickle.load(f) 64 | with open(f"{data_dir}/object/{modelId}/faces.pkl", "rb") as f: 65 | self.faces = pickle.load(f) 66 | 67 | self.bbox_min = np.min(self.vertices, 0) 68 | self.bbox_max = np.max(self.vertices, 0) 69 | 70 | if mirror: 71 | t = np.asarray([[-1, 0, 0, 0], \ 72 | [0, 1, 0, 0], \ 73 | [0, 0, 1, 0], \ 74 | [0, 0, 0, 1]]) 75 | self.transform(t) 76 | self.modelId = modelId+"_mirror" 77 | else: 78 | self.modelId = modelId 79 | 80 | def save(self): 81 | data_dir = utils.get_data_root_dir() 82 | dest_dir = f"{data_dir}/object/{self.modelId}" 83 | if not os.path.exists(dest_dir): 84 | os.makedirs(dest_dir) 85 | with open(f"{dest_dir}/vertices.pkl", "wb") as f: 86 | pickle.dump(self.vertices, f, pickle.HIGHEST_PROTOCOL) 87 | with open(f"{dest_dir}/faces.pkl", "wb") as f: 88 | pickle.dump(self.faces, f, pickle.HIGHEST_PROTOCOL) 89 | 90 | 91 | def transform(self, t): 92 | self.vertices = np.dot(self.vertices, t) 93 | 94 | def get_triangles(self): 95 | for face in self.faces: 96 | yield (self.vertices[face[0]][:3], \ 97 | self.vertices[face[1]][:3], \ 98 | self.vertices[face[2]][:3],) 99 | 100 | def xmax(self): 101 | return np.amax(self.vertices, axis = 0)[0] 102 | 103 | def xmin(self): 104 | return np.amin(self.vertices, axis = 0)[0] 105 | 106 | def ymax(self): 107 | return np.amax(self.vertices, axis = 0)[2] 108 | 109 | def ymin(self): 110 | return np.amin(self.vertices, axis = 0)[2] 111 | 112 | def zmax(self): 113 | return np.amax(self.vertices, axis = 0)[1] 114 | 115 | def zmin(self): 116 | return np.amin(self.vertices, axis = 0)[1] 117 | 118 | def parse_objects(): 119 | """ 120 | parse .obj objects and save them to pickle files 121 | """ 122 | data_dir = utils.get_data_root_dir() 123 | obj_dir = data_dir + "/suncg_data/object/" 124 | print("Parsing SUNCG object files...") 125 | l = len(os.listdir(obj_dir)) 126 | for (i, modelId) in enumerate(os.listdir(obj_dir)): 127 | print(f"{i+1} of {l}...", end="\r") 128 | if not modelId in ["mgcube", ".DS_Store"]: 129 | o = Obj(modelId, from_source = True) 130 | o.save() 131 | o = Obj(modelId, from_source = True, mirror = True) 132 | o.save() 133 | print() 134 | 135 | if __name__ == "__main__": 136 | parse_objects() 137 | 138 | 139 | 140 | -------------------------------------------------------------------------------- /scene-synth/data/object_data.py: -------------------------------------------------------------------------------- 1 | import csv 2 | import os 3 | import numpy as np 4 | import utils 5 | 6 | class ObjectCategories(): 7 | """ 8 | Determine which categories does each object belong to 9 | """ 10 | def __init__(self): 11 | print("Initializing Category Map...") #Debug purposes 12 | fname = "ModelCategoryMapping.csv" 13 | self.model_to_categories = {} 14 | 15 | root_dir = os.path.dirname(os.path.abspath(__file__)) 16 | model_cat_file = f"{root_dir}/{fname}" 17 | 18 | with open(model_cat_file, "r") as f: 19 | categories = csv.reader(f) 20 | for l in categories: 21 | self.model_to_categories[l[1]] = [l[2],l[3],l[5]] 22 | 23 | fname = "ModelCategoryMapping_grains.csv" 24 | self.model_to_categories2 = {} 25 | 26 | root_dir = os.path.dirname(os.path.abspath(__file__)) 27 | model_cat_file = f"{root_dir}/{fname}" 28 | 29 | with open(model_cat_file, "r") as f: 30 | categories = csv.reader(f) 31 | for l in categories: 32 | self.model_to_categories2[l[1]] = [l[8]] 33 | 34 | def get_fine_category(self, model_id): 35 | model_id = model_id.replace("_mirror","") 36 | return self.model_to_categories[model_id][0] 37 | 38 | def get_coarse_category(self, model_id): 39 | model_id = model_id.replace("_mirror","") 40 | return self.model_to_categories[model_id][1] 41 | 42 | def get_final_category(self, model_id): 43 | """ 44 | Final categories used in the generated dataset 45 | Minor tweaks from fine categories 46 | """ 47 | model_id = model_id.replace("_mirror","") 48 | 49 | #return self.model_to_categories2[model_id][0] 50 | 51 | category = self.model_to_categories[model_id][0] 52 | if model_id == "199": 53 | category = "dressing_table_with_stool" 54 | if model_id in ["150", "453", "s__1138", "s__1251"]: 55 | category = "bidet" 56 | if category == "nightstand": 57 | category = "stand" 58 | if category == "bookshelf": 59 | category = "shelving" 60 | if category == "books": 61 | category = "book" 62 | if category == "xbox" or category == "playstation": 63 | category = "console" 64 | return category 65 | 66 | class ObjectCategoriesGrains(): 67 | """ 68 | Determine which categories does each object belong to 69 | """ 70 | def __init__(self): 71 | print("Initializing Category Map...") #Debug purposes 72 | fname = "ModelCategoryMapping_grains.csv" 73 | self.model_to_categories = {} 74 | 75 | root_dir = os.path.dirname(os.path.abspath(__file__)) 76 | model_cat_file = f"{root_dir}/{fname}" 77 | 78 | with open(model_cat_file, "r") as f: 79 | categories = csv.reader(f) 80 | for l in categories: 81 | self.model_to_categories[l[1]] = [l[8]] 82 | 83 | def get_grains_category(self, model_id): 84 | model_id = model_id.replace("_mirror","") 85 | return self.model_to_categories[model_id][0] 86 | 87 | class ObjectData(): 88 | """ 89 | Various information associated with the objects 90 | """ 91 | def __init__(self): 92 | print("Initializing Object Data") #Debug purposes 93 | self.model_to_data = {} 94 | 95 | root_dir = os.path.dirname(os.path.abspath(__file__)) 96 | model_data_file = f"{root_dir}/Models.csv" 97 | 98 | with open(model_data_file, "r") as f: 99 | data = csv.reader(f) 100 | for l in data: 101 | if l[0] != 'id': # skip header row 102 | self.model_to_data[l[0]] = l[1:] 103 | 104 | def get_front(self, model_id): 105 | model_id = model_id.replace("_mirror","") 106 | # TODO compensate for mirror (can have effect if not axis-aligned in model space) 107 | return [float(a) for a in self.model_to_data[model_id][0].split(",")] 108 | 109 | def get_aligned_dims(self, model_id): 110 | """Return canonical alignment dimensions of model *in meters*""" 111 | model_id = model_id.replace('_mirror', '') # NOTE dims don't change since mirroring is symmetric on yz plane 112 | return [float(a)/100.0 for a in self.model_to_data[model_id][4].split(',')] 113 | 114 | def get_model_semantic_frame_matrix(self, model_id): 115 | """Return canonical semantic frame matrix for model. 116 | Transforms from semantic frame [0,1]^3, [x,y,z] = [right,up,back] to raw model coordinates.""" 117 | up = np.array([0, 1, 0]) # NOTE: up is assumed to always be +Y for SUNCG objects 118 | front = np.array(self.get_front(model_id)) 119 | has_mirror = '_mirror' in model_id 120 | model_id = model_id.replace('_mirror', '') 121 | hdims = np.array(self.get_aligned_dims(model_id)) * 0.5 122 | p_min = np.array([float(a) for a in self.model_to_data[model_id][2].split(',')]) 123 | p_max = np.array([float(a) for a in self.model_to_data[model_id][3].split(',')]) 124 | if has_mirror: 125 | p_max[0] = -p_max[0] 126 | p_min[0] = -p_min[0] 127 | model_space_center = (p_max + p_min) * 0.5 128 | m = np.identity(4) 129 | m[:3, 0] = np.cross(front, up) * hdims[0] # +x = right 130 | m[:3, 1] = np.array(up) * hdims[1] # +y = up 131 | m[:3, 2] = -front * hdims[2] # +z = back = -front 132 | m[:3, 3] = model_space_center # origin = center 133 | # r = np.identity(3) 134 | # r[:3, 0] = np.cross(front, up) # +x = right 135 | # r[:3, 1] = np.array(up) # +y = up 136 | # r[:3, 2] = -front # +z = back = -front 137 | # s = np.identity(3) 138 | # s[0, 0] = hdims[0] 139 | # s[1, 1] = hdims[1] 140 | # s[2, 2] = hdims[2] 141 | # sr = np.matmul(s, r) 142 | # m = np.identity(4) 143 | # m[:3, :3] = sr 144 | # m[:3, 3] = model_space_center 145 | return m 146 | 147 | def get_alignment_matrix(self, model_id): 148 | """ 149 | Since some models in the dataset are not aligned in the way we want 150 | Generate matrix that realign them 151 | """ 152 | #alignment happens BEFORE mirror, so make sure no mirrored 153 | #object will ever call this! 154 | #model_id = model_id.replace("_mirror","") 155 | if self.get_front(model_id) == [0,0,1]: 156 | return None 157 | else: 158 | #Let's just do case by case enumeration!!! 159 | if model_id in ["106", "114", "142", "323", "333", "363", "364", 160 | "s__1782", "s__1904"]: 161 | M = [[-1,0,0,0], 162 | [0,1,0,0], 163 | [0,0,-1,0], 164 | [0,0,0,1]] 165 | elif model_id in ["s__1252", "s__400", "s__885"]: 166 | M = [[0,0,-1,0], 167 | [0,1,0,0], 168 | [1,0,0,0], 169 | [0,0,0,1]] 170 | elif model_id in ["146", "190", "s__404", "s__406"]: 171 | M = [[0,0,1,0], 172 | [0,1,0,0], 173 | [-1,0,0,0], 174 | [0,0,0,1]] 175 | else: 176 | print(model_id) 177 | raise NotImplementedError 178 | 179 | return np.asarray(M) 180 | 181 | def get_setIds(self, model_id): 182 | model_id = model_id.replace("_mirror","") 183 | return [a for a in self.model_to_data[model_id][8].split(",")] 184 | 185 | if __name__ == "__main__": 186 | b = ObjectData() 187 | #print(b.get_front("40")) 188 | #print(b.get_setIds("s__2240")) 189 | -------------------------------------------------------------------------------- /scene-synth/data/projection.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import math 3 | import scipy.misc as m 4 | 5 | class ProjectionGenerator(): 6 | """ 7 | Generates projection between original 3D space and rendered 2D image space 8 | Given the position of the room 9 | """ 10 | def __init__(self, room_size_cap=(6.05,4.05,6.05), zpad=0.5, img_size=256): 11 | """ 12 | See top_down.TopDownView for explanation of room_size, zpad and img_size 13 | """ 14 | self.room_size_cap = room_size_cap 15 | self.zpad = zpad 16 | self.img_size = img_size 17 | self.xscale = self.img_size / self.room_size_cap[0] 18 | self.yscale = self.img_size / self.room_size_cap[2] 19 | self.zscale = 1.0 / (self.room_size_cap[1] + self.zpad) 20 | 21 | def get_projection(self, room): 22 | """ 23 | Generates projection matrices specific to a room, 24 | need to be room-specific since every room is located in a different position, 25 | but they are all rendered centered in the image """ 26 | xscale, yscale, zscale = self.xscale, self.yscale, self.zscale 27 | 28 | xshift = -(room.xmin * 0.5 + room.xmax * 0.5 - self.room_size_cap[0] / 2.0) 29 | yshift = -(room.ymin * 0.5 + room.ymax * 0.5 - self.room_size_cap[2] / 2.0) 30 | zshift = -room.zmin + self.zpad 31 | 32 | t_scale = np.asarray([[xscale, 0, 0, 0], \ 33 | [0, zscale, 0, 0], \ 34 | [0, 0, yscale, 0], \ 35 | [0, 0, 0, 1]]) 36 | 37 | t_shift = np.asarray([[1, 0, 0, 0], \ 38 | [0, 1, 0, 0], \ 39 | [0, 0, 1, 0], \ 40 | [xshift, zshift, yshift, 1]]) 41 | 42 | t_3To2 = np.dot(t_shift, t_scale) 43 | 44 | t_scale = np.asarray([[1/xscale, 0, 0, 0], \ 45 | [0, 1/zscale, 0, 0], \ 46 | [0, 0, 1/yscale, 0], \ 47 | [0, 0, 0, 1]]) 48 | 49 | t_shift = np.asarray([[1, 0, 0, 0], \ 50 | [0, 1, 0, 0], \ 51 | [0, 0, 1, 0], \ 52 | [-xshift, -zshift, -yshift, 1]]) 53 | 54 | t_2To3 = np.dot(t_scale, t_shift) 55 | 56 | return Projection(t_3To2, t_2To3, self.img_size) 57 | 58 | class Projection(): 59 | def __init__(self, t_2d, t_3d, img_size): 60 | self.t_2d = t_2d 61 | self.t_3d = t_3d 62 | self.img_size = img_size 63 | 64 | def to_2d(self, t=None): 65 | """ 66 | Parameters 67 | ---------- 68 | t(Matrix or None): transformation matrix of the object 69 | if None, then returns room projection 70 | """ 71 | if t is None: 72 | return self.t_2d 73 | else: 74 | return np.dot(t, self.t_2d) 75 | 76 | def to_3d(self, t=None): 77 | if t is None: 78 | return self.t_3d 79 | else: 80 | return np.dot(t, self.t_3d) 81 | 82 | def get_ortho_parameters(self): 83 | bottom_left = np.asarray([0,0,0,1]) 84 | top_right = np.asarray([self.img_size,1,self.img_size,1]) 85 | bottom_left = np.dot(bottom_left, self.t_3d) 86 | top_right = np.dot(top_right, self.t_3d) 87 | return (bottom_left[0], top_right[0], bottom_left[2], top_right[2], bottom_left[1], top_right[1]) 88 | 89 | if __name__ == "__main__": 90 | pass 91 | #from . import * 92 | #a = ProjectionGenerator() 93 | #house = House(include_support_information=False, file_dir="/home/kai/data/3d-scene/3d-scene/data/bedroom_final/json/4200.json") 94 | #room = house.rooms[0] 95 | #p = a.get_projection(room) 96 | #print(p.to_2d()) 97 | #print(p.to_3d()) 98 | #print(p.get_ortho_parameters()) 99 | -------------------------------------------------------------------------------- /scene-synth/data/rendered.py: -------------------------------------------------------------------------------- 1 | """ 2 | Handles pre-rendered rooms generated by data.top_down 3 | RenderedScene loads the pre-rendered data 4 | along with other precomputed information, 5 | and creates RenderedComposite, which combines those to 6 | create the multi-channel top-down view used in the pipeline 7 | """ 8 | from torch.utils import data 9 | from data import ObjectCategories, House, Obj 10 | import random 11 | import numpy as np 12 | import math 13 | import pickle 14 | import os 15 | import json 16 | import copy 17 | import torch 18 | import utils 19 | 20 | 21 | class RenderedScene(): 22 | """ 23 | Loading a rendered room 24 | Attributes 25 | ---------- 26 | category_map (ObjectCategories): object category mapping 27 | that should be the same across all instances of the class 28 | categories (list[string]): all categories present in this room type. 29 | Loaded once when the first room is loaded to reduce disk access. 30 | cat_to_index (dict[string, int]): maps a category to corresponding index 31 | current_data_dir (string): keep track of the current data directory, if 32 | it changes, then categories and cat_to_index should be recomputed 33 | """ 34 | category_map = ObjectCategories() 35 | categories = None 36 | cat_to_index = None 37 | current_data_dir = None 38 | 39 | def __init__(self, index, data_dir, data_root_dir=None, \ 40 | shuffle=True, load_objects=True, seed=None, rotation=0): 41 | """ 42 | Load a rendered scene from file 43 | Parameters 44 | ---------- 45 | index (int): room number 46 | data_dir (string): location of the pre-rendered rooms 47 | data_root_dir (string or None, optional): if specified, 48 | use this as the root directory 49 | shuffle (bool, optional): If true, randomly order the objects 50 | in the room. Otherwise use the default order as written 51 | in the original dataset 52 | load_objects (bool, optional): If false, only load the doors 53 | and windows. Otherwise load all objects in the room 54 | seed (int or None, optional): if set, use a fixed random seed 55 | so we can replicate a particular experiment 56 | """ 57 | if seed: 58 | random.seed(seed) 59 | 60 | if not data_root_dir: 61 | data_root_dir = utils.get_data_root_dir() 62 | 63 | if RenderedScene.categories is None or RenderedScene.current_data_dir != data_dir: 64 | with open(f"{data_root_dir}/{data_dir}/final_categories_frequency", "r") as f: 65 | lines = f.readlines() 66 | cats = [line.split()[0] for line in lines] 67 | 68 | RenderedScene.categories = [cat for cat in cats if cat not in set(['window', 'door'])] 69 | RenderedScene.cat_to_index = {RenderedScene.categories[i]:i for i in range(len(RenderedScene.categories))} 70 | RenderedScene.current_data_dir = data_dir 71 | 72 | #print(index, rotation) 73 | if rotation != 0: 74 | fname = f"{index}_{rotation}" 75 | else: 76 | fname = index 77 | 78 | with open(f"{data_root_dir}/{data_dir}/{fname}.pkl", "rb") as f: 79 | (self.floor, self.wall, nodes), self.room = pickle.load(f) 80 | 81 | self.index = index 82 | self.rotation = rotation 83 | 84 | self.object_nodes = [] 85 | self.door_window_nodes = [] 86 | for node in nodes: 87 | category = RenderedScene.category_map.get_final_category(node["modelId"]) 88 | if category in ["door", "window"]: 89 | node["category"] = category 90 | self.door_window_nodes.append(node) 91 | elif load_objects: 92 | node["category"] = RenderedScene.cat_to_index[category] 93 | self.object_nodes.append(node) 94 | 95 | if shuffle: 96 | random.shuffle(self.object_nodes) 97 | 98 | self.size = self.floor.shape[0] 99 | 100 | def create_composite(self): 101 | """ 102 | Create a initial composite that only contains the floor, 103 | wall, doors and windows. See RenderedComposite for how 104 | to add more objects 105 | """ 106 | r = RenderedComposite(RenderedScene.categories, self.floor, self.wall, self.door_window_nodes) 107 | return r 108 | 109 | class RenderedComposite(): 110 | """ 111 | Multi-channel top-down composite, used as input to NN 112 | """ 113 | def __init__(self, categories, floor, wall, door_window_nodes=None): 114 | #Optional door_window just in case 115 | self.size = floor.shape[0] 116 | 117 | self.categories = categories 118 | 119 | self.room_mask = (floor + wall) 120 | self.room_mask[self.room_mask != 0] = 1 121 | 122 | self.wall_mask = wall.clone() 123 | self.wall_mask[self.wall_mask != 0] = 0.5 124 | 125 | self.height_map = torch.max(floor, wall) 126 | self.cat_map = torch.zeros((len(self.categories),self.size,self.size)) 127 | 128 | self.sin_map = torch.zeros((self.size,self.size)) 129 | self.cos_map = torch.zeros((self.size,self.size)) 130 | 131 | self.door_map = torch.zeros((self.size, self.size)) 132 | self.window_map = torch.zeros((self.size, self.size)) 133 | 134 | if door_window_nodes: 135 | for node in door_window_nodes: 136 | h = node["height_map"] 137 | xsize, ysize = h.size() 138 | xmin = math.floor(node["bbox_min"][0]) 139 | ymin = math.floor(node["bbox_min"][2]) 140 | if xmin < 0: xmin = 0 141 | if ymin < 0: ymin = 0 142 | if xsize==256: xmin = 0 143 | if ysize==256: ymin = 0 144 | to_add = torch.zeros((self.size, self.size)) 145 | to_add[xmin:xmin+xsize,ymin:ymin+ysize] = h 146 | update = to_add > self.height_map 147 | self.height_map[update] = to_add[update] 148 | self.wall_mask[to_add>0] = 1 149 | to_add[to_add>0] = 0.5 150 | if node["category"] == "door": 151 | self.door_map = self.door_map + to_add 152 | else: 153 | self.window_map = self.window_map + to_add 154 | 155 | def get_transformation(self, transform): 156 | """ 157 | Bad naming, really just getting the sin and cos of the 158 | angle of rotation. 159 | """ 160 | a = transform[0] 161 | b = transform[8] 162 | scale = (a**2+b**2)**0.5 163 | return (b/scale, a/scale) 164 | 165 | def add_height_map(self, to_add, category, sin, cos): 166 | """ 167 | Add a new object to the composite. 168 | Height map, category, and angle of rotation are 169 | all the information required. 170 | """ 171 | update = to_add>self.height_map 172 | self.height_map[update] = to_add[update] 173 | mask = torch.zeros(to_add.size()) 174 | mask[to_add>0] = 0.5 175 | self.cat_map[category] = self.cat_map[category] + mask 176 | self.sin_map[update] = (sin + 1) / 2 177 | self.cos_map[update] = (cos + 1) / 2 178 | 179 | def add_node(self, node): 180 | """ 181 | Add a new object to the composite. 182 | Computes the necessary information and calls 183 | add_height_map 184 | """ 185 | h = node["height_map"] 186 | category = node["category"] 187 | xsize, ysize = h.shape 188 | xmin = math.floor(node["bbox_min"][0]) 189 | ymin = math.floor(node["bbox_min"][2]) 190 | if xmin < 0: 191 | xmin = 0 192 | if ymin < 0: 193 | ymin = 0 194 | if xsize==256: xmin = 0 195 | if ysize==256: ymin = 0 196 | to_add = torch.zeros((self.size, self.size)) 197 | to_add[xmin:xmin+xsize,ymin:ymin+ysize] = h 198 | sin, cos = self.get_transformation(node["transform"]) 199 | self.add_height_map(to_add, category, sin, cos) 200 | 201 | def add_nodes(self, nodes): 202 | for node in nodes: 203 | self.add_node(node) 204 | 205 | # Use these to build a composite that renders OBBs insted of full object geometry 206 | def add_node_obb(self, node): 207 | h = node["height_map_obb"] 208 | category = node["category"] 209 | xsize, ysize = h.shape 210 | # xmin = math.floor(node["bbox_min_obb"][0]) 211 | # ymin = math.floor(node["bbox_min_obb"][2]) 212 | xmin = max(0, math.floor(node["bbox_min_obb"][0])) 213 | ymin = max(0, math.floor(node["bbox_min_obb"][2])) 214 | to_add = torch.zeros((self.size, self.size)) 215 | # ##### TEST 216 | # tmp = to_add[xmin:xmin+xsize,ymin:ymin+ysize] 217 | # if tmp.size()[0] != h.shape[0] or tmp.size()[1] != h.shape[1]: 218 | # print('------------') 219 | # print(tmp.size()) 220 | # print(h.shape) 221 | # print(f'{xmin}, {ymin}') 222 | # ##### 223 | to_add[xmin:xmin+xsize,ymin:ymin+ysize] = h 224 | sin, cos = self.get_transformation(node["transform"]) 225 | self.add_height_map(to_add, category, sin, cos) 226 | def add_nodes_obb(self, nodes): 227 | for node in nodes: 228 | self.add_node_obb(node) 229 | 230 | def get_cat_map(self): 231 | return self.cat_map.clone() 232 | 233 | def add_and_get_composite(self, to_add, category, sin, cos, \ 234 | num_extra_channels=1, temporary=True): 235 | """ 236 | Sometimes we need to create a composite to test 237 | if some objects should be added, without actually 238 | fixing the object to the scene. This method allows doing so. 239 | See get_composite. 240 | """ 241 | if not temporary: 242 | raise NotImplementedError 243 | update = to_add>self.height_map 244 | mask = torch.zeros(to_add.size()) 245 | mask[to_add>0] = 0.5 246 | composite = torch.zeros((len(self.categories)+num_extra_channels+8, self.size, self.size)) 247 | composite[0] = self.room_mask 248 | composite[1] = self.wall_mask 249 | composite[2] = self.cat_map.sum(0) + mask 250 | composite[3] = self.height_map 251 | composite[3][update] = to_add[update] 252 | composite[4] = self.sin_map 253 | composite[4][update] = (sin + 1) / 2 254 | composite[5] = self.cos_map 255 | composite[5][update] = (cos + 1) / 2 256 | composite[6] = self.door_map 257 | composite[7] = self.window_map 258 | for i in range(len(self.categories)): 259 | composite[i+8] = self.cat_map[i] 260 | composite[8+category] += mask 261 | 262 | return composite 263 | 264 | 265 | def get_composite(self, num_extra_channels=1, ablation=None): 266 | """ 267 | Create the actual multi-channel representation. 268 | Which is a N x img_size x img_size tensor. 269 | See the paper for more information. 270 | Current channel order: 271 | -0: room mask 272 | -1: wall mask 273 | -2: object mask 274 | -3: height map 275 | -4, 5: sin and cos of the angle of rotation 276 | -6, 7: single category channel for door and window 277 | -8~8+C: single category channel for all other categories 278 | Parameters 279 | ---------- 280 | num_extra_channels (int, optional): number of extra empty 281 | channels at the end. 1 for most tasks, 0 for should continue 282 | ablation (string or None, optional): if set, return a subset of all 283 | the channels for ablation study, see the paper for more details 284 | """ 285 | if ablation is None: 286 | composite = torch.zeros((len(self.categories)+num_extra_channels+8, self.size, self.size)) 287 | composite[0] = self.room_mask 288 | composite[1] = self.wall_mask 289 | composite[2] = self.cat_map.sum(0) 290 | composite[3] = self.height_map 291 | composite[4] = self.sin_map 292 | composite[5] = self.cos_map 293 | composite[6] = self.door_map 294 | composite[7] = self.window_map 295 | for i in range(len(self.categories)): 296 | composite[i+8] = self.cat_map[i] 297 | elif ablation == "depth": 298 | composite = torch.zeros((1+num_extra_channels, self.size, self.size)) 299 | composite[0] = self.height_map 300 | elif ablation == "basic": 301 | composite = torch.zeros((6+num_extra_channels, self.size, self.size)) 302 | composite[0] = self.room_mask 303 | composite[1] = self.wall_mask 304 | composite[2] = self.cat_map.sum(0) 305 | composite[3] = self.height_map 306 | composite[4] = self.sin_map 307 | composite[5] = self.cos_map 308 | else: 309 | raise NotImplementedError 310 | 311 | return composite 312 | 313 | 314 | if __name__ == "__main__": 315 | from . import ProjectionGenerator 316 | import scipy.misc as m 317 | pgen = ProjectionGenerator() 318 | a = RenderedScene(5, load_objects=True, data_dir="dining_final") 319 | c = a.create_composite() 320 | 321 | for node in a.object_nodes: 322 | c.add_node(node) 323 | 324 | img = c.get_composite() 325 | img = img[7].numpy() 326 | img = m.toimage(img, cmin=0, cmax=1) 327 | img.show() 328 | 329 | -------------------------------------------------------------------------------- /scene-synth/data/top_down.py: -------------------------------------------------------------------------------- 1 | from data import Obj, ProjectionGenerator, Projection, Node 2 | import numpy as np 3 | import math 4 | import scipy.misc as m 5 | from numba import jit 6 | import torch 7 | 8 | # For rendering OBBS 9 | from math_utils.OBB import OBB 10 | from math_utils import Transform 11 | 12 | class TopDownView(): 13 | """ 14 | Take a room, pre-render top-down views 15 | Of floor, walls and individual objects 16 | That can be used to generate the multi-channel views used in our pipeline 17 | """ 18 | #Padding to avoid problems with boundary cases 19 | def __init__(self, height_cap=4.05, length_cap=6.05, size=512): 20 | #Padding by 0.05m to avoid problems with boundary cases 21 | """ 22 | Parameters 23 | ---------- 24 | height_cap (int): the maximum height (in meters) of rooms allowed, which will be rendered with 25 | a value of 1 in the depth channel. To separate the floor from empty spaces, 26 | floors will have a height of 0.5m. See zpad below 27 | length_cap (int): the maximum length/width of rooms allowed. 28 | size (int): size of the rendered top-down image 29 | Returns 30 | ------- 31 | visualization (Image): a visualization of the rendered room, which is 32 | simply the superimposition of all the rendered parts 33 | (floor, wall, nodes) (Triple[torch.Tensor, torch.Tensor, list[torch.Tensor]): 34 | rendered invidiual parts of the room, as 2D torch tensors 35 | this is the part used by the pipeline 36 | """ 37 | self.size = size 38 | self.pgen = ProjectionGenerator(room_size_cap=(length_cap, height_cap, length_cap), \ 39 | zpad=0.5, img_size=size) 40 | 41 | def render(self, room): 42 | projection = self.pgen.get_projection(room) 43 | 44 | visualization = np.zeros((self.size,self.size)) 45 | nodes = [] 46 | 47 | for node in room.nodes: 48 | modelId = node.modelId #Camelcase due to original json 49 | 50 | t = np.asarray(node.transform).reshape(4,4) 51 | 52 | o = Obj(modelId) 53 | t = projection.to_2d(t) 54 | o.transform(t) 55 | 56 | save_t = t 57 | t = projection.to_2d() 58 | bbox_min = np.dot(np.asarray([node.xmin, node.zmin, node.ymin, 1]), t) 59 | bbox_max = np.dot(np.asarray([node.xmax, node.zmax, node.ymax, 1]), t) 60 | xmin = math.floor(bbox_min[0]) 61 | ymin = math.floor(bbox_min[2]) 62 | xsize = math.ceil(bbox_max[0]) - xmin + 1 63 | ysize = math.ceil(bbox_max[2]) - ymin + 1 64 | 65 | description = {} 66 | description["modelId"] = modelId 67 | description["transform"] = node.transform 68 | description["bbox_min"] = bbox_min 69 | description["bbox_max"] = bbox_max 70 | description["id"] = node.id 71 | description["child"] = [c.id for c in node.child] if node.child else None 72 | description["parent"] = node.parent.id if isinstance(node.parent, Node) else node.parent 73 | #if description["parent"] is None or description["parent"] == "Ceiling": 74 | # print(description["modelId"]) 75 | # print(description["parent"]) 76 | # print(node.zmin - room.zmin) 77 | #print("FATAL ERROR") 78 | 79 | #Since it is possible that the bounding box information of a room 80 | #Was calculated without some doors/windows 81 | #We need to handle these cases 82 | if ymin < 0: 83 | ymin = 0 84 | if xmin < 0: 85 | xmin = 0 86 | 87 | #xmin = 0 88 | #ymin = 0 89 | #xsize = 256 90 | #ysize = 256 91 | 92 | #print(list(bbox_min), list(bbox_max)) 93 | #print(xmin, ymin, xsize, ysize) 94 | rendered = self.render_object(o, xmin, ymin, xsize, ysize, self.size) 95 | description["height_map"] = torch.from_numpy(rendered).float() 96 | 97 | #tmp = np.zeros((self.size, self.size)) 98 | #tmp[xmin:xmin+rendered.shape[0],ymin:ymin+rendered.shape[1]] = rendered 99 | #visualization += tmp 100 | 101 | # Compute the pixel-space dimensions of the object before it has been 102 | # transformed (i.e. in object space) 103 | objspace_bbox_min = np.dot(o.bbox_min, t) 104 | objspace_bbox_max = np.dot(o.bbox_max, t) 105 | description['objspace_dims'] = np.array([ 106 | objspace_bbox_max[0] - objspace_bbox_min[0], 107 | objspace_bbox_max[2] - objspace_bbox_min[2] 108 | ]) 109 | 110 | # Render an OBB height map as well 111 | bbox_dims = o.bbox_max - o.bbox_min 112 | model_matrix = Transform(scale=bbox_dims[:3], translation=o.bbox_min[:3]).as_mat4() 113 | full_matrix = np.matmul(np.transpose(save_t), model_matrix) 114 | obb = OBB.from_local2world_transform(full_matrix) 115 | obb_tris = np.asarray(obb.get_triangles(), dtype=np.float32) 116 | bbox_min = np.min(np.min(obb_tris, 0), 0) 117 | bbox_max = np.max(np.max(obb_tris, 0), 0) 118 | xmin, ymin = math.floor(bbox_min[0]), math.floor(bbox_min[2]) 119 | xsize, ysize = math.ceil(bbox_max[0]) - xmin + 1, math.ceil(bbox_max[2]) - ymin + 1 120 | rendered_obb = self.render_object_helper(obb_tris, xmin, ymin, xsize, ysize, self.size) 121 | description["height_map_obb"] = torch.from_numpy(rendered_obb).float() 122 | description['bbox_min_obb'] = bbox_min 123 | description['bbox_max_obb'] = bbox_max 124 | 125 | tmp = np.zeros((self.size, self.size)) 126 | tmp[xmin:xmin+rendered.shape[0],ymin:ymin+rendered.shape[1]] = rendered 127 | visualization += tmp 128 | 129 | nodes.append(description) 130 | 131 | if hasattr(room, "transform"): 132 | t = projection.to_2d(np.asarray(room.transform).reshape(4,4)) 133 | else: 134 | t = projection.to_2d() 135 | #Render the floor 136 | o = Obj(room.modelId+"f", room.house_id, is_room=True) 137 | o.transform(t) 138 | floor = self.render_object(o, 0, 0, self.size, self.size, self.size) 139 | visualization += floor 140 | floor = torch.from_numpy(floor).float() 141 | 142 | #Render the walls 143 | o = Obj(room.modelId+"w", room.house_id, is_room=True) 144 | o.transform(t) 145 | wall = self.render_object(o, 0, 0, self.size, self.size, self.size) 146 | visualization += wall 147 | wall = torch.from_numpy(wall).float() 148 | return (visualization, (floor, wall, nodes)) 149 | 150 | @staticmethod 151 | @jit(nopython=True) 152 | def render_object_helper(triangles, xmin, ymin, xsize, ysize, img_size): 153 | result = np.zeros((img_size, img_size), dtype=np.float32) 154 | N, _, _ = triangles.shape 155 | 156 | for triangle in range(N): 157 | x0,z0,y0 = triangles[triangle][0] 158 | x1,z1,y1 = triangles[triangle][1] 159 | x2,z2,y2 = triangles[triangle][2] 160 | a = -y1*x2 + y0*(-x1+x2) + x0*(y1-y2) + x1*y2 161 | if a != 0: 162 | for i in range(max(0,math.floor(min(x0,x1,x2))), \ 163 | min(img_size,math.ceil(max(x0,x1,x2)))): 164 | for j in range(max(0,math.floor(min(y0,y1,y2))), \ 165 | min(img_size,math.ceil(max(y0,y1,y2)))): 166 | x = i+0.5 167 | y = j+0.5 168 | s = (y0*x2 - x0*y2 + (y2-y0)*x + (x0-x2)*y)/a 169 | t = (x0*y1 - y0*x1 + (y0-y1)*x + (x1-x0)*y)/a 170 | if s < 0 and t < 0: 171 | s = -s 172 | t = -t 173 | if 0 < s < 1 and 0 < t < 1 and s + t <= 1: 174 | height = z0 *(1-s-t) + z1*s + z2*t 175 | result[i][j] = max(result[i][j], height) 176 | 177 | return result[xmin:xmin+xsize, ymin:ymin+ysize] 178 | 179 | @staticmethod 180 | def render_object(o, xmin, ymin, xsize, ysize, img_size): 181 | """ 182 | Render a cropped top-down view of object 183 | 184 | Parameters 185 | ---------- 186 | o (list[triple]): object to be rendered, represented as a triangle mesh 187 | xmin, ymin (int): min coordinates of the bounding box containing the object, 188 | with respect to the full image 189 | xsize, ysze (int); size of the bounding box containing the object 190 | img_size (int): size of the full image 191 | """ 192 | triangles = np.asarray(list(o.get_triangles()), dtype=np.float32) 193 | return TopDownView.render_object_helper(triangles, xmin, ymin, xsize, ysize, img_size) 194 | 195 | @staticmethod 196 | @jit(nopython=True) 197 | def render_object_full_size_helper(triangles, size): 198 | result = np.zeros((size, size), dtype=np.float32) 199 | N, _, _ = triangles.shape 200 | 201 | for triangle in range(N): 202 | x0,z0,y0 = triangles[triangle][0] 203 | x1,z1,y1 = triangles[triangle][1] 204 | x2,z2,y2 = triangles[triangle][2] 205 | a = -y1*x2 + y0*(-x1+x2) + x0*(y1-y2) + x1*y2 206 | if a != 0: 207 | for i in range(max(0,math.floor(min(x0,x1,x2))), \ 208 | min(size,math.ceil(max(x0,x1,x2)))): 209 | for j in range(max(0,math.floor(min(y0,y1,y2))), \ 210 | min(size,math.ceil(max(y0,y1,y2)))): 211 | x = i+0.5 212 | y = j+0.5 213 | s = (y0*x2 - x0*y2 + (y2-y0)*x + (x0-x2)*y)/a 214 | t = (x0*y1 - y0*x1 + (y0-y1)*x + (x1-x0)*y)/a 215 | if s < 0 and t < 0: 216 | s = -s 217 | t = -t 218 | if 0 < s < 1 and 0 < t < 1 and s + t <= 1: 219 | height = z0 *(1-s-t) + z1*s + z2*t 220 | result[i][j] = max(result[i][j], height) 221 | 222 | return result 223 | 224 | @staticmethod 225 | def render_object_full_size(o, size): 226 | """ 227 | Render a full-sized top-down view of the object, see render_object 228 | """ 229 | triangles = np.asarray(list(o.get_triangles()), dtype=np.float32) 230 | return TopDownView.render_object_full_size_helper(triangles, size) 231 | 232 | if __name__ == "__main__": 233 | from .house import House 234 | h = House(id_="51515da17cd4b575775cea4f5546737a") 235 | r = h.rooms[0] 236 | renderer = TopDownView() 237 | img = renderer.render(r)[0] 238 | img = m.toimage(img, cmin=0, cmax=1) 239 | img.show() 240 | -------------------------------------------------------------------------------- /scene-synth/filters/bedroom.py: -------------------------------------------------------------------------------- 1 | from data.house import * 2 | from data.dataset import DatasetFilter 3 | from data.object_data import ObjectCategories 4 | from .global_category_filter import * 5 | import utils 6 | 7 | """ 8 | Bedroom filter 9 | """ 10 | 11 | def bedroom_filter(version, room_size, second_tier, source): 12 | data_dir = utils.get_data_root_dir() 13 | with open(f"{data_dir}/{source}/coarse_categories_frequency", "r") as f: 14 | coarse_categories_frequency = ([s[:-1] for s in f.readlines()]) 15 | coarse_categories_frequency = [s.split(" ") for s in coarse_categories_frequency] 16 | coarse_categories_frequency = dict([(a,int(b)) for (a,b) in coarse_categories_frequency]) 17 | category_map = ObjectCategories() 18 | 19 | if version == "final": 20 | filtered, rejected, door_window = GlobalCategoryFilter.get_filter() 21 | with open(f"{data_dir}/{source}/final_categories_frequency", "r") as f: 22 | frequency = ([s[:-1] for s in f.readlines()]) 23 | frequency = [s.split(" ") for s in frequency] 24 | frequency = dict([(a,int(b)) for (a,b) in frequency]) 25 | 26 | def node_criteria(node, room): 27 | category = category_map.get_final_category(node.modelId) 28 | if category in filtered: return False 29 | return True 30 | 31 | def room_criteria(room, house): 32 | node_count = 0 33 | bed_count = 0 #Must have one bed 34 | for node in room.nodes: 35 | category = category_map.get_final_category(node.modelId) 36 | if category in rejected: 37 | return False 38 | if not category in door_window: 39 | node_count += 1 40 | 41 | t = np.asarray(node.transform).reshape((4,4)).transpose() 42 | a = t[0][0] 43 | b = t[0][2] 44 | c = t[2][0] 45 | d = t[2][2] 46 | 47 | xscale = (a**2 + c**2)**0.5 48 | yscale = (b**2 + d**2)**0.5 49 | zscale = t[1][1] 50 | 51 | if not 0.8 20: return False 64 | if bed_count < 1: return False 65 | 66 | return True 67 | elif version == "latent": 68 | filtered, rejected, door_window, second_tier_include = \ 69 | GlobalCategoryFilter.get_filter_latent() 70 | rejected += ["gym_equipment", "fireplace", "dining_table", "short_kitchen_cabinet",\ 71 | "hanging_kitchen_cabinet", "refrigerator"] 72 | filtered += ["pillow", "glass", "cup"] 73 | filtered, rejected, door_window, second_tier_include = \ 74 | set(filtered), set(rejected), set(door_window), set(second_tier_include) 75 | with open(f"{data_dir}/{source}/final_categories_frequency", "r") as f: 76 | frequency = ([s[:-1] for s in f.readlines()]) 77 | frequency = [s.split(" ") for s in frequency] 78 | frequency = dict([(a,int(b)) for (a,b) in frequency]) 79 | 80 | def node_criteria(node, room): 81 | category = category_map.get_final_category(node.modelId) 82 | if category in door_window: 83 | return True 84 | if category in filtered: return False 85 | 86 | if second_tier: 87 | if node.zmin - room.zmin > 0.1 and \ 88 | (category not in second_tier_include or node.parent is None): 89 | return False 90 | 91 | if node.parent: 92 | if isinstance(node.parent, Node) and node.zmin < node.parent.zmax-0.1: 93 | return False 94 | node_now = node 95 | while isinstance(node_now, Node) and node_now.parent: 96 | node_now = node_now.parent 97 | if node_now != "Floor": 98 | return False 99 | else: 100 | if node.zmin - room.zmin > 0.1: 101 | return False 102 | if category in ["vase", "console", "book", "television", "table_lamp", "fishbowl"]: 103 | return False 104 | 105 | #Quick filter for second-tier non ceiling mount 106 | #if node.zmin - room.zmin < 0.1: 107 | # return False 108 | #else: 109 | # if node.zmax - room.zmax > -0.2: 110 | # return False 111 | return True 112 | 113 | def room_criteria(room, house): 114 | if room.height > 4: return False 115 | if room.length > room_size: return False 116 | if room.width > room_size: return False 117 | floor_node_count = 0 118 | node_count = 0 119 | bed_count = 0 #Must have one bed 120 | scaled = False 121 | #dirty fix! 122 | for i in range(5): 123 | room.nodes = [node for node in room.nodes if not \ 124 | ((node.parent and isinstance(node.parent, Node) and \ 125 | (node.parent) not in room.nodes)) 126 | ] 127 | for node in room.nodes: 128 | category = category_map.get_final_category(node.modelId) 129 | if category in rejected: 130 | return False 131 | if not category in door_window: 132 | node_count += 1 133 | 134 | if node.zmin - room.zmin < 0.1: 135 | floor_node_count += 1 136 | 137 | t = np.asarray(node.transform).reshape((4,4)).transpose() 138 | a = t[0][0] 139 | b = t[0][2] 140 | c = t[2][0] 141 | d = t[2][2] 142 | 143 | xscale = (a**2 + c**2)**0.5 144 | yscale = (b**2 + d**2)**0.5 145 | zscale = t[1][1] 146 | 147 | if not 0.9 20: return False 175 | if bed_count < 1: return False 176 | 177 | return True 178 | 179 | else: 180 | raise NotImplementedError 181 | 182 | dataset_f = DatasetFilter(room_filters = [room_criteria], node_filters = [node_criteria]) 183 | 184 | return dataset_f 185 | 186 | 187 | -------------------------------------------------------------------------------- /scene-synth/filters/collision.py: -------------------------------------------------------------------------------- 1 | from data.house import * 2 | from data.dataset import DatasetFilter 3 | from data.object_data import ObjectCategories 4 | import copy 5 | import os 6 | from utils import stdout_redirected 7 | 8 | """ 9 | Filters out rooms with large collisions 10 | """ 11 | def collision_filter(oc): 12 | category_map = ObjectCategories() 13 | def room_criteria(room, house): 14 | with stdout_redirected(): 15 | oc.reset() 16 | oc.init_from_room(copy.deepcopy(house), room.id) 17 | collisions = oc.get_collisions() 18 | 19 | has_contact = False 20 | for (collision_pair, contact_record) in collisions.items(): 21 | if contact_record.distance < -0.15: 22 | idA = contact_record.idA 23 | idB = contact_record.idB 24 | if idA == idB: 25 | node = next(node for node in room.nodes if node.id==idA) 26 | if category_map.get_final_category(node.modelId) not in ["window", "door"]: 27 | has_contact = True 28 | else: 29 | has_contact = True 30 | 31 | return not has_contact 32 | 33 | dataset_f = DatasetFilter(room_filters = [room_criteria]) 34 | return dataset_f 35 | -------------------------------------------------------------------------------- /scene-synth/filters/floor_node.py: -------------------------------------------------------------------------------- 1 | from data.house import * 2 | from data.dataset import DatasetFilter 3 | from data.object_data import ObjectCategories 4 | 5 | """ 6 | Floor node if either: 7 | Parent is floor (from house_relations, not loaded by default in the released version of the code, 8 | so parent is always None and this gets ignored ) 9 | Distance from floor is less than 0.1 meters 10 | Is a door/window (since those have to be included) 11 | """ 12 | 13 | def floor_node_filter(): 14 | category = ObjectCategories() 15 | def node_criteria(node, room): 16 | return (node.parent and node.parent == "Floor") or (node.zmin - room.zmin < 0.1) or (hasattr(node, "modelId") and category.get_final_category(node.modelId) in ["door", "window"]) 17 | 18 | dataset_f = DatasetFilter(node_filters = [node_criteria]) 19 | return dataset_f 20 | -------------------------------------------------------------------------------- /scene-synth/filters/global_category_filter.py: -------------------------------------------------------------------------------- 1 | """ 2 | Just a list of categories that should be filtered out 3 | Shared by all room specific filters 4 | """ 5 | class GlobalCategoryFilter(): 6 | 7 | def get_filter(): 8 | 9 | door_window = ["window", 10 | "door"] 11 | 12 | second_tier = ["table_lamp", 13 | "chandelier", 14 | "guitar", 15 | "amplifier", 16 | "keyboard", 17 | "drumset", 18 | "microphone", 19 | "accordion", 20 | "toy", 21 | "xbox", 22 | "playstation", 23 | "fishbowl", 24 | "chessboard", 25 | "iron", 26 | "helmet", 27 | "telephone", 28 | "stationary_container", 29 | "ceiling_fan", 30 | "bottle", 31 | "fruit_bowl", 32 | "glass", 33 | "knife_rack", 34 | "plates", 35 | "books", 36 | "book", 37 | "television", 38 | "wood_board", 39 | "switch", 40 | "pillow", 41 | "laptop", 42 | "clock", 43 | "helmet", 44 | "bottle", 45 | "trinket", 46 | "glass", 47 | "range_hood", 48 | "candle", 49 | "soap_dish"] 50 | 51 | wall_objects = ["wall_lamp", 52 | "mirror", 53 | "curtain", 54 | "blind"] 55 | 56 | unimportant = ["toy", 57 | "fish_tank", 58 | "tricycle", 59 | "vacuum_cleaner", 60 | "weight_scale", 61 | "heater", 62 | "picture_frame", 63 | "beer", 64 | "shoes", 65 | "weight_scale", 66 | "decoration", 67 | "ladder", 68 | "tripod", 69 | "air_conditioner", 70 | "cart", 71 | "fireplace_tools", 72 | "vase"] 73 | 74 | inhabitants = ["person", 75 | "cat", 76 | "bird", 77 | "dog", 78 | "pet"] 79 | 80 | special_filter = ["rug",] 81 | 82 | filtered = second_tier + unimportant + inhabitants + special_filter + wall_objects 83 | 84 | unwanted_complex_structure = ["partition", 85 | "column", 86 | "arch", 87 | "stairs"] 88 | set_items = ["chair_set", 89 | "stereo_set", 90 | "table_and_chair", 91 | "slot_machine_and_chair", 92 | "kitchen_set", 93 | "double_desk", 94 | "double_desk_with_chairs", 95 | "dressing_table_with_stool", 96 | "kitchen_island_with_range_hood_and_table"] 97 | 98 | outdoor = ["lawn_mower", 99 | "car", 100 | "motorcycle", 101 | "bicycle", 102 | "garage_door", 103 | "outdoor_seating", 104 | "fence"] 105 | 106 | rejected = unwanted_complex_structure + set_items + outdoor 107 | 108 | return filtered, rejected, door_window 109 | 110 | def get_filter_latent(): 111 | 112 | door_window = ["window", 113 | "door"] 114 | 115 | second_tier_include = ["table_lamp", 116 | "television", 117 | "picture_frame", 118 | "books", 119 | "book", 120 | "laptop", 121 | "floor_lamp", 122 | "vase", 123 | "plant", 124 | "console", 125 | "stereo_set", 126 | "toy", 127 | "fish_tank", 128 | "cup", 129 | "glass", 130 | "fruit_bowl", 131 | "bottle", 132 | "fishbowl", 133 | "pillow", 134 | ] 135 | 136 | second_tier = ["chandelier", 137 | "guitar", 138 | "amplifier", 139 | "keyboard", 140 | "drumset", 141 | "microphone", 142 | "accordion", 143 | "chessboard", 144 | "iron", 145 | "helmet", 146 | "stationary_container", 147 | "ceiling_fan", 148 | "knife_rack", 149 | "plates", 150 | "wood_board", 151 | "switch", 152 | "clock", 153 | "helmet", 154 | "trinket", 155 | "range_hood", 156 | "candle", 157 | "soap_dish"] 158 | 159 | wall_objects = ["wall_lamp", 160 | "mirror", 161 | "curtain", 162 | "wall_shelf", 163 | "blinds", 164 | "blind"] 165 | 166 | unimportant = ["tricycle", 167 | "fish_tank", 168 | "vacuum_cleaner", 169 | "weight_scale", 170 | "heater", 171 | "picture_frame", 172 | "beer", 173 | "shoes", 174 | "weight_scale", 175 | "decoration", 176 | "ladder", 177 | "tripod", 178 | "air_conditioner", 179 | "cart", 180 | "fireplace_tools", 181 | "ironing_board", 182 | ] 183 | 184 | inhabitants = ["person", 185 | "cat", 186 | "bird", 187 | "dog", 188 | "pet"] 189 | 190 | special_filter = ["rug",] 191 | 192 | filtered = second_tier + unimportant + inhabitants + special_filter + wall_objects 193 | 194 | unwanted_complex_structure = ["partition", 195 | "column", 196 | "arch", 197 | "stairs"] 198 | set_items = ["chair_set", 199 | "stereo_set", 200 | "table_and_chair", 201 | "slot_machine_and_chair", 202 | "kitchen_set", 203 | "double_desk", 204 | "double_desk_with_chairs", 205 | "desk_with_shelves", 206 | "dressing_table_with_stool", 207 | "armchair_with_ottoman", 208 | "kitchen_island_with_range_hood_and_table"] 209 | 210 | outdoor = ["lawn_mower", 211 | "car", 212 | "motorcycle", 213 | "bicycle", 214 | "garage_door", 215 | "outdoor_seating", 216 | "fence"] 217 | 218 | rejected = unwanted_complex_structure + set_items + outdoor 219 | 220 | return filtered, rejected, door_window, second_tier_include 221 | 222 | 223 | 224 | 225 | 226 | 227 | 228 | -------------------------------------------------------------------------------- /scene-synth/filters/good_house.py: -------------------------------------------------------------------------------- 1 | from data.house import * 2 | import utils 3 | 4 | """ 5 | Good houses 6 | The list of good houses should be included in filters/ already 7 | In addition, if a house is not scaled to meters, we reject those 8 | """ 9 | def good_house_criteria(house): 10 | data_dir = os.path.dirname(os.path.abspath(__file__)) 11 | with open(f"{data_dir}/good_houses", 'r') as f: 12 | good_houses = [s[:-1] for s in f.readlines()[1:]] 13 | if house.scaleToMeters != 1: return False 14 | if not house.id in good_houses: return False 15 | return True 16 | 17 | -------------------------------------------------------------------------------- /scene-synth/filters/livingroom.py: -------------------------------------------------------------------------------- 1 | from data.house import * 2 | from data.dataset import DatasetFilter 3 | from data.object_data import ObjectCategories 4 | from .global_category_filter import * 5 | import utils 6 | 7 | """ 8 | Living room filter 9 | """ 10 | 11 | def livingroom_filter(version, room_size, second_tier, source): 12 | data_dir = utils.get_data_root_dir() 13 | with open(f"{data_dir}/{source}/coarse_categories_frequency", "r") as f: 14 | coarse_categories_frequency = ([s[:-1] for s in f.readlines()]) 15 | coarse_categories_frequency = [s.split(" ") for s in coarse_categories_frequency] 16 | coarse_categories_frequency = dict([(a,int(b)) for (a,b) in coarse_categories_frequency]) 17 | category_map = ObjectCategories() 18 | if version == "final": 19 | filtered, rejected, door_window = GlobalCategoryFilter.get_filter() 20 | with open(f"{data_dir}/{source}/final_categories_frequency", "r") as f: 21 | frequency = ([s[:-1] for s in f.readlines()]) 22 | frequency = [s.split(" ") for s in frequency] 23 | frequency = dict([(a,int(b)) for (a,b) in frequency]) 24 | 25 | def node_criteria(node, room): 26 | category = category_map.get_final_category(node.modelId) 27 | if category in filtered: return False 28 | return True 29 | 30 | def room_criteria(room, house): 31 | node_count = 0 32 | for node in room.nodes: 33 | category = category_map.get_final_category(node.modelId) 34 | if category in rejected: 35 | return False 36 | if not category in door_window: 37 | node_count += 1 38 | 39 | t = np.asarray(node.transform).reshape((4,4)).transpose() 40 | a = t[0][0] 41 | b = t[0][2] 42 | c = t[2][0] 43 | d = t[2][2] 44 | 45 | xscale = (a**2 + c**2)**0.5 46 | yscale = (b**2 + d**2)**0.5 47 | zscale = t[1][1] 48 | 49 | if not 0.8 20: return False 59 | return True 60 | 61 | elif version == "latent": 62 | filtered, rejected, door_window, second_tier_include = \ 63 | GlobalCategoryFilter.get_filter_latent() 64 | rejected += ["chair", "outdoor_lamp", "water_dispenser", "office_chair", "game_table", \ 65 | "dining_table"] 66 | filtered, rejected, door_window, second_tier_include = \ 67 | set(filtered), set(rejected), set(door_window), set(second_tier_include) 68 | with open(f"{data_dir}/{source}/final_categories_frequency", "r") as f: 69 | frequency = ([s[:-1] for s in f.readlines()]) 70 | frequency = [s.split(" ") for s in frequency] 71 | frequency = dict([(a,int(b)) for (a,b) in frequency]) 72 | 73 | def node_criteria(node, room): 74 | category = category_map.get_final_category(node.modelId) 75 | if category in door_window: 76 | return True 77 | if category in filtered: return False 78 | 79 | if second_tier: 80 | if node.zmin - room.zmin > 0.1 and \ 81 | (category not in second_tier_include or node.parent is None): 82 | return False 83 | 84 | if node.parent: 85 | if isinstance(node.parent, Node) and node.zmin < node.parent.zmax-0.1: 86 | return False 87 | node_now = node 88 | while isinstance(node_now, Node) and node_now.parent: 89 | node_now = node_now.parent 90 | if node_now != "Floor": 91 | return False 92 | else: 93 | if node.zmin - room.zmin > 0.1: 94 | return False 95 | if category in ["book", "console", "television", "table_lamp", "fishbowl"]: 96 | return False 97 | 98 | #Quick filter for second-tier non ceiling mount 99 | #if node.zmin - room.zmin < 0.1: 100 | # return False 101 | #else: 102 | # if node.zmax - room.zmax > -0.2: 103 | # return False 104 | return True 105 | 106 | def room_criteria(room, house): 107 | if room.height > 4: return False 108 | if room.length > room_size: return False 109 | if room.width > room_size: return False 110 | floor_node_count = 0 111 | node_count = 0 112 | scaled = False 113 | #dirty fix! 114 | for i in range(5): 115 | room.nodes = [node for node in room.nodes if not \ 116 | ((node.parent and isinstance(node.parent, Node) and \ 117 | (node.parent) not in room.nodes)) 118 | ] 119 | for node in room.nodes: 120 | category = category_map.get_final_category(node.modelId) 121 | if category in rejected: 122 | return False 123 | if not category in door_window: 124 | node_count += 1 125 | 126 | if node.zmin - room.zmin < 0.1: 127 | floor_node_count += 1 128 | 129 | t = np.asarray(node.transform).reshape((4,4)).transpose() 130 | a = t[0][0] 131 | b = t[0][2] 132 | c = t[2][0] 133 | d = t[2][2] 134 | 135 | xscale = (a**2 + c**2)**0.5 136 | yscale = (b**2 + d**2)**0.5 137 | zscale = t[1][1] 138 | 139 | if not 0.9 20: return False 165 | 166 | return True 167 | else: 168 | raise NotImplementedError 169 | 170 | dataset_f = DatasetFilter(room_filters = [room_criteria], node_filters = [node_criteria]) 171 | 172 | return dataset_f 173 | -------------------------------------------------------------------------------- /scene-synth/filters/office.py: -------------------------------------------------------------------------------- 1 | from data.house import * 2 | from data.dataset import DatasetFilter 3 | from data.object_data import ObjectCategories 4 | from .global_category_filter import * 5 | import utils 6 | 7 | """ 8 | Office filter 9 | """ 10 | 11 | def office_filter(version, room_size, second_tier, source): 12 | data_dir = utils.get_data_root_dir() 13 | with open(f"{data_dir}/{source}/coarse_categories_frequency", "r") as f: 14 | coarse_categories_frequency = ([s[:-1] for s in f.readlines()]) 15 | coarse_categories_frequency = [s.split(" ") for s in coarse_categories_frequency] 16 | coarse_categories_frequency = dict([(a,int(b)) for (a,b) in coarse_categories_frequency]) 17 | category_map = ObjectCategories() 18 | if version == "final": 19 | filtered, rejected, door_window = GlobalCategoryFilter.get_filter() 20 | rejected += ["gym_equipment"] 21 | with open(f"{data_dir}/{source}/final_categories_frequency", "r") as f: 22 | frequency = ([s[:-1] for s in f.readlines()]) 23 | frequency = [s.split(" ") for s in frequency] 24 | frequency = dict([(a,int(b)) for (a,b) in frequency]) 25 | 26 | def node_criteria(node, room): 27 | category = category_map.get_final_category(node.modelId) 28 | if category in filtered: return False 29 | return True 30 | 31 | def room_criteria(room, house): 32 | node_count = 0 33 | for node in room.nodes: 34 | category = category_map.get_final_category(node.modelId) 35 | if category in rejected: 36 | return False 37 | if not category in door_window: 38 | node_count += 1 39 | 40 | t = np.asarray(node.transform).reshape((4,4)).transpose() 41 | a = t[0][0] 42 | b = t[0][2] 43 | c = t[2][0] 44 | d = t[2][2] 45 | 46 | xscale = (a**2 + c**2)**0.5 47 | yscale = (b**2 + d**2)**0.5 48 | zscale = t[1][1] 49 | 50 | if not 0.8 20: return False 60 | return True 61 | 62 | elif version == "latent": 63 | filtered, rejected, door_window, second_tier_include = \ 64 | GlobalCategoryFilter.get_filter_latent() 65 | rejected += ["gym_equipment", "small_refrigerator", "glass", "bottle", "computer", \ 66 | "trash_can", "chair", "dining_table"] 67 | filtered, rejected, door_window, second_tier_include = \ 68 | set(filtered), set(rejected), set(door_window), set(second_tier_include) 69 | with open(f"{data_dir}/{source}/final_categories_frequency", "r") as f: 70 | frequency = ([s[:-1] for s in f.readlines()]) 71 | frequency = [s.split(" ") for s in frequency] 72 | frequency = dict([(a,int(b)) for (a,b) in frequency]) 73 | 74 | def node_criteria(node, room): 75 | category = category_map.get_final_category(node.modelId) 76 | if category in door_window: 77 | return True 78 | if category in filtered: return False 79 | 80 | if second_tier: 81 | if node.zmin - room.zmin > 0.1 and \ 82 | (category not in second_tier_include or node.parent is None): 83 | return False 84 | 85 | if node.parent: 86 | if isinstance(node.parent, Node) and node.zmin < node.parent.zmax-0.1: 87 | return False 88 | node_now = node 89 | while isinstance(node_now, Node) and node_now.parent: 90 | node_now = node_now.parent 91 | if node_now != "Floor": 92 | return False 93 | else: 94 | if node.zmin - room.zmin > 0.1: 95 | return False 96 | if category in ["vase", "console", "book"]: 97 | return False 98 | 99 | #Quick filter for second-tier non ceiling mount 100 | #if node.zmin - room.zmin < 0.1: 101 | # return False 102 | #else: 103 | # if node.zmax - room.zmax > -0.2: 104 | # return False 105 | return True 106 | 107 | def room_criteria(room, house): 108 | if room.height > 4: return False 109 | if room.length > room_size: return False 110 | if room.width > room_size: return False 111 | floor_node_count = 0 112 | node_count = 0 113 | scaled = False 114 | #dirty fix! 115 | for i in range(5): 116 | room.nodes = [node for node in room.nodes if not \ 117 | ((node.parent and isinstance(node.parent, Node) and \ 118 | (node.parent) not in room.nodes)) 119 | ] 120 | for node in room.nodes: 121 | category = category_map.get_final_category(node.modelId) 122 | if category in rejected: 123 | return False 124 | if not category in door_window: 125 | node_count += 1 126 | 127 | if node.zmin - room.zmin < 0.1: 128 | floor_node_count += 1 129 | 130 | t = np.asarray(node.transform).reshape((4,4)).transpose() 131 | a = t[0][0] 132 | b = t[0][2] 133 | c = t[2][0] 134 | d = t[2][2] 135 | 136 | xscale = (a**2 + c**2)**0.5 137 | yscale = (b**2 + d**2)**0.5 138 | zscale = t[1][1] 139 | 140 | if not 0.9 20: return False 166 | 167 | return True 168 | 169 | else: 170 | raise NotImplementedError 171 | 172 | dataset_f = DatasetFilter(room_filters = [room_criteria], node_filters = [node_criteria]) 173 | 174 | return dataset_f 175 | 176 | -------------------------------------------------------------------------------- /scene-synth/filters/renderable.py: -------------------------------------------------------------------------------- 1 | from data.house import * 2 | from data.dataset import DatasetFilter 3 | from data.object_data import ObjectCategories 4 | import utils 5 | 6 | """ 7 | Rooms renderable under current setting 8 | Filters out rooms that are too large 9 | and rooms that miss floor/wall objs 10 | """ 11 | 12 | def renderable_room_filter(height_cap=4, length_cap=6, width_cap=6): 13 | def room_criteria(room, house): 14 | data_dir = utils.get_data_root_dir() 15 | #if room.height > height_cap: return False 16 | #if room.length > length_cap: return False 17 | #if room.width > width_cap: return False 18 | if not os.path.isfile(f"{data_dir}/suncg_data/room/{room.house_id}/{room.modelId}f.obj"): 19 | return False 20 | if not os.path.isfile(f"{data_dir}/suncg_data/room/{room.house_id}/{room.modelId}w.obj"): 21 | return False 22 | for node in room.nodes: 23 | if node.type == "Box": 24 | return False 25 | return True 26 | 27 | dataset_f = DatasetFilter(room_filters = [room_criteria]) 28 | return dataset_f 29 | -------------------------------------------------------------------------------- /scene-synth/filters/room_type.py: -------------------------------------------------------------------------------- 1 | from data.house import * 2 | 3 | """ 4 | Just filters by room type 5 | A room may have multiple types, all must match 6 | This removes dual function rooms e.g. bedroom+kitchen 7 | """ 8 | def room_type_criteria(types): 9 | if not isinstance(types, list): types = [types] 10 | def wrapper(room, _): 11 | return all(t in types for t in room.roomTypes) and len(room.roomTypes)>0 12 | return wrapper 13 | 14 | -------------------------------------------------------------------------------- /scene-synth/filters/toilet.py: -------------------------------------------------------------------------------- 1 | from data.house import * 2 | from data.dataset import DatasetFilter 3 | from data.object_data import ObjectCategories 4 | from .global_category_filter import * 5 | import utils 6 | 7 | """ 8 | Verions: 9 | set1: baseline 10 | """ 11 | 12 | def toilet_filter(version, room_size, second_tier, source): 13 | data_dir = utils.get_data_root_dir() 14 | with open(f"{data_dir}/{source}/coarse_categories_frequency", "r") as f: 15 | coarse_categories_frequency = ([s[:-1] for s in f.readlines()]) 16 | coarse_categories_frequency = [s.split(" ") for s in coarse_categories_frequency] 17 | coarse_categories_frequency = dict([(a,int(b)) for (a,b) in coarse_categories_frequency]) 18 | category_map = ObjectCategories() 19 | if version == "set1": 20 | rejected_object_categories = ["bathroom_stuff", "partition", "kitchenware", \ 21 | "sofa", "kitchen_appliance", "column", "bed", \ 22 | "table", "hanger"] 23 | filtered_object_categories = ["person", "pet", "rug", "curtain", "trash_can", \ 24 | "toy", "mirror", "picture_frame"] 25 | no_count_object_categories = ["window", "door"] 26 | freq_threshold = 1000 27 | 28 | def node_criteria(node, room): 29 | coarse_category = category_map.get_coarse_category(node.modelId) 30 | #Recursively check if all children are filtered out or not 31 | #This is done because if there is a node that we do not want 32 | #But it has a child that isn't filtered out 33 | #Then we have to reject the room completely, see room_criteria 34 | if node.child: 35 | if any(node_criteria(c, room) for c in node.child): 36 | return True 37 | #Still nodes without parent, so can't just check if floor parent 38 | if node.parent and node.parent !="Floor" and \ 39 | not coarse_category in no_count_object_categories: 40 | return False 41 | if coarse_category in rejected_object_categories: return True 42 | if coarse_categories_frequency[coarse_category] < freq_threshold: return False 43 | if coarse_category in filtered_object_categories: return False 44 | if node.length * node.width < 0.05: return False 45 | return True 46 | 47 | def room_criteria(room, house): 48 | if room.height > 4: return False 49 | if room.length > 6: return False 50 | if room.width > 6: return False 51 | true_node_count = 0 52 | category_dict = {} 53 | for node in room.nodes: 54 | coarse_category = category_map.get_coarse_category(node.modelId) 55 | if coarse_categories_frequency[coarse_category] < freq_threshold: return False 56 | if coarse_category in rejected_object_categories: return False 57 | if coarse_category in filtered_object_categories: return False 58 | if not coarse_category in no_count_object_categories: 59 | if node.zmin - room.zmin > 0.1: 60 | return False 61 | category_dict[coarse_category] = category_dict.get(coarse_category,0) + 1 62 | true_node_count += 1 63 | 64 | if true_node_count < 4 or true_node_count > 20: return False 65 | 66 | if not os.path.isfile(f"./data/suncg_data/room/{room.house_id}/{room.modelId}f.obj"): 67 | return False 68 | if not os.path.isfile(f"./data/suncg_data/room/{room.house_id}/{room.modelId}w.obj"): 69 | return False 70 | 71 | return True 72 | 73 | elif version == "final": 74 | filtered, rejected, door_window = GlobalCategoryFilter.get_filter() 75 | rejected += ["short_kitchen_cabinet", "coffee_table", "straight_chair"] 76 | with open(f"./data/{source}/final_categories_frequency", "r") as f: 77 | frequency = ([s[:-1] for s in f.readlines()]) 78 | frequency = [s.split(" ") for s in frequency] 79 | frequency = dict([(a,int(b)) for (a,b) in frequency]) 80 | 81 | def node_criteria(node, room): 82 | category = category_map.get_final_category(node.modelId) 83 | if category in filtered: return False 84 | return True 85 | 86 | def room_criteria(room, house): 87 | node_count = 0 88 | toilet_count = 0 89 | bathtub_count = 0 90 | for node in room.nodes: 91 | category = category_map.get_final_category(node.modelId) 92 | if category in rejected: 93 | return False 94 | if not category in door_window: 95 | node_count += 1 96 | 97 | t = np.asarray(node.transform).reshape((4,4)).transpose() 98 | a = t[0][0] 99 | b = t[0][2] 100 | c = t[2][0] 101 | d = t[2][2] 102 | 103 | xscale = (a**2 + c**2)**0.5 104 | yscale = (b**2 + d**2)**0.5 105 | zscale = t[1][1] 106 | 107 | if not 0.8 10: return False 122 | if toilet_count > 2: return False 123 | if bathtub_count > 3: return False 124 | 125 | return True 126 | 127 | elif version == "latent": 128 | filtered, rejected, door_window, second_tier_include = \ 129 | GlobalCategoryFilter.get_filter_latent() 130 | rejected += ["desk", "television", "towel_rack", "sofa"] 131 | filtered, rejected, door_window, second_tier_include = \ 132 | set(filtered), set(rejected), set(door_window), set(second_tier_include) 133 | with open(f"{data_dir}/{source}/final_categories_frequency", "r") as f: 134 | frequency = ([s[:-1] for s in f.readlines()]) 135 | frequency = [s.split(" ") for s in frequency] 136 | frequency = dict([(a,int(b)) for (a,b) in frequency]) 137 | 138 | def node_criteria(node, room): 139 | category = category_map.get_final_category(node.modelId) 140 | if category in door_window: 141 | return True 142 | if category in filtered: return False 143 | 144 | if second_tier: 145 | if node.zmin - room.zmin > 0.1 and \ 146 | (category not in second_tier_include or node.parent is None): 147 | return False 148 | 149 | if node.parent: 150 | if isinstance(node.parent, Node) and node.zmin < node.parent.zmax-0.1: 151 | return False 152 | node_now = node 153 | while isinstance(node_now, Node) and node_now.parent: 154 | node_now = node_now.parent 155 | if node_now != "Floor": 156 | return False 157 | else: 158 | if node.zmin - room.zmin > 0.1: 159 | return False 160 | 161 | #Quick filter for second-tier non ceiling mount 162 | #if node.zmin - room.zmin < 0.1: 163 | # return False 164 | #else: 165 | # if node.zmax - room.zmax > -0.2: 166 | # return False 167 | return True 168 | 169 | def room_criteria(room, house): 170 | if room.height > 4: return False 171 | if room.length > room_size: return False 172 | if room.width > room_size: return False 173 | floor_node_count = 0 174 | node_count = 0 175 | scaled = False 176 | #dirty fix! 177 | for i in range(5): 178 | room.nodes = [node for node in room.nodes if not \ 179 | ((node.parent and isinstance(node.parent, Node) and \ 180 | (node.parent) not in room.nodes)) 181 | ] 182 | for node in room.nodes: 183 | category = category_map.get_final_category(node.modelId) 184 | if category in rejected: 185 | return False 186 | if not category in door_window: 187 | node_count += 1 188 | 189 | if node.zmin - room.zmin < 0.1: 190 | floor_node_count += 1 191 | 192 | t = np.asarray(node.transform).reshape((4,4)).transpose() 193 | a = t[0][0] 194 | b = t[0][2] 195 | c = t[2][0] 196 | d = t[2][2] 197 | 198 | xscale = (a**2 + c**2)**0.5 199 | yscale = (b**2 + d**2)**0.5 200 | zscale = t[1][1] 201 | 202 | if not 0.9 20: return False 228 | 229 | return True 230 | else: 231 | raise NotImplementedError 232 | 233 | dataset_f = DatasetFilter(room_filters = [room_criteria], node_filters = [node_criteria]) 234 | 235 | return dataset_f 236 | -------------------------------------------------------------------------------- /scene-synth/loc.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import torch.optim as optim 6 | from torchvision import datasets, transforms 7 | from torch.autograd import Variable 8 | from loc_dataset import * 9 | from PIL import Image 10 | import scipy.misc as m 11 | import numpy as np 12 | import math 13 | import utils 14 | 15 | """ 16 | Module that predicts the location of the next object 17 | """ 18 | 19 | def conv3x3(in_planes, out_planes, stride=1): 20 | "3x3 convolution with padding" 21 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 22 | padding=1, bias=False) 23 | 24 | 25 | class BasicBlock(nn.Module): 26 | expansion = 1 27 | 28 | def __init__(self, inplanes, planes, stride=1, downsample=None): 29 | super(BasicBlock, self).__init__() 30 | self.conv1 = conv3x3(inplanes, planes, stride) 31 | self.bn1 = nn.BatchNorm2d(planes) 32 | self.relu = nn.LeakyReLU(inplace=True) 33 | self.conv2 = conv3x3(planes, planes) 34 | self.bn2 = nn.BatchNorm2d(planes) 35 | self.downsample = downsample 36 | self.stride = stride 37 | 38 | def forward(self, x): 39 | residual = x 40 | 41 | out = self.conv1(x) 42 | out = self.bn1(out) 43 | out = self.relu(out) 44 | 45 | out = self.conv2(out) 46 | out = self.bn2(out) 47 | 48 | if self.downsample is not None: 49 | residual = self.downsample(x) 50 | 51 | out += residual 52 | out = self.relu(out) 53 | 54 | return out 55 | 56 | class ResNet(nn.Module): 57 | 58 | def __init__(self, block, layers, num_classes=2, num_input_channels=17, use_fc=False): 59 | self.inplanes = 64 60 | super(ResNet, self).__init__() 61 | self.conv1 = nn.Conv2d(num_input_channels, 64, kernel_size=7, stride=4, padding=3, 62 | bias=False) 63 | self.bn1 = nn.BatchNorm2d(64) 64 | self.relu = nn.LeakyReLU(inplace=True) 65 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 66 | self.layer1 = self._make_layer(block, 64, layers[0]) 67 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 68 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 69 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 70 | 71 | self.avgpool = nn.AdaptiveAvgPool2d(1) 72 | 73 | self.use_fc = use_fc 74 | if use_fc: 75 | self.fc = nn.Linear(512*block.expansion, num_classes) 76 | 77 | for m in self.modules(): 78 | if isinstance(m, nn.Conv2d): 79 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 80 | m.weight.data.normal_(0, math.sqrt(2. / n)) 81 | elif isinstance(m, nn.BatchNorm2d): 82 | m.weight.data.fill_(1) 83 | m.bias.data.zero_() 84 | 85 | def _make_layer(self, block, planes, blocks, stride=1): 86 | downsample = None 87 | if stride != 1 or self.inplanes != planes * block.expansion: 88 | downsample = nn.Sequential( 89 | nn.Conv2d(self.inplanes, planes * block.expansion, 90 | kernel_size=1, stride=stride, bias=False), 91 | nn.BatchNorm2d(planes * block.expansion), 92 | ) 93 | 94 | layers = [] 95 | layers.append(block(self.inplanes, planes, stride, downsample)) 96 | self.inplanes = planes * block.expansion 97 | for i in range(1, blocks): 98 | layers.append(block(self.inplanes, planes)) 99 | 100 | return nn.Sequential(*layers) 101 | 102 | def forward(self, x): 103 | x = self.conv1(x) 104 | x = self.bn1(x) 105 | x = self.relu(x) 106 | x = self.maxpool(x) 107 | 108 | x = self.layer1(x) 109 | x = self.layer2(x) 110 | x = self.layer3(x) 111 | x = self.layer4(x) 112 | 113 | x = self.avgpool(x) 114 | 115 | return x 116 | 117 | def resnet34(pretrained=False, **kwargs): 118 | """Constructs a ResNet-34 model. 119 | Args: 120 | pretrained (bool): If True, returns a model pre-trained on ImageNet 121 | """ 122 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 123 | return model 124 | 125 | class DownConvBlock(nn.Module): 126 | def __init__(self, inplanes, outplanes): 127 | super(DownConvBlock, self).__init__() 128 | self.conv = nn.Conv2d(inplanes, outplanes, stride=2, kernel_size=4, padding=1) 129 | self.bn = nn.BatchNorm2d(outplanes) 130 | self.act = nn.LeakyReLU() 131 | def forward(self, x): 132 | return self.act(self.bn(self.conv(x))) 133 | 134 | class UpConvBlock(nn.Module): 135 | def __init__(self, inplanes, outplanes): 136 | super(UpConvBlock, self).__init__() 137 | self.conv = nn.Conv2d(inplanes, outplanes, stride=1, kernel_size=3, padding=1) 138 | self.bn = nn.BatchNorm2d(outplanes) 139 | self.act = nn.LeakyReLU() 140 | def forward(self, x): 141 | x = F.upsample(x, mode='nearest', scale_factor=2) 142 | return self.act(self.bn(self.conv(x))) 143 | 144 | class Model(nn.Module): 145 | def __init__(self, num_classes, num_input_channels): 146 | super(Model, self).__init__() 147 | 148 | self.model = nn.Sequential( 149 | nn.Dropout(p=0.2), 150 | resnet34(num_input_channels=num_input_channels), 151 | nn.Dropout(p=0.1), 152 | UpConvBlock(512, 256), 153 | UpConvBlock(256, 128), 154 | UpConvBlock(128, 64), 155 | UpConvBlock(64, 32), 156 | UpConvBlock(32, 16), 157 | UpConvBlock(16, 8), 158 | nn.Dropout(p=0.1), 159 | nn.Conv2d(8,num_classes,1,1) 160 | ) 161 | 162 | def forward(self, x): 163 | x = self.model(x) 164 | return x 165 | 166 | if __name__ == '__main__': 167 | parser = argparse.ArgumentParser(description='Location Training with Auxillary Tasks') 168 | parser.add_argument('--data-folder', type=str, default="bedroom_6x6", metavar='S') 169 | parser.add_argument('--num-workers', type=int, default=6, metavar='N') 170 | parser.add_argument('--last-epoch', type=int, default=-1, metavar='N') 171 | parser.add_argument('--train-size', type=int, default=6000, metavar='N') 172 | parser.add_argument('--save-dir', type=str, default="loc_test", metavar='S') 173 | parser.add_argument('--ablation', type=str, default=None, metavar='S') 174 | parser.add_argument('--lr', type=float, default=0.001, metavar='N') 175 | parser.add_argument('--eps', type=float, default=1e-6, metavar='N') 176 | parser.add_argument('--centroid-weight', type=float, default=10, metavar="N") 177 | args = parser.parse_args() 178 | 179 | save_dir = args.save_dir 180 | utils.ensuredir(save_dir) 181 | batch_size = 16 182 | 183 | with open(f"data/{args.data_folder}/final_categories_frequency", "r") as f: 184 | lines = f.readlines() 185 | num_categories = len(lines)-2 186 | 187 | num_input_channels = num_categories+8 188 | 189 | logfile = open(f"{save_dir}/log_location.txt", 'w') 190 | def LOG(msg): 191 | print(msg) 192 | logfile.write(msg + '\n') 193 | logfile.flush() 194 | 195 | LOG('Building model...') 196 | model = Model(num_classes=num_categories+1, num_input_channels=num_input_channels) 197 | 198 | weight = [args.centroid_weight for i in range(num_categories+1)] 199 | weight[0] = 1 200 | print(weight) 201 | 202 | weight = torch.from_numpy(np.asarray(weight)).float().cuda() 203 | cross_entropy = nn.CrossEntropyLoss(weight=weight) 204 | softmax = nn.Softmax() 205 | 206 | LOG('Converting to CUDA...') 207 | model.cuda() 208 | cross_entropy.cuda() 209 | 210 | LOG('Building dataset...') 211 | train_dataset = LocDataset( 212 | data_root_dir = 'data', 213 | data_folder = args.data_folder, 214 | scene_indices = (0, args.train_size), 215 | ) 216 | 217 | LOG('Building data loader...') 218 | train_loader = torch.utils.data.DataLoader( 219 | train_dataset, 220 | batch_size = batch_size, 221 | num_workers = args.num_workers, 222 | shuffle = True 223 | ) 224 | 225 | LOG('Building optimizer...') 226 | optimizer = optim.Adam(model.parameters(), 227 | lr = args.lr, 228 | weight_decay = 2e-6, 229 | ) 230 | 231 | if args.last_epoch < 0: 232 | load = False 233 | starting_epoch = 0 234 | else: 235 | load = True 236 | last_epoch = args.last_epoch 237 | 238 | if load: 239 | LOG('Loading saved models...') 240 | model.load_state_dict(torch.load(f"{save_dir}/location_{last_epoch}.pt")) 241 | optimizer.load_state_dict(torch.load(f"{save_dir}/location_optim_backup.pt")) 242 | starting_epoch = last_epoch + 1 243 | 244 | current_epoch = starting_epoch 245 | num_seen = 0 246 | 247 | model.train() 248 | LOG(f'=========================== Epoch {current_epoch} ===========================') 249 | 250 | def train(): 251 | global num_seen, current_epoch 252 | for batch_idx, (data, target) \ 253 | in enumerate(train_loader): 254 | 255 | data, target = data.cuda(), target.cuda() 256 | 257 | optimizer.zero_grad() 258 | output = model(data) 259 | loss = cross_entropy(output,target) 260 | print(loss) 261 | 262 | loss.backward() 263 | optimizer.step() 264 | 265 | num_seen += batch_size 266 | if num_seen % 800 == 0: 267 | LOG(f'Examples {num_seen}/10000') 268 | if num_seen % 10000 == 0: 269 | LOG('Validating') 270 | num_seen = 0 271 | current_epoch += 1 272 | LOG(f'=========================== Epoch {current_epoch} ===========================') 273 | if current_epoch % 10 == 0: 274 | torch.save(model.state_dict(), f"{save_dir}/location_{current_epoch}.pt") 275 | torch.save(optimizer.state_dict(), f"{save_dir}/location_optim_backup.pt") 276 | 277 | while True: 278 | train() 279 | -------------------------------------------------------------------------------- /scene-synth/loc_dataset.py: -------------------------------------------------------------------------------- 1 | from torch.utils import data 2 | import torch 3 | from PIL import Image 4 | from torch.autograd import Variable 5 | import numpy as np 6 | import scipy.misc as m 7 | import random 8 | import math 9 | import pickle 10 | from data import ObjectCategories, RenderedScene, RenderedComposite, House, ProjectionGenerator, DatasetToJSON, ObjectData 11 | import copy 12 | import utils 13 | from collections import defaultdict 14 | import os 15 | from functools import cmp_to_key 16 | 17 | """ 18 | Dataset for predicting location 19 | """ 20 | 21 | class LocDataset(): 22 | def __init__(self, scene_indices=(0,4000), data_folder="bedroom", data_root_dir=None, seed=None): 23 | super(LocDataset, self).__init__() 24 | self.category_map = ObjectCategories() 25 | self.seed = seed 26 | self.data_folder = data_folder 27 | self.data_root_dir = data_root_dir 28 | self.scene_indices = scene_indices 29 | 30 | data_root_dir = utils.get_data_root_dir() 31 | with open(f"{data_root_dir}/{data_folder}/final_categories_frequency", "r") as f: 32 | lines = f.readlines() 33 | self.n_categories = len(lines)-2 # -2 for 'window' and 'door' 34 | 35 | def __len__(self): 36 | return self.scene_indices[1]-self.scene_indices[0] 37 | 38 | def __getitem__(self,index): 39 | if self.seed: 40 | random.seed(self.seed) 41 | 42 | i = index+self.scene_indices[0] 43 | scene = RenderedScene(i, self.data_folder, self.data_root_dir) 44 | composite = scene.create_composite() 45 | 46 | object_nodes = scene.object_nodes 47 | 48 | random.shuffle(object_nodes) 49 | 50 | if 'parent' in object_nodes[0]: 51 | #print([a["category"] for a in object_nodes]) 52 | # Make sure that all second-tier objects come *after* first tier ones 53 | def is_second_tier(node): 54 | return (node['parent'] != 'Wall') and \ 55 | (node['parent'] != 'Floor') 56 | object_nodes.sort(key = lambda node: int(is_second_tier(node))) 57 | 58 | # Make sure that all children come after their parents 59 | def cmp_parent_child(node1, node2): 60 | # Less than (negative): node1 is the parent of node2 61 | if node2['parent'] == node1['id']: 62 | return -1 63 | # Greater than (postive): node2 is the parent of node1 64 | elif node1['parent'] == node2['id']: 65 | return 1 66 | # Equal (zero): all other cases 67 | else: 68 | return 0 69 | object_nodes.sort(key = cmp_to_key(cmp_parent_child)) 70 | #print([a["category"] for a in object_nodes]) 71 | #print("________________") 72 | 73 | num_objects = random.randint(0, len(object_nodes)) 74 | #num_objects = len(object_nodes) 75 | num_categories = len(scene.categories) 76 | 77 | centroids = [] 78 | parent_ids = ["Floor", "Wall"] 79 | for i in range(num_objects): 80 | node = object_nodes[i] 81 | if node["parent"] == "Wall": 82 | print("Massive messup!") 83 | composite.add_node(node) 84 | xsize, ysize = node["height_map"].shape 85 | xmin, _, ymin, _ = node["bbox_min"] 86 | xmax, _, ymax, _ = node["bbox_max"] 87 | parent_ids.append(node["id"]) 88 | 89 | inputs = composite.get_composite(num_extra_channels=0) 90 | size = inputs.shape[1] 91 | 92 | for i in range(num_objects, len(object_nodes)): 93 | node = object_nodes[i] 94 | if node["parent"] == "Wall": 95 | print("Massive messup!") 96 | if node["parent"] in parent_ids: 97 | xsize, ysize = node["height_map"].shape 98 | xmin, _, ymin, _ = node["bbox_min"] 99 | xmax, _, ymax, _ = node["bbox_max"] 100 | centroids.append(((xmin+xmax)/2, (ymin+ymax)/2, node["category"])) 101 | 102 | output = torch.zeros((64,64)).long() 103 | 104 | for (x,y,label) in centroids: 105 | output[math.floor(x/4),math.floor(y/4)] = label+1 106 | 107 | return inputs, output 108 | 109 | -------------------------------------------------------------------------------- /scene-synth/math_utils/OBB.py: -------------------------------------------------------------------------------- 1 | import random 2 | import numpy as np 3 | 4 | from pyquaternion import Quaternion 5 | from math_utils import Transform 6 | 7 | 8 | class OBB(object): 9 | def __init__(self, center, half_widths, rotation_matrix): 10 | self._c = np.squeeze(center) 11 | self._h = np.squeeze(half_widths) 12 | self._R = rotation_matrix 13 | self._recompute_transforms() 14 | 15 | def _recompute_transforms(self): 16 | # local-to-world transform: takes [0,1]^3 points in OBB to world space points 17 | self._local_to_world = np.identity(4) 18 | for i in range(3): 19 | self._local_to_world[:3, i] = self._R[:, i] * self._h[i] 20 | self._local_to_world[:3, 3] = self._c 21 | 22 | # world-to-local transform: world space points within OBB to [0,1]^3 23 | self._world_to_local = np.identity(4) 24 | for i in range(3): 25 | self._world_to_local[i, :3] = self._R[:, i] * (1.0 / self._h[i]) 26 | t_inv = - np.matmul(self._world_to_local[:3, :3], np.transpose(self._c)) 27 | self._world_to_local[:3, 3] = np.squeeze(t_inv) 28 | 29 | @classmethod 30 | def from_local2world_transform(cls, transform): 31 | xform = Transform.from_mat4(transform) 32 | return cls(center=xform.translation, half_widths=xform.scale, rotation_matrix=xform.rotation.rotation_matrix) 33 | 34 | @classmethod 35 | def from_node(cls, node, aligned_dims): 36 | xform = Transform.from_mat4x4_flat_row_major(node.transform) 37 | return cls(center=xform.translation, half_widths=aligned_dims * xform.scale * 0.5, 38 | rotation_matrix=xform.rotation.rotation_matrix) 39 | 40 | @property 41 | def half_extents(self): 42 | return self._h 43 | 44 | @property 45 | def rotation_matrix(self): 46 | return self._R 47 | 48 | @property 49 | def rotation_quat(self): 50 | return Quaternion(matrix=self._R) 51 | 52 | @property 53 | def dimensions(self): 54 | return 2.0 * self._h 55 | 56 | @property 57 | def half_dimensions(self): 58 | return self._h 59 | 60 | @property 61 | def centroid(self): 62 | return self._c 63 | 64 | @property 65 | def world2local(self): 66 | return self._world_to_local 67 | 68 | @property 69 | def local2world(self): 70 | return self._local_to_world 71 | 72 | def __repr__(self): 73 | return 'OBB: {c:' + str(self._c) + ',h:' + str(self._h) + ',R:' + str(self._R.tolist()) + '}' 74 | 75 | def transform_point(self, p): 76 | return np.matmul(self._world_to_local, np.append(p, [1], axis=0))[:3] 77 | 78 | def transform_point_toworld(self, p): 79 | return np.matmul(self._local_to_world, np.append(p, [1], axis=0))[:3] 80 | 81 | def transform_direction(self, d): 82 | return np.matmul(np.transpose(self._R), d) 83 | 84 | def distance_to_point(self, p): 85 | if self.contains_point(p): 86 | return 0.0 87 | closest = self.closest_point(p) 88 | return np.linalg.norm(p - closest) 89 | 90 | def contains_point(self, p): 91 | p_local = np.matmul(self._world_to_local, np.append(p, [1], axis=0))[:3] 92 | bound = 1.0 93 | for i in range(3): 94 | if abs(p_local[i]) > bound: 95 | return False 96 | return True # here only if all three coord within bounds 97 | 98 | def closest_point(self, p): 99 | d = p - self._c 100 | closest = np.copy(self._c) 101 | for i in range(3): 102 | closest += np.clip(self._R[:, i] * d, -self._h[i], self._h[i]) * self._R[:, i] 103 | return closest 104 | 105 | def sample(self): 106 | p = np.copy(self._c) 107 | for i in range(3): 108 | r = random.random() * 2.0 - 1.0 109 | p += r * self._h[i] * self._R[:, i] 110 | return p 111 | 112 | def project_to_axis(self, direction): 113 | """ 114 | Projects this OBB onto the given 1D direction vector and returns projection interval [min, max]. 115 | If vector is unnormalized, the output is scaled by the length of the vector 116 | :param direction: the axis on which to project this OBB 117 | :return: out_min - the minimum extent of this OBB along direction, out_max - the maximum extent 118 | """ 119 | x = abs(np.dot(direction, self._R[0]) * self._h[0]) 120 | y = abs(np.dot(direction, self._R[1]) * self._h[1]) 121 | z = abs(np.dot(direction, self._R[2]) * self._h[2]) 122 | p = np.dot(direction, self._c) 123 | out_min = p - x - y - z 124 | out_max = p + x + y + z 125 | return out_min, out_max 126 | 127 | def signed_distance_to_plane(self, plane): 128 | p_min, p_max = self.project_to_axis(plane.normal) 129 | p_min -= plane.d 130 | p_max -= plane.d 131 | if p_min * p_max <= 0.0: 132 | return 0.0 133 | return p_min if abs(p_min) < abs(p_max) else p_max 134 | 135 | def to_aabb(self): 136 | h_size = abs(self._R[0] * self._h[0]) + abs(self._R[1] * self._h[1]) + abs(self._R[2] * self._h[2]) 137 | p_min = self._c - h_size 138 | p_max = self._c + h_size 139 | return p_min, p_max 140 | 141 | # Return a list of vertex position tuples which represent this OBB as a triangle mesh 142 | def get_triangles(self): 143 | # For interpretability, let's call x = left/right, y = bottom/top, z = back/front 144 | tris = [ 145 | # BOTTOM FACE 146 | ([0, 0, 0], [0, 0, 1], [1, 0, 1]), 147 | ([1, 0, 1], [1, 0, 0], [0, 0, 0]), 148 | # TOP FACE 149 | ([0, 1, 0], [0, 1, 1], [1, 1, 1]), 150 | ([1, 1, 1], [1, 1, 0], [0, 1, 0]), 151 | # LEFT FACE 152 | ([0, 0, 0], [0, 0, 1], [0, 1, 1]), 153 | ([0, 1, 1], [0, 1, 0], [0, 0, 0]), 154 | # RIGHT FACE 155 | ([1, 0, 0], [1, 0, 1], [1, 1, 1]), 156 | ([1, 1, 1], [1, 1, 0], [1, 0, 0]), 157 | # BACK FACE 158 | ([0, 0, 0], [0, 1, 0], [1, 1, 0]), 159 | ([1, 1, 0], [1, 0, 0], [0, 0, 0]), 160 | # FRONT FACE 161 | ([0, 0, 1], [0, 1, 1], [1, 1, 1]), 162 | ([1, 1, 1], [1, 0, 1], [0, 0, 1]) 163 | ] 164 | return [ 165 | tuple([ self.transform_point_toworld(np.array(vert)) for vert in tri ]) 166 | for tri in tris 167 | ] 168 | -------------------------------------------------------------------------------- /scene-synth/math_utils/Simulator.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import pybullet as p 4 | import subprocess as sp 5 | import time 6 | 7 | from pyquaternion import Quaternion 8 | from collections import namedtuple 9 | from itertools import groupby 10 | from math_utils import Transform 11 | import utils 12 | 13 | # handle to a simulated rigid body 14 | Body = namedtuple('Body', ['id', 'bid', 'vid', 'cid', 'static']) 15 | 16 | # a body-body contact record 17 | Contact = namedtuple( 18 | 'Contact', ['flags', 'idA', 'idB', 'linkIndexA', 'linkIndexB', 19 | 'positionOnAInWS', 'positionOnBInWS', 20 | 'contactNormalOnBInWS', 'distance', 'normalForce'] 21 | ) 22 | 23 | # ray intersection record 24 | Intersection = namedtuple('Intersection', ['id', 'linkIndex', 'ray_fraction', 'position', 'normal']) 25 | 26 | 27 | class Simulator: 28 | def __init__(self, mode='direct', bullet_server_binary=None, data_dir_base=None, verbose=False): 29 | self._mode = mode 30 | self._verbose = verbose 31 | module_dir = os.path.dirname(os.path.abspath(__file__)) 32 | data_root_dir = utils.get_data_root_dir() 33 | if data_dir_base: 34 | self._data_dir_base = data_dir_base 35 | else: 36 | self._data_dir_base = os.path.join(data_root_dir, 'suncg_data') 37 | if bullet_server_binary: 38 | self._bullet_server_binary = bullet_server_binary 39 | else: 40 | self._bullet_server_binary = os.path.join(module_dir, '..', 'bullet_shared_memory_server') 41 | self._obj_id_to_body = {} 42 | self._bid_to_body = {} 43 | self._pid = None 44 | self._bullet_server = None 45 | self.connect() 46 | 47 | def connect(self): 48 | # disconnect and kill existing servers 49 | if self._pid: 50 | p.disconnect(physicsClientId=self._pid) 51 | self._pid = None 52 | if self._bullet_server: 53 | print(f'Restarting by killing bullet server pid={self._bullet_server.pid}...') 54 | self._bullet_server.kill() 55 | time.sleep(1) # seems necessary to prevent deadlock on re-connection attempt 56 | self._bullet_server = None 57 | 58 | # reconnect to appropriate server type 59 | if self._mode == 'gui': 60 | self._pid = p.connect(p.GUI) 61 | elif self._mode == 'direct': 62 | self._pid = p.connect(p.DIRECT) 63 | elif self._mode == 'shared_memory': 64 | print(f'Restarting bullet server process...') 65 | self._bullet_server = sp.Popen([self._bullet_server_binary]) 66 | time.sleep(1) # seems necessary to prevent deadlock on connection attempt 67 | self._pid = p.connect(p.SHARED_MEMORY) 68 | else: 69 | raise RuntimeError(f'Unknown simulator server mode={self._mode}') 70 | 71 | # reset and initialize gui if needed 72 | p.resetSimulation(physicsClientId=self._pid) 73 | if self._mode == 'gui': 74 | p.configureDebugVisualizer(p.COV_ENABLE_Y_AXIS_UP, 1, physicsClientId=self._pid) 75 | p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self._pid) 76 | p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self._pid) 77 | p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self._pid) 78 | # disable rendering during loading -> much faster 79 | p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 0, physicsClientId=self._pid) 80 | 81 | def __del__(self): 82 | if self._bullet_server: 83 | print(f'Process terminating. Killing bullet server pid={self._bullet_server.pid}...') 84 | self._bullet_server.kill() 85 | 86 | def run(self): 87 | if self._mode == 'gui': 88 | # infinite ground plane and gravity 89 | # plane_cid = p.createCollisionShape(p.GEOM_PLANE, planeNormal=[0, 1, 0]) 90 | # plane_bid = p.createMultiBody(baseMass=0, baseCollisionShapeIndex=plane_cid) 91 | p.setGravity(0, -10, 0, physicsClientId=self._pid) 92 | p.setRealTimeSimulation(1, physicsClientId=self._pid) 93 | self.set_gui_rendering(enabled=True) 94 | while True: 95 | contacts = self.get_contacts(include_collision_with_static=True) 96 | if self._verbose: 97 | print(f'#contacts={len(contacts)}, contact_pairs={contacts.keys()}') 98 | else: 99 | self.step() 100 | contacts = self.get_contacts(include_collision_with_static=True) 101 | if self._verbose: 102 | for (k, c) in contacts.items(): 103 | cp = self.get_closest_point(obj_id_a=c.idA, obj_id_b=c.idB) 104 | print(f'contact pair={k} record={cp}') 105 | print(f'#contacts={len(contacts)}, contact_pairs={contacts.keys()}') 106 | 107 | def set_gui_rendering(self, enabled): 108 | if not self._mode == 'gui': 109 | return False 110 | center = np.array([0.0, 0.0, 0.0]) 111 | num_obj = 0 112 | for obj_id in self._obj_id_to_body.keys(): 113 | pos, _ = self.get_state(obj_id) 114 | if not np.allclose(pos, [0, 0, 0]): # try to ignore room object "identity" transform 115 | num_obj += 1 116 | center += pos 117 | center /= num_obj 118 | p.resetDebugVisualizerCamera(cameraDistance=10.0, 119 | cameraYaw=45.0, 120 | cameraPitch=-30.0, 121 | cameraTargetPosition=center, 122 | physicsClientId=self._pid) 123 | p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 1 if enabled else 0, physicsClientId=self._pid) 124 | return enabled 125 | 126 | def add_mesh(self, obj_id, obj_file, transform, vis_mesh_file=None, static=False): 127 | if static: 128 | cid = p.createCollisionShape(p.GEOM_MESH, fileName=obj_file, meshScale=transform.scale, 129 | flags=p.GEOM_FORCE_CONCAVE_TRIMESH, physicsClientId=self._pid) 130 | else: 131 | cid = p.createCollisionShape(p.GEOM_MESH, fileName=obj_file, meshScale=transform.scale, 132 | physicsClientId=self._pid) 133 | vid = -1 134 | if vis_mesh_file: 135 | vid = p.createVisualShape(p.GEOM_MESH, fileName=vis_mesh_file, meshScale=transform.scale, 136 | physicsClientId=self._pid) 137 | rot_q = np.roll(transform.rotation.elements, -1) # w,x,y,z -> x,y,z,w (which pybullet expects) 138 | mass = 0 if static else 1 139 | bid = p.createMultiBody(baseMass=mass, 140 | baseCollisionShapeIndex=cid, 141 | baseVisualShapeIndex=vid, 142 | basePosition=transform.translation, 143 | baseOrientation=rot_q, 144 | physicsClientId=self._pid) 145 | body = Body(id=obj_id, bid=bid, vid=vid, cid=cid, static=static) 146 | self._obj_id_to_body[obj_id] = body 147 | self._bid_to_body[bid] = body 148 | return body 149 | 150 | def add_box(self, obj_id, half_extents, transform, static=False): 151 | cid = p.createCollisionShape(p.GEOM_BOX, halfExtents=half_extents, physicsClientId=self._pid) 152 | rot_q = np.roll(transform.rotation.elements, -1) # w,x,y,z -> x,y,z,w (which pybullet expects) 153 | mass = 0 if static else 1 154 | bid = p.createMultiBody(baseMass=mass, 155 | baseCollisionShapeIndex=cid, 156 | basePosition=transform.translation, 157 | baseOrientation=rot_q, 158 | physicsClientId=self._pid) 159 | body = Body(id=obj_id, bid=bid, vid=-1, cid=cid, static=static) 160 | self._obj_id_to_body[obj_id] = body 161 | self._bid_to_body[bid] = body 162 | return body 163 | 164 | # House-specific functions 165 | def add_object(self, node, create_vis_mesh=False, static=False): 166 | model_id = node.modelId.replace('_mirror', '') # TODO: need to otherwise account for mirror? 167 | object_dir = os.path.join(self._data_dir_base, 'object') 168 | basename = f'{object_dir}/{model_id}/{model_id}' 169 | vis_obj_filename = f'{basename}.obj' if create_vis_mesh else None 170 | col_obj_filename = f'{basename}.vhacd.obj' 171 | if not os.path.exists(col_obj_filename): 172 | print('WARNING: collision mesh {col_obj_filename} unavailable, using visual mesh instead.') 173 | col_obj_filename = f'{basename}.obj' 174 | return self.add_mesh(obj_id=node.id, obj_file=col_obj_filename, transform=Transform.from_node(node), 175 | vis_mesh_file=vis_obj_filename, static=static) 176 | 177 | def add_wall(self, node): 178 | h = node['height'] 179 | p0 = np.transpose(np.matrix(node['points'][0])) 180 | p1 = np.transpose(np.matrix(node['points'][1])) 181 | c = (p0 + p1) * 0.5 182 | c[1] = h * 0.5 183 | dp = p1 - p0 184 | dp_l = np.linalg.norm(dp) 185 | dp = dp / dp_l 186 | angle = np.arccos(dp[0]) 187 | rot_q = Quaternion(axis=[0, 1, 0], radians=angle) 188 | half_extents = np.array([dp_l, h, node['depth']]) * 0.5 189 | return self.add_box(obj_id=node['id'], half_extents=half_extents, 190 | transform=Transform(translation=c, rotation=rot_q), static=True) 191 | 192 | def add_room(self, node, wall=True, floor=True, ceiling=False): 193 | def add_architecture(n, obj_file, suffix): 194 | return self.add_mesh(obj_id=n.id + suffix, obj_file=obj_file, transform=Transform(), vis_mesh_file=None, 195 | static=True) 196 | room_id = node.modelId 197 | room_dir = os.path.join(self._data_dir_base, 'room') 198 | basename = f'{room_dir}/{node.house_id}/{room_id}' 199 | body_ids = [] 200 | if wall: 201 | body_wall = add_architecture(node, f'{basename}w.obj', '') # treat walls as room (=room.id, no suffix) 202 | body_ids.append(body_wall) 203 | if floor: 204 | body_floor = add_architecture(node, f'{basename}f.obj', 'f') 205 | body_ids.append(body_floor) 206 | if ceiling: 207 | body_ceiling = add_architecture(node, f'{basename}c.obj', 'c') 208 | body_ids.append(body_ceiling) 209 | return body_ids 210 | 211 | def add_house(self, house, no_walls=False, no_ceil=False, no_floor=False, use_separate_walls=False, static=False): 212 | for node in house.nodes: 213 | if not node.valid: 214 | continue 215 | if not hasattr(node, 'body'): 216 | node.body = None 217 | if node.type == 'Object': 218 | node.body = self.add_object(node, static=static) 219 | if node.type == 'Room': 220 | ceil = False if no_ceil else not (hasattr(node, 'hideCeiling') and node.hideCeiling == 1) 221 | wall = False if (no_walls or use_separate_walls) else not (hasattr(node, 'hideWalls') and node.hideWalls == 1) 222 | floor = False if no_floor else not (hasattr(node, 'hideFloor') and node.hideFloor == 1) 223 | node.body = self.add_room(node, wall=wall, floor=floor, ceiling=ceil) 224 | if node.type == 'Box': 225 | half_widths = list(map(lambda x: 0.5 * x, node.dimensions)) 226 | node.body = self.add_box(obj_id=node.id, half_extents=half_widths, transform=Transform.from_node(node), 227 | static=static) 228 | if use_separate_walls and not no_walls: 229 | for wall in house.walls: 230 | wall['body'] = self.add_wall(wall) 231 | 232 | def add_house_room_only(self, house, room, no_walls=False, no_ceil=True, no_floor=False, use_separate_walls=False, 233 | only_architecture=False, static=False): 234 | #walls, ceil, floor logic not fleshed out due to current limited use case 235 | room_node = [node for node in house.nodes if node.id == room.id] 236 | if len(room_node) < 1: 237 | raise Exception("Missing Room") 238 | if only_architecture: 239 | house.nodes = room_node 240 | else: 241 | house.nodes = [node for node in room.nodes] 242 | house.nodes.append(room_node[0]) 243 | 244 | for node in house.nodes: 245 | if not node.valid: 246 | continue 247 | if not hasattr(node, 'body'): 248 | node.body = None 249 | if node.type == 'Object': 250 | node.body = self.add_object(node, static=static) 251 | if node.type == 'Room': 252 | ceil = False if no_ceil else not (hasattr(node, 'hideCeiling') and node.hideCeiling == 1) 253 | wall = False if (no_walls or use_separate_walls) else not (hasattr(node, 'hideWalls') and node.hideWalls == 1) 254 | floor = False if no_floor else not (hasattr(node, 'hideFloor') and node.hideFloor == 1) 255 | node.body = self.add_room(node, wall=wall, floor=floor, ceiling=ceil) 256 | if node.type == 'Box': 257 | half_widths = list(map(lambda x: 0.5 * x, node.dimensions)) 258 | node.body = self.add_box(obj_id=node.id, half_extents=half_widths, transform=Transform.from_node(node), 259 | static=static) 260 | if use_separate_walls and not no_walls: 261 | for wall in house.walls: 262 | wall['body'] = self.add_wall(wall) 263 | 264 | def remove(self, obj_id): 265 | body = self._obj_id_to_body[obj_id] 266 | p.removeBody(bodyUniqueId=body.bid, physicsClientId=self._pid) 267 | del self._obj_id_to_body[obj_id] 268 | del self._bid_to_body[body.bid] 269 | 270 | def set_state(self, obj_id, position, rotation_q): 271 | body = self._obj_id_to_body[obj_id] 272 | rot_q = np.roll(rotation_q.elements, -1) # w,x,y,z -> x,y,z,w (which pybullet expects) 273 | p.resetBasePositionAndOrientation(bodyUniqueId=body.bid, posObj=position, ornObj=rot_q, 274 | physicsClientId=self._pid) 275 | 276 | def get_state(self, obj_id): 277 | body = self._obj_id_to_body[obj_id] 278 | pos, q = p.getBasePositionAndOrientation(bodyUniqueId=body.bid, physicsClientId=self._pid) 279 | rotation = Quaternion(x=q[0], y=q[1], z=q[2], w=q[3]) 280 | return pos, rotation 281 | 282 | def step(self): 283 | p.stepSimulation(physicsClientId=self._pid) 284 | 285 | def reset(self): 286 | p.resetSimulation(physicsClientId=self._pid) 287 | self._obj_id_to_body = {} 288 | self._bid_to_body = {} 289 | 290 | def get_closest_point(self, obj_id_a, obj_id_b, max_distance=np.inf): 291 | """ 292 | Return record with distance between closest points between pair of nodes if within max_distance or None. 293 | """ 294 | bid_a = self._obj_id_to_body[obj_id_a].bid 295 | bid_b = self._obj_id_to_body[obj_id_b].bid 296 | cps = p.getClosestPoints(bodyA=bid_a, bodyB=bid_b, distance=max_distance, physicsClientId=self._pid) 297 | cp = None 298 | if len(cps) > 0: 299 | closest_points = self._convert_contacts(cps) 300 | cp = min(closest_points, key=lambda x: x.distance) 301 | del cps # NOTE force garbage collection of pybullet objects 302 | return cp 303 | 304 | def get_contacts(self, obj_id_a=None, obj_id_b=None, only_closest_contact_per_pair=True, 305 | include_collision_with_static=True): 306 | """ 307 | Return all current contacts. When include_collision_with_statics is true, include contacts with static bodies 308 | """ 309 | bid_a = self._obj_id_to_body[obj_id_a].bid if obj_id_a else -1 310 | bid_b = self._obj_id_to_body[obj_id_b].bid if obj_id_b else -1 311 | cs = p.getContactPoints(bodyA=bid_a, bodyB=bid_b, physicsClientId=self._pid) 312 | contacts = self._convert_contacts(cs) 313 | del cs # NOTE force garbage collection of pybullet objects 314 | 315 | if not include_collision_with_static: 316 | def not_contact_with_static(c): 317 | static_a = self._obj_id_to_body[c.idA].static 318 | static_b = self._obj_id_to_body[c.idB].static 319 | return not static_a and not static_b 320 | contacts = filter(not_contact_with_static, contacts) 321 | # print(f'#all_contacts={len(all_contacts)} to #non_static_contacts={len(non_static_contacts)}') 322 | 323 | if only_closest_contact_per_pair: 324 | def bid_pair_key(x): 325 | return str(x.idA) + '_' + str(x.idB) 326 | contacts = sorted(contacts, key=bid_pair_key) 327 | min_dist_contact_by_pair = {} 328 | for k, g in groupby(contacts, key=bid_pair_key): 329 | min_dist_contact = min(g, key=lambda x: x.distance) 330 | min_dist_contact_by_pair[k] = min_dist_contact 331 | contacts = min_dist_contact_by_pair.values() 332 | 333 | # convert into dictionary of form (id_a, id_b) -> Contact 334 | contacts_dict = {} 335 | for c in contacts: 336 | key = (c.idA, c.idB) 337 | contacts_dict[key] = c 338 | 339 | return contacts_dict 340 | 341 | def _convert_contacts(self, contacts): 342 | out = [] 343 | for c in contacts: 344 | bid_a = c[1] 345 | bid_b = c[2] 346 | if bid_a not in self._bid_to_body or bid_b not in self._bid_to_body: 347 | continue 348 | id_a = self._bid_to_body[bid_a].id 349 | id_b = self._bid_to_body[bid_b].id 350 | o = Contact(flags=c[0], idA=id_a, idB=id_b, linkIndexA=c[3], linkIndexB=c[4], 351 | positionOnAInWS=c[5], positionOnBInWS=c[6], contactNormalOnBInWS=c[7], 352 | distance=c[8], normalForce=c[9]) 353 | out.append(o) 354 | return out 355 | 356 | def ray_test(self, from_pos, to_pos): 357 | hit = p.rayTest(rayFromPosition=from_pos, rayToPosition=to_pos, physicsClientId=self._pid) 358 | intersection = Intersection._make(*hit) 359 | del hit # NOTE force garbage collection of pybullet objects 360 | if intersection.id >= 0: # if intersection, replace bid with id 361 | intersection = intersection._replace(id=self._bid_to_body[intersection.id].id) 362 | return intersection 363 | 364 | 365 | if __name__ == '__main__': 366 | from data import House 367 | sim = Simulator(mode='gui', verbose=True) 368 | h = House(id_='d119e6e0bd567d923aea774c2a984bf0', include_arch_information=False) 369 | # h = House(index=5) 370 | sim.add_house(h, no_walls=False, no_ceil=True, use_separate_walls=False, static=False) 371 | sim.run() 372 | -------------------------------------------------------------------------------- /scene-synth/math_utils/__init__.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | import scipy 3 | import math 4 | import matplotlib as mpl 5 | import numpy as np 6 | 7 | from scipy.stats import vonmises 8 | from pyquaternion import Quaternion 9 | 10 | try: 11 | import matplotlib.pyplot as plt 12 | except RuntimeError: 13 | print('ERROR importing pyplot, plotting functions unavailable') 14 | 15 | 16 | class Transform: 17 | """ 18 | Represents an affine transformation composed of translation, rotation, scale 19 | """ 20 | def __init__(self, translation=np.asarray((0, 0, 0)), rotation=Quaternion(), scale=np.asarray((1, 1, 1))): 21 | self._t = translation 22 | self._r = rotation 23 | self._s = scale 24 | 25 | @property 26 | def rotation(self): 27 | return self._r 28 | 29 | @property 30 | def translation(self): 31 | return self._t 32 | 33 | @property 34 | def scale(self): 35 | return self._s 36 | 37 | @classmethod 38 | def from_mat4(cls, m): 39 | """ 40 | Construct Transform from matrix m 41 | :param m: 4x4 affine transformation matrix 42 | :return: rotation (quaternion), translation, scale 43 | """ 44 | translation = m[:3, 3] 45 | scale_rotation = m[:3, :3] 46 | u, _ = scipy.linalg.polar(scale_rotation) 47 | rotation_q = Quaternion(matrix=u) 48 | r_axis = rotation_q.get_axis(undefined=[0, 1, 0]) 49 | # ensure that rotation axis is +y 50 | if np.array(r_axis).dot(np.array([0, 1, 0])) < 0: 51 | rotation_q = -rotation_q 52 | scale = np.linalg.norm(scale_rotation, axis=0) 53 | return cls(translation=translation, rotation=rotation_q, scale=scale) 54 | 55 | @classmethod 56 | def from_mat4x4_flat_row_major(cls, mat): 57 | return cls.from_mat4(np.array(mat).reshape((4, 4)).transpose()) 58 | 59 | @classmethod 60 | def from_node(cls, node): 61 | if hasattr(node, 'transform'): 62 | # get rotation and translation out of 4x4 xform 63 | return cls.from_mat4x4_flat_row_major(node.transform) 64 | elif hasattr(node, 'type') and hasattr(node, 'bbox'): 65 | p_min = np.array(node.bbox['min']) 66 | p_max = np.array(node.bbox['max']) 67 | center = (p_max + p_min) * 0.5 68 | half_dims = (p_max - p_min) * 0.5 69 | return cls(translation=center, scale=half_dims) 70 | else: 71 | return None 72 | 73 | def as_mat4(self): 74 | m = np.identity(4) 75 | m[:3, 3] = self._t 76 | r = self._r.rotation_matrix 77 | for i in range(3): 78 | m[:3, i] = r[:, i] * self._s[i] 79 | return m 80 | 81 | def as_mat4_flat_row_major(self): 82 | return list(self.as_mat4().transpose().flatten()) 83 | 84 | def inverse(self): 85 | return self.from_mat4(np.linalg.inv(self.as_mat4())) 86 | 87 | def set_translation(self, t): 88 | self._t = t 89 | 90 | def translate(self, t): 91 | self._t += t 92 | 93 | def rotate(self, radians, axis): 94 | q = Quaternion(axis=axis, radians=radians) 95 | self._r = q * Quaternion(matrix=self._r) 96 | 97 | def set_rotation(self, radians, axis=np.asarray((0, 1, 0))): 98 | self._r = Quaternion(axis=axis, radians=radians) 99 | 100 | def rotate_y(self, radians): 101 | self.rotate(radians, [0, 1, 0]) 102 | 103 | def rescale(self, s): 104 | self._s = np.multiply(self._s, s) 105 | 106 | def transform_point(self, p): 107 | # TODO cache optimization 108 | return np.matmul(self.as_mat4(), np.append(p, [1], axis=0))[:3] 109 | 110 | def transform_direction(self, d): 111 | # TODO cache optimization 112 | return np.matmul(np.transpose(self._r.rotation_matrix), d) 113 | 114 | 115 | def relative_pos_to_xz_distance_angle(p): 116 | ang = math.atan2(p[2], p[0]) 117 | dist = math.sqrt(p[0] * p[0] + p[2] * p[2]) 118 | return dist, ang 119 | 120 | 121 | def relative_dir_to_xz_angle(d): 122 | return math.atan2(d[2], d[0]) 123 | 124 | 125 | def nparr2str_compact(a): 126 | # drop '[ and ]' at ends and condense whitespace 127 | v = [] 128 | for i in range(a.shape[0]): 129 | v.append('%.6g' % a[i]) 130 | return ' '.join(v) 131 | # return re.sub(' +', ' ', np.array2string(a, precision=5, suppress_small=True)[1:-1]).strip() 132 | 133 | 134 | def str2nparr(sa): 135 | return np.fromstring(sa, dtype=float, sep=' ') 136 | 137 | 138 | def plot_gmm_fit(x, y_, means, covariances, index, title): 139 | color_iter = itertools.cycle(['navy', 'c', 'cornflowerblue', 'gold', 'darkorange']) 140 | subplot = plt.subplot(3, 1, 1 + index) 141 | for i, (mean, covariance, color) in enumerate(zip(means, covariances, color_iter)): 142 | v, w = scipy.linalg.eigh(covariances) 143 | v = 2. * np.sqrt(2.) * np.sqrt(v) 144 | u = w[0] / scipy.linalg.norm(w[0]) 145 | # DP GMM will not use every component it has access to unless it needs it, don't plot redundant components 146 | if not np.any(y_ == i): 147 | continue 148 | plt.scatter(x[y_ == i, 0], x[y_ == i, 1], .8, color=color) 149 | # Plot an ellipse to show the Gaussian component 150 | angle = np.arctan(u[1] / u[0]) 151 | angle = 180. * angle / np.pi # convert to degrees 152 | ell = mpl.patches.Ellipse(mean, v[0], v[1], 180. + angle, color=color) 153 | ell.set_clip_box(subplot.bbox) 154 | ell.set_alpha(0.5) 155 | subplot.add_artist(ell) 156 | plt.title(title) 157 | 158 | 159 | def plot_vmf_fit(x, mu, kappa, scale, index, title): 160 | subplot = plt.subplot(3, 1, 1 + index) 161 | plt.hist(x, bins=8, normed=True, histtype='stepfilled') 162 | domain = np.linspace(vonmises.ppf(0.01, kappa), vonmises.ppf(0.99, kappa), 100) 163 | plt.plot(domain, vonmises.pdf(domain, kappa=kappa, loc=mu, scale=scale)) 164 | plt.title(title) 165 | 166 | 167 | def plot_hist_fit(x, hist_dist, index, title): 168 | subplot = plt.subplot(3, 1, 1 + index) 169 | plt.hist(x, bins=16, normed=True, histtype='stepfilled') 170 | domain = np.linspace(-math.pi, math.pi, 16) 171 | plt.plot(domain, hist_dist.pdf(domain)) 172 | plt.title(title) 173 | -------------------------------------------------------------------------------- /scene-synth/model_prior.py: -------------------------------------------------------------------------------- 1 | from data import RenderedScene, ObjectCategories 2 | import os 3 | import pickle 4 | import numpy as np 5 | import utils 6 | 7 | 8 | """ 9 | Simple bigram model 10 | """ 11 | class ModelPrior(): 12 | 13 | def __init__(self): 14 | pass 15 | 16 | def learn(self, data_folder="bedroom_final", data_root_dir=None): 17 | if not data_root_dir: 18 | data_root_dir = utils.get_data_root_dir() 19 | data_dir = f"{data_root_dir}/{data_folder}" 20 | self.data_dir = data_dir 21 | self.category_map = ObjectCategories() 22 | 23 | files = os.listdir(data_dir) 24 | files = [f for f in files if ".pkl" in f and not "domain" in f and not "_" in f] 25 | 26 | with open(f"{data_dir}/final_categories_frequency", "r") as f: 27 | lines = f.readlines() 28 | cats = [line.split()[0] for line in lines] 29 | 30 | self.categories = [cat for cat in cats if cat not in set(['window', 'door'])] 31 | self.cat_to_index = {self.categories[i]:i for i in range(len(self.categories))} 32 | 33 | with open(f"{data_dir}/model_frequency", "r") as f: 34 | lines = f.readlines() 35 | models = [line.split()[0] for line in lines] 36 | self.model_freq = [int(l[:-1].split()[1]) for l in lines] 37 | 38 | self.models = [model for model in models if self.category_map.get_final_category(model) not in set(['window', 'door'])] 39 | self.model_to_index = {self.models[i]:i for i in range(len(self.models))} 40 | 41 | N = len(self.models) 42 | self.num_categories = len(self.categories) 43 | 44 | self.model_index_to_cat = [self.cat_to_index[self.category_map.get_final_category(self.models[i])] for i in range(N)] 45 | 46 | self.count = [[0 for i in range(N)] for j in range(N)] 47 | 48 | for index in range(len(files)): 49 | #for index in range(100): 50 | with open(f"{data_dir}/{index}.pkl", "rb") as f: 51 | (_, _, nodes), _ = pickle.load(f) 52 | 53 | object_nodes = [] 54 | for node in nodes: 55 | modelId = node["modelId"] 56 | category = self.category_map.get_final_category(modelId) 57 | if not category in ["door", "window"]: 58 | object_nodes.append(node) 59 | 60 | for i in range(len(object_nodes)): 61 | for j in range(i+1, len(object_nodes)): 62 | a = self.model_to_index[object_nodes[i]["modelId"]] 63 | b = self.model_to_index[object_nodes[j]["modelId"]] 64 | self.count[a][b] += 1 65 | self.count[b][a] += 1 66 | print(index) 67 | 68 | self.N = N 69 | 70 | def save(self, dest=None): 71 | if dest == None: 72 | dest = f"{self.data_dir}/model_prior.pkl" 73 | with open(dest, "wb") as f: 74 | pickle.dump(self.__dict__, f, pickle.HIGHEST_PROTOCOL) 75 | 76 | def load(self, data_dir): 77 | source = f"{data_dir}/model_prior.pkl" 78 | with open(source, "rb") as f: 79 | self.__dict__ = pickle.load(f) 80 | 81 | 82 | def sample(self, category, models): 83 | N = self.N 84 | indices = [i for i in range(N) if self.model_index_to_cat[i]==category] 85 | p = [self.model_freq[indices[i]] for i in range(len(indices))] 86 | p = np.asarray(p) 87 | for model in models: 88 | i = self.model_to_index[model] 89 | p1 = [self.count[indices[j]][i] for j in range(len(indices))] 90 | p1 = np.asarray(p1) 91 | p1 = p1/p1.sum() 92 | p = p * p1 93 | 94 | p = p/sum(p) 95 | numbers = np.asarray([i for i in range(len(indices))]) 96 | return self.models[indices[np.random.choice(numbers, p=p)]] 97 | 98 | def get_models(self, category, important, others): 99 | N = self.N 100 | indices = [i for i in range(N) if self.model_index_to_cat[i]==category] 101 | to_remove = [] 102 | 103 | freq = [self.model_freq[indices[i]] for i in range(len(indices))] 104 | total_freq = sum(freq) 105 | for j in range(len(indices)): 106 | if freq[j] / total_freq < 0.01: 107 | if not indices[j] in to_remove: 108 | to_remove.append(indices[j]) 109 | 110 | for model in important: 111 | i = self.model_to_index[model] 112 | freq = [self.count[indices[j]][i] for j in range(len(indices))] 113 | total_freq = sum(freq) 114 | if total_freq > 0: 115 | for j in range(len(indices)): 116 | if freq[j] / total_freq < 0.1: 117 | if not indices[j] in to_remove: 118 | to_remove.append(indices[j]) 119 | 120 | for model in others: 121 | i = self.model_to_index[model] 122 | freq = [self.count[indices[j]][i] for j in range(len(indices))] 123 | total_freq = sum(freq) 124 | if total_freq > 0: 125 | for j in range(len(indices)): 126 | if freq[j] / total_freq < 0.05: 127 | if not indices[j] in to_remove: 128 | to_remove.append(indices[j]) 129 | 130 | for item in to_remove: 131 | if len(indices) > 1: 132 | indices.remove(item) 133 | 134 | return [self.models[index] for index in indices] 135 | 136 | 137 | if __name__ == "__main__": 138 | a = ModelPrior() 139 | # a.learn("toilet_6x6") 140 | a.learn("bedroom_6x6") 141 | # a.learn("living_6x6_aug") 142 | #a.learn("office_6x6_aug") 143 | a.save() 144 | #a.load() 145 | #print(a.get_models(1, ["415"])) 146 | -------------------------------------------------------------------------------- /scene-synth/models/__init__.py: -------------------------------------------------------------------------------- 1 | from models.resnet import * 2 | 3 | -------------------------------------------------------------------------------- /scene-synth/models/resnet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import math 3 | import torch.utils.model_zoo as model_zoo 4 | 5 | 6 | __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 7 | 'resnet152'] 8 | 9 | 10 | model_urls = { 11 | 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 12 | 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 13 | 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 14 | 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 15 | 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 16 | } 17 | 18 | 19 | def conv3x3(in_planes, out_planes, stride=1): 20 | "3x3 convolution with padding" 21 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 22 | padding=1, bias=False) 23 | 24 | 25 | class BasicBlock(nn.Module): 26 | expansion = 1 27 | 28 | def __init__(self, inplanes, planes, stride=1, downsample=None): 29 | super(BasicBlock, self).__init__() 30 | self.conv1 = conv3x3(inplanes, planes, stride) 31 | self.bn1 = nn.BatchNorm2d(planes) 32 | self.relu = nn.LeakyReLU(inplace=True) 33 | self.conv2 = conv3x3(planes, planes) 34 | self.bn2 = nn.BatchNorm2d(planes) 35 | self.downsample = downsample 36 | self.stride = stride 37 | 38 | def forward(self, x): 39 | residual = x 40 | 41 | out = self.conv1(x) 42 | out = self.bn1(out) 43 | out = self.relu(out) 44 | 45 | out = self.conv2(out) 46 | out = self.bn2(out) 47 | 48 | if self.downsample is not None: 49 | residual = self.downsample(x) 50 | 51 | out += residual 52 | out = self.relu(out) 53 | 54 | return out 55 | 56 | 57 | class Bottleneck(nn.Module): 58 | expansion = 4 59 | 60 | def __init__(self, inplanes, planes, stride=1, downsample=None): 61 | super(Bottleneck, self).__init__() 62 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 63 | self.bn1 = nn.BatchNorm2d(planes) 64 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, 65 | padding=1, bias=False) 66 | self.bn2 = nn.BatchNorm2d(planes) 67 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 68 | self.bn3 = nn.BatchNorm2d(planes * 4) 69 | self.relu = nn.LeakyReLU(inplace=True) 70 | self.downsample = downsample 71 | self.stride = stride 72 | 73 | def forward(self, x): 74 | residual = x 75 | 76 | out = self.conv1(x) 77 | out = self.bn1(out) 78 | out = self.relu(out) 79 | 80 | out = self.conv2(out) 81 | out = self.bn2(out) 82 | out = self.relu(out) 83 | 84 | out = self.conv3(out) 85 | out = self.bn3(out) 86 | 87 | if self.downsample is not None: 88 | residual = self.downsample(x) 89 | 90 | out += residual 91 | out = self.relu(out) 92 | 93 | return out 94 | 95 | 96 | class ResNet(nn.Module): 97 | 98 | def __init__(self, block, layers, num_classes=2, num_input_channels=17, use_fc=True): 99 | self.inplanes = 64 100 | super(ResNet, self).__init__() 101 | self.conv1 = nn.Conv2d(num_input_channels, 64, kernel_size=7, stride=2, padding=3, 102 | bias=False) 103 | self.bn1 = nn.BatchNorm2d(64) 104 | self.relu = nn.LeakyReLU(inplace=True) 105 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 106 | self.layer1 = self._make_layer(block, 64, layers[0]) 107 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 108 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 109 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 110 | 111 | self.avgpool = nn.AdaptiveAvgPool2d(1) 112 | 113 | # self.layer5 = self._make_layer(block, 1024, layers[3], stride=2) 114 | # self.layer6 = self._make_layer(block, 1024, layers[3], stride=2) 115 | # self.layer7 = self._make_layer(block, 1024, layers[3], stride=2) 116 | 117 | self.use_fc = use_fc 118 | if use_fc: 119 | self.fc = nn.Linear(512*block.expansion, num_classes) 120 | # self.fc = nn.Linear(1024, num_classes) 121 | 122 | for m in self.modules(): 123 | if isinstance(m, nn.Conv2d): 124 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 125 | m.weight.data.normal_(0, math.sqrt(2. / n)) 126 | elif isinstance(m, nn.BatchNorm2d): 127 | m.weight.data.fill_(1) 128 | m.bias.data.zero_() 129 | 130 | def _make_layer(self, block, planes, blocks, stride=1): 131 | downsample = None 132 | if stride != 1 or self.inplanes != planes * block.expansion: 133 | downsample = nn.Sequential( 134 | nn.Conv2d(self.inplanes, planes * block.expansion, 135 | kernel_size=1, stride=stride, bias=False), 136 | nn.BatchNorm2d(planes * block.expansion), 137 | ) 138 | 139 | layers = [] 140 | layers.append(block(self.inplanes, planes, stride, downsample)) 141 | self.inplanes = planes * block.expansion 142 | for i in range(1, blocks): 143 | layers.append(block(self.inplanes, planes)) 144 | 145 | return nn.Sequential(*layers) 146 | 147 | def forward(self, x): 148 | x = self.conv1(x) 149 | x = self.bn1(x) 150 | x = self.relu(x) 151 | x = self.maxpool(x) 152 | # print(f'Size after maxpool: {x.size()}') 153 | 154 | x = self.layer1(x) 155 | x = self.layer2(x) 156 | x = self.layer3(x) 157 | x = self.layer4(x) 158 | 159 | # x = self.layer5(x) 160 | # x = self.layer6(x) 161 | 162 | # print(f'Size before avgpool: {x.size()}') 163 | x = self.avgpool(x) 164 | 165 | x = x.view(x.size(0), -1) 166 | if self.use_fc: 167 | x = self.fc(x) 168 | 169 | return x 170 | 171 | 172 | def resnet18(pretrained=False, **kwargs): 173 | """Constructs a ResNet-18 model. 174 | Args: 175 | pretrained (bool): If True, returns a model pre-trained on ImageNet 176 | """ 177 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 178 | if pretrained: 179 | model.load_state_dict(model_zoo.load_url(model_urls['resnet18'])) 180 | return model 181 | 182 | 183 | def resnet34(pretrained=False, **kwargs): 184 | """Constructs a ResNet-34 model. 185 | Args: 186 | pretrained (bool): If True, returns a model pre-trained on ImageNet 187 | """ 188 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 189 | if pretrained: 190 | model.load_state_dict(model_zoo.load_url(model_urls['resnet34'])) 191 | return model 192 | 193 | 194 | def resnet50(pretrained=False, **kwargs): 195 | """Constructs a ResNet-50 model. 196 | Args: 197 | pretrained (bool): If True, returns a model pre-trained on ImageNet 198 | """ 199 | model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 200 | if pretrained: 201 | model.load_state_dict(model_zoo.load_url(model_urls['resnet50'])) 202 | return model 203 | 204 | 205 | def resnet101(pretrained=False, **kwargs): 206 | """Constructs a ResNet-101 model. 207 | Args: 208 | pretrained (bool): If True, returns a model pre-trained on ImageNet 209 | """ 210 | model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 211 | if pretrained: 212 | model.load_state_dict(model_zoo.load_url(model_urls['resnet101'])) 213 | return model 214 | 215 | 216 | def resnet152(pretrained=False, **kwargs): 217 | """Constructs a ResNet-152 model. 218 | Args: 219 | pretrained (bool): If True, returns a model pre-trained on ImageNet 220 | """ 221 | model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) 222 | if pretrained: 223 | model.load_state_dict(model_zoo.load_url(model_urls['resnet152'])) 224 | return model 225 | -------------------------------------------------------------------------------- /scene-synth/models/utils.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | 7 | ## ------------------------------------------------------------------------------------------------ 8 | # Make reshape into a layer, so we can put it in nn.Sequential 9 | class Reshape(nn.Module): 10 | def __init__(self, *args): 11 | super(Reshape, self).__init__() 12 | self.shape = args 13 | 14 | def forward(self, x): 15 | return x.view(self.shape) 16 | 17 | # Render a signed distance field image for a given oriented box 18 | # img_size: a 2-tuple of image dimensions 19 | # box_dims: a 2D tensor of x,y box lengths 20 | # loc: a 2D tensor location 21 | # orient: a 2D tensor (sin, cos) orientation 22 | # NOTE: This operates on batches 23 | # NOTE: Since box_dims are expressed in percentage of image size, everything is set 24 | # up to work in that coordinate system. In particular, this means that locations are 25 | # rescaled from [-1, 1] to [0, 1] 26 | def render_obb_sdf(img_size, box_dims, loc, orient): 27 | batch_size = loc.shape[0] 28 | 29 | # I don't know why, but the orient dims need to be reversed... 30 | orient = torch.stack([orient[:, 1], orient[:, 0]], dim=-1) 31 | 32 | # Rescale location 33 | loc = (loc + 1) * 0.5 34 | 35 | # Rescale the box dims to be half-widths 36 | r = box_dims * 0.5 37 | 38 | # Create coordinate grid 39 | linear_points_x = torch.linspace(0, 1, img_size[0]) 40 | linear_points_y = torch.linspace(0, 1, img_size[1]) 41 | coords = torch.Tensor(batch_size, img_size[0], img_size[1], 3) 42 | coords[:, :, :, 0] = torch.ger(linear_points_x, torch.ones(img_size[0])) 43 | coords[:, :, :, 1] = torch.ger(torch.ones(img_size[1]), linear_points_y) 44 | coords[:, :, :, 2] = 1 # Homogeneous coords 45 | # (Non-standard ordering of dimensions (x, y, b, 3) in order to support broadcasting) 46 | coords = coords.permute(1, 2, 0, 3) 47 | if loc.is_cuda: 48 | coords = coords.cuda() 49 | 50 | # Attempting to do this faster by inverse-transforming the coordinate points and 51 | # then evaluating the SDF for a standard axis-aligned box 52 | inv_rot_matrices = torch.zeros(batch_size, 3, 3) 53 | if loc.is_cuda: 54 | inv_rot_matrices = inv_rot_matrices.cuda() 55 | cos = orient[:, 0] 56 | sin = orient[:, 1] 57 | inv_rot_matrices[:, 0, 0] = cos 58 | inv_rot_matrices[:, 1, 1] = cos 59 | inv_rot_matrices[:, 0, 1] = sin 60 | inv_rot_matrices[:, 1, 0] = -sin 61 | inv_rot_matrices[:, 2, 2] = 1 62 | inv_trans_matrices = [torch.eye(3, 3) for i in range(0, batch_size)] 63 | inv_trans_matrices = torch.stack(inv_trans_matrices) 64 | if loc.is_cuda: 65 | inv_trans_matrices = inv_trans_matrices.cuda() 66 | inv_trans_matrices[:, 0, 2] = -loc[:, 0] 67 | inv_trans_matrices[:, 1, 2] = -loc[:, 1] 68 | inv_matrices = torch.matmul(inv_rot_matrices, inv_trans_matrices) 69 | # Discard the last row (to get 2x3 matrices) 70 | inv_matrices = inv_matrices[:, 0:2, :] 71 | # Multiply the coords by these matrices 72 | # Batch matrix multiply of (b, 2, 3) by (x, y, b, 3, 1) -> (x, y, b, 2, 1) 73 | coords = coords.view(img_size[0], img_size[1], batch_size, 3, 1) 74 | coords = torch.matmul(inv_matrices, coords) 75 | coords = coords.view(img_size[0], img_size[1], batch_size, 2) 76 | 77 | # Now evaluate the SDF of an axis-aligned box centered at the origin 78 | xdist = coords[:, :, :, 0].abs() - r[:, 0] 79 | ydist = coords[:, :, :, 1].abs() - r[:, 1] 80 | dist = torch.max(xdist, ydist) 81 | dist = dist.permute(2, 0, 1) 82 | return dist.view(batch_size, 1, img_size[0], img_size[1]) 83 | 84 | def log(x): 85 | return torch.log(torch.max(x, 1e-5)) 86 | 87 | def set_requires_grad(net, requires_grad=False): 88 | for param in net.parameters(): 89 | param.requires_grad = requires_grad 90 | 91 | def make_walls_img(dims): 92 | wall_img = torch.zeros(1, img_size, img_size) 93 | w = int((dims[0] / 2) * img_size) 94 | h = int((dims[1] / 2) * img_size) 95 | mid = img_size // 2 96 | wall_img[0, (mid+h):(mid+h+wall_thickness), 0:(mid+w+wall_thickness)] = 1 # Horiz wall 97 | wall_img[0, 0:(mid+h+wall_thickness), (mid+w):(mid+w+wall_thickness)] = 1 # Vert wall 98 | return wall_img 99 | def make_walls_img_batch(dims_batch): 100 | batch_size = dims_batch.shape[0] 101 | return torch.stack([make_walls_img(dims_batch[i]) for i in range(batch_size)], dim=0) 102 | 103 | def unitnormal_normal_kld(mu, logvar, size_average=True): 104 | # Always reduce along the data dimensionality axis 105 | # output = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp(), dim=-1) 106 | output = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=-1) 107 | # Average along the batch axis, if requested 108 | if size_average: 109 | output = torch.mean(output) 110 | return output 111 | 112 | def inverse_xform_img(img, loc, orient, output_size): 113 | batch_size = img.shape[0] 114 | matrices = torch.zeros(batch_size, 2, 3).cuda() 115 | cos = orient[:, 0] 116 | sin = orient[:, 1] 117 | matrices[:, 0, 0] = cos 118 | matrices[:, 1, 1] = cos 119 | matrices[:, 0, 1] = -sin 120 | matrices[:, 1, 0] = sin 121 | matrices[:, 0, 2] = loc[:, 1] 122 | matrices[:, 1, 2] = loc[:, 0] 123 | out_size = torch.Size((batch_size, img.shape[1], output_size, output_size)) 124 | grid = F.affine_grid(matrices, out_size) 125 | return F.grid_sample(img, grid) 126 | 127 | def forward_xform_img(img, loc, orient, output_size): 128 | # First, build the inverse rotation matrices 129 | batch_size = img.shape[0] 130 | inv_rot_matrices = torch.zeros(batch_size, 3, 3).cuda() 131 | cos = orient[:, 0] 132 | sin = orient[:, 1] 133 | inv_rot_matrices[:, 0, 0] = cos 134 | inv_rot_matrices[:, 1, 1] = cos 135 | inv_rot_matrices[:, 0, 1] = sin 136 | inv_rot_matrices[:, 1, 0] = -sin 137 | inv_rot_matrices[:, 2, 2] = 1 138 | # Then, build the inverse translation matrices 139 | # (Apparently, x and y need to be swapped. I don't know why...) 140 | inv_trans_matrices = [torch.eye(3, 3) for i in range(0, batch_size)] 141 | inv_trans_matrices = torch.stack(inv_trans_matrices).cuda() 142 | inv_trans_matrices[:, 0, 2] = -loc[:, 1] 143 | inv_trans_matrices[:, 1, 2] = -loc[:, 0] 144 | # Multiply them to get the full affine matrix 145 | inv_matrices = torch.matmul(inv_rot_matrices, inv_trans_matrices) 146 | # Discard the last row (affine_grid expects 2x3 matrices) 147 | inv_matrices = inv_matrices[:, 0:2, :] 148 | # Finalize 149 | out_size = torch.Size((batch_size, img.shape[1], output_size, output_size)) 150 | grid = F.affine_grid(inv_matrices, out_size) 151 | return F.grid_sample(img, grid) 152 | 153 | def default_loc_orient(batch_size): 154 | loc = torch.zeros(batch_size, 2).cuda() 155 | orient = torch.stack([torch.Tensor([math.cos(0), math.sin(0)]) for i in range(batch_size)], dim=0).cuda() 156 | return loc, orient 157 | 158 | class DownConvBlock(nn.Module): 159 | def __init__(self, inplanes, outplanes): 160 | super(DownConvBlock, self).__init__() 161 | self.conv = nn.Conv2d(inplanes, outplanes, stride=2, kernel_size=3, padding=1, bias=True) 162 | self.bn = nn.BatchNorm2d(outplanes) 163 | self.activation = nn.ReLU() 164 | def forward(self, x): 165 | # return self.activation(self.conv(x)) 166 | return self.activation(self.bn(self.conv(x))) 167 | 168 | def make_walls_img(dims): 169 | wall_img = torch.zeros(1, img_size, img_size) 170 | w = int((dims[0] / 2) * img_size) 171 | h = int((dims[1] / 2) * img_size) 172 | mid = img_size // 2 173 | wall_img[0, (mid+h):(mid+h+wall_thickness), 0:(mid+w+wall_thickness)] = 1 # Horiz wall 174 | wall_img[0, 0:(mid+h+wall_thickness), (mid+w):(mid+w+wall_thickness)] = 1 # Vert wall 175 | # wall_img[0, :, (mid+w):(mid+w+wall_thickness)] = 1 # Vert wall 176 | return wall_img 177 | 178 | def render_orientation_sdf(img_size, dims, loc, orient): 179 | batch_size = loc.shape[0] 180 | 181 | # I don't know why, but the orient dims need to be reversed... 182 | orient = torch.stack([orient[:, 1], orient[:, 0]], dim=-1) 183 | 184 | # Rescale location 185 | loc = (loc + 1) * 0.5 186 | 187 | # Rescale the box dims to be half-widths 188 | r = dims * 0.5 189 | 190 | # Create coordinate grid 191 | linear_points_x = torch.linspace(0, 1, img_size[0]) 192 | linear_points_y = torch.linspace(0, 1, img_size[1]) 193 | coords = torch.Tensor(batch_size, img_size[0], img_size[1], 3) 194 | coords[:, :, :, 0] = torch.ger(linear_points_x, torch.ones(img_size[0])) 195 | coords[:, :, :, 1] = torch.ger(torch.ones(img_size[1]), linear_points_y) 196 | coords[:, :, :, 2] = 1 # Homogeneous coords 197 | # (Non-standard ordering of dimensions (x, y, b, 3) in order to support broadcasting) 198 | coords = coords.permute(1, 2, 0, 3) 199 | if loc.is_cuda: 200 | coords = coords.cuda() 201 | 202 | # Attempting to do this faster by inverse-transforming the coordinate points and 203 | # then evaluating the SDF for a standard axis-aligned plane 204 | inv_rot_matrices = torch.zeros(batch_size, 3, 3) 205 | if loc.is_cuda: 206 | inv_rot_matrices = inv_rot_matrices.cuda() 207 | cos = orient[:, 0] 208 | sin = orient[:, 1] 209 | inv_rot_matrices[:, 0, 0] = cos 210 | inv_rot_matrices[:, 1, 1] = cos 211 | inv_rot_matrices[:, 0, 1] = sin 212 | inv_rot_matrices[:, 1, 0] = -sin 213 | inv_rot_matrices[:, 2, 2] = 1 214 | inv_trans_matrices = [torch.eye(3, 3) for i in range(0, batch_size)] 215 | inv_trans_matrices = torch.stack(inv_trans_matrices) 216 | if loc.is_cuda: 217 | inv_trans_matrices = inv_trans_matrices.cuda() 218 | inv_trans_matrices[:, 0, 2] = -loc[:, 0] 219 | inv_trans_matrices[:, 1, 2] = -loc[:, 1] 220 | inv_matrices = torch.matmul(inv_rot_matrices, inv_trans_matrices) 221 | # Discard the last row (to get 2x3 matrices) 222 | inv_matrices = inv_matrices[:, 0:2, :] 223 | # Multiply the coords by these matrices 224 | # Batch matrix multiply of (b, 2, 3) by (x, y, b, 3, 1) -> (x, y, b, 2, 1) 225 | coords = coords.view(img_size[0], img_size[1], batch_size, 3, 1) 226 | coords = torch.matmul(inv_matrices, coords) 227 | coords = coords.view(img_size[0], img_size[1], batch_size, 2) 228 | 229 | # Now evaluate the SDF of an axis-aligned plane 230 | # (Signed dist to plane is just the x coordinate) 231 | dist = coords[:, :, :, 0] 232 | dist = dist.permute(2, 0, 1) 233 | return dist.view(batch_size, 1, img_size[0], img_size[1]) 234 | 235 | # xdist = (coords[:, :, :, 0].abs() - r[:, 0]).max(torch.Tensor([0]).cuda()) * coords[:, :, :, 0].sign() 236 | # ydist = (coords[:, :, :, 1].abs() - r[:, 1]).max(torch.Tensor([0]).cuda()) * coords[:, :, :, 1].sign() 237 | # dist = torch.stack([xdist, ydist], dim=-1) 238 | # dist = dist.permute(2, 3, 0, 1) 239 | # return dist 240 | 241 | # Like stack the regular bbox SDF with the orientation plane SDF 242 | def render_oriented_sdf(img_sizes, dims, loc, orient): 243 | sdf = render_obb_sdf(img_sizes, dims, loc, orient) 244 | osdf = render_orientation_sdf(img_sizes, dims, loc, orient) 245 | # sdf = render_obb_sdf_slow(img_sizes, dims, loc, orient) 246 | # osdf = render_orientation_sdf_slow(img_sizes, dims, loc, orient) 247 | return torch.cat([sdf, osdf], dim=1) 248 | 249 | CARDINAL_ANGLES = torch.Tensor([0, math.pi/2, math.pi, 3*math.pi/2]) 250 | CARDINAL_DIRECTIONS = torch.stack([CARDINAL_ANGLES.cos(), CARDINAL_ANGLES.sin()], dim=1) 251 | 252 | # Snap an orientation to its nearest cardinal direction 253 | def snap_orient(orient): 254 | sims = [F.cosine_similarity(orient, cdir.unsqueeze(0).cuda()) for cdir in CARDINAL_DIRECTIONS] 255 | sims = torch.stack(sims, dim=1) 256 | maxvals, indices = sims.max(dim=1) 257 | return CARDINAL_DIRECTIONS[indices].cuda() 258 | 259 | def should_snap(orient): 260 | snap_sims = [F.cosine_similarity(orient, cdir.unsqueeze(0)) for cdir in CARDINAL_DIRECTIONS] 261 | snap_sims = torch.stack(snap_sims, dim=1) 262 | snap_sim, _ = snap_sims.max(dim=1) 263 | snap = (snap_sim > (1 - 1e-4)).float() 264 | return snap 265 | 266 | def index_to_onehot(indices, numclasses): 267 | """ 268 | Turn index into one-hot vector 269 | indices = torch.LongTensor of indices (batched) 270 | """ 271 | b_size = indices.size()[0] 272 | one_hot = torch.zeros(b_size, numclasses) 273 | for i in range(0, b_size): 274 | cat_index = indices[i][0] 275 | one_hot[i][cat_index] = 1 276 | return one_hot 277 | 278 | def index_to_onehot_fast(indices, numclasses): 279 | """A better version of index to one-hot that uses scatter""" 280 | indices = indices.unsqueeze(-1) 281 | onehot = torch.zeros(indices.shape[0], numclasses) 282 | if indices.is_cuda: 283 | onehot = onehot.cuda() 284 | onehot.scatter_(1, indices, 1) 285 | return onehot 286 | 287 | def nearest_downsample(img, factor): 288 | """ 289 | Do nearest-neighbor downsampling on an image tensor with 290 | arbitrary channels 291 | Img must be a 4D B x C x H x W tensor 292 | Factor must be an integer (is converted to one) 293 | """ 294 | assert (factor > 0 and factor == int(factor)), 'Downsample factor must be positive integer' 295 | # We do this by convolving with a strided convolutional kernel, whose kernel 296 | # is size one and is the identity matrix between input and output channels 297 | nchannels = img.shape[1] 298 | kernel = torch.eye(nchannels).view(nchannels, nchannels, 1, 1) 299 | if img.is_cuda: 300 | kernel = kernel.cuda() 301 | return F.conv2d(img, weight=kernel, stride=factor) 302 | 303 | def softmax2d(img): 304 | """Softmax across the pixels of an image batch""" 305 | size = img.shape[2] 306 | img = img.view(img.shape[0], -1) 307 | img = F.softmax(img, dim=-1) 308 | img = img.view(-1, size, size) 309 | return img 310 | 311 | def mask_to_outline(img): 312 | num_channels = img.shape[1] 313 | efx = torch.zeros(num_channels, num_channels, 3, 3) 314 | efy = torch.zeros(num_channels, num_channels, 3, 3) 315 | for i in range(num_channels): 316 | efx[i, i, :, :] = edge_filters_x 317 | efy[i, i, :, :] = edge_filters_y 318 | if img.is_cuda: 319 | efx = efx.cuda() 320 | efy = efy.cuda() 321 | edges_x = F.conv2d(img, efx, padding=1) 322 | edges_y = F.conv2d(img, efy, padding=1) 323 | edges = edges_x + edges_y 324 | return (edges != 0).float() 325 | 326 | def mask_to_sdf(img): 327 | outline = mask_to_outline(img) 328 | outline_neg = 1 - outline 329 | dists = torch.tensor(distance_transform_edt(outline_neg)).float() 330 | diag_len = math.sqrt(2*img.shape[2]*img.shape[2]) 331 | dists = dists / diag_len # normalize 332 | return torch.where(img > 0, -dists, dists) 333 | 334 | # Converting image batch tensors that represent binary masks into outlines 335 | # and signed distance functions 336 | edge_filters_x = torch.tensor([ 337 | [0, 0, 0], 338 | [-1, 0, 1], 339 | [0, 0, 0], 340 | ]).float() 341 | edge_filters_y = torch.tensor([ 342 | [0, -1, 0], 343 | [0, 0, 0], 344 | [0, 1, 0], 345 | ]).float() 346 | 347 | -------------------------------------------------------------------------------- /scene-synth/obj_dims.py: -------------------------------------------------------------------------------- 1 | from data import * 2 | from utils import * 3 | import pickle 4 | 5 | # data_dir = "toilet_6x6" 6 | # data_dir = "bedroom_6x6" 7 | # data_dir = "living_6x6_aug" 8 | data_dir = "office_6x6_aug" 9 | data_root_dir = utils.get_data_root_dir() 10 | img_size = 256 11 | room_dim = 6.05 12 | 13 | with open(f"{data_root_dir}/{data_dir}/model_frequency", "r") as f: 14 | models = f.readlines() 15 | 16 | models = [l[:-1].split(" ") for l in models] 17 | models = [l[0] for l in models] 18 | #print(models) 19 | 20 | model_dims = dict() 21 | 22 | for model in models: 23 | o = Obj(model) 24 | dims = [(o.bbox_max[0] - o.bbox_min[0])/room_dim, \ 25 | (o.bbox_max[2] - o.bbox_min[2])/room_dim] 26 | 27 | model_dims[model] = dims 28 | 29 | with open(f"{data_root_dir}/{data_dir}/model_dims.pkl", 'wb') as f: 30 | pickle.dump(model_dims, f, pickle.HIGHEST_PROTOCOL) 31 | 32 | -------------------------------------------------------------------------------- /scene-synth/priors/arrangement.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import argparse 4 | import math 5 | import os 6 | import utils 7 | 8 | import numpy as np 9 | 10 | 11 | from data import House, DatasetToJSON 12 | from math_utils import Transform 13 | from priors.observations import ObjectCollection 14 | from priors.pairwise import PairwiseArrangementPrior # needed for pickle laod 15 | from pyquaternion import Quaternion 16 | from random import random 17 | 18 | 19 | class ArrangementPriors: 20 | """ 21 | A collection of object arrangement priors for estimating probability of an arrangement and sampling arrangements 22 | """ 23 | def __init__(self, priors_file, w_dist=1.0, w_clearance=1.0, w_same_category=1.0, w_closest=1.0, w_orientation=1.0): 24 | self._priors = utils.pickle_load_compressed(priors_file) 25 | self._num_priors = 0. 26 | self._num_observations = 0. 27 | self._w_dist = w_dist 28 | self._w_clearance = w_clearance 29 | self._w_same_category = w_same_category 30 | self._w_closest = w_closest 31 | self._w_orientation = w_orientation 32 | for p_key, p in self._priors.items(): 33 | self._priors[p_key] = p 34 | self._num_priors += 1 35 | self._num_observations += p.num_observations 36 | p.trim() 37 | 38 | @property 39 | def num_observations(self): 40 | """ 41 | Return total number of observations in this ArrangementPriors set 42 | """ 43 | return self._num_observations 44 | 45 | @property 46 | def pairwise_priors(self): 47 | """ 48 | Return all pairwise priors 49 | """ 50 | return self._priors 51 | 52 | def get(self, category_key): 53 | """ 54 | Return pairwise prior for given category_key 55 | """ 56 | return self._priors[category_key] 57 | 58 | def pairwise_occurrence_prob(self, category_key): 59 | """ 60 | Return probability of pairwise observation with given category_key 61 | """ 62 | norm = self._num_observations + self._num_priors # add one smoothing 63 | if category_key in self._priors: 64 | return (self._priors[category_key].num_observations + 1) / norm 65 | else: 66 | return 1. / norm 67 | 68 | def log_prob(self, pairwise_observations_by_key): 69 | """ 70 | Return log probability of given set of pairwise relative observations under arrangement priors 71 | """ 72 | sum_lp = 0 73 | 74 | for key in pairwise_observations_by_key: 75 | # get observations and lp for this pairwise prior 76 | observations = pairwise_observations_by_key[key] 77 | 78 | # parameterize observations 79 | records = map(lambda o: o.parameterize(scheme='offsets_angles'), observations) 80 | x = np.stack(records, axis=0) 81 | # print(observations) 82 | # print(x) 83 | 84 | if key not in self._priors: 85 | lp = math.log(0.5) 86 | else: # evaluate each observation against relevant prior 87 | prior = self._priors[key] 88 | if key.obj_category == 'room': # distance and angle to closest point on wall 89 | lp = self._w_clearance * prior.log_prob_closest(x) 90 | elif key.obj_category == key.ref_obj_category: # same category instance 91 | lp = self._w_same_category * self._w_closest * prior.log_prob_closest(x) 92 | else: # all other categories 93 | lp = self._w_closest * prior.log_prob_closest(x) 94 | lp += self._w_orientation * prior.log_prob_orientation(x) 95 | 96 | if self._w_dist > 0: 97 | dist = np.sum(np.square(x[:, 2:4]), axis=1, keepdims=True) 98 | lp = lp - (self._w_dist * dist/2) # log(exp(-x^2/2)) 99 | 100 | pairwise_occ_p = self.pairwise_occurrence_prob(key) 101 | lp = lp * pairwise_occ_p 102 | 103 | # sum up lps 104 | sum_lp += np.sum(lp, axis=0) 105 | 106 | if hasattr(sum_lp, 'shape'): 107 | sum_lp = np.sum(sum_lp) 108 | return sum_lp 109 | 110 | 111 | class ArrangementGreedySampler: 112 | """ 113 | Iterative optimization of object arrangements using greedy sampling of ArrangementPriors 114 | """ 115 | def __init__(self, arrangement_priors, num_angle_divisions=8, num_pairwise_priors=-1, sim_mode='direct'): 116 | self._objects = ObjectCollection(sim_mode=sim_mode) 117 | self.room_id = None 118 | self._priors = arrangement_priors 119 | self._num_angle_divisions = num_angle_divisions 120 | self._num_pairwise_priors = num_pairwise_priors 121 | 122 | @property 123 | def objects(self): 124 | return self._objects 125 | 126 | def init(self, house, only_architecture=True, room_id=None): 127 | if not room_id: 128 | room_id = house.rooms[0].id 129 | self._objects.init_from_room(house, room_id, only_architecture=only_architecture) 130 | self.room_id = room_id 131 | 132 | def log_prob(self, filter_ref_obj=None, ignore_categories=list()): 133 | observations = self._objects.get_relative_observations(self.room_id, filter_ref_obj=filter_ref_obj, 134 | ignore_categories=ignore_categories) 135 | observations_by_key = {} 136 | 137 | # top_k_prior_categories = None 138 | # if filter_ref_obj and num_pairwise_priors specified, filter observations to only those in top k priors 139 | # if filter_ref_obj and self._num_pairwise_priors > 0: 140 | # category = self._objects.category(filter_ref_obj.modelId, scheme='final') 141 | # priors = list(filter(lambda p: p.ref_obj_category == category, self._priors.pairwise_priors)) 142 | # k = min(self._num_pairwise_priors, len(priors)) 143 | # priors = list(sorted(priors, key=lambda p: self._priors.pairwise_occurrence_log_prob(p)))[-k:] 144 | # top_k_prior_categories = set(map(lambda p: p.obj_category, priors)) 145 | 146 | for o in observations.values(): 147 | # only pairwise observations in which filter_ref_obj is the reference 148 | if filter_ref_obj and o.ref_id != filter_ref_obj.id: 149 | continue 150 | key = self._objects.get_observation_key(o) 151 | os_key = observations_by_key.get(key, []) 152 | os_key.append(o) 153 | observations_by_key[key] = os_key 154 | return self._priors.log_prob(observations_by_key) 155 | 156 | def get_candidate_transform(self, node, max_iterations=100): 157 | num_checks = 0 158 | zmin = self._objects.room.zmin 159 | while True: 160 | num_checks += 1 161 | p = self._objects.room.obb.sample() 162 | ray_from = [p[0], zmin - .5, p[2]] 163 | ray_to = [p[0], zmin + .5, p[2]] 164 | intersection = ags._objects.simulator.ray_test(ray_from, ray_to) 165 | if intersection.id == self.room_id + 'f' or num_checks > max_iterations: 166 | break 167 | xform = Transform() 168 | xform.set_translation([p[0], zmin + .1, p[2]]) 169 | angle = random() * 2 * math.pi 170 | angular_resolution = 2 * math.pi / self._num_angle_divisions 171 | angle = round(angle / angular_resolution) * angular_resolution 172 | xform.set_rotation(radians=angle) 173 | return xform 174 | 175 | def sample_placement(self, node, n_samples, houses_log=None, max_attempts_per_sample=10, ignore_categories=list(), 176 | collision_threshold=0): 177 | """ 178 | Sample placement for given node 179 | """ 180 | self._objects.add_object(node) 181 | max_lp = -np.inf 182 | max_xform = None 183 | max_house = None 184 | num_noncolliding_samples = 0 185 | for i in range(max_attempts_per_sample*n_samples): 186 | xform = self.get_candidate_transform(node) 187 | self._objects.update(node, xform=xform, update_sim=True) 188 | collisions = self._objects.get_collisions(obj_id_a=node.id) 189 | # print(f'i={i}, samples_so_far={num_noncolliding_samples}, n_samples={n_samples},' 190 | # f'max_attempts_per_sample={max_attempts_per_sample}') 191 | if collision_threshold > 0: 192 | if min(collisions.values(), key=lambda c: c.distance).distance < -collision_threshold: 193 | continue 194 | elif len(collisions) > 0: 195 | continue 196 | lp = self.log_prob(filter_ref_obj=node, ignore_categories=ignore_categories) 197 | print(f'lp={lp}') 198 | if lp > max_lp: 199 | max_xform = xform 200 | max_lp = lp 201 | if houses_log is not None: 202 | max_house = self._objects.as_house() 203 | num_noncolliding_samples += 1 204 | if num_noncolliding_samples == n_samples: 205 | break 206 | if houses_log is not None: 207 | houses_log.append(max_house) 208 | self._objects.update(node, xform=max_xform, update_sim=True) 209 | 210 | def placeable_objects_sorted_by_size(self, house): 211 | objects = [] 212 | fixed_objects = [] 213 | for n in house.levels[0].nodes: 214 | if n.type != 'Object': 215 | continue 216 | category = self._objects.category(n.modelId, scheme='final') 217 | if category == 'door' or category == 'window': 218 | fixed_objects.append(n) 219 | continue 220 | objects.append(n) 221 | 222 | for o in objects: 223 | # to_delete = [] 224 | # for k in o.__dict__: 225 | # if k not in ['id', 'modelId', 'transform', 'type', 'valid', 'bbox']: 226 | # to_delete.append(k) 227 | # for k in to_delete: 228 | # delattr(o, k) 229 | dims = self._objects._object_data.get_aligned_dims(o.modelId) 230 | o.volume = dims[0] * dims[1] * dims[2] 231 | objects = list(sorted(objects, key=lambda x: x.volume, reverse=True)) 232 | return objects, fixed_objects 233 | 234 | 235 | if __name__ == '__main__': 236 | module_dir = os.path.dirname(os.path.abspath(__file__)) 237 | parser = argparse.ArgumentParser(description='Arrangement model') 238 | parser.add_argument('--task', type=str, default='arrange', help='task [arrange]') 239 | parser.add_argument('--input', type=str, required=True, help='input scene to arrange') 240 | parser.add_argument('--output_dir', type=str, required=True, help='output directory to save arranged house states') 241 | parser.add_argument('--priors_dir', type=str, required=True, help='base directory to load priors') 242 | parser.add_argument('--only_final_state', action='store_true', help='save only final') 243 | parser.add_argument('--sim_mode', dest='sim_mode', default='direct', help='sim server [gui|direct|shared_memory]') 244 | parser.add_argument('--restart_sim_every_round', action='store_true', help='restart sim after every object placed') 245 | parser.add_argument('--max_objects', type=int, default=np.inf, help='maximum number of objects to place') 246 | parser.add_argument('--num_samples_per_object', type=int, default=100, help='placement samples per object') 247 | parser.add_argument('--ignore_doors_windows', action='store_true', help='ignore door and window priors') 248 | parser.add_argument('--collision_threshold', type=float, default=0, help='how much collision penetration to allow') 249 | args = parser.parse_args() 250 | 251 | if args.task == 'arrange': 252 | # initialize arrangement priors and sampler 253 | filename = os.path.join(args.priors_dir, f'priors.pkl.gz') 254 | ap = ArrangementPriors(priors_file=filename, 255 | w_dist=0.0, 256 | w_clearance=1.0, 257 | w_closest=1.0, 258 | w_orientation=1.0, 259 | w_same_category=1.0) 260 | ags = ArrangementGreedySampler(arrangement_priors=ap, sim_mode=args.sim_mode) 261 | 262 | # load house and set architecture 263 | original_house = House(file_dir=args.input, include_support_information=False) 264 | ags.init(original_house, only_architecture=True) 265 | 266 | # split into placeable objects and fixed portals (doors, windows) 267 | placeables, fixed_objects = ags.placeable_objects_sorted_by_size(original_house) 268 | 269 | # add fixed objects (not to be arranged) 270 | for n in fixed_objects: 271 | ags.objects.add_object(n) 272 | 273 | # place objects one-by-one 274 | houses_states = [] 275 | num_placeables = len(placeables) 276 | for (i, n) in enumerate(placeables): 277 | print(f'Placing {i+1}/{num_placeables}') 278 | if args.sim_mode == 'gui' and i < 2: # center view on first couple of placements if gui_mode 279 | ags.objects.simulator.set_gui_rendering(enabled=True) 280 | houses_log = None if args.only_final_state else houses_states 281 | ags.sample_placement(node=n, n_samples=args.num_samples_per_object, houses_log=houses_log, 282 | ignore_categories=['window', 'door'] if args.ignore_doors_windows else []) 283 | if args.restart_sim_every_round: 284 | ags.objects.simulator.connect() 285 | ags.objects.reinit_simulator() 286 | if i+1 >= args.max_objects: 287 | break 288 | 289 | # append final house state 290 | house_final = ags._objects.as_house() 291 | houses_states.append(house_final) 292 | 293 | # save out states 294 | input_id = os.path.splitext(os.path.basename(args.input))[0] 295 | output_dir = os.path.join(args.output_dir, input_id) 296 | print(f'Saving results to {output_dir}...') 297 | dj = DatasetToJSON(dest=output_dir) 298 | empty_is = [] 299 | for i, h in enumerate(houses_states): 300 | if h: 301 | h.trim() 302 | else: 303 | print(f'Warning: empty house state for step {i}, ignoring...') 304 | houses_states = list(filter(lambda h: h is not None, houses_states)) 305 | for (i, s) in enumerate(dj.step(houses_states)): 306 | print(f'Saved {i}') 307 | 308 | print('DONE') 309 | -------------------------------------------------------------------------------- /scene-synth/priors/pairwise.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import argparse 4 | import os 5 | import utils 6 | from math_utils import * 7 | from sklearn import mixture 8 | 9 | from priors.observations import RelativeObservationsDatabase, RelativeObservation, ObservationCategory 10 | 11 | 12 | class PairwiseArrangementPrior: 13 | """Encapsulates arrangement priors for a specific category key (room_type, obj_category, ref_obj_category)""" 14 | def __init__(self, category_key): 15 | self._category_key = category_key 16 | self._records = [] 17 | self._gmm_centroid = None # GMM for offset to centroid of reference object 18 | self._gmm_closest = None # GMM for offset to closest point on reference object 19 | self._hist_orientation = None # histogram distribution for orientation of object front in reference frame 20 | 21 | @property 22 | def category_key(self): 23 | return self._category_key 24 | 25 | @property 26 | def centroid(self): 27 | return self._gmm_centroid 28 | 29 | @property 30 | def closest(self): 31 | return self._gmm_closest 32 | 33 | @property 34 | def orientation(self): 35 | return self._hist_orientation 36 | # return self._vmf 37 | # return self._vmf_orientation 38 | 39 | def add_observations(self, observations, parameterization_scheme): 40 | for o in observations: 41 | self._records.append(o.parameterize(scheme=parameterization_scheme)) 42 | 43 | @property 44 | def num_observations(self): 45 | return len(self._records) 46 | 47 | def fit_model(self, gmm_opts=None): 48 | if gmm_opts is None: 49 | gmm_opts = {'n_components': 4, 'covariance_type': 'tied', 'n_init': 5, 'init_params': 'random'} 50 | X = np.stack(self._records, axis=0) 51 | try: 52 | self._gmm_centroid = mixture.BayesianGaussianMixture(**gmm_opts).fit(X[:, 0:2]) 53 | self._gmm_closest = mixture.BayesianGaussianMixture(**gmm_opts).fit(X[:, 2:4]) 54 | angles = X[:, 4][:, np.newaxis] 55 | # avoid dirac delta-like peaks that cause fit issues by adding noise 56 | # angles = np.clip(np.random.normal(angles, scale=math.pi/16), a_min=-np.pi, a_max=np.pi) 57 | low = -math.pi * 33. / 32. 58 | high = math.pi * 33. / 32. 59 | hist = np.histogram(angles, bins=np.linspace(low, high, num=17)) 60 | h = hist[0] + 1 # add one (i.e. laplace) smoothing on angle histogram 61 | h = h / np.sum(h) 62 | self._hist_orientation = scipy.stats.rv_histogram((h, hist[1])) 63 | # kappa, mu, scale = vonmises.fit(angles, fscale=1) 64 | # self._vmf = {'kappa': kappa, 'mu': mu, 'scale': scale} 65 | # cos_sin_angles = np.concatenate((np.cos(angles), np.sin(angles)), axis=1) 66 | # self._vmf_orientation = VonMisesFisherMixture(**vmf_opts).fit(cos_sin_angles) 67 | # print(self._vmf_orientation.cluster_centers_, self._vmf_orientation.concentrations_) 68 | print(f'GMM+VMF fit obtained for {self._category_key} with {len(self._records)} observations') 69 | except Exception: 70 | print(f'Error fitting priors for {self.category_key} with {len(self._records)} samples.') 71 | 72 | def plot_model(self): 73 | X = np.stack(self._records, axis=0) 74 | plot_gmm_fit(X[:, 0:2], self._gmm_centroid.predict(X[:, 0:2]), self._gmm_centroid.means_, 75 | self._gmm_centroid.covariances_, 0, str(self._category_key) + ':centroid') 76 | plot_gmm_fit(X[:, 2:4], self._gmm_closest.predict(X[:, 2:4]), self._gmm_closest.means_, 77 | self._gmm_closest.covariances_, 1, str(self._category_key) + ':closest') 78 | ang = X[:, 4].reshape(-1, 1) 79 | plot_hist_fit(ang, self._hist_orientation, 2, title=str(self._category_key) + ':orientation') 80 | # plot_vmf_fit(ang, mu=self._vmf['mu'], kappa=self._vmf['kappa'], scale=self._vmf['scale'], 81 | # index=2, title=str(self._category_key) + ':orientation') 82 | # samples = self.sample(n_samples=1000) 83 | # # print(samples) 84 | # # print(prior.log_prob(samples)) 85 | # c = samples[:, 0:2] 86 | # plot_gmm_fit(c, prior.centroid.predict(c), prior.centroid.means_, prior.centroid.covariances_, 2, 87 | # str(prior.category_key) + ':samples') 88 | plt.show() 89 | 90 | def log_prob(self, x): 91 | lp_centroid = self._gmm_centroid.score_samples(x[:, 0:2])[:, np.newaxis] 92 | lp_closest = self._gmm_closest.score_samples(x[:, 2:4])[:, np.newaxis] 93 | angles = x[:, 4][:, np.newaxis] 94 | lp_orientation = self._hist_orientation.logpdf(angles) 95 | # lp_vmf = vonmises.logpdf(angles, kappa=self._vmf['kappa'], loc=self._vmf['mu'], scale=self._vmf['scale']) 96 | # cos_sin_angles = np.concatenate((np.cos(angles), np.sin(angles)), axis=1) 97 | # lp_orientation = self._vmf_orientation.log_likelihood(cos_sin_angles)[0][:, np.newaxis] 98 | # print(lp_centroid) 99 | # print(lp_closest) 100 | # print(np.concatenate((angles, lp_orientation), axis=1)) 101 | # print(lp_vmf) 102 | return lp_centroid + lp_closest + lp_orientation 103 | 104 | def log_prob_offset(self, x): 105 | if not self._gmm_centroid: 106 | print(f'Warning: tried log_prob_offset with no prior for {self.category_key}') 107 | return math.log(0.5) 108 | return self._gmm_centroid.score_samples(x[:, 0:2])[:, np.newaxis] 109 | 110 | def log_prob_closest(self, x): 111 | if not self._gmm_closest: 112 | print(f'Warning: tried log_prob_closest with no prior for {self.category_key}') 113 | return math.log(0.5) 114 | return self._gmm_closest.score_samples(x[:, 2:4])[:, np.newaxis] 115 | 116 | def log_prob_orientation(self, x): 117 | if not self._hist_orientation: 118 | print(f'Warning: tried log_prob_orientation with no prior for {self.category_key}') 119 | return math.log(0.5) 120 | return self._hist_orientation.logpdf(x[:, 4][:, np.newaxis]) 121 | 122 | def sample(self, n_samples): 123 | # self._vmf_orientation.sample(n_samples=n_samples) 124 | centroid_samples = self._gmm_centroid.sample(n_samples=n_samples)[0] 125 | closest_samples = self._gmm_closest.sample(n_samples=n_samples)[0] 126 | orientation_samples = self._hist_orientation.rvs(size=(n_samples, 1)) 127 | # orientation_samples = np.random.vonmises(self._vmf['mu'], self._vmf['kappa'], size=(n_samples, 1)) 128 | # print(np.mean(orientation_samples)) 129 | return np.concatenate((centroid_samples, closest_samples, orientation_samples), axis=1) 130 | 131 | def trim(self): 132 | self._records = [None] * len(self._records) # NOTE force garbage collection of records 133 | 134 | 135 | if __name__ == '__main__': 136 | parser = argparse.ArgumentParser(description='Compute arrangement priors') 137 | parser.add_argument('--task', type=str, required=True, help=' task [fit|load]') 138 | parser.add_argument('--priors_dir', type=str, required=True, help='directory to save priors') 139 | parser.add_argument('--room_type', type=str, help='room type to filter priors') 140 | parser.add_argument('--min_num_observations', type=int, default=100, help='ignore priors with < this observations') 141 | args = parser.parse_args() 142 | 143 | rod = RelativeObservationsDatabase(name='suncg_priors', priors_dir=args.priors_dir, verbose=True) 144 | 145 | if args.task == 'fit': 146 | rod.load() 147 | priors = {} 148 | for obs_category in rod.grouped_observations: 149 | if args.room_type and obs_category.room_types != args.room_type: 150 | continue 151 | pap = PairwiseArrangementPrior(obs_category) 152 | observations = rod.grouped_observations[obs_category] 153 | if len(observations) > args.min_num_observations: 154 | pap.add_observations(observations, parameterization_scheme='offsets_angles') 155 | pap.fit_model() 156 | priors[obs_category] = pap 157 | filename = os.path.join(args.priors_dir, f'priors.pkl.gz') 158 | utils.pickle_dump_compressed(priors, filename) 159 | 160 | if args.task == 'load': 161 | filename = os.path.join(args.priors_dir, f'priors.pkl.gz') 162 | priors = utils.pickle_load_compressed(filename) 163 | for (cat_key, prior) in priors.items(): 164 | print(f'Loaded PairwiseArrangementPrior for {cat_key} with {prior.num_observations} observations') 165 | if prior.num_observations > args.min_num_observations: 166 | print(f'centroid: mu={prior.centroid.means_}, w={prior.centroid.weights_}') 167 | print(f'closest: mu={prior.closest.means_}, w={prior.closest.weights_}') 168 | # print(f'orientation: mu={prior.orientation["mu"]}, kappa={prior.orientation["kappa"]}') 169 | # samples = prior.sample(n_samples=1000) 170 | # lp = prior.log_prob(samples) 171 | # i_max = np.argmax(lp, axis=0) 172 | # print(f'max sample={samples[i_max, :]} lp={lp[i_max, :]}') 173 | # x = np.stack(prior._records, axis=0) 174 | # lp_x = prior.log_prob(x) 175 | # i_max_x = np.argmax(lp_x, axis=0) 176 | # print(f'max observ={x[i_max_x, :]} lp={lp_x[i_max_x, :]}') 177 | if cat_key.obj_category == 'office_chair' and cat_key.ref_obj_category == 'desk': 178 | prior.plot_model() 179 | 180 | print('DONE') 181 | -------------------------------------------------------------------------------- /scene-synth/requirements.txt: -------------------------------------------------------------------------------- 1 | matplotlib 2 | numpy 3 | numba 4 | pybullet>=1.7.5 5 | pyquaternion 6 | scikit-learn 7 | scipy 8 | -------------------------------------------------------------------------------- /scene-synth/scene_filter.py: -------------------------------------------------------------------------------- 1 | from data import * 2 | import os 3 | import pickle 4 | import math 5 | import pickle 6 | 7 | """ 8 | Performs various kind of operations to the dataset 9 | """ 10 | filter_description = [] 11 | #==================EDIT HERE==================== 12 | #filter_description should be a list of tuples in the form of (filter_description, *params for filter) 13 | #Here are some examples 14 | 15 | room_size = 6 16 | #filter_description = [("floor_node",)] 17 | #filter_description = [("room_type", ["Office"]), ("floor_node",), ("renderable",)] 18 | #filter_description = [("room_type", ["Bathroom", "Toilet"]), ("renderable",)] 19 | #filter_description = [("renderable",)] 20 | #filter_description = [("bedroom", "final"), ("collision",)] 21 | filter_description = [("living", "latent", room_size), ("collision",)] 22 | #filter_description = [("good_house",)] 23 | source = "living_temp" #Source and dest are relative to utils.get_data_root_dir() 24 | dest = "living1111" 25 | #I hate typing True and False 26 | stats = 1 #If true, print stats about the dataset 27 | save_freq = 1 #If true, save category frequency count to the dest directory 28 | render = 1 #If true, render the dataset 29 | save = 0 #If true, save the filtered dataset back as pkl files 30 | json = 1 #If true, save the json files 31 | #There seems to be a KNOWN BUG (I can't remember) that prevents using render/json together with save 32 | #So avoid using them together to be save 33 | render_size = 256 34 | augment = False 35 | #=============================================== 36 | 37 | def get_filter(source, filter_type, *args): 38 | #Just brute force enumerate possible implemented filters 39 | #See filters/ 40 | if filter_type == "good_house": 41 | from filters.good_house import good_house_criteria 42 | house_f = good_house_criteria 43 | dataset_f = DatasetFilter(house_filters = [house_f]) 44 | elif filter_type == "room_type": 45 | from filters.room_type import room_type_criteria 46 | room_f = room_type_criteria(*args) 47 | dataset_f = DatasetFilter(room_filters = [room_f]) 48 | elif filter_type == "bedroom": 49 | from filters.bedroom import bedroom_filter 50 | dataset_f = bedroom_filter(*args, source) 51 | elif filter_type == "office": 52 | from filters.office import office_filter 53 | dataset_f = office_filter(*args, source) 54 | elif filter_type == "living": 55 | from filters.livingroom import livingroom_filter 56 | dataset_f = livingroom_filter(*args, source) 57 | elif filter_type == "toilet": 58 | from filters.toilet import toilet_filter 59 | dataset_f = toilet_filter(*args, source) 60 | elif filter_type == "floor_node": 61 | from filters.floor_node import floor_node_filter 62 | dataset_f = floor_node_filter(*args) 63 | elif filter_type == "renderable": 64 | from filters.renderable import renderable_room_filter 65 | dataset_f = renderable_room_filter(*args) 66 | elif filter_type == "collision": 67 | from filters.collision import collision_filter 68 | from priors.observations import ObjectCollection 69 | oc = ObjectCollection() 70 | dataset_f = collision_filter(oc) 71 | else: 72 | raise NotImplementedError 73 | return dataset_f 74 | 75 | def run_filter(filter_description, source, dest, stats, save_freq, render, save, json, render_size=render_size, num_threads=1, augment=False, room_size=room_size): 76 | """ 77 | Parameters 78 | ---------- 79 | source, dest (String): location of source/dest relative (should be) to utils.get_data_root_dir() 80 | stats (bool): If true, print stats about the dataset 81 | save_freq (bool): If true, save category frequency count to the dest directory 82 | render (bool): If true, render the dataset 83 | save (bool): If true, save the filtered dataset back as pkl files 84 | json (bool): If true, save the json files 85 | render_size (int): size of the top down render 86 | """ 87 | actions = [] 88 | for description in filter_description: 89 | actions.append(get_filter(source, *description)) 90 | if stats: 91 | actions.append(DatasetStats(save_freq=save_freq, details=True, model_details=False, save_dest=dest)) 92 | if save: 93 | actions.append(DatasetSaver(dest=dest, batch_size=1000)) 94 | if augment: 95 | print("Augmenting the data as well") 96 | actions.append(DatasetRotAugmentation()) 97 | if render: 98 | actions.append(DatasetRenderer(dest=dest, size=render_size, room_size=room_size+0.05)) 99 | if json: 100 | actions.append(DatasetToJSON(dest=dest)) 101 | 102 | d = Dataset(actions, source=source) 103 | d.run(num_threads=num_threads) 104 | 105 | if __name__ == "__main__": 106 | run_filter(filter_description, source, dest, stats, save_freq, render, save, json, render_size, augment=augment) 107 | -------------------------------------------------------------------------------- /scene-synth/support_prior.py: -------------------------------------------------------------------------------- 1 | from data import RenderedScene, ObjectCategories 2 | import os 3 | import pickle 4 | import numpy as np 5 | import utils 6 | 7 | 8 | """ 9 | Simple bigram model 10 | """ 11 | class SupportPrior(): 12 | 13 | def __init__(self): 14 | pass 15 | 16 | def learn(self, data_folder="bedroom_final", data_root_dir=None): 17 | if not data_root_dir: 18 | data_root_dir = utils.get_data_root_dir() 19 | data_dir = f"{data_root_dir}/{data_folder}" 20 | self.data_dir = data_dir 21 | self.category_map = ObjectCategories() 22 | 23 | files = os.listdir(data_dir) 24 | files = [f for f in files if ".pkl" in f and not "domain" in f and not "_" in f] 25 | 26 | with open(f"{data_dir}/final_categories_frequency", "r") as f: 27 | lines = f.readlines() 28 | cats = [line.split()[0] for line in lines] 29 | self.category_count = [int(line.split()[1]) for line in lines if line.split()[0] not in ["window", "door"]] 30 | 31 | self.categories = [cat for cat in cats if cat not in set(['window', 'door'])] 32 | self.cat_to_index = {self.categories[i]:i for i in range(len(self.categories))} 33 | self.num_categories = len(self.categories) 34 | self.categories.append("floor") 35 | N = self.num_categories 36 | 37 | self.support_count = [[0 for i in range(N+1)] for j in range(N)] 38 | 39 | for index in range(len(files)): 40 | print(index) 41 | with open(f"{data_dir}/{index}.pkl", "rb") as f: 42 | (_, _, nodes), _ = pickle.load(f) 43 | 44 | object_nodes = [] 45 | id_to_cat = {} 46 | for node in nodes: 47 | modelId = node["modelId"] 48 | category = self.category_map.get_final_category(modelId) 49 | if not category in ["door", "window"]: 50 | object_nodes.append(node) 51 | id_to_cat[node["id"]] = self.cat_to_index[category] 52 | node["category"] = self.cat_to_index[category] 53 | 54 | for node in object_nodes: 55 | parent = node["parent"] 56 | category = node["category"] 57 | if parent == "Floor" or parent is None: 58 | self.support_count[category][-1] += 1 59 | else: 60 | self.support_count[category][id_to_cat[parent]] += 1 61 | #quit() 62 | 63 | self.possible_supports={} 64 | for i in range(self.num_categories): 65 | print(f"Support for {self.categories[i]}:") 66 | supports = [(c, self.support_count[i][c]/self.category_count[i]) for c in range(N+1)] 67 | supports = sorted(supports, key = lambda x:-x[1]) 68 | supports = [s for s in supports if s[1] > 0.01] 69 | for s in supports: 70 | print(f" {self.categories[s[0]]}:{s[1]:4f}") 71 | self.possible_supports[i] = [s[0] for s in supports] 72 | 73 | print(self.possible_supports) 74 | self.N = N 75 | 76 | def save(self, dest=None): 77 | if dest == None: 78 | dest = f"{self.data_dir}/support_prior.pkl" 79 | with open(dest, "wb") as f: 80 | pickle.dump(self.__dict__, f, pickle.HIGHEST_PROTOCOL) 81 | 82 | def load(self, data_dir): 83 | source = f"{data_dir}/support_prior.pkl" 84 | with open(source, "rb") as f: 85 | self.__dict__ = pickle.load(f) 86 | 87 | 88 | if __name__ == "__main__": 89 | a = SupportPrior() 90 | # a.learn("toilet_6x6") 91 | # a.learn("bedroom_6x6") 92 | # a.learn("living_6x6_aug") 93 | a.learn("toilet_6x6") 94 | a.save() 95 | #a.load() 96 | #print(a.get_models(1, ["415"])) 97 | -------------------------------------------------------------------------------- /scene-synth/utils.py: -------------------------------------------------------------------------------- 1 | import gzip 2 | import math 3 | import os 4 | import os.path 5 | import pickle 6 | import torch 7 | import torch.nn.functional as F 8 | from contextlib import contextmanager 9 | from scipy.ndimage import distance_transform_edt 10 | import sys 11 | 12 | def ensuredir(dirname): 13 | """Ensure a directory exists""" 14 | if not os.path.exists(dirname): 15 | os.makedirs(dirname) 16 | 17 | ''' 18 | Turn a number into a string that is zero-padded up to length n 19 | ''' 20 | def zeropad(num, n): 21 | sn = str(num) 22 | while len(sn) < n: 23 | sn = '0' + sn 24 | return sn 25 | 26 | def pickle_dump_compressed(object, filename, protocol=pickle.HIGHEST_PROTOCOL): 27 | """Pickles + compresses an object to file""" 28 | file = gzip.GzipFile(filename, 'wb') 29 | file.write(pickle.dumps(object, protocol)) 30 | file.close() 31 | 32 | def pickle_load_compressed(filename): 33 | """Loads a compressed pickle file and returns reconstituted object""" 34 | file = gzip.GzipFile(filename, 'rb') 35 | buffer = b"" 36 | while True: 37 | data = file.read() 38 | if data == b"": 39 | break 40 | buffer += data 41 | object = pickle.loads(buffer) 42 | file.close() 43 | return object 44 | 45 | def get_data_root_dir(): 46 | """ 47 | Get root dir of the data, defaults to /data if env viariable is not set 48 | """ 49 | env_path = os.environ.get("SCENESYNTH_DATA_PATH") 50 | if env_path: 51 | return env_path 52 | else: 53 | root_dir = os.path.dirname(os.path.abspath(__file__)) 54 | return f"{root_dir}/data" 55 | 56 | def memoize(func): 57 | """ 58 | Decorator to memoize a function 59 | https://medium.com/@nkhaja/memoization-and-decorators-with-python-32f607439f84 60 | """ 61 | cache = func.cache = {} 62 | 63 | @functools.wraps(func) 64 | def memoized_func(*args, **kwargs): 65 | key = str(args) + str(kwargs) 66 | if key not in cache: 67 | cache[key] = func(*args, **kwargs) 68 | return cache[key] 69 | 70 | return memoized_func 71 | 72 | @contextmanager 73 | def stdout_redirected(to=os.devnull): 74 | """ 75 | From https://stackoverflow.com/questions/5081657/how-do-i-prevent-a-c-shared-library-to-print-on-stdout-in-python 76 | Suppress C warnings 77 | """ 78 | fd = sys.stdout.fileno() 79 | 80 | ##### assert that Python and C stdio write using the same file descriptor 81 | ####assert libc.fileno(ctypes.c_void_p.in_dll(libc, "stdout")) == fd == 1 82 | 83 | def _redirect_stdout(to): 84 | sys.stdout.close() # + implicit flush() 85 | os.dup2(to.fileno(), fd) # fd writes to 'to' file 86 | sys.stdout = os.fdopen(fd, 'w') # Python writes to fd 87 | 88 | with os.fdopen(os.dup(fd), 'w') as old_stdout: 89 | with open(to, 'w') as file: 90 | _redirect_stdout(to=file) 91 | try: 92 | yield # allow code to be run with the redirected stdout 93 | finally: 94 | _redirect_stdout(to=old_stdout) # restore stdout. 95 | # buffering and flags such as 96 | # CLOEXEC may be different 97 | --------------------------------------------------------------------------------