├── .gitignore ├── Data ├── GraphConstructor.py ├── LargeFilenameDataset.py ├── VolumeDesignGraph.py ├── get_uniform_voxel.py ├── process_data.py └── process_variation_data.py ├── LICENSE.md ├── Model ├── losses.py └── models.py ├── README.md ├── images ├── Building-GAN.png └── Volumetric_Design.png ├── inference.py ├── runs └── iccv2021 │ └── checkpoints │ └── 70.ckpt ├── train.py ├── train_args.py ├── util.py ├── util_eval.py └── util_graph.py /.gitignore: -------------------------------------------------------------------------------- 1 | /Data/6types-raw_data/ 2 | /Data/6types-processed_data/ 3 | /Data/6types-raw_data_/ 4 | /Data/6types-processed_data_/ 5 | /Data/6types-raw_data__/ 6 | /Data/6types-processed_data__/ 7 | /Data/__pycache__/ 8 | /Data/get_uniform_voxel_slices.py 9 | /Data/uniform_voxel_dataset_class.py 10 | /Data/6types-raw_data.zip 11 | /Model/__pycache__/ 12 | /runs/* 13 | !/runs/iccv2021/ 14 | /.idea/ 15 | /inference*/ 16 | /__pycache__/ 17 | inference_10k.py 18 | -------------------------------------------------------------------------------- /Data/GraphConstructor.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | from Data.VolumeDesignGraph import VolumeDesignGraph 4 | import torch 5 | from collections import Counter 6 | 7 | 8 | class GraphConstructor: 9 | number_of_class = 6 10 | dimension_norm_factor = 100 11 | 12 | global_graph_prefix = "graph_global_" 13 | local_graph_prefix = "graph_local_" 14 | voxel_graph_prefix = "voxel_" 15 | data_id_length = 6 16 | 17 | def __init__(self, ): 18 | pass 19 | 20 | @staticmethod 21 | def get_program_unicode(n): 22 | return n["floor"], n["type"], n["type_id"] 23 | 24 | @staticmethod 25 | def one_hot_vector(_id): 26 | assert _id < GraphConstructor.number_of_class 27 | return [1 if i == _id else 0 for i in range(GraphConstructor.number_of_class)] 28 | 29 | @staticmethod 30 | def load_graph_jsons(data_id, data_dir): 31 | data_id_str = str(data_id).zfill(GraphConstructor.data_id_length) 32 | with open(os.path.join(data_dir, "global_graph_data", GraphConstructor.global_graph_prefix + data_id_str+".json")) as f: 33 | global_graph = json.load(f) # Keys: FAR, site area, Global node (Room Type, target_ratio, connections) 34 | with open(os.path.join(data_dir, "local_graph_data", GraphConstructor.local_graph_prefix + data_id_str+".json")) as f: 35 | local_graph = json.load(f) 36 | with open(os.path.join(data_dir, "voxel_data", GraphConstructor.voxel_graph_prefix + data_id_str+".json")) as f: 37 | raw_voxel_graph = json.load(f) 38 | 39 | program_graph = GraphConstructor.construct_program_graph(local_graph, global_graph) 40 | voxel_graph = GraphConstructor.construct_voxel_graph(raw_voxel_graph) 41 | 42 | return GraphConstructor.tensorfy(program_graph, voxel_graph, data_id_str) 43 | 44 | @staticmethod # not used 45 | def load_graph_from_UI(global_graph, local_graph, raw_voxel_graph, data_id): 46 | data_id_str = str(data_id).zfill(GraphConstructor.data_id_length) 47 | 48 | program_graph = GraphConstructor.construct_program_graph(local_graph, global_graph) 49 | voxel_graph = GraphConstructor.construct_voxel_graph(raw_voxel_graph) 50 | 51 | return GraphConstructor.tensorfy(program_graph, voxel_graph, data_id_str) 52 | 53 | @staticmethod 54 | def tensorfy(program_grpah, voxel_graph, data_id_str): 55 | program_floor_index = torch.tensor(program_grpah["program_floor_index"], dtype=torch.long) 56 | story_level_feature = torch.FloatTensor(program_grpah["story_level_feature"]) # N x 1 57 | program_class_feature = torch.FloatTensor(program_grpah["program_class_feature"]) 58 | program_edge = torch.tensor(program_grpah["program_edge"], dtype=torch.long) 59 | neg_same_story_program_edge = torch.tensor(program_grpah["neg_same_story_program_edge"], dtype=torch.long) 60 | neg_diff_story_program_edge = torch.tensor(program_grpah["neg_diff_story_program_edge"], dtype=torch.long) 61 | program_class_cluster = torch.tensor(program_grpah["program_class_cluster"], dtype=torch.long) 62 | program_target_ratio = torch.FloatTensor(program_grpah["program_target_ratio"]) 63 | FAR = torch.FloatTensor([program_grpah["FAR"]]) 64 | voxel_floor_index = torch.tensor(voxel_graph["voxel_floor_index"], dtype=torch.long) 65 | voxel_projection_cluster = torch.tensor(voxel_graph["voxel_projection_cluster"], dtype=torch.long) 66 | voxel_feature = torch.FloatTensor(voxel_graph["voxel_feature"]) 67 | voxel_bool = torch.FloatTensor(voxel_graph["voxel_bool"]) 68 | voxel_label = torch.FloatTensor(voxel_graph["voxel_label"]) 69 | voxel_edge = torch.tensor(voxel_graph["voxel_edge"], dtype=torch.long) 70 | 71 | # add index_select for cross-edges 72 | pfc = program_grpah["program_floor_index"] 73 | program_node_id_in_story_group = [] 74 | voxel_node_counter = Counter(voxel_graph["voxel_floor_index"]) 75 | for program_node_id, story_id in enumerate(pfc): 76 | if program_node_id-1 >= 0 and story_id == pfc[program_node_id-1]: 77 | program_node_id_in_story_group[-1].append(program_node_id) 78 | else: 79 | program_node_id_in_story_group.append([program_node_id]) 80 | 81 | cross_edge_program_index_select, cross_edge_voxel_index_select, voxel_id = [], [], 0 82 | for global_story_id, program_node_id_in_this_story_group in enumerate(program_node_id_in_story_group): 83 | number_of_voxel_in_this_story = voxel_node_counter[global_story_id] 84 | number_of_program_in_this_story = len(program_node_id_in_this_story_group) 85 | cross_edge_program_index_select += program_node_id_in_this_story_group * number_of_voxel_in_this_story 86 | for _ in range(number_of_voxel_in_this_story): 87 | cross_edge_voxel_index_select += [voxel_id] * number_of_program_in_this_story 88 | voxel_id += 1 89 | 90 | assert len(cross_edge_program_index_select) == len(cross_edge_voxel_index_select) 91 | 92 | cross_edge_program_index_select = torch.tensor(cross_edge_program_index_select, dtype=torch.long) 93 | cross_edge_voxel_index_select = torch.tensor(cross_edge_voxel_index_select, dtype=torch.long) 94 | 95 | # add label_program_index. This is for the discriminator to point the voxel label back to its corresponding program node 96 | voxel_label_program_node_index = [] 97 | 98 | for voxel_id, story_id in enumerate(voxel_graph["voxel_floor_index"]): 99 | program_unicode = (story_id, ) + voxel_graph["node_id_to_program_unicode"][voxel_id] # (floor, -1, 0) if type=-1 100 | program_node_id = program_grpah["unicode_to_node_id_map"].get(program_unicode, -1) 101 | if program_node_id != -1: 102 | assert(program_floor_index[program_node_id]) == story_id 103 | voxel_label_program_node_index.append(program_node_id) 104 | 105 | voxel_label_program_node_index = torch.tensor(voxel_label_program_node_index, dtype=torch.long) 106 | 107 | # add voxel_edge_mask to extract edges across same story in voxel graph. This is not used. 108 | voxel_edge_mask = [] 109 | for p1, p2 in zip(voxel_edge[0], voxel_edge[1]): 110 | voxel_edge_mask.append(1.0 if voxel_graph["voxel_floor_index"][p1] == voxel_graph["voxel_floor_index"][p2] else 0.0) 111 | voxel_edge_mask = torch.FloatTensor(voxel_edge_mask) 112 | 113 | graph = VolumeDesignGraph(data_id_str, program_floor_index, story_level_feature, program_class_feature, 114 | program_edge, neg_same_story_program_edge, neg_diff_story_program_edge, program_class_cluster, program_target_ratio, FAR, 115 | voxel_floor_index, voxel_projection_cluster, voxel_feature, voxel_bool, voxel_label, 116 | voxel_edge, voxel_edge_mask, cross_edge_program_index_select, cross_edge_voxel_index_select, voxel_label_program_node_index) 117 | return graph 118 | 119 | @staticmethod 120 | def construct_voxel_graph(voxel_graph): 121 | # Extract Node Feature 122 | unicode_to_node_id_map, node_id_to_program_unicode, projection_map = {}, {}, {} 123 | voxel_floor_index, voxel_feature, voxel_bool, voxel_label = [], [], [], [] 124 | for i, n in enumerate(voxel_graph["voxel_node"]): 125 | voxel_bool.append(0 if n["type"] < 0 else 1) 126 | voxel_label.append(GraphConstructor.one_hot_vector(n["type"])) 127 | voxel_feature.append([*[x / GraphConstructor.dimension_norm_factor for x in n["dimension"] + n["coordinate"]], n["weight"]]) 128 | unicode_to_node_id_map[tuple(n["location"])] = i 129 | node_id_to_program_unicode[i] = (n["type"], n["type_id"]) 130 | voxel_floor_index.append(n["location"][0]) 131 | 132 | horizontal_location = (n["location"][2], n["location"][1]) # x, y 133 | if projection_map.get(horizontal_location, None): 134 | projection_map.get(horizontal_location).append(i) 135 | else: 136 | projection_map[horizontal_location] = [i] 137 | 138 | voxel_projection_cluster = [-1] * len(voxel_graph["voxel_node"]) 139 | for i, v in enumerate(projection_map.values()): 140 | for n_i in v: 141 | voxel_projection_cluster[n_i] = i 142 | assert(-1 not in voxel_projection_cluster) 143 | 144 | # Extract Edge Feature <- Already bi-directional 145 | voxel_edge_src, voxel_edge_dst = [], [] 146 | for i, n_i in enumerate(voxel_graph["voxel_node"]): 147 | for unicode_j in n_i["neighbors"]: 148 | voxel_edge_src.append(i) 149 | voxel_edge_dst.append(unicode_to_node_id_map[tuple(unicode_j)]) 150 | voxel_edge = [voxel_edge_src, voxel_edge_dst] 151 | 152 | return { 153 | "voxel_floor_index": voxel_floor_index, # dict (F --> nodes in each floor) 154 | "voxel_projection_cluster": voxel_projection_cluster, # N x 1 155 | "voxel_feature": voxel_feature, # N x 7 156 | "voxel_bool": voxel_bool, # N x 1 157 | "voxel_label": voxel_label, # N x C 158 | "voxel_edge": voxel_edge, # 2 x E 159 | "node_id_to_program_unicode": node_id_to_program_unicode # voxel node id -> type, type id 160 | } 161 | 162 | @staticmethod 163 | def construct_program_graph(local_graph, global_graph): 164 | # Extract Node Feature 165 | unicode_to_node_id_map = {} 166 | program_floor_index, story_level_feature, program_class_feature, program_class_cluster = [], [], [], [] 167 | for i, n in enumerate(local_graph["node"]): 168 | # n is a dict. Keys: floor, type, type_id 169 | unicode = GraphConstructor.get_program_unicode(n) 170 | unicode_to_node_id_map[unicode] = i 171 | # n["region_far"] and n["center"] are not used 172 | 173 | story_level_feature.append(n["floor"]) 174 | program_class_feature.append(GraphConstructor.one_hot_vector(n["type"])) 175 | program_class_cluster.append(n["type"]) 176 | program_floor_index.append(n["floor"]) 177 | 178 | # Extract Edge Feature <- Already bi-directional 179 | program_edge_src, program_edge_dst, adj_matrix = [], [], [[0 for _ in range(len(local_graph["node"]))] for _ in range(len(local_graph["node"]))] 180 | for i, n_i in enumerate(local_graph["node"]): 181 | for unicode_j in n_i["neighbors"]: 182 | program_edge_src.append(i) 183 | program_edge_dst.append(unicode_to_node_id_map[tuple(unicode_j)]) 184 | adj_matrix[i][unicode_to_node_id_map[tuple(unicode_j)]] = 1 185 | program_edge = [program_edge_src, program_edge_dst] 186 | 187 | # Construct Negative edges on the same story 188 | neg_same_story_program_edge_src, neg_same_story_program_edge_dst, neg_diff_story_program_edge_src, neg_diff_story_program_edge_dst = [], [], [], [] 189 | for i in range(len(local_graph["node"])): 190 | for j in range(i+1, len(local_graph["node"])): 191 | if adj_matrix[i][j] == 1: 192 | continue 193 | if story_level_feature[i] == story_level_feature[j]: 194 | neg_same_story_program_edge_src += [i, j] 195 | neg_same_story_program_edge_dst += [j, i] 196 | else: 197 | neg_diff_story_program_edge_src += [i, j] 198 | neg_diff_story_program_edge_dst += [j, i] 199 | neg_same_story_program_edge = [neg_same_story_program_edge_src, neg_same_story_program_edge_dst] 200 | neg_diff_story_program_edge = [neg_diff_story_program_edge_src, neg_diff_story_program_edge_dst] 201 | 202 | # Global Feature and Edge 203 | program_target_ratio = [0] * GraphConstructor.number_of_class 204 | FAR = global_graph["far"] 205 | for c in global_graph["global_node"]: 206 | program_target_ratio[c["type"]] = c["proportion"] 207 | 208 | return { 209 | "program_floor_index": program_floor_index, # N x 1 210 | "story_level_feature": story_level_feature, # N x 1 211 | "program_class_feature": program_class_feature, # N x C 212 | "program_edge": program_edge, # 2 x E 213 | "neg_same_story_program_edge": neg_same_story_program_edge, # 2 x E 214 | "neg_diff_story_program_edge": neg_diff_story_program_edge, # 2 x E 215 | "program_class_cluster": program_class_cluster, # N x 1 216 | "program_target_ratio": program_target_ratio, # C 217 | "FAR": FAR, # 1 218 | "unicode_to_node_id_map": unicode_to_node_id_map # floor, type, type_id -> program node id 219 | } 220 | 221 | -------------------------------------------------------------------------------- /Data/LargeFilenameDataset.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import os 3 | from torch.utils.data import Dataset 4 | 5 | 6 | class LargeFilenameDataset(Dataset): 7 | def __init__(self, data_dir, fname_list): 8 | self.data_dir = data_dir 9 | self.fname_list = fname_list 10 | 11 | def __len__(self): 12 | return len(self.fname_list) 13 | 14 | def __getitem__(self, i): 15 | return torch.load(os.path.join(self.data_dir, self.fname_list[i])) 16 | -------------------------------------------------------------------------------- /Data/VolumeDesignGraph.py: -------------------------------------------------------------------------------- 1 | from torch_geometric.data import Data 2 | 3 | 4 | class VolumeDesignGraph(Data): 5 | def __init__(self, data_id_str=None, program_floor_cluster=None, story_level_feature=None, program_class_feature=None, 6 | program_edge=None, neg_same_story_program_edge=None, neg_diff_story_program_edge=None, program_class_cluster=None, program_target_ratio=None, FAR=None, 7 | voxel_floor_cluster=None, voxel_projection_cluster=None, voxel_feature=None, voxel_bool=None, voxel_label=None, 8 | voxel_edge=None, voxel_edge_mask=None, cross_edge_program_index_select=None, cross_edge_voxel_index_select=None, voxel_label_program_node_index=None): 9 | """ 10 | data_id_str: (str) data ID 11 | 12 | program_edge: (e_program x 2) program edges (include same_story and diff_story) 13 | story_level_feature: (n_program_node x 1) the story level for each program nodes 14 | program_class_feature: (n_program_node x d) one-hot vector of program type 15 | program_class_cluster: (n_program_type x n_element) list of program node IDs on the same program type 16 | program_floor_cluster: (n_cluster x n_element) list of program node IDs on the same story 17 | program_target_ratio: (n_program_type x 1) the target ratio between all program types 18 | FAR: (1) target floor area ratio 19 | 20 | voxel_feature: (n_voxel x d) voxel feature 21 | voxel_floor_cluster: (n_cluster x n_element) list of voxel node IDs on the same story 22 | voxel_projection_cluster:　(n_cluster x n_element) list of voxel node IDs on the same x,y coordinates 23 | voxel_bool: (n_voxel x 1) if this voxel is void 24 | voxel_label: (n_voxel x program_type) which program type is this voxel 25 | voxel_edge: (e_voxel x 2) voxel edges 26 | voxel_edge_mask: (e_voxel x 1) Not used! if the two voxel nodes on the voxel edge is on the same story 27 | 28 | cross_edge_program_index_select: (e_cross x 1) the program node IDs of the cross edges 29 | cross_edge_voxel_index_select: (e_cross x 1) the voxel node IDS of the cross edges 30 | voxel_label_program_node_index: (n_voxel x 1) which program node ID is assigned to this voxel 31 | """ 32 | 33 | super(VolumeDesignGraph, self).__init__() 34 | self.data_id_str = data_id_str 35 | self.program_floor_cluster = program_floor_cluster 36 | self.story_level_feature, self.program_class_feature = story_level_feature, program_class_feature 37 | self.program_edge, self.neg_same_story_program_edge, self.neg_diff_story_program_edge = program_edge, neg_same_story_program_edge, neg_diff_story_program_edge 38 | self.program_class_cluster = program_class_cluster 39 | self.program_target_ratio = program_target_ratio 40 | self.FAR = FAR 41 | 42 | self.voxel_floor_cluster = voxel_floor_cluster 43 | self.voxel_projection_cluster = voxel_projection_cluster 44 | self.voxel_feature = voxel_feature 45 | self.voxel_bool, self.voxel_label = voxel_bool, voxel_label 46 | self.voxel_edge = voxel_edge 47 | self.voxel_edge_mask = voxel_edge_mask 48 | 49 | self.cross_edge_program_index_select = cross_edge_program_index_select 50 | self.cross_edge_voxel_index_select = cross_edge_voxel_index_select 51 | self.voxel_label_program_node_index = voxel_label_program_node_index 52 | 53 | self._batch_program_inc_keys = ["program_edge", "cross_edge_program_index_select", "voxel_label_program_node_index", "neg_same_story_program_edge", "neg_diff_story_program_edge"] 54 | self._batch_voxel_inc_keys = ["voxel_edge", "cross_edge_voxel_index_select"] 55 | self._floor_inc_keys = ["program_floor_cluster", "voxel_floor_cluster"] 56 | self._program_class_inc_keys = ["program_class_cluster"] 57 | self._voxel_projection_inc_keys = ["voxel_projection_cluster"] 58 | self._cat_dim_keys = ["program_edge", "voxel_edge", "neg_same_story_program_edge", "neg_diff_story_program_edge"] 59 | 60 | def __inc__(self, key, value): 61 | if key in self._batch_program_inc_keys: 62 | return self.program_class_feature.size(0) # number of edges in program graph 63 | if key in self._batch_voxel_inc_keys: 64 | return self.voxel_feature.size(0) # number of nodes in voxel graph 65 | if key in self._floor_inc_keys: 66 | return int(max(self.story_level_feature) + 1) # number of stories in each volume design graph 67 | if key in self._program_class_inc_keys: 68 | return self.program_class_feature.size(1) # number of program class 69 | if key in self._voxel_projection_inc_keys: 70 | return int(max(self.voxel_projection_cluster) + 1) # number of voxels on XY plane 71 | else: 72 | return super(VolumeDesignGraph, self).__inc__(key, value) 73 | 74 | def __cat_dim__(self, key, value): 75 | # `*index*` and `*face*` should be concatenated in the last dimension, 76 | # everything else in the first dimension. 77 | return -1 if key in self._cat_dim_keys else 0 78 | 79 | 80 | 81 | -------------------------------------------------------------------------------- /Data/get_uniform_voxel.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from pathlib import Path 4 | import json 5 | 6 | voxel_dir = "6types-raw_data/voxel_data" 7 | uni_voxel_label_dir = "6types-raw_data_uniform/uni_voxel_label" 8 | Path(uni_voxel_label_dir).mkdir(parents=True, exist_ok=True) 9 | uni_voxel_id_dir = "6types-raw_data_uniform/uni_voxel_id" 10 | Path(uni_voxel_id_dir).mkdir(parents=True, exist_ok=True) 11 | 12 | k = 120000 13 | 14 | for i in range(k): 15 | 16 | print(i, "/120000") 17 | 18 | with open(voxel_dir + '/voxel_{}.json'.format(str(i).zfill(5))) as f: 19 | voxels = json.load(f)["voxel_node"] 20 | 21 | uni_voxels_label = np.ones((50, 40, 40)) 22 | uni_voxels_label = - uni_voxels_label 23 | 24 | uni_voxels_id = np.zeros_like(uni_voxels_label) 25 | 26 | for id, voxel in enumerate(voxels): 27 | 28 | sz = np.array(range(int(voxel["coordinate"][0]), int(voxel["coordinate"][0] + voxel["dimension"][0]))) 29 | sy = np.array(range(int(voxel["coordinate"][1]), int(voxel["coordinate"][1] + voxel["dimension"][1]))) 30 | sx = np.array(range(int(voxel["coordinate"][2]), int(voxel["coordinate"][2] + voxel["dimension"][2]))) 31 | uni_voxels_label[np.ix_(sz, sy, sx)] = voxel["type"] 32 | uni_voxels_id[np.ix_(sz, sy, sx)] = id 33 | 34 | uni_voxels_label_tensor = torch.from_numpy(uni_voxels_label) 35 | uni_voxels_label_tensor.type(torch.DoubleTensor) 36 | 37 | uni_voxels_id_tensor = torch.from_numpy(uni_voxels_id) 38 | uni_voxels_id_tensor.type(torch.long) 39 | 40 | torch.save(uni_voxels_label_tensor, uni_voxel_label_dir + "/uni_voxel_label_{}.pt".format(str(i).zfill(6))) 41 | torch.save(uni_voxels_id_tensor, uni_voxel_id_dir + "/uni_voxel_id_{}.pt".format(str(i).zfill(6))) 42 | 43 | -------------------------------------------------------------------------------- /Data/process_data.py: -------------------------------------------------------------------------------- 1 | """ 2 | Folder structure: 3 | raw_data 4 | - global_graph_data 5 | - local_graph_data 6 | - voxel_data 7 | 8 | """ 9 | import torch 10 | from Data.GraphConstructor import GraphConstructor 11 | import os 12 | 13 | raw_data_dir = "6types-raw_data" 14 | output_dir = "6types-processed_data" 15 | global_graph_dir = os.path.join(raw_data_dir, "global_graph_data") 16 | local_graph_dir = os.path.join(raw_data_dir, "local_graph_data") 17 | voxel_graph_dir = os.path.join(raw_data_dir, "voxel_data") 18 | 19 | 20 | for fname in os.listdir(global_graph_dir): 21 | data_id = int(''.join(filter(str.isdigit, fname))) 22 | try: 23 | g = GraphConstructor.load_graph_jsons(data_id, raw_data_dir) 24 | output_fname = "data" + str(data_id).zfill(GraphConstructor.data_id_length) + ".pt" 25 | torch.save(g, os.path.join(output_dir, output_fname)) 26 | except: 27 | print("Error loading data " + str(data_id).zfill(GraphConstructor.data_id_length)) 28 | 29 | print("Data processing completed") 30 | 31 | -------------------------------------------------------------------------------- /Data/process_variation_data.py: -------------------------------------------------------------------------------- 1 | """ 2 | This code generates variation samples by fixing the graph data and varying the random variable z. 3 | """ 4 | 5 | from Data.GraphConstructor import GraphConstructor 6 | import os 7 | import json 8 | import torch 9 | from shutil import copyfile 10 | 11 | 12 | num_VG = 1 13 | num_z = 10 14 | 15 | num_of_data = 1000 16 | data_dir = "./6types-raw_data/sum" 17 | output_dir = "./6types-processed-variation-v1-z10_data" 18 | os.mkdir(output_dir) 19 | 20 | raw_dir = "./6types-raw-variation-v1-z10_data" 21 | raw_voxel = raw_dir + "/voxel_data" 22 | raw_global = raw_dir + "/global_graph_data" 23 | raw_local = raw_dir + "/local_graph_data" 24 | os.mkdir(raw_dir) 25 | os.mkdir(raw_voxel) 26 | os.mkdir(raw_local) 27 | os.mkdir(raw_global) 28 | 29 | voxel_dir = "./variation_test" 30 | 31 | cnt = 0 32 | for data_id in range(num_of_data): 33 | print("generate " + str(num_VG*num_z) + "data from Data" + str(data_id)) 34 | data_id_str = str(96000 + data_id).zfill(GraphConstructor.data_id_length) 35 | with open(os.path.join(data_dir, "global_graph_data", GraphConstructor.global_graph_prefix + data_id_str + ".json")) as f: 36 | global_graph = json.load(f) # Keys: FAR, site area, Global node (Room Type, target_ratio, connections) 37 | with open(os.path.join(data_dir, "local_graph_data", GraphConstructor.local_graph_prefix + data_id_str + ".json")) as f: 38 | local_graph = json.load(f) 39 | program_graph = GraphConstructor.construct_program_graph(local_graph, global_graph) 40 | 41 | num_of_story = max(program_graph["story_level_feature"]) + 1 # starting from 1 ~ 11 42 | voxel_graph_dir = os.path.join(voxel_dir, str(num_of_story), "voxel_data") 43 | for i in range(num_VG): 44 | with open(os.path.join(voxel_graph_dir, GraphConstructor.voxel_graph_prefix + str(i).zfill(GraphConstructor.data_id_length) + ".json")) as f: 45 | raw_voxel_graph = json.load(f) 46 | voxel_graph = GraphConstructor.construct_voxel_graph(raw_voxel_graph) 47 | 48 | for _ in range(num_z): 49 | g = GraphConstructor.tensorfy(program_graph, voxel_graph, str(cnt).zfill(GraphConstructor.data_id_length)) 50 | output_fname = "data" + str(cnt).zfill(GraphConstructor.data_id_length) + ".pt" 51 | torch.save(g, os.path.join(output_dir, output_fname)) 52 | 53 | # copy json files and rename them based on current cnt 54 | global_copy_f = os.path.join(data_dir, "global_graph_data", GraphConstructor.global_graph_prefix + data_id_str + ".json") 55 | local_copy_f = os.path.join(data_dir, "local_graph_data", GraphConstructor.local_graph_prefix + data_id_str + ".json") 56 | voxel_copy_f = os.path.join(voxel_graph_dir, GraphConstructor.voxel_graph_prefix + str(i).zfill(GraphConstructor.data_id_length) + ".json") 57 | 58 | global_paste_f = os.path.join(raw_global, "graph_global_"+str(cnt).zfill(GraphConstructor.data_id_length)+".json") 59 | local_paste_f = os.path.join(raw_local, "graph_local_"+str(cnt).zfill(GraphConstructor.data_id_length)+".json") 60 | voxel_paste_f = os.path.join(raw_voxel, "voxel_" + str(cnt).zfill(GraphConstructor.data_id_length) + ".json") 61 | 62 | copyfile(global_copy_f, global_paste_f) 63 | copyfile(local_copy_f, local_paste_f) 64 | copyfile(voxel_copy_f, voxel_paste_f) 65 | 66 | cnt += 1 67 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | Attribution-NonCommercial-ShareAlike 4.0 International 2 | 3 | ======================================================================= 4 | 5 | Creative Commons Corporation ("Creative Commons") is not a law firm and 6 | does not provide legal services or legal advice. Distribution of 7 | Creative Commons public licenses does not create a lawyer-client or 8 | other relationship. Creative Commons makes its licenses and related 9 | information available on an "as-is" basis. Creative Commons gives no 10 | warranties regarding its licenses, any material licensed under their 11 | terms and conditions, or any related information. Creative Commons 12 | disclaims all liability for damages resulting from their use to the 13 | fullest extent possible. 14 | 15 | Using Creative Commons Public Licenses 16 | 17 | Creative Commons public licenses provide a standard set of terms and 18 | conditions that creators and other rights holders may use to share 19 | original works of authorship and other material subject to copyright 20 | and certain other rights specified in the public license below. The 21 | following considerations are for informational purposes only, are not 22 | exhaustive, and do not form part of our licenses. 23 | 24 | Considerations for licensors: Our public licenses are 25 | intended for use by those authorized to give the public 26 | permission to use material in ways otherwise restricted by 27 | copyright and certain other rights. Our licenses are 28 | irrevocable. Licensors should read and understand the terms 29 | and conditions of the license they choose before applying it. 30 | Licensors should also secure all rights necessary before 31 | applying our licenses so that the public can reuse the 32 | material as expected. Licensors should clearly mark any 33 | material not subject to the license. This includes other CC- 34 | licensed material, or material used under an exception or 35 | limitation to copyright. More considerations for licensors: 36 | wiki.creativecommons.org/Considerations_for_licensors 37 | 38 | Considerations for the public: By using one of our public 39 | licenses, a licensor grants the public permission to use the 40 | licensed material under specified terms and conditions. If 41 | the licensor's permission is not necessary for any reason--for 42 | example, because of any applicable exception or limitation to 43 | copyright--then that use is not regulated by the license. Our 44 | licenses grant only permissions under copyright and certain 45 | other rights that a licensor has authority to grant. Use of 46 | the licensed material may still be restricted for other 47 | reasons, including because others have copyright or other 48 | rights in the material. A licensor may make special requests, 49 | such as asking that all changes be marked or described. 50 | Although not required by our licenses, you are encouraged to 51 | respect those requests where reasonable. More considerations 52 | for the public: 53 | wiki.creativecommons.org/Considerations_for_licensees 54 | 55 | ======================================================================= 56 | 57 | Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 58 | Public License 59 | 60 | By exercising the Licensed Rights (defined below), You accept and agree 61 | to be bound by the terms and conditions of this Creative Commons 62 | Attribution-NonCommercial-ShareAlike 4.0 International Public License 63 | ("Public License"). To the extent this Public License may be 64 | interpreted as a contract, You are granted the Licensed Rights in 65 | consideration of Your acceptance of these terms and conditions, and the 66 | Licensor grants You such rights in consideration of benefits the 67 | Licensor receives from making the Licensed Material available under 68 | these terms and conditions. 69 | 70 | 71 | Section 1 -- Definitions. 72 | 73 | a. Adapted Material means material subject to Copyright and Similar 74 | Rights that is derived from or based upon the Licensed Material 75 | and in which the Licensed Material is translated, altered, 76 | arranged, transformed, or otherwise modified in a manner requiring 77 | permission under the Copyright and Similar Rights held by the 78 | Licensor. For purposes of this Public License, where the Licensed 79 | Material is a musical work, performance, or sound recording, 80 | Adapted Material is always produced where the Licensed Material is 81 | synched in timed relation with a moving image. 82 | 83 | b. Adapter's License means the license You apply to Your Copyright 84 | and Similar Rights in Your contributions to Adapted Material in 85 | accordance with the terms and conditions of this Public License. 86 | 87 | c. BY-NC-SA Compatible License means a license listed at 88 | creativecommons.org/compatiblelicenses, approved by Creative 89 | Commons as essentially the equivalent of this Public License. 90 | 91 | d. Copyright and Similar Rights means copyright and/or similar rights 92 | closely related to copyright including, without limitation, 93 | performance, broadcast, sound recording, and Sui Generis Database 94 | Rights, without regard to how the rights are labeled or 95 | categorized. For purposes of this Public License, the rights 96 | specified in Section 2(b)(1)-(2) are not Copyright and Similar 97 | Rights. 98 | 99 | e. Effective Technological Measures means those measures that, in the 100 | absence of proper authority, may not be circumvented under laws 101 | fulfilling obligations under Article 11 of the WIPO Copyright 102 | Treaty adopted on December 20, 1996, and/or similar international 103 | agreements. 104 | 105 | f. Exceptions and Limitations means fair use, fair dealing, and/or 106 | any other exception or limitation to Copyright and Similar Rights 107 | that applies to Your use of the Licensed Material. 108 | 109 | g. License Elements means the license attributes listed in the name 110 | of a Creative Commons Public License. The License Elements of this 111 | Public License are Attribution, NonCommercial, and ShareAlike. 112 | 113 | h. Licensed Material means the artistic or literary work, database, 114 | or other material to which the Licensor applied this Public 115 | License. 116 | 117 | i. Licensed Rights means the rights granted to You subject to the 118 | terms and conditions of this Public License, which are limited to 119 | all Copyright and Similar Rights that apply to Your use of the 120 | Licensed Material and that the Licensor has authority to license. 121 | 122 | j. Licensor means the individual(s) or entity(ies) granting rights 123 | under this Public License. 124 | 125 | k. NonCommercial means not primarily intended for or directed towards 126 | commercial advantage or monetary compensation. For purposes of 127 | this Public License, the exchange of the Licensed Material for 128 | other material subject to Copyright and Similar Rights by digital 129 | file-sharing or similar means is NonCommercial provided there is 130 | no payment of monetary compensation in connection with the 131 | exchange. 132 | 133 | l. Share means to provide material to the public by any means or 134 | process that requires permission under the Licensed Rights, such 135 | as reproduction, public display, public performance, distribution, 136 | dissemination, communication, or importation, and to make material 137 | available to the public including in ways that members of the 138 | public may access the material from a place and at a time 139 | individually chosen by them. 140 | 141 | m. Sui Generis Database Rights means rights other than copyright 142 | resulting from Directive 96/9/EC of the European Parliament and of 143 | the Council of 11 March 1996 on the legal protection of databases, 144 | as amended and/or succeeded, as well as other essentially 145 | equivalent rights anywhere in the world. 146 | 147 | n. You means the individual or entity exercising the Licensed Rights 148 | under this Public License. Your has a corresponding meaning. 149 | 150 | 151 | Section 2 -- Scope. 152 | 153 | a. License grant. 154 | 155 | 1. Subject to the terms and conditions of this Public License, 156 | the Licensor hereby grants You a worldwide, royalty-free, 157 | non-sublicensable, non-exclusive, irrevocable license to 158 | exercise the Licensed Rights in the Licensed Material to: 159 | 160 | a. reproduce and Share the Licensed Material, in whole or 161 | in part, for NonCommercial purposes only; and 162 | 163 | b. produce, reproduce, and Share Adapted Material for 164 | NonCommercial purposes only. 165 | 166 | 2. Exceptions and Limitations. For the avoidance of doubt, where 167 | Exceptions and Limitations apply to Your use, this Public 168 | License does not apply, and You do not need to comply with 169 | its terms and conditions. 170 | 171 | 3. Term. The term of this Public License is specified in Section 172 | 6(a). 173 | 174 | 4. Media and formats; technical modifications allowed. The 175 | Licensor authorizes You to exercise the Licensed Rights in 176 | all media and formats whether now known or hereafter created, 177 | and to make technical modifications necessary to do so. The 178 | Licensor waives and/or agrees not to assert any right or 179 | authority to forbid You from making technical modifications 180 | necessary to exercise the Licensed Rights, including 181 | technical modifications necessary to circumvent Effective 182 | Technological Measures. For purposes of this Public License, 183 | simply making modifications authorized by this Section 2(a) 184 | (4) never produces Adapted Material. 185 | 186 | 5. Downstream recipients. 187 | 188 | a. Offer from the Licensor -- Licensed Material. Every 189 | recipient of the Licensed Material automatically 190 | receives an offer from the Licensor to exercise the 191 | Licensed Rights under the terms and conditions of this 192 | Public License. 193 | 194 | b. Additional offer from the Licensor -- Adapted Material. 195 | Every recipient of Adapted Material from You 196 | automatically receives an offer from the Licensor to 197 | exercise the Licensed Rights in the Adapted Material 198 | under the conditions of the Adapter's License You apply. 199 | 200 | c. No downstream restrictions. You may not offer or impose 201 | any additional or different terms or conditions on, or 202 | apply any Effective Technological Measures to, the 203 | Licensed Material if doing so restricts exercise of the 204 | Licensed Rights by any recipient of the Licensed 205 | Material. 206 | 207 | 6. No endorsement. Nothing in this Public License constitutes or 208 | may be construed as permission to assert or imply that You 209 | are, or that Your use of the Licensed Material is, connected 210 | with, or sponsored, endorsed, or granted official status by, 211 | the Licensor or others designated to receive attribution as 212 | provided in Section 3(a)(1)(A)(i). 213 | 214 | b. Other rights. 215 | 216 | 1. Moral rights, such as the right of integrity, are not 217 | licensed under this Public License, nor are publicity, 218 | privacy, and/or other similar personality rights; however, to 219 | the extent possible, the Licensor waives and/or agrees not to 220 | assert any such rights held by the Licensor to the limited 221 | extent necessary to allow You to exercise the Licensed 222 | Rights, but not otherwise. 223 | 224 | 2. Patent and trademark rights are not licensed under this 225 | Public License. 226 | 227 | 3. To the extent possible, the Licensor waives any right to 228 | collect royalties from You for the exercise of the Licensed 229 | Rights, whether directly or through a collecting society 230 | under any voluntary or waivable statutory or compulsory 231 | licensing scheme. In all other cases the Licensor expressly 232 | reserves any right to collect such royalties, including when 233 | the Licensed Material is used other than for NonCommercial 234 | purposes. 235 | 236 | 237 | Section 3 -- License Conditions. 238 | 239 | Your exercise of the Licensed Rights is expressly made subject to the 240 | following conditions. 241 | 242 | a. Attribution. 243 | 244 | 1. If You Share the Licensed Material (including in modified 245 | form), You must: 246 | 247 | a. retain the following if it is supplied by the Licensor 248 | with the Licensed Material: 249 | 250 | i. identification of the creator(s) of the Licensed 251 | Material and any others designated to receive 252 | attribution, in any reasonable manner requested by 253 | the Licensor (including by pseudonym if 254 | designated); 255 | 256 | ii. a copyright notice; 257 | 258 | iii. a notice that refers to this Public License; 259 | 260 | iv. a notice that refers to the disclaimer of 261 | warranties; 262 | 263 | v. a URI or hyperlink to the Licensed Material to the 264 | extent reasonably practicable; 265 | 266 | b. indicate if You modified the Licensed Material and 267 | retain an indication of any previous modifications; and 268 | 269 | c. indicate the Licensed Material is licensed under this 270 | Public License, and include the text of, or the URI or 271 | hyperlink to, this Public License. 272 | 273 | 2. You may satisfy the conditions in Section 3(a)(1) in any 274 | reasonable manner based on the medium, means, and context in 275 | which You Share the Licensed Material. For example, it may be 276 | reasonable to satisfy the conditions by providing a URI or 277 | hyperlink to a resource that includes the required 278 | information. 279 | 3. If requested by the Licensor, You must remove any of the 280 | information required by Section 3(a)(1)(A) to the extent 281 | reasonably practicable. 282 | 283 | b. ShareAlike. 284 | 285 | In addition to the conditions in Section 3(a), if You Share 286 | Adapted Material You produce, the following conditions also apply. 287 | 288 | 1. The Adapter's License You apply must be a Creative Commons 289 | license with the same License Elements, this version or 290 | later, or a BY-NC-SA Compatible License. 291 | 292 | 2. You must include the text of, or the URI or hyperlink to, the 293 | Adapter's License You apply. You may satisfy this condition 294 | in any reasonable manner based on the medium, means, and 295 | context in which You Share Adapted Material. 296 | 297 | 3. You may not offer or impose any additional or different terms 298 | or conditions on, or apply any Effective Technological 299 | Measures to, Adapted Material that restrict exercise of the 300 | rights granted under the Adapter's License You apply. 301 | 302 | 303 | Section 4 -- Sui Generis Database Rights. 304 | 305 | Where the Licensed Rights include Sui Generis Database Rights that 306 | apply to Your use of the Licensed Material: 307 | 308 | a. for the avoidance of doubt, Section 2(a)(1) grants You the right 309 | to extract, reuse, reproduce, and Share all or a substantial 310 | portion of the contents of the database for NonCommercial purposes 311 | only; 312 | 313 | b. if You include all or a substantial portion of the database 314 | contents in a database in which You have Sui Generis Database 315 | Rights, then the database in which You have Sui Generis Database 316 | Rights (but not its individual contents) is Adapted Material, 317 | including for purposes of Section 3(b); and 318 | 319 | c. You must comply with the conditions in Section 3(a) if You Share 320 | all or a substantial portion of the contents of the database. 321 | 322 | For the avoidance of doubt, this Section 4 supplements and does not 323 | replace Your obligations under this Public License where the Licensed 324 | Rights include other Copyright and Similar Rights. 325 | 326 | 327 | Section 5 -- Disclaimer of Warranties and Limitation of Liability. 328 | 329 | a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE 330 | EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS 331 | AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF 332 | ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS, 333 | IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION, 334 | WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR 335 | PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS, 336 | ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT 337 | KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT 338 | ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU. 339 | 340 | b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE 341 | TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION, 342 | NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT, 343 | INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES, 344 | COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR 345 | USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN 346 | ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR 347 | DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR 348 | IN PART, THIS LIMITATION MAY NOT APPLY TO YOU. 349 | 350 | c. The disclaimer of warranties and limitation of liability provided 351 | above shall be interpreted in a manner that, to the extent 352 | possible, most closely approximates an absolute disclaimer and 353 | waiver of all liability. 354 | 355 | 356 | Section 6 -- Term and Termination. 357 | 358 | a. This Public License applies for the term of the Copyright and 359 | Similar Rights licensed here. However, if You fail to comply with 360 | this Public License, then Your rights under this Public License 361 | terminate automatically. 362 | 363 | b. Where Your right to use the Licensed Material has terminated under 364 | Section 6(a), it reinstates: 365 | 366 | 1. automatically as of the date the violation is cured, provided 367 | it is cured within 30 days of Your discovery of the 368 | violation; or 369 | 370 | 2. upon express reinstatement by the Licensor. 371 | 372 | For the avoidance of doubt, this Section 6(b) does not affect any 373 | right the Licensor may have to seek remedies for Your violations 374 | of this Public License. 375 | 376 | c. For the avoidance of doubt, the Licensor may also offer the 377 | Licensed Material under separate terms or conditions or stop 378 | distributing the Licensed Material at any time; however, doing so 379 | will not terminate this Public License. 380 | 381 | d. Sections 1, 5, 6, 7, and 8 survive termination of this Public 382 | License. 383 | 384 | 385 | Section 7 -- Other Terms and Conditions. 386 | 387 | a. The Licensor shall not be bound by any additional or different 388 | terms or conditions communicated by You unless expressly agreed. 389 | 390 | b. Any arrangements, understandings, or agreements regarding the 391 | Licensed Material not stated herein are separate from and 392 | independent of the terms and conditions of this Public License. 393 | 394 | 395 | Section 8 -- Interpretation. 396 | 397 | a. For the avoidance of doubt, this Public License does not, and 398 | shall not be interpreted to, reduce, limit, restrict, or impose 399 | conditions on any use of the Licensed Material that could lawfully 400 | be made without permission under this Public License. 401 | 402 | b. To the extent possible, if any provision of this Public License is 403 | deemed unenforceable, it shall be automatically reformed to the 404 | minimum extent necessary to make it enforceable. If the provision 405 | cannot be reformed, it shall be severed from this Public License 406 | without affecting the enforceability of the remaining terms and 407 | conditions. 408 | 409 | c. No term or condition of this Public License will be waived and no 410 | failure to comply consented to unless expressly agreed to by the 411 | Licensor. 412 | 413 | d. Nothing in this Public License constitutes or may be interpreted 414 | as a limitation upon, or waiver of, any privileges and immunities 415 | that apply to the Licensor or You, including from the legal 416 | processes of any jurisdiction or authority. 417 | 418 | ======================================================================= 419 | 420 | Creative Commons is not a party to its public 421 | licenses. Notwithstanding, Creative Commons may elect to apply one of 422 | its public licenses to material it publishes and in those instances 423 | will be considered the “Licensor.” The text of the Creative Commons 424 | public licenses is dedicated to the public domain under the CC0 Public 425 | Domain Dedication. Except for the limited purpose of indicating that 426 | material is shared under a Creative Commons public license or as 427 | otherwise permitted by the Creative Commons policies published at 428 | creativecommons.org/policies, Creative Commons does not authorize the 429 | use of the trademark "Creative Commons" or any other trademark or logo 430 | of Creative Commons without its prior written consent including, 431 | without limitation, in connection with any unauthorized modifications 432 | to any of its public licenses or any other arrangements, 433 | understandings, or agreements concerning use of licensed material. For 434 | the avoidance of doubt, this paragraph does not form part of the 435 | public licenses. 436 | 437 | Creative Commons may be contacted at creativecommons.org. 438 | -------------------------------------------------------------------------------- /Model/losses.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch_geometric as tg 4 | from util import * 5 | from util_graph import get_program_ratio 6 | 7 | 8 | class VolumetricDesignLoss_D(nn.Module): 9 | def __init__(self, gan_loss, gp_lambda=10): 10 | super(VolumetricDesignLoss_D, self).__init__() 11 | self.gan_loss = gan_loss 12 | self.gp_lambda = gp_lambda 13 | 14 | def forward(self, real_validity_voxel, fake_validity_voxel, gp=None): 15 | if self.gan_loss == "WGANGP": 16 | Dv_loss = -torch.mean(real_validity_voxel[0]) - torch.mean(real_validity_voxel[1]) + torch.mean(fake_validity_voxel[0]) + torch.mean(fake_validity_voxel[1]) 17 | return Dv_loss + (self.gp_lambda * gp if gp is not None else 0) 18 | 19 | device = real_validity_voxel[0].get_device() 20 | valid0 = torch.FloatTensor(real_validity_voxel[0].shape[0], 1).fill_(1.0).to(device) 21 | valid1 = torch.FloatTensor(real_validity_voxel[1].shape[0], 1).fill_(1.0).to(device) 22 | fake0 = torch.FloatTensor(fake_validity_voxel[0].shape[0], 1).fill_(0.0).to(device) 23 | fake1 = torch.FloatTensor(fake_validity_voxel[1].shape[0], 1).fill_(0.0).to(device) 24 | 25 | if self.gan_loss == "NSGAN": # NS GAN log(D(x))+log(1-D(G(z))) 26 | loss = nn.BCELoss() 27 | return loss(real_validity_voxel[0], valid0) + loss(real_validity_voxel[1], valid1) + loss(fake_validity_voxel[0], fake0) + loss(fake_validity_voxel[1], fake1) 28 | elif self.gan_loss == "LSGAN": # LS GAN (D(x)-1)^2 + (D(G(z)))^2 29 | loss = nn.MSELoss() 30 | return 0.5 * (loss(real_validity_voxel[0], valid0) + loss(real_validity_voxel[1], valid1) + loss(fake_validity_voxel[0], fake0) + loss(fake_validity_voxel[1], fake1)) 31 | elif self.gan_loss == "hinge": # SA GAN 32 | loss = nn.ReLU() 33 | return loss(1.0 - real_validity_voxel[0]).mean() + loss(1.0 - real_validity_voxel[1]).mean() + loss(fake_validity_voxel[0] + 1.0).mean() + loss(fake_validity_voxel[1] + 1.0).mean() 34 | else: 35 | raise TypeError("self.gan_loss is not valid") 36 | 37 | 38 | class VolumetricDesignLoss_G(nn.Module): 39 | def __init__(self, lp_weight, tr_weight, far_weight, embedding_dim, sample_size, similarity_fun, gan_loss, lp_loss, hinge_margin): # hidden_dim = 128 40 | super(VolumetricDesignLoss_G, self).__init__() 41 | self.gan_loss = gan_loss 42 | self.tr_weight = tr_weight 43 | self.lp_weight = lp_weight 44 | self.far_weight = far_weight 45 | 46 | def forward(self, fake_validity_voxel, graph, att, mask, area_index_in_voxel_feature): 47 | device = att.get_device() if att.is_cuda else "cpu" 48 | target0 = torch.FloatTensor(fake_validity_voxel[0].shape[0], 1).fill_(1.0).to(device) 49 | target1 = torch.FloatTensor(fake_validity_voxel[1].shape[0], 1).fill_(1.0).to(device) 50 | 51 | # adversarial loss 52 | if self.gan_loss == "WGANGP": 53 | adversarial_loss_voxel = -torch.mean(fake_validity_voxel[0])-torch.mean(fake_validity_voxel[1]) 54 | adversarial_loss = adversarial_loss_voxel # + adversarial_loss_program 55 | elif self.gan_loss == "NSGAN": 56 | loss = nn.BCELoss() 57 | adversarial_loss = loss(fake_validity_voxel[0], target0) + loss(fake_validity_voxel[1], target1) 58 | elif self.gan_loss == "LSGAN": 59 | loss = nn.MSELoss() 60 | adversarial_loss = loss(fake_validity_voxel[0], target0) + loss(fake_validity_voxel[1], target1) 61 | elif self.gan_loss == "hinge": 62 | adversarial_loss = -torch.mean(fake_validity_voxel[0])-torch.mean(fake_validity_voxel[1]) 63 | 64 | # auxiliary loss 65 | normalized_program_class_weight, _, FAR = get_program_ratio(graph, att, mask, area_index_in_voxel_feature) 66 | target_ratio_loss = torch.Tensor([0]).to(device) if self.tr_weight == 0 else self.tr_weight * nn.functional.smooth_l1_loss(normalized_program_class_weight.flatten(), graph.program_target_ratio) 67 | far_loss = torch.Tensor([0]).to(device) if self.far_weight == 0 else self.far_weight * nn.functional.smooth_l1_loss(FAR.view(FAR.shape[0]), graph.FAR) 68 | total_loss = adversarial_loss + target_ratio_loss + far_loss 69 | 70 | return total_loss, adversarial_loss, target_ratio_loss, far_loss 71 | 72 | 73 | def compute_gradient_penalty(Dv, batch, label, out): 74 | # Interpolated sample 75 | device = out.get_device() 76 | u = torch.FloatTensor(label.shape[0], 1).uniform_(0, 1).to(device) # weight between model and gt label 77 | mixed_sample = torch.autograd.Variable(label * u + out * (1 - u), requires_grad=True).to(device) # Nv x C 78 | mask = (mixed_sample.max(dim=-1)[0] != 0).type(torch.float32).view(-1, 1) 79 | sample = softmax_to_hard(mixed_sample, -1) * mask 80 | 81 | # compute gradient penalty 82 | dv_loss = Dv(batch, sample) 83 | grad_b = torch.autograd.grad(outputs=dv_loss[0], inputs=sample, grad_outputs=torch.ones(dv_loss[0].shape).to(device), retain_graph=True, create_graph=True, only_inputs=True)[0] 84 | grad_s = torch.autograd.grad(outputs=dv_loss[1], inputs=sample, grad_outputs=torch.ones(dv_loss[1].shape).to(device), retain_graph=True, create_graph=True, only_inputs=True)[0] 85 | dv_gp_b = ((grad_b.norm(2, 1) - 1) ** 2).mean() 86 | dv_gp_s = ((grad_s.norm(2, 1) - 1) ** 2).mean() 87 | dv_gp = dv_gp_b + dv_gp_s 88 | 89 | return dv_gp, dv_gp_b, dv_gp_s 90 | 91 | 92 | -------------------------------------------------------------------------------- /Model/models.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch_geometric.nn import MessagePassing, inits 4 | import torch_geometric as tg 5 | from Data.GraphConstructor import GraphConstructor 6 | from torch_scatter import scatter, scatter_max 7 | from util import gumbel_softmax, softmax_to_hard 8 | from util_graph import find_max_out_program_index 9 | import math 10 | 11 | 12 | # ---- POSITION ENCODING ------------------------------- 13 | def position_encoder(d_model, max_len=20): 14 | pe = torch.zeros(max_len, d_model) 15 | position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) 16 | div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) 17 | pe[:, 0::2] = torch.sin(position * div_term) 18 | pe[:, 1::2] = torch.cos(position * div_term) 19 | return pe 20 | 21 | 22 | def get_voxel_floor_level(vfc, vbi): 23 | # floor ids are serialized in batch, need to subtract the accumulated num of node in each graph 24 | pool = -tg.nn.max_pool_x(cluster=vbi, x=-vfc, batch=torch.cuda.LongTensor([0] * vfc.shape[0]))[0] 25 | return vfc-pool.index_select(0, vbi) 26 | # --------------------------------------------------------- 27 | 28 | 29 | def MLP(dims, act=""): 30 | assert len(dims) >= 2 31 | nn_list = [] 32 | for i in range(len(dims)-1): 33 | nn_list.append(nn.Linear(dims[i], dims[i+1], bias=True)) 34 | if "leaky" in act: 35 | nn_list.append(nn.LeakyReLU(negative_slope=0.01)) 36 | elif "relu" in act: 37 | nn_list.append(nn.ReLU()) 38 | elif "tanh" in act: 39 | nn_list.append(nn.Tanh()) 40 | elif "sigmoid" in act: 41 | nn_list.append(nn.Sigmoid()) 42 | return nn.Sequential(*nn_list) 43 | 44 | 45 | def extract_pos(feature): 46 | """ Split position features and the rest""" 47 | pos = feature[:, 3:6] # 3, 4, 5 are position coordinates 48 | non_pos = torch.cat((feature[:, 0:3], feature[:, 6:]), dim=-1) 49 | return pos, non_pos 50 | 51 | 52 | class ProgramBlock(MessagePassing): 53 | def __init__(self, hidden_dim): 54 | super(ProgramBlock, self).__init__() 55 | self.aggr = 'mean' 56 | self.messageMLP = MLP([2 * hidden_dim, hidden_dim]) 57 | self.updateMLP = MLP([3 * hidden_dim, hidden_dim], act="leaky") 58 | 59 | def forward(self, feature, edge_index, class_index, class_feature, batch_index): 60 | """ 61 | feature=x, edge_index=graph.program_edge, class_index=graph.program_class_cluster, 62 | class_feature=graph.program_target_ratio, batch_index=graph.program_class_feature_batch 63 | """ 64 | kwargs = {"feature": feature, "class_index": class_index, "class_feature": class_feature, "batch_index": batch_index} 65 | return self.propagate(edge_index, size=None, **kwargs) 66 | 67 | def message(self, feature_j, feature_i=None): 68 | return self.messageMLP(torch.cat((feature_i, feature_j), dim=-1)) 69 | 70 | def update(self, aggr_out, feature=None, class_index=None, class_feature=None, batch_index=None): 71 | agg_class_feature = tg.nn.avg_pool_x(cluster=class_index, x=feature, batch=batch_index, size=GraphConstructor.number_of_class)[0] * class_feature.view(-1, 1) # r_{Cl(i)}c_{i} in equation 4 72 | pool_class_feature_on_node = agg_class_feature.index_select(0, class_index) 73 | return self.updateMLP(torch.cat((feature, aggr_out, pool_class_feature_on_node), dim=-1)) 74 | 75 | 76 | class ProgramGNN(nn.Module): 77 | def __init__(self, input_dim, hidden_dim, noise_dim, layer_num): # hidden_dim = 128 78 | super(ProgramGNN, self).__init__() 79 | self.layer_num = layer_num 80 | self.enc = MLP([input_dim + noise_dim, hidden_dim]) 81 | self.block = ProgramBlock(hidden_dim) 82 | 83 | def forward(self, cf, slf, e, cc, tr, cfb, z): 84 | """ 85 | cf=graph.program_class_feature, slf=graph.story_level_feature, e=graph.program_edge, cc=graph.program_class_cluster, 86 | tr=graph.program_target_ratio, cfb=graph.program_class_feature_batch, z=noise 87 | """ 88 | x = torch.cat([cf, slf.view(-1, 1), z], dim=-1) 89 | x = self.enc(x) 90 | for i in range(self.layer_num): 91 | x = x + self.block(feature=x, edge_index=e, class_index=cc, class_feature=tr, batch_index=cfb) 92 | return x 93 | 94 | 95 | class Attention(nn.Module): 96 | def __init__(self, hidden_dim): 97 | super(Attention, self).__init__() 98 | self.dec = MLP([hidden_dim, hidden_dim // 2, 2]) 99 | self.NN_v = nn.Linear(hidden_dim, hidden_dim) 100 | self.NN_p = nn.Linear(hidden_dim, hidden_dim) 101 | self.theta = nn.Parameter(torch.Tensor(1, hidden_dim)) 102 | inits.glorot(self.theta) 103 | 104 | def forward(self, program_graph_feature, voxel_feature, cross_edge_program_index, cross_edge_voxel_index): 105 | """ 106 | program_graph_feature=program_node_feature, program_class_feature=graph.program_class_feature, voxel_feature=voxel_node_feature, 107 | cross_edge_program_index=graph.cross_edge_program_index_select, cross_edge_voxel_index=graph.cross_edge_voxel_index_select 108 | """ 109 | soft_mask = torch.nn.functional.gumbel_softmax(self.dec(voxel_feature)) 110 | hard_mask = softmax_to_hard(soft_mask, dim=-1) 111 | out_mask = {"hard": hard_mask[:, 0].view(-1, 1), "soft": soft_mask[:, 0].view(-1, 1)} 112 | 113 | program_feature_j = program_graph_feature.index_select(0, cross_edge_program_index) 114 | voxel_feature_i = voxel_feature.index_select(0, cross_edge_voxel_index) 115 | # compute attention; equation 10 116 | att = (self.theta.view(1, -1) * torch.tanh(self.NN_v(voxel_feature_i) + self.NN_p(program_feature_j))).sum(-1) 117 | # equation 11 118 | soft_att, hard_att = gumbel_softmax(att, cross_edge_voxel_index) 119 | out_att = {"hard": hard_att.view(-1, 1), "soft": soft_att.view(-1, 1)} 120 | 121 | # equation 12 122 | weighted_program_feature = out_att["soft"] * program_graph_feature[cross_edge_program_index] 123 | new_voxel_feature = voxel_feature + out_mask["soft"] * scatter(weighted_program_feature, cross_edge_voxel_index, dim=0, dim_size=voxel_feature.shape[0], reduce="sum") 124 | 125 | return out_mask, out_att, new_voxel_feature 126 | 127 | @staticmethod 128 | def construct_output(program_class_feature, num_voxel, att, mask, cross_edge_program_index, cross_edge_voxel_index): 129 | # program node index that each voxel node has the max attention 130 | max_out_program_index = find_max_out_program_index(att["hard"] * mask["hard"].index_select(0, cross_edge_voxel_index), cross_edge_voxel_index, cross_edge_program_index, num_voxel) 131 | 132 | # If a voxel node is not masked, paste the selected program class feature (i.e. which class) to the voxel node as labels. 133 | hard_out_label = att["hard"] * program_class_feature[cross_edge_program_index] 134 | hard_out_label = scatter(hard_out_label, cross_edge_voxel_index, dim=0, dim_size=num_voxel, reduce="sum") 135 | masked_hard_label = mask["hard"] * hard_out_label 136 | 137 | soft_label = att["soft"] * program_class_feature[cross_edge_program_index] 138 | soft_label = scatter(soft_label, cross_edge_voxel_index, dim=0, dim_size=num_voxel, reduce="sum") 139 | masked_soft_label = mask["hard"] * soft_label 140 | 141 | return masked_hard_label, masked_soft_label, max_out_program_index.view(-1) 142 | 143 | 144 | class VoxelBlock(MessagePassing): 145 | def __init__(self, hidden_dim, if_agg_story_in_update=True): 146 | super(VoxelBlock, self).__init__() 147 | self.aggr = 'mean' 148 | self.if_agg_story_in_update = if_agg_story_in_update # True in VoxelGNN_G, False in VoxelGNN_D 149 | self.messageMLP = MLP([2 * hidden_dim + 3, hidden_dim]) 150 | self.updateMLP = MLP([4 * hidden_dim, hidden_dim], act="leaky") if if_agg_story_in_update else MLP([2 * hidden_dim, hidden_dim]) 151 | 152 | def forward(self, voxel_feature, edge_index, batch_index, position, floor_index): 153 | """ 154 | voxel_feature=v, edge_index=graph.voxel_edge, position=pos, floor_index=graph.voxel_floor_cluster, batch_index=graph.voxel_feature_batch 155 | cross_edge_voxel_index=graph.cross_edge_program_index_select, cross_edge_program_index=graph.cross_edge_voxel_index_select, 156 | """ 157 | kwargs = {"voxel_feature": voxel_feature, "batch_index": batch_index, "position": position, "floor_index": floor_index} 158 | return self.propagate(edge_index, size=None, **kwargs) 159 | 160 | def message(self, voxel_feature_j, voxel_feature_i=None, position_i=None, position_j=None): 161 | return self.messageMLP(torch.cat((voxel_feature_i, voxel_feature_j, position_i-position_j), dim=-1)) 162 | 163 | def update(self, aggr_out, voxel_feature=None, floor_index=None, batch_index=None): 164 | if self.if_agg_story_in_update: 165 | # concatenating story feature and building feature to the input of update MLP 166 | agg_story_feature, agg_story_feature_batch = tg.nn.avg_pool_x(cluster=floor_index, x=voxel_feature, batch=batch_index) # aggregate features on the same story 167 | agg_building_feature = tg.nn.avg_pool_x(cluster=agg_story_feature_batch, x=agg_story_feature, batch=torch.cuda.LongTensor([0] * agg_story_feature.shape[0]))[0] # aggregate all features in the graph 168 | agg_story_feature = agg_story_feature.index_select(0, floor_index) # populate the feature to each voxel node 169 | agg_building_feature = agg_building_feature.index_select(0, batch_index) # populate the feature to each voxel node 170 | return self.updateMLP(torch.cat((voxel_feature, aggr_out, agg_story_feature, agg_building_feature-agg_story_feature), dim=-1)) 171 | else: 172 | return self.updateMLP(torch.cat((voxel_feature, aggr_out), dim=-1)) 173 | 174 | 175 | class VoxelGNN_G(nn.Module): 176 | def __init__(self, input_dim, hidden_dim, noise_dim, layer_num): 177 | super(VoxelGNN_G, self).__init__() 178 | self.layer_num = layer_num 179 | self.noise_dim = noise_dim 180 | self.register_buffer('pe', position_encoder(hidden_dim, 20)) 181 | self.enc = MLP([input_dim - 3 + noise_dim, hidden_dim]) 182 | self.block = VoxelBlock(hidden_dim, if_agg_story_in_update=True) 183 | self.attention = Attention(hidden_dim) 184 | 185 | def forward(self, x, v, z, e, cep, cev, vbi, vfc, vfb): 186 | """ 187 | x=program_node_feature, v=graph.voxel_feature, e=graph.voxel_edge, cep=graph.cross_edge_program_index_select, 188 | cev=graph.cross_edge_voxel_index_select, vbi=graph.voxel_feature_batch, vfc=graph.voxel_floor_cluster, vfb=graph.voxel_feature_batch 189 | """ 190 | pos, v = extract_pos(v) 191 | v = self.enc(torch.cat((v, z), dim=-1)) if self.noise_dim != 0 else self.enc(v) 192 | v = v + self.pe[get_voxel_floor_level(vfc, vbi), :] 193 | 194 | _, _, v = self.attention(program_graph_feature=x, voxel_feature=v, cross_edge_program_index=cep, cross_edge_voxel_index=cev) 195 | 196 | for i in range(self.layer_num): 197 | # voxel GNN propagation 198 | v = v + self.block(voxel_feature=v, edge_index=e, position=pos, floor_index=vfc, batch_index=vfb) 199 | # pointer-based cross-modal module 200 | if i % 2 == 0 and i != 0: 201 | out_mask, out_att, v = self.attention(program_graph_feature=x, voxel_feature=v, cross_edge_program_index=cep, cross_edge_voxel_index=cev) 202 | 203 | return out_mask, out_att, v 204 | 205 | 206 | class VoxelGNN_D(nn.Module): 207 | def __init__(self, input_feature_dim, input_label_dim, hidden_dim, layer_num): 208 | super(VoxelGNN_D, self).__init__() 209 | self.layer_num = layer_num 210 | self.register_buffer('pe', position_encoder(hidden_dim, 20)) 211 | self.feature_enc = MLP([input_feature_dim-3, hidden_dim]) 212 | self.label_enc = MLP([input_label_dim, hidden_dim]) 213 | self.block = VoxelBlock(2*hidden_dim, if_agg_story_in_update=False) 214 | 215 | def forward(self, v, l, e, vbi, vfc, vfb): 216 | """ 217 | v=graph.voxel_feature, l=out_label or graph.voxel_label, 218 | e=graph.voxel_edge, vbi=graph.voxel_feature_batch, vfc=graph.voxel_floor_cluster, vfb=graph.voxel_feature_batch 219 | """ 220 | pos, v = extract_pos(v) 221 | v = self.feature_enc(v) 222 | v = v + self.pe[get_voxel_floor_level(vfc, vbi), :] 223 | 224 | l = self.label_enc(l) 225 | v = torch.cat([v, l], dim=-1) 226 | 227 | for i in range(self.layer_num): 228 | v = v + self.block(voxel_feature=v, edge_index=e, position=pos, floor_index=vfc, batch_index=vfb) 229 | return v 230 | 231 | 232 | class Generator(nn.Module): 233 | def __init__(self, program_input_dim, voxel_input_dim, hidden_dim, noise_dim, program_layer, voxel_layer, device): # hidden = 128 234 | super(Generator, self).__init__() 235 | self.device = device 236 | self.hidden_dim = hidden_dim 237 | self.noise_dim = noise_dim 238 | self.programGNN = ProgramGNN(program_input_dim, hidden_dim, noise_dim, program_layer) 239 | self.voxelGNN = VoxelGNN_G(voxel_input_dim, hidden_dim, 0, voxel_layer) 240 | 241 | def forward(self, graph, program_z, voxel_z): 242 | program_node_feature = self.programGNN(cf=graph.program_class_feature, slf=graph.story_level_feature, e=graph.program_edge, cc=graph.program_class_cluster, 243 | tr=graph.program_target_ratio, cfb=graph.program_class_feature_batch, z=program_z) 244 | mask, att, voxel_node_feature = self.voxelGNN(x=program_node_feature, v=graph.voxel_feature, z=voxel_z, e=graph.voxel_edge, cep=graph.cross_edge_program_index_select, 245 | cev=graph.cross_edge_voxel_index_select, vbi=graph.voxel_feature_batch, vfc=graph.voxel_floor_cluster, vfb=graph.voxel_feature_batch) 246 | 247 | out, soft_out, max_out_program_index = Attention.construct_output(program_class_feature=graph.program_class_feature, num_voxel=graph.voxel_feature.shape[0], att=att, mask=mask, 248 | cross_edge_program_index=graph.cross_edge_program_index_select, cross_edge_voxel_index=graph.cross_edge_voxel_index_select) 249 | 250 | return out, soft_out, mask, att, max_out_program_index 251 | 252 | 253 | class DiscriminatorVoxel(nn.Module): 254 | def __init__(self, voxel_input_feature_dim, voxel_input_label_dim, hidden_dim, voxel_layer, act=""): 255 | super(DiscriminatorVoxel, self).__init__() 256 | self.voxelGNN_D = VoxelGNN_D(voxel_input_feature_dim, voxel_input_label_dim, hidden_dim, voxel_layer) 257 | self.building_dec = MLP([2 * hidden_dim, hidden_dim, 1], act) 258 | self.story_dec = MLP([2 * hidden_dim, hidden_dim, 1], act) 259 | 260 | def forward(self, graph, out_label): 261 | voxel_node_feature = self.voxelGNN_D(v=graph.voxel_feature, l=out_label, e=graph.voxel_edge, vbi=graph.voxel_feature_batch, vfc=graph.voxel_floor_cluster, vfb=graph.voxel_feature_batch) 262 | story_feature = tg.nn.global_max_pool(voxel_node_feature, graph.voxel_floor_cluster) 263 | graph_feature = tg.nn.global_max_pool(voxel_node_feature, graph.voxel_feature_batch.to(out_label.get_device())) 264 | return self.building_dec(graph_feature), self.story_dec(story_feature) 265 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Building-GAN 2 | ====== 3 | 4 | Code and instructions for our paper: 5 | 6 | [Building-GAN: Graph-Conditioned Architectural Volumetric Design Generation](https://arxiv.org/abs/2104.13316), ICCV 2021. 7 | 8 | 9 | 10 | Volumetric Design Process 11 | ----- 12 | 13 | 14 | Data 15 | ------ 16 | - Download the dataset [here](https://d271velnk8wqmg.cloudfront.net/6types-raw_data.zip). 17 | - Put the subfolders and files in `raw-data` under the folder `6types-raw_data`. 18 | - Run `Data/process_data.py` to process the raw data. 19 | 20 | For the detail about how the raw data are processed, please refer the `Data/process_data.py`. 21 | 22 | In the dataset, each volumetric design comprises three json files: 23 | - Global Graph: contains the FAR, program ratios, and the associated rooms for each program type. 24 | - Local Graph: contains the bubble diagram--the type and size of each room and the connectivity between rooms 25 | - Voxel: contains the voxel graph 26 | 27 | Running pretrained models 28 | ------ 29 | 30 | For running a pre-trained model, please follow the steps below: 31 | - The pre-trained model is located at `runs/iccv2021/checkpoints/` 32 | - Run ```python inference.py``` 33 | - Check out the results in the `inference/{model}/{epch_current_time}/output` folder. 34 | - Check out the variation results from the same program graph in the `inference/{model}/{epch_current_time}/var_output*` folders. 35 | 36 | Training models 37 | ------ 38 | 39 | For training a model from scratch, please follow the steps below: 40 | - Follow the steps in Data section. 41 | - run ```python train.py ```. Customized arguments can be set according to ```train_args.py```. 42 | - Check out ```output``` and ```checkpoints``` folders for intermediate outputs and checkpoints, respectively. They are under the ```runs/run_id/``` where run_id is the serial number of the 43 | experiment. 44 | 45 | Requirements 46 | ------ 47 | - PyTorch >= 1.7.0 48 | - PyTorch Geometric 1.6.2 49 | 50 | Citation 51 | ------ 52 | ``` 53 | @article{chang2021building, 54 | title={Building-GAN: Graph-Conditioned Architectural Volumetric Design Generation}, 55 | author={Chang, Kai-Hung and Cheng, Chin-Yi and Luo, Jieliang and Murata, Shingo and Nourbakhsh, Mehdi and Tsuji, Yoshito}, 56 | booktitle={International Conference on Computer Vision}, 57 | year={2021} 58 | } 59 | ``` 60 | 61 | Contact 62 | ------ 63 | Unfortunately this repo is no longer actively maintained. 64 | If you have any question, feel free to contact Chin-Yi Cheng @chinyich or Kai-Hung Chang @kaihungc1993 65 | 66 | 67 | ## License 68 | Shield: [![CC BY-NC-SA 4.0][cc-by-nc-sa-shield]][cc-by-nc-sa] 69 | 70 | This work is licensed under a 71 | [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License][cc-by-nc-sa]. 72 | 73 | [![CC BY-NC-SA 4.0][cc-by-nc-sa-image]][cc-by-nc-sa] 74 | 75 | [cc-by-nc-sa]: http://creativecommons.org/licenses/by-nc-sa/4.0/ 76 | [cc-by-nc-sa-image]: https://licensebuttons.net/l/by-nc-sa/4.0/88x31.png 77 | [cc-by-nc-sa-shield]: https://img.shields.io/badge/License-CC%20BY--NC--SA%204.0-lightgrey.svg 78 | -------------------------------------------------------------------------------- /images/Building-GAN.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutodeskAILab/Building-GAN/bb8948cf1f80504a35c09de9d8593e45b35e6f5b/images/Building-GAN.png -------------------------------------------------------------------------------- /images/Volumetric_Design.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutodeskAILab/Building-GAN/bb8948cf1f80504a35c09de9d8593e45b35e6f5b/images/Volumetric_Design.png -------------------------------------------------------------------------------- /inference.py: -------------------------------------------------------------------------------- 1 | from train_args import make_args 2 | args = make_args() 3 | 4 | gpu_pci_str = str(args.cuda) # single-gpu version 5 | import os 6 | gpu_pci_ids = [int(i) for i in gpu_pci_str.split(',')] 7 | cuda = len(gpu_pci_ids) > 0 8 | if_multi_gpu = len(gpu_pci_ids) > 1 9 | if cuda: 10 | os.environ['CUDA_VISIBLE_DEVICES'] = gpu_pci_str 11 | 12 | from datetime import datetime 13 | current_time = datetime.now().strftime('%b%d_%H-%M-%S') 14 | 15 | from torch_geometric.data import DataLoader, Batch 16 | 17 | from Model.models import Generator 18 | from util_eval import * 19 | from util_graph import * 20 | 21 | 22 | device_ids = list(range(torch.cuda.device_count())) if cuda else [] 23 | assert(args.batch_size % len(device_ids) == 0) 24 | print('Using GPU {}'.format(os.environ['CUDA_VISIBLE_DEVICES']) if cuda else "Using CPU") 25 | if cuda: 26 | print([torch.cuda.get_device_name(device_id) for device_id in device_ids]) 27 | print(args) 28 | device = torch.device('cuda:'+str(device_ids[0]) if cuda else 'cpu') 29 | 30 | inference_dir = "inference" 31 | trained_id = "iccv2021" 32 | epoch_id = '70' 33 | 34 | trained_file = 'runs/{}/checkpoints/{}.ckpt'.format(trained_id, epoch_id) 35 | 36 | viz_dir = os.path.join(inference_dir, trained_id, epoch_id + "_" + current_time, "output") 37 | var_viz_dir1 = os.path.join(inference_dir, trained_id, epoch_id + "_" + current_time, "var_output1") 38 | var_viz_dir2 = os.path.join(inference_dir, trained_id, epoch_id + "_" + current_time, "var_output2") 39 | 40 | mkdirs = [inference_dir, os.path.join(inference_dir, trained_id), os.path.join(inference_dir, trained_id, epoch_id + "_" + current_time), viz_dir, var_viz_dir1, var_viz_dir2] 41 | for mkdir in mkdirs: 42 | if not os.path.exists(mkdir): 43 | os.mkdir(mkdir) 44 | 45 | truncated = False 46 | if not truncated: 47 | trunc_num = 1.0 48 | else: 49 | trunc_num = 0.7 50 | 51 | # Load Data 52 | follow_batch = ['program_target_ratio', 'program_class_feature', 'voxel_feature'] 53 | train_size = args.train_size//args.batch_size * args.batch_size 54 | test_size = args.test_size//args.batch_size * args.batch_size 55 | 56 | data_fname_list = sorted(os.listdir(args.train_data_dir)) 57 | print("Total %d data: %d train / %d test" % (len(data_fname_list), train_size, test_size)) 58 | 59 | variation_test_data1 = torch.load(os.path.join(args.train_data_dir, data_fname_list[args.variation_eval_id1])) 60 | variation_test_data2 = torch.load(os.path.join(args.train_data_dir, data_fname_list[args.variation_eval_id2])) 61 | variation_test_batch1 = Batch.from_data_list([variation_test_data1 for _ in range(args.variation_num)], follow_batch) 62 | variation_test_batch2 = Batch.from_data_list([variation_test_data2 for _ in range(args.variation_num)], follow_batch) 63 | 64 | test_data_list = [] 65 | for fname in data_fname_list[train_size:train_size + test_size]: 66 | test_data_list.append(torch.load(os.path.join(args.train_data_dir, fname))) 67 | 68 | test_data_loader = DataLoader(test_data_list, follow_batch=follow_batch, batch_size=args.batch_size, shuffle=False, num_workers=args.n_cpu) 69 | program_input_dim = test_data_list[0].program_class_feature.size(-1) + 1 # data_list[0].story_level_feature.size(-1) 70 | voxel_input_dim = test_data_list[0].voxel_feature.size(-1) 71 | voxel_label_dim = test_data_list[0].voxel_label.size(-1) 72 | 73 | # Load Model 74 | generator = Generator(program_input_dim, voxel_input_dim, args.latent_dim, args.noise_dim, args.program_layer, args.voxel_layer, device).to(device) 75 | generator.load_state_dict(torch.load(trained_file), strict=False) 76 | generator.eval() 77 | 78 | # evaluate 79 | n_batches = 20 # total number of generated samples = n_batches * args.batch_size 80 | evaluate(test_data_loader, generator, args.raw_dir, viz_dir, follow_batch, device_ids, number_of_batches=n_batches,trunc=trunc_num) 81 | generate_multiple_outputs_from_batch(variation_test_batch1, args.variation_num, generator, args.raw_dir, var_viz_dir1, follow_batch, device_ids, trunc=trunc_num) 82 | generate_multiple_outputs_from_batch(variation_test_batch2, args.variation_num, generator, args.raw_dir, var_viz_dir2, follow_batch, device_ids, trunc=trunc_num) 83 | -------------------------------------------------------------------------------- /runs/iccv2021/checkpoints/70.ckpt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutodeskAILab/Building-GAN/bb8948cf1f80504a35c09de9d8593e45b35e6f5b/runs/iccv2021/checkpoints/70.ckpt -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | from train_args import make_args 2 | import os 3 | args = make_args() 4 | if str(args.cuda) == -1: 5 | cuda = False 6 | else: 7 | cuda = True 8 | os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda) 9 | 10 | import torch 11 | from datetime import datetime 12 | from tensorboardX import SummaryWriter 13 | from torch_geometric.data import DataLoader, Batch 14 | from Model.models import Generator, DiscriminatorVoxel 15 | from Model.losses import VolumetricDesignLoss_D, VolumetricDesignLoss_G, compute_gradient_penalty 16 | from util_eval import * 17 | from util_graph import * 18 | from Data.LargeFilenameDataset import LargeFilenameDataset 19 | 20 | current_time = datetime.now().strftime('%b%d_%H-%M-%S') 21 | run_id = current_time + "_" + args.comment 22 | device_ids = list(range(torch.cuda.device_count())) if cuda else [] 23 | device = torch.device('cuda:'+str(device_ids[0]) if cuda else 'cpu') 24 | print('Using GPU {}'.format(os.environ['CUDA_VISIBLE_DEVICES']) if cuda else "Using CPU") 25 | print([torch.cuda.get_device_name(device_id) for device_id in device_ids] if cuda else "Using CPU") 26 | print(run_id) 27 | print(args) 28 | 29 | 30 | run_dir = "runs" 31 | viz_dir = os.path.join(run_dir, run_id, "output") 32 | var_viz_dir1 = os.path.join(run_dir, run_id, "var_output1") 33 | var_viz_dir2 = os.path.join(run_dir, run_id, "var_output2") 34 | model_dir = os.path.join(run_dir, run_id, "checkpoints") 35 | mkdirs = [os.path.join(run_dir, run_id), viz_dir, var_viz_dir1, var_viz_dir2, model_dir] 36 | for mkdir in mkdirs: 37 | os.mkdir(mkdir) 38 | writer_train = SummaryWriter(log_dir=os.path.join(run_dir, run_id)) 39 | 40 | 41 | # Load Data 42 | data_fname_list = sorted(os.listdir(args.train_data_dir)) 43 | follow_batch = ['program_target_ratio', 'program_class_feature', 'voxel_feature'] 44 | train_size = args.train_size//args.batch_size * args.batch_size 45 | test_size = args.test_size//args.batch_size * args.batch_size 46 | # print("Total %d data: %d train / %d test" % (len(data_fname_list), train_size, len(data_fname_list)-train_size)) 47 | print("Total %d data: %d train / %d test" % (len(data_fname_list), train_size, test_size)) 48 | 49 | train_data_list = LargeFilenameDataset(args.train_data_dir, data_fname_list[:train_size]) 50 | train_data_loader = DataLoader(train_data_list, follow_batch=follow_batch, batch_size=args.batch_size, shuffle=True, num_workers=args.n_cpu) 51 | test_data_list = [torch.load(os.path.join(args.train_data_dir, fname)) for fname in data_fname_list[train_size:train_size + test_size]] 52 | test_data_loader = DataLoader(test_data_list, follow_batch=follow_batch, batch_size=20, shuffle=False, num_workers=args.n_cpu) 53 | 54 | variation_test_data1 = torch.load(os.path.join(args.train_data_dir, data_fname_list[args.variation_eval_id1])) 55 | variation_test_data2 = torch.load(os.path.join(args.train_data_dir, data_fname_list[args.variation_eval_id2])) 56 | variation_test_batch1 = Batch.from_data_list([variation_test_data1 for _ in range(args.variation_num)], follow_batch) 57 | variation_test_batch2 = Batch.from_data_list([variation_test_data2 for _ in range(args.variation_num)], follow_batch) 58 | 59 | # Load Model 60 | program_input_dim = test_data_list[0].program_class_feature.size(-1) + 1 61 | voxel_input_dim = test_data_list[0].voxel_feature.size(-1) 62 | voxel_label_dim = test_data_list[0].voxel_label.size(-1) 63 | 64 | generator = Generator(program_input_dim, voxel_input_dim, args.latent_dim, args.noise_dim, args.program_layer, args.voxel_layer, device).to(device) 65 | discriminator = DiscriminatorVoxel(voxel_input_dim, voxel_label_dim, args.latent_dim, args.voxel_layer, act="sigmoid" if args.gan_loss == "NSGAN" else "").to(device) 66 | d_loss_func = VolumetricDesignLoss_D(gan_loss=args.gan_loss, gp_lambda=args.gp_lambda).to(device) 67 | g_loss_func = VolumetricDesignLoss_G(lp_weight=args.lp_weight, tr_weight=args.tr_weight, far_weight=args.far_weight, embedding_dim=args.latent_dim, sample_size=args.lp_sample_size, similarity_fun=args.lp_similarity_fun, 68 | gan_loss=args.gan_loss, lp_loss=args.lp_loss_fun, hinge_margin=args.lp_hinge_margin).to(device) 69 | optimizer_G = torch.optim.Adam(generator.parameters(), lr=args.g_lr, betas=(args.b1, args.b2)) 70 | optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=args.d_lr, betas=(args.b1, args.b2)) 71 | 72 | 73 | for epoch in range(args.n_epochs): 74 | for i, batch in enumerate(train_data_loader): 75 | batch = batch.to(device) 76 | 77 | # --------------------- 78 | # Train Discriminator 79 | # --------------------- 80 | for p in list(discriminator.parameters()): 81 | p.requires_grad = True 82 | optimizer_D.zero_grad() 83 | 84 | # random variables 85 | program_z = torch.rand(tuple([batch.program_class_feature.shape[0], args.noise_dim])).to(device) 86 | voxel_z = torch.rand(tuple([batch.voxel_feature.shape[0], args.noise_dim])).to(device) 87 | out, soft_out, mask, att, max_out_program_index = generator(batch, program_z, voxel_z) 88 | 89 | if i % args.n_critic_d == 0: 90 | fake_validity_voxel = discriminator(detach_batch(batch), out.detach()) 91 | real_validity_voxel = discriminator(batch, batch.voxel_label) 92 | 93 | if args.gan_loss == "WGANGP" and d_loss_func.gp_lambda != 0: 94 | gp, gp_b, gp_s = compute_gradient_penalty(discriminator, detach_batch(batch), batch.voxel_label.data, soft_out.data) 95 | else: 96 | gp, gp_b, gp_s = torch.tensor(0, device=device), torch.tensor(0, device=device), torch.tensor(0, device=device) 97 | 98 | d_loss = d_loss_func(real_validity_voxel, fake_validity_voxel, gp) 99 | d_loss.backward() 100 | if i % args.n_critic_d == 0: 101 | optimizer_D.step() 102 | 103 | # ----------------- 104 | # Train Generator 105 | # ----------------- 106 | for p in list(discriminator.parameters()): 107 | p.requires_grad = False 108 | optimizer_G.zero_grad() 109 | 110 | if i % args.n_critic_g == 0: 111 | validity_G_voxel = discriminator(batch, out) 112 | g_loss, gan_loss, tr_loss, far_loss = g_loss_func(validity_G_voxel, batch, att["hard"], mask["hard"], area_index_in_voxel_feature=6) 113 | g_loss.backward() 114 | optimizer_G.step() 115 | 116 | if epoch % args.plot_period == 0: 117 | avg_fake_validity_b, avg_fake_validity_s = torch.mean(fake_validity_voxel[0]), torch.mean(fake_validity_voxel[1]) 118 | avg_real_validity_b, avg_real_validity_s = torch.mean(real_validity_voxel[0]), torch.mean(real_validity_voxel[1]) 119 | print("[Epoch %d/%d] [Batch %d/%d] [D loss: %.4f, gp: %.4f] [G loss: %.4f (%.4f, %.4f, %.4f)] [Validity Real: (%.4f, %.4f), Fake (%.4f, %.4f)]" % 120 | (epoch, args.n_epochs, i + 1, len(train_data_loader), d_loss.item(), gp.item(), g_loss.item(), gan_loss.item(), tr_loss.item(), far_loss.item(), 121 | avg_real_validity_b.item(), avg_real_validity_s.item(), avg_fake_validity_b.item(), avg_fake_validity_s.item())) 122 | writer_train.add_scalar('d_loss', d_loss, epoch) 123 | writer_train.add_scalar('_gp', gp.item(), epoch) 124 | writer_train.add_scalar('g_loss', g_loss, epoch) 125 | writer_train.add_scalar('gan_loss', gan_loss, epoch) 126 | writer_train.add_scalar('target_ratio_loss', tr_loss, epoch) 127 | writer_train.add_scalar('D(fake)', avg_fake_validity_b, epoch) 128 | writer_train.add_scalar('D(real)', avg_real_validity_b, epoch) 129 | writer_train.add_scalar('D(fake)_story', avg_fake_validity_s, epoch) 130 | writer_train.add_scalar('D(real)_story', avg_real_validity_s, epoch) 131 | writer_train.close() 132 | 133 | if epoch % args.eval_period == 0: 134 | os.mkdir(os.path.join(viz_dir, str(epoch))) 135 | os.mkdir(os.path.join(var_viz_dir1, str(epoch))) 136 | os.mkdir(os.path.join(var_viz_dir2, str(epoch))) 137 | evaluate(test_data_loader, generator, args.raw_dir, os.path.join(viz_dir, str(epoch)), follow_batch, device_ids, number_of_batches=1) 138 | generate_multiple_outputs_from_batch(variation_test_batch1, args.variation_num, generator, args.raw_dir, os.path.join(var_viz_dir1, str(epoch)), follow_batch, device_ids) 139 | generate_multiple_outputs_from_batch(variation_test_batch2, args.variation_num, generator, args.raw_dir, os.path.join(var_viz_dir2, str(epoch)), follow_batch, device_ids) 140 | save_model = os.path.join(model_dir, '{}.ckpt'.format(epoch)) 141 | torch.save(generator.state_dict(), save_model) 142 | 143 | -------------------------------------------------------------------------------- /train_args.py: -------------------------------------------------------------------------------- 1 | from argparse import ArgumentParser 2 | 3 | 4 | def add_bool_argument(arg_parser, true_arg, false_arg, dest_var, help_str): 5 | arg_parser.add_argument(true_arg, dest=dest_var, action='store_true', help=help_str) 6 | arg_parser.add_argument(false_arg, dest=dest_var, action='store_false', help=help_str) 7 | 8 | 9 | def make_args(): 10 | parser = ArgumentParser() 11 | parser.add_argument('--cuda', dest='cuda', default='0', type=str) # -1 if not using GPU 12 | parser.add_argument('--comment', dest='comment', default='0', type=str, help='comment') 13 | 14 | # Data related 15 | parser.add_argument('--batch_size', default=8, type=int, help='size of graph batching') 16 | # add_bool_argument(parser, "--if_curriculum", "--if_curriculum_no", dest_var="if_curriculum", help_str="if use curriculum") 17 | parser.add_argument('--train_data_dir', default='Data/6types-processed_data', type=str, help='where to load the training data if not preload') 18 | parser.add_argument('--raw_dir', default='Data/6types-raw_data', type=str, help='For evaluation to copy and paste') # normal runs 19 | # parser.add_argument('--preload_dir', default='Data/6types_preload_data.pkl', type=str, help='where to load the training data') 20 | 21 | # 96000, 96005, 96003 22 | parser.add_argument('--train_size', default=96000, type=int, help='how many data are train data') # 23 | parser.add_argument('--test_size', default=4000, type=int, help='how many data are test data') # 24 | parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation") # 8, if RAM OOM -> decrease to 4 25 | parser.add_argument("--variation_eval_id1", type=int, default=96018, help="Data index for variation test") # 96008 26 | parser.add_argument("--variation_eval_id2", type=int, default=96010, help="Data index for variation test") 27 | parser.add_argument("--variation_num", type=int, default=25, help="How many variation to generate for the variation test data") 28 | 29 | # ------------------------------------------------------------------------------------------------- 30 | 31 | # Model related 32 | parser.add_argument("--latent_dim", type=int, default=128, help="dimensionality of the latent space") 33 | parser.add_argument("--noise_dim", type=int, default=32, help="dimensionality of the noise space") 34 | parser.add_argument("--program_layer", type=int, default=4, help="numbers of message passing in program graph") 35 | parser.add_argument("--voxel_layer", type=int, default=12, help="numbers of message passing in voxel graph") 36 | 37 | # Loss related 38 | parser.add_argument('--gan_loss', default='WGANGP', type=str, help='WGANGP/NSGAN/LSGAN/hinge') 39 | parser.add_argument("--gp_lambda", type=float, default=10.0, help="gradient penalty coefficient in D loss") 40 | parser.add_argument("--lp_weight", type=float, default=0.0, help="link prediction coefficient in G loss") 41 | parser.add_argument("--tr_weight", type=float, default=0.0, help="program target ratio coefficient in G loss") 42 | parser.add_argument("--far_weight", type=float, default=0.0, help="far coefficient in G loss") 43 | 44 | # Link prediction related 45 | # link prediction is not used in this implementation 46 | parser.add_argument("--lp_sample_size", type=int, default=20, help="link prediction sample size") 47 | parser.add_argument('--lp_similarity_fun', default='cos', type=str, help='link prediction similarity type: cos/dot/l2/mlp') 48 | parser.add_argument('--lp_loss_fun', default='hinge', type=str, help='link prediction loss type: hinge/BCE/skipgram') 49 | parser.add_argument("--lp_hinge_margin", type=float, default=1.0, help="link prediction hinge loss margin") 50 | 51 | # Training parameter 52 | parser.add_argument("--n_epochs", type=int, default=1000, help="number of epochs of training") 53 | parser.add_argument("--n_critic_d", type=int, default=1, help="number of training steps for discriminator per iter") 54 | parser.add_argument("--n_critic_g", type=int, default=5, help="number of training steps for discriminator per iter") 55 | parser.add_argument("--n_critic_p", type=int, default=5, help="number of training steps for discriminator per iter") 56 | parser.add_argument("--plot_period", type=int, default=10, help="number of epochs between each result plot") # 10 57 | parser.add_argument("--eval_period", type=int, default=20, help="number of epochs between each evaluation and save model") # 20 58 | 59 | # Optimizer parameter 60 | parser.add_argument("--g_lr", type=float, default=0.0001, help="adam: learning rate") 61 | parser.add_argument("--d_lr", type=float, default=0.0001, help="adam: learning rate") 62 | parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient") 63 | parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient") 64 | 65 | parser.set_defaults(if_curriculum=False) 66 | args = parser.parse_args() 67 | return args 68 | -------------------------------------------------------------------------------- /util.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch_scatter import scatter_max, scatter_add 3 | from torch_geometric.utils.num_nodes import maybe_num_nodes 4 | 5 | 6 | def softmax_to_hard(y_soft, dim): 7 | """the function uses the hard version trick as in gumbel softmax""" 8 | index = y_soft.max(dim, keepdim=True)[1] 9 | y_hard = torch.zeros_like(y_soft, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0) 10 | return y_hard - y_soft.detach() + y_soft 11 | 12 | 13 | def gumbel_softmax(src, index, tau=-1, num_nodes=None): 14 | """modified from torch_geometric.utils import softmax""" 15 | num_nodes = maybe_num_nodes(index, num_nodes) 16 | 17 | gumbels = -torch.empty_like(src, memory_format=torch.legacy_contiguous_format).exponential_().log() # ~Gumbel(0,1) 18 | gumbels = (src + gumbels) / tau # ~Gumbel(logits,tau) 19 | 20 | out = gumbels - scatter_max(gumbels, index, dim=0, dim_size=num_nodes)[0][index] 21 | out = out.exp() 22 | out = out / (scatter_add(out, index, dim=0, dim_size=num_nodes)[index] + 1e-16) 23 | 24 | argmax = scatter_max(out, index, dim=0, dim_size=num_nodes)[1] 25 | out_hard = torch.zeros_like(out, memory_format=torch.legacy_contiguous_format).scatter_(0, argmax, 1.0) 26 | out_hard = out_hard - out.detach() + out 27 | 28 | return out, out_hard 29 | 30 | 31 | -------------------------------------------------------------------------------- /util_eval.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import os 3 | import json 4 | from Data.GraphConstructor import GraphConstructor 5 | from torch_geometric.data import Batch 6 | from util_graph import get_program_ratio, data_parallel, rebatch_for_multi_gpu 7 | 8 | 9 | def save_output(batch_size, batch, class_weights, program_weights, FAR, max_out_program_index, out, follow_batch, raw_dir, output_dir, new_data_id_strs=None): 10 | """ 11 | Save the evaluation results 12 | """ 13 | if not os.path.exists(os.path.join(output_dir, "global_graph_data")): 14 | os.mkdir(os.path.join(output_dir, "global_graph_data")) 15 | os.mkdir(os.path.join(output_dir, "local_graph_data")) 16 | os.mkdir(os.path.join(output_dir, "voxel_data")) 17 | 18 | num_of_program_node_accum = 0 19 | batch_all_g = [] 20 | data = rebatch_for_multi_gpu(batch, list(range(batch_size)), follow_batch, out, class_weights, program_weights, FAR, max_out_program_index) 21 | 22 | """ 23 | --- data --- 24 | g: graph 25 | o: voxel label (n[type]) 26 | cw: (n[new_proportion] in global graph) -- program class ratio/weight 27 | pw: (n[region_far] in local graph) 28 | far: (g[far] in global graph) 29 | pid: the selected program node id for each voxel node 30 | """ 31 | 32 | for i, (g, o, cw, pw, far, pid) in enumerate(data): 33 | data_id_str = g["data_id_str"][0] 34 | new_data_id_str = g["data_id_str"][0] if new_data_id_strs is None else str(new_data_id_strs[i]).zfill(GraphConstructor.data_id_length) 35 | o = o.cpu().data.numpy().tolist() 36 | cw, pw, far = cw.cpu().data.numpy(), pw.cpu().data.numpy(), far.item() 37 | 38 | # Modify Global data 39 | with open(os.path.join(raw_dir, "global_graph_data", GraphConstructor.global_graph_prefix + data_id_str + ".json")) as f: 40 | global_graph = json.load(f) 41 | global_graph["new_far"] = far 42 | for n in global_graph["global_node"]: 43 | n["new_proportion"] = float(cw[n['type']]) 44 | with open(os.path.join(output_dir, "global_graph_data", GraphConstructor.global_graph_prefix + new_data_id_str + ".json"), 'w') as f: 45 | json.dump(global_graph, f) 46 | 47 | # Modify Local data 48 | d = {} # program id to its type and type id 49 | with open(os.path.join(raw_dir, "local_graph_data", GraphConstructor.local_graph_prefix + data_id_str + ".json")) as f: 50 | local_graph = json.load(f) 51 | for i, (n, c) in enumerate(zip(local_graph["node"], pw)): 52 | n["region_far"] = float(c) 53 | d[i] = [n["type"], n["type_id"]] 54 | with open(os.path.join(output_dir, "local_graph_data", GraphConstructor.local_graph_prefix + new_data_id_str + ".json"), 'w') as f: 55 | json.dump(local_graph, f) 56 | 57 | # Modify Voxel data 58 | with open(os.path.join(raw_dir, "voxel_data", GraphConstructor.voxel_graph_prefix + data_id_str + ".json")) as f: 59 | voxel_graph = json.load(f) 60 | for n, label, _pid in zip(voxel_graph["voxel_node"], o, pid): 61 | query = d[_pid.item() - num_of_program_node_accum] 62 | n["type"] = query[0] if 1.0 in label else -1 # # label.index(1.0) 63 | n["type_id"] = query[1] if 1.0 in label else 0 64 | num_of_program_node_accum += pw.shape[0] 65 | with open(os.path.join(output_dir, "voxel_data", GraphConstructor.voxel_graph_prefix + new_data_id_str + ".json"),'w') as f: 66 | json.dump(voxel_graph, f) 67 | 68 | all_graphs = [global_graph, local_graph, voxel_graph] 69 | batch_all_g.append(all_graphs) 70 | return batch_all_g 71 | 72 | 73 | def evaluate(data_loader, generator, raw_dir, output_dir, follow_batch, device_ids, number_of_batches=0, trunc=1.0): 74 | number_of_batches = min(number_of_batches, len(data_loader)) 75 | device = device_ids[0] 76 | with torch.no_grad(): 77 | total_inter, total_program_edge = 0, 0 78 | for i, g in enumerate(data_loader): 79 | if i >= number_of_batches: 80 | break 81 | program_z_shape = [g.program_class_feature.shape[0], generator.noise_dim] 82 | program_z = torch.rand(tuple(program_z_shape)).to(device) 83 | voxel_z_shape = [g.voxel_feature.shape[0], generator.noise_dim] 84 | voxel_z = torch.rand(tuple(voxel_z_shape)).to(device) 85 | if trunc < 1.0: 86 | program_z.clamp_(min=-trunc, max=trunc) 87 | voxel_z.clamp_(min=-trunc, max=trunc) 88 | 89 | g.to(device) 90 | out, soft_out, mask, att, max_out_program_index = generator(g, program_z, voxel_z) 91 | inter_edges, missing_edges, gen_edges = check_connectivity(g, max_out_program_index, mask['hard']) 92 | total_inter += inter_edges.shape[1] 93 | total_program_edge += g.program_edge.shape[1] 94 | normalized_program_class_weight, normalized_program_weight, FAR = get_program_ratio(g, att["hard"], mask["hard"], area_index_in_voxel_feature=6) 95 | all_g = save_output(data_loader.batch_size, g, normalized_program_class_weight, normalized_program_weight,FAR, max_out_program_index, out, follow_batch, raw_dir, output_dir) 96 | 97 | acc = total_inter/total_program_edge 98 | print('acc=', acc) 99 | return all_g 100 | 101 | 102 | def check_connectivity(g, max_out_program_index, mask): 103 | """ 104 | Extract connectivity from the generated design 105 | inter_edges: program edge observed in the generated output 106 | missing_edges: program edges only in the input program graph 107 | gen_edges: program edges only in the generated output 108 | """ 109 | # Look at the g.voxel_edge and see if the two voxel nodes are masked 110 | voxel_edge_out_mask = mask.reshape([-1])[g.voxel_edge] # Ev x 2 111 | sums = torch.sum(voxel_edge_out_mask, dim=0) # Ev x 1 112 | masked_edges = g.voxel_edge[:, sums == 2] # Ev x 2 sums ==2 means voxel edges observed in the generated output 113 | 114 | if masked_edges.shape[1] != 0: 115 | # Now put program index onto the voxel edge and delete duplicates 116 | predicted_program_edges = torch.unique(max_out_program_index[masked_edges], dim=1) 117 | # union of program edges and program edges from the generated output 118 | mixed_edges = torch.cat((g.program_edge, predicted_program_edges), dim=1) 119 | unique_mix_edges, mix_counts = mixed_edges.unique(return_counts=True, dim=1) 120 | inter_edges = unique_mix_edges[:, mix_counts > 1] 121 | 122 | # program edges only in the input program graph 123 | mixed_gt_edges = torch.cat((g.program_edge, inter_edges), dim=1) 124 | unique_gt_edges, mix_gt_counts = mixed_gt_edges.unique(return_counts=True, dim=1) 125 | missing_edges = unique_gt_edges[:, mix_gt_counts == 1] 126 | 127 | # program edges only in the generated output 128 | mixed_gen_edges = torch.cat((predicted_program_edges, inter_edges), dim=1) 129 | unique_gen_edges, mix_gen_counts = mixed_gen_edges.unique(return_counts=True, dim=1) 130 | gen_edges = unique_gen_edges[:, mix_gen_counts == 1] 131 | else: # there is no voxel edge 132 | inter_edges = masked_edges 133 | missing_edges = g.program_edge 134 | gen_edges = masked_edges 135 | 136 | return inter_edges, missing_edges, gen_edges 137 | 138 | 139 | def generate_multiple_outputs_from_batch(batch, variation_num, generator, raw_dir, output_dir, follow_batch, device_ids, trunc=1.0): 140 | device = device_ids[0] 141 | batch.to(device) 142 | with torch.no_grad(): 143 | 144 | program_z_shape = [batch.program_class_feature.shape[0], generator.noise_dim] 145 | program_z = torch.rand(tuple(program_z_shape)).to(device) 146 | voxel_z_shape = [batch.voxel_feature.shape[0], generator.noise_dim] 147 | voxel_z = torch.rand(tuple(voxel_z_shape)).to(device) 148 | if trunc < 1.0: 149 | program_z.clamp_(min=-trunc, max=trunc) 150 | voxel_z.clamp_(min=-trunc, max=trunc) 151 | 152 | batch.to(device) 153 | out, soft_out, mask, att, max_out_program_index = generator(batch, program_z, voxel_z) 154 | 155 | normalized_program_class_weight, normalized_program_weight, FAR = get_program_ratio(batch, att["hard"], mask["hard"], area_index_in_voxel_feature=6) 156 | save_output(variation_num, batch, normalized_program_class_weight, normalized_program_weight, FAR, max_out_program_index, out, follow_batch, 157 | raw_dir, output_dir, new_data_id_strs=list(range(variation_num))) 158 | 159 | 160 | def generate_multiple_outputs_from_data(data, variation_num, generator, raw_dir, output_dir, follow_batch, device_ids): 161 | batch = Batch.from_data_list([data for _ in range(variation_num)], follow_batch) 162 | generate_multiple_outputs_from_batch(batch, variation_num, generator, raw_dir, output_dir, follow_batch, device_ids) 163 | 164 | -------------------------------------------------------------------------------- /util_graph.py: -------------------------------------------------------------------------------- 1 | from Data.VolumeDesignGraph import VolumeDesignGraph 2 | from torch_geometric.data import Batch 3 | from torch_scatter import scatter_add, scatter, scatter_max 4 | import torch 5 | import copy 6 | 7 | 8 | def detach_batch(batch): 9 | detached_batch = Batch() 10 | detached_batch.__data_class__ = VolumeDesignGraph 11 | detached_batch.__slices__ = copy.deepcopy(batch.__slices__) 12 | 13 | for key in batch.keys: 14 | if torch.is_tensor(batch[key]): 15 | detached_batch[key] = batch[key].detach() 16 | else: 17 | detached_batch[key] = copy.deepcopy(batch[key]) 18 | return detached_batch 19 | 20 | 21 | def get_program_ratio(graph, att, mask, area_index_in_voxel_feature): 22 | """ 23 | For each program type, we sum all the areas from the corresponding voxel nodes and obtain program weight. 24 | We can then normalize it and compute FAR. 25 | """ 26 | device = att.get_device() if att.is_cuda else "cpu" 27 | # Nv x 1 if voxel node is masked, area = 0; otherwise, area. 28 | masked_voxel_weight = mask * graph.voxel_feature[:, area_index_in_voxel_feature].view(-1, 1) 29 | # E x 1 put area values on the cross edge 30 | painted_voxel_weight = (att * torch.index_select(masked_voxel_weight, 0, graph.cross_edge_voxel_index_select)) 31 | # Np x 1 sum of voxel node areas on the program node 32 | program_weight = scatter(src=painted_voxel_weight, index=graph.cross_edge_program_index_select, dim=0, dim_size=graph.program_class_feature.shape[0], reduce="sum") 33 | # Sums the areas of program nodes for each type 34 | program_class_weight = scatter(src=program_weight, index=graph.program_class_cluster, dim=0, dim_size=graph.program_target_ratio.shape[0], reduce='sum') 35 | # Sums the total area for each graph 36 | batch_sum = scatter_add(program_class_weight, graph.program_target_ratio_batch.to(device), dim=0, dim_size=graph.FAR.shape[0])[graph.program_target_ratio_batch] 37 | # Normalize the program ratio in each graph 38 | normalized_program_class_weight = program_class_weight / (batch_sum + 1e-16) 39 | # Compute FAR 40 | FAR = scatter(src=program_class_weight, index=graph.program_target_ratio_batch, dim=0, dim_size=graph.FAR.shape[0], reduce="sum") 41 | return normalized_program_class_weight, program_weight, FAR 42 | 43 | 44 | def find_max_out_program_index(logit, cross_edge_voxel_index, cross_edge_program_index, num_of_voxels): 45 | """ max_out_program_index (Nv x 1) is the program node index that each voxel node has the max attention. We also compute voxel nodes that are masked (mask["hard"])""" 46 | _, out_cross_edge_index = scatter_max(logit, index=cross_edge_voxel_index, dim=0, dim_size=num_of_voxels) 47 | max_out_program_index = cross_edge_program_index[out_cross_edge_index] 48 | return max_out_program_index 49 | 50 | 51 | def unbatch_data_to_data_list(batch): 52 | """ 53 | Modified by torch_geometric.data.batch.to_data_list() since the keys are not in order, so when calling __inc__ and number of nodes are queried, 54 | error occurs because the node features might not be populated yet. 55 | The "xxx_batch" will be dropped in the output as in the data_list. You recreate them when making batches 56 | """ 57 | if batch.__slices__ is None: 58 | raise RuntimeError('Cannot reconstruct data list from batch because the batch object was not created using Batch.from_data_list()') 59 | keys = [key for key in batch.keys if key[-5:] != 'batch'] 60 | cumsum = {key: 0 for key in keys} 61 | data_list = [] 62 | for i in range(len(batch.__slices__[keys[0]]) - 1): 63 | data = batch.__data_class__() 64 | for key in keys: 65 | if data[key] is not None: 66 | continue 67 | if torch.is_tensor(batch[key]): 68 | data[key] = batch[key].narrow(data.__cat_dim__(key, batch[key]), batch.__slices__[key][i], batch.__slices__[key][i + 1] - batch.__slices__[key][i]) 69 | # if batch[key].dtype != torch.bool: data[key] = data[key] - cumsum[key] 70 | else: 71 | data[key] = batch[key][batch.__slices__[key][i]:batch.__slices__[key][i + 1]] 72 | # cumsum[key] = cumsum[key] + data.__inc__(key, data[key]) 73 | for key in keys: 74 | if torch.is_tensor(batch[key]) and batch[key].dtype != torch.bool: 75 | data[key] = data[key] - cumsum[key] 76 | cumsum[key] = cumsum[key] + data.__inc__(key, data[key]) 77 | 78 | data["data_id_str"] = data["data_id_str"][0] # This is added, otherwise this will be list not str 79 | data_list.append(data) 80 | return data_list 81 | 82 | 83 | def rebatch_graph_for_multi_gpu(batch, device_ids, follow_batch): 84 | """ 85 | Given a batch of data, split to multiple mini-batches. 86 | """ 87 | data_list = unbatch_data_to_data_list(batch) 88 | mini_batch_size = len(data_list)//len(device_ids) 89 | mini_batch_list, mini_batch_slices = [], [0] 90 | for i in range(len(device_ids)): 91 | mini_batch_list.append(Batch.from_data_list(data_list[i * mini_batch_size: (i+1) * mini_batch_size], follow_batch=follow_batch)) 92 | mini_batch_slices.append((i+1) * mini_batch_size) 93 | return mini_batch_list, mini_batch_slices 94 | 95 | 96 | def rebatch_for_multi_gpu(batch, device_ids, follow_batch, *args): 97 | """ 98 | This function rebatches batched graphs(batch) and other information(*args) based on the given device/new_batch ids. 99 | Return dimension [M batches x N features (list of batch and args)] 100 | 101 | split_by_graph: FAR, class_weights 102 | split_by_program_node: z noise, program_weights 103 | split_by_voxel_node: out, mask_hard, max_out_program_index 104 | 105 | example args: batch = batch, args = out, class_weights, program_weights, FAR, max_out_program_index 106 | return 107 | g: graph 108 | ------------------ 109 | o: voxel label (n[type]) 110 | cw: (n[new_proportion] in global graph) -- program class ratio/weight 111 | pw: (n[region_far] in local graph) 112 | far: (g[far] in global graph) 113 | pid: the selected program node id for each voxel node 114 | """ 115 | 116 | # First rebatch the batched graphs 117 | mini_batch_list, mini_batch_slices = rebatch_graph_for_multi_gpu(batch, device_ids, follow_batch) 118 | 119 | # Then identify batching method for the args 120 | rebatch_type_key = [] 121 | for arg in args: 122 | if arg.shape[0] == batch["FAR"].shape[0]: 123 | rebatch_type_key.append("FAR") 124 | elif arg.shape[0] == batch["program_class_feature"].shape[0]: 125 | rebatch_type_key.append("program_class_feature") 126 | elif arg.shape[0] == batch["voxel_feature"].shape[0]: 127 | rebatch_type_key.append("voxel_feature") 128 | elif arg.shape[0] == batch["program_target_ratio"].shape[0]: 129 | rebatch_type_key.append("program_target_ratio") 130 | else: 131 | raise ValueError("unknown input") 132 | 133 | # Save to the output data structure 134 | ret, data = [], batch.__data_class__() 135 | for i, (device_id, mini_batch) in enumerate(zip(device_ids, mini_batch_list)): 136 | placeholder = [mini_batch] 137 | start, end = mini_batch_slices[i], mini_batch_slices[i+1] 138 | for arg, key in zip(args, rebatch_type_key): 139 | mini_arg = arg.narrow(data.__cat_dim__(key, None), batch.__slices__[key][start], batch.__slices__[key][end] - batch.__slices__[key][start]) 140 | placeholder.append(mini_arg) 141 | try: 142 | placeholder = [ele.to(device_id) for ele in placeholder] 143 | except: 144 | pass 145 | ret.append(tuple(placeholder)) 146 | return ret 147 | 148 | 149 | def data_parallel(module, batch, _input, follow_batch, device_ids): 150 | """ 151 | Reference code for multi-gpu setups. Not used in the this code repo 152 | """ 153 | output_device = device_ids[0] 154 | replicas = torch.nn.parallel.replicate(module, device_ids) 155 | inputs = rebatch_for_multi_gpu(batch, device_ids, follow_batch, *_input) 156 | replicas = replicas[:len(inputs)] 157 | outputs = torch.nn.parallel.parallel_apply(replicas, inputs) 158 | return torch.nn.parallel.gather(outputs, output_device) 159 | --------------------------------------------------------------------------------