-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
133 lines (103 loc) · 3.74 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import os
import random
import pickle
import numpy as np
import matplotlib.pyplot as plt
import torch_geometric
# loading and unloading pickled files
def save_pickle(object, filename):
"""
Saves a pickled file.
"""
with open(filename, 'wb') as f:
pickle.dump(object, f)
f.close()
def load_pickle(filename):
"""
Loads a pickled file.
"""
with open(filename, 'rb') as f:
object = pickle.load(f)
f.close()
return object
# RPLAN image loading
def load_image_rplan(id, data_path):
img_path = os.path.join(data_path, 'original', f'{id}.png')
img = (255*plt.imread(img_path)[..., 1]).astype(int)
return img
# Floor plan colorization
def colorize_floorplan(img, classes, cmap):
"""
Colorizes an integer-valued image (multi-class segmentation mask)
based on a map of classes and colormap.
"""
h, w = np.shape(img)
img_c = (np.ones((h, w, 3)) * 255).astype(int)
for cat in classes:
color = np.array(cmap(cat))[:3] * 255
img_c[img == cat, :] = color.astype(int)
return img_c
# Graph-based utilities
def pyg_to_nx(graph, node_attrs=None, edge_attrs=None, graph_attrs=None):
"""
Creates a networkx graph from a pytorch geometric graph data structure
:param edge_attrs: transferred edge attributes
:param node_attrs: transferred node attributes
:param graph: a pytorch geometric graph structure
:return: graph with fewer node and edge categories
"""
if graph_attrs is None:
graph_attrs = []
if node_attrs is None:
node_attrs = []
if edge_attrs is None:
edge_attrs = []
return torch_geometric.utils.to_networkx(graph, to_undirected=True,
node_attrs=node_attrs,
edge_attrs=edge_attrs,
graph_attrs=graph_attrs)
def simple_graph(graph):
"""
Creates simple graph with reduced node and edge attributes.
Remaining node attributes: categorical room types. Should be called 'category'.
Remaining edge attributes: door types. Should be called 'door'
"""
# Creates simple Networkx graph with only 'category' and 'door' attributes.
graph = torch_geometric.utils.to_networkx(graph,
to_undirected=True,
node_attrs=['category'],
edge_attrs=['door'])
return graph
def remove_attributes_from_graph(graph, list_attr=None):
"""
Removes attributes from graph.
:param graph: Input topological graph.
:param list_attr: Attributes to-be removed.
:return: Output topological graph with removed attributes.
"""
if list_attr is None:
list_attr = ['polygons']
for attr in list_attr:
for n in graph.nodes(): # delete irrelevant nodes
del graph.nodes[n][attr]
return graph
# gather IDs
def gather_ids(ids_path='./data/valid_ids.pickle', shuffle=True):
ids = load_pickle(ids_path)['full'].tolist()
if shuffle: random.Random(4).shuffle(ids)
return ids
# tensor to image conversion
def tens2img(img):
img = img.cpu().clone().detach().numpy() #make sure works always
return img.transpose(1,2,0)
# remaining image helper functions
def hex_to_rgb(hex_code):
hex_code = hex_code.lstrip('#')
return tuple(int(hex_code[i:i+2], 16) for i in (0, 2, 4))
def pad_image(img, pad, color):
# initialize padded image
img_pad = np.ndarray((img.shape[0] + pad * 2, img.shape[1] + pad * 2, 3))
# create padded image
for i, c in enumerate(color):
img_pad[:,:,i] = np.pad(img[:,:,i], ((pad, pad), (pad, pad)), constant_values=c)
return img_pad