1import torch 2import numpy as np 3 4class PLCDataset(torch.utils.data.Dataset): 5 def __init__(self, 6 feature_file, 7 loss_file, 8 sequence_length=1000, 9 nb_features=20, 10 nb_burg_features=36, 11 lpc_order=16): 12 13 self.features_in = nb_features + nb_burg_features 14 self.nb_burg_features = nb_burg_features 15 total_features = self.features_in + lpc_order 16 self.sequence_length = sequence_length 17 self.nb_features = nb_features 18 19 self.features = np.memmap(feature_file, dtype='float32', mode='r') 20 self.lost = np.memmap(loss_file, dtype='int8', mode='r') 21 self.lost = self.lost.astype('float32') 22 23 self.nb_sequences = self.features.shape[0]//self.sequence_length//total_features 24 25 self.features = self.features[:self.nb_sequences*self.sequence_length*total_features] 26 self.features = self.features.reshape((self.nb_sequences, self.sequence_length, total_features)) 27 self.features = self.features[:,:,:self.features_in] 28 29 #self.lost = self.lost[:(len(self.lost)//features.shape[1]-1)*features.shape[1]] 30 #self.lost = self.lost.reshape((-1, self.sequence_length)) 31 32 def __len__(self): 33 return self.nb_sequences 34 35 def __getitem__(self, index): 36 features = self.features[index, :, :] 37 burg_lost = (np.random.rand(features.shape[0]) > .1).astype('float32') 38 burg_lost = np.reshape(burg_lost, (features.shape[0], 1)) 39 burg_mask = np.tile(burg_lost, (1,self.nb_burg_features)) 40 41 lost_offset = np.random.randint(0, high=self.lost.shape[0]-self.sequence_length) 42 lost = self.lost[lost_offset:lost_offset+self.sequence_length] 43 #randomly add a few 10-ms losses so that the model learns to handle them 44 lost = lost * (np.random.rand(lost.shape[-1]) > .02).astype('float32') 45 #randomly break long consecutive losses so we don't try too hard to predict them 46 lost = 1 - ((1-lost) * (np.random.rand(lost.shape[-1]) > .1).astype('float32')) 47 lost = np.reshape(lost, (features.shape[0], 1)) 48 lost_mask = np.tile(lost, (1,features.shape[-1])) 49 in_features = features*lost_mask 50 in_features[:,:self.nb_burg_features] = in_features[:,:self.nb_burg_features]*burg_mask 51 52 #For the first frame after a loss, we don't have valid features, but the Burg estimate is valid. 53 #in_features[:,1:,self.nb_burg_features:] = in_features[:,1:,self.nb_burg_features:]*lost_mask[:,:-1,self.nb_burg_features:] 54 out_lost = np.copy(lost) 55 #out_lost[:,1:,:] = out_lost[:,1:,:]*out_lost[:,:-1,:] 56 57 out_features = np.concatenate([features[:,self.nb_burg_features:], 1.-out_lost], axis=-1) 58 burg_sign = 2*burg_lost - 1 59 # last dim is 1 for received packet, 0 for lost packet, and -1 when just the Burg info is missing 60 return in_features*lost_mask, lost*burg_sign, out_features 61