-
Notifications
You must be signed in to change notification settings - Fork 10
/
prep_data_struct.py
196 lines (182 loc) · 8.45 KB
/
prep_data_struct.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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
"""
Script could be used as a template for preparing the image-beam data structure required to train
the modified ResNet on beam or blockage prediction. NOTE: script might need modification based on the
dataset you generated and where it is saved. PLEASE READ THROUGH.
"""
import numpy as np
import h5py as h5
import pickle
import scipy.io as sciio
import skimage.io as imio
import os
from shutil import copyfile
root_img_dir = '' # Where ViWi images are
codebook = sciio.loadmat('DFT_codebook64')['W']# Every column is a beam (codebook must be generated beforehand, see UPA_codebook_generator_DFT.m)
def getMATLAB(matlab_file=None, save_mode=False,pickle_name=None):
'''
It converts MAT data into a python dictionary of numpy arrays.
:param matlab_file: path to the matlab data structure to read. The structure is expected
to have two fields: wireless channels and user locations. The channel field
should be a 4D array: #of antennas, # of subcarriers, # of user positions, # of BS
user location field should be a 2D array: 3, # of user positions
:param save_mode: whether the numpy data needs to be saved or not
:param pickle_name: name of the pickle file where the data is stored
:return: dictionary of numpy arrays containing the raw channel and location data.
-----------------------------------------------------------------------------------
NOTE:
The MATLAB data structure has to be prepared following the generation of the wireless data using the ViWi data-generation
script.
'''
# Read MATLAB structure:
f = h5.File(matlab_file, 'r')
key1 = list(f.keys())
raw_data = f[key1[1]]
key2 = list(raw_data.keys())
channels = raw_data[key2[0]][:] # Wireless data field
loc = raw_data[key2[1]][:] # Loc data field
s1 = channels.shape
s2 = loc.shape
# Construct and store numpy dictionary
X = channels.view(np.double).reshape(s1+(2,))
X = X.astype(np.float32) # This is necessary to reduce the precision
if len(s1) == 4:
X_comp = X[:, :, :, :, 0] + X[:, :, :, :, 1] * 1j # size: # of BSs X # of users X sub-carriers X # of antennas
else:
X_comp = X[:, :, :, 0] + X[:, :, :, 1] * 1j # size: # of users X sub-carriers X # of antennas
# Normalize channels
rec_pow = np.mean(np.abs(X_comp)**2)
X_comp = X_comp/np.sqrt(rec_pow)
raw_data = {'ch': X_comp,
'loc': loc,
'norm_fact': rec_pow}
print(raw_data['ch'].shape)
if save_mode:
f = open(pickle_name, 'wb')
pickle.dump(raw_data, f, protocol=4)
f.close()
return raw_data
def beamPredStruct(raw_data,codebook,val_per,image_path=None):
'''
This function prepares an image data structure for training and testing
a CNN on mmWave beam prediction. The function is designed for the
direct distributed-camera scenario.
:param raw_data: Wireless data dictionary with the keys: ch, loc, and
norm_fact.
:param codebook: Beamforming DFT matrix
:param val_per: Precentage of validation (test) data
:param image_path: Path to the ViWi IMAGE folder
:return:
'''
image_names = os.listdir(image_path)
image_names = sorted(image_names)
shuf_ind = np.random.permutation(len(image_names))
loc = raw_data['loc'][:, 0:2]# User coordinates as output by ViWi
num_train = len(image_names) - np.ceil( val_per*len(image_names) )
count = 0
train_list = []
test_list = []
for i in shuf_ind:
# Find the coordinates in the image: (NOTE an image is tagged with the coordinates of its single user)
split_name = image_names[i].split('_')
x_axis = float( split_name[2] )
y_axis = float( split_name[3][:-4] )
coord = np.array([x_axis, y_axis])
cam_num = int( split_name[1] )-1
# Find the channel of those coordinates:
diff = np.sum( np.abs(loc - coord), axis=1 )
user = np.argmin(diff)
# print('coord {} and locatio {}'.format( coord,loc[user] ))
h = raw_data['ch'][cam_num,user,:,:] # Channel for img_name
# Finding the best beamforming vector:
codebook_H = codebook.conj()
rec_pow_sub = np.power(np.abs(np.matmul(h, codebook_H)), 2) # per subcarrier
rate_per_sub = np.log2( 1+rec_pow_sub )
print(rec_pow_sub.shape)
ave_rate_per_beam = np.mean(rate_per_sub, axis=0) # averaged over subcarriers
beam_ind = np.argmax(ave_rate_per_beam)+1
print('image name {} and beam index {}'.format(image_names[i], beam_ind))
# Dividing images into folders
count += 1
if count <= num_train:
if not os.path.exists('train_images'):
os.mkdir('train_images')
sub_dir_name = beam_ind
sub_dir_path = os.getcwd()+'/train_images/'+str(sub_dir_name)
if not os.path.exists( sub_dir_path ):
os.mkdir( sub_dir_path )
copyfile( image_path+'/'+image_names[i], sub_dir_path+'/'+image_names[i] )
else:
copyfile( image_path+'/'+image_names[i], sub_dir_path+'/'+image_names[i] )
train_list.append(image_names[i])
else:
if not os.path.exists('test_images'):
os.mkdir('test_images')
sub_dir_name = beam_ind
sub_dir_path = os.getcwd()+'/test_images/'+str(sub_dir_name)
if not os.path.exists( sub_dir_path ):
os.mkdir( sub_dir_path )
copyfile( image_path+'/'+image_names[i], sub_dir_path+'/'+image_names[i] )
else:
copyfile( image_path+'/'+image_names[i], sub_dir_path+'/'+image_names[i] )
test_list.append(image_names[i])
# return [train_list,test_list]
def blockagePredStruct(raw_data,codebook,val_per,image_path=None):
'''
This function prepares an image data structure for training and testing
a CNN on blockage prediction. The function is designed for the
blocked colocated-camera scenario.
:param raw_data: Wireless data dictionary with the keys: ch, loc, and
abs_max.
:param codebook: Beamforming DFT matrix
:param val_per: Precentage of validation (test) data
:param image_path: Path to the ViWi IMAGE folder
:return:
'''
image_names = os.listdir(image_path)
image_names = sorted(image_names)
shuf_ind = np.random.permutation(len(image_names))
loc = raw_data['loc'][:, 0:2] # User coordinates as output by ViWi
num_train = len(image_names) - np.ceil(val_per * len(image_names))
count = 0
train_list = []
test_list = []
for i in shuf_ind:
# Find the coordinates in the image:
split_name = image_names[i].split('_')
x_axis = float(split_name[2])
y_axis = float(split_name[3])
blk_status = int(split_name[4][:-4])
coord = np.array([x_axis, y_axis])
cam_num = int(split_name[1]) - 1
# Dividing images into folders
count += 1
if count <= num_train:
if not os.path.exists('train_images_blk'):
os.mkdir('train_images_blk')
sub_dir_name = blk_status
sub_dir_path = os.getcwd() + '/train_images_blk/' + str(sub_dir_name)
if not os.path.exists(sub_dir_path):
os.mkdir(sub_dir_path)
copyfile(image_path + '/' + image_names[i], sub_dir_path + '/' + image_names[i])
else:
copyfile(image_path + '/' + image_names[i], sub_dir_path + '/' + image_names[i])
train_list.append(image_names[i])
else:
if not os.path.exists('test_images_blk'):
os.mkdir('test_images_blk')
sub_dir_name = blk_status
sub_dir_path = os.getcwd() + '/test_images_blk/' + str(sub_dir_name)
if not os.path.exists(sub_dir_path):
os.mkdir(sub_dir_path)
copyfile(image_path + '/' + image_names[i], sub_dir_path + '/' + image_names[i])
else:
copyfile(image_path + '/' + image_names[i], sub_dir_path + '/' + image_names[i])
test_list.append(image_names[i])
if __name__ == '__main__':
raw_data = getMATLAB(matlab_file='raw_data_dist_D_64ULA_5p_64s.mat',
save_mode=False,
pickle_name=None)
beamPredStruct(raw_data,codebook,0.3,root_img_dir)
# blockagePredStruct(raw_data,codebook,0.3,root_img_dir)
# print('Training samples: {} and testing samples: {}'.format( len(train_list), len(test_list) ))
print('break')