-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathnesting_states_dataset.py
84 lines (70 loc) · 3.26 KB
/
nesting_states_dataset.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
import os
import numpy as np
from torch.utils.data import Dataset, DataLoader
def prepare_state(sample_data, max_num_of_polys):
# state = np.expand_dims(sample_data["state"], 0)
state = sample_data["state"]
poly_idxs = sample_data["poly_idxs"]
num_polys = poly_idxs.shape[0]
poly_idxs = np.expand_dims(np.pad(poly_idxs, (0, max_num_of_polys - num_polys), 'constant'), axis=-1)
mask = np.zeros(max_num_of_polys)
mask[:num_polys] = 1
mask = np.array(mask, dtype=np.bool)
state, poly_idxs = np.array(state, dtype=np.float32), np.array(poly_idxs, dtype=np.float32)
return state, poly_idxs, mask
def prepare_sample(sample_data, max_grid_width, max_grid_height, max_num_of_polys):
state, poly_idxs, mask = prepare_state(sample_data, max_num_of_polys)
num_polys = mask.sum()
ideal_action = np.array(sample_data["ideal_action"], dtype=np.float32)
ideal_action[:, 0] /= max_grid_width
ideal_action[:, 1] /= max_grid_height
ideal_action = np.pad(ideal_action, ((0, max_num_of_polys - num_polys), (0, 0)), 'constant',
constant_values=(np.array([0, 0])))
return state, poly_idxs, mask, ideal_action
class NestingStatesDataset(Dataset):
def __init__(self, data_dirs, max_num_of_polys, max_grid_width, max_grid_height):
self.max_grid_height = max_grid_height
self.max_grid_width = max_grid_width
self.max_num_of_polys = max_num_of_polys
self.all_files = []
for dd in data_dirs:
dir_files = os.listdir(dd)
dir_files = [os.path.join(dd, f) for f in dir_files]
self.all_files.extend(dir_files)
def __len__(self):
return len(self.all_files)
def __getitem__(self, idx):
data = np.load(self.all_files[idx], allow_pickle=True)
state, poly_idxs, mask, ideal_action = prepare_sample(data, self.max_grid_width, self.max_grid_height, self.max_num_of_polys)
# return
# state = np.expand_dims(data["state"], 0)
# poly_idxs = data["poly_idxs"]
# ideal_action = np.array(data["ideal_action"])
# ideal_action[:,0] /= self.max_grid_width
# ideal_action[:,1] /= self.max_grid_height
#
# num_polys = poly_idxs.shape[0]
#
# poly_idxs = np.expand_dims(np.pad(poly_idxs, (0, self.max_num_of_polys - num_polys), 'constant'), axis=-1)
# ideal_action = np.pad(ideal_action, ((0, self.max_num_of_polys - num_polys), (0,0)), 'constant', constant_values=(np.array([np.inf, np.inf])))
#
# mask = np.zeros(self.max_num_of_polys)
# mask[:num_polys] = 1
# mask = np.array(mask, dtype=np.bool)
# state, poly_idxs = np.array(state, dtype=np.float32), np.array(poly_idxs, dtype=np.float32)
return state, poly_idxs, ideal_action, mask
if __name__ == "__main__":
ds = NestingStatesDataset(r"C:\Proiecte\AiNesterPlusPlus\scratch\1", 1000, 1000, 1000)
dl = DataLoader(ds, batch_size=16)
from models.cross_attention_vit import CrossAttentionVit
m = CrossAttentionVit()
for state, poly_idxs, ideal_action, mask in dl:
print(state.shape)
x = (state, poly_idxs, mask)
y = m(x)
print(y)
print(len(ds))
print(ds[0])
a = ds[0]
poly_rasters = a["rasters"]
print(poly_rasters)