1*a58d3d2aSXin Li#!/usr/bin/python3 2*a58d3d2aSXin Li'''Copyright (c) 2021-2022 Amazon 3*a58d3d2aSXin Li 4*a58d3d2aSXin Li Redistribution and use in source and binary forms, with or without 5*a58d3d2aSXin Li modification, are permitted provided that the following conditions 6*a58d3d2aSXin Li are met: 7*a58d3d2aSXin Li 8*a58d3d2aSXin Li - Redistributions of source code must retain the above copyright 9*a58d3d2aSXin Li notice, this list of conditions and the following disclaimer. 10*a58d3d2aSXin Li 11*a58d3d2aSXin Li - Redistributions in binary form must reproduce the above copyright 12*a58d3d2aSXin Li notice, this list of conditions and the following disclaimer in the 13*a58d3d2aSXin Li documentation and/or other materials provided with the distribution. 14*a58d3d2aSXin Li 15*a58d3d2aSXin Li THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 16*a58d3d2aSXin Li ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 17*a58d3d2aSXin Li LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 18*a58d3d2aSXin Li A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE FOUNDATION OR 19*a58d3d2aSXin Li CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20*a58d3d2aSXin Li EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21*a58d3d2aSXin Li PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22*a58d3d2aSXin Li PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF 23*a58d3d2aSXin Li LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING 24*a58d3d2aSXin Li NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 25*a58d3d2aSXin Li SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26*a58d3d2aSXin Li''' 27*a58d3d2aSXin Li 28*a58d3d2aSXin Liimport numpy as np 29*a58d3d2aSXin Lifrom tensorflow.keras.utils import Sequence 30*a58d3d2aSXin Li 31*a58d3d2aSXin Liclass PLCLoader(Sequence): 32*a58d3d2aSXin Li def __init__(self, features, lost, nb_burg_features, batch_size): 33*a58d3d2aSXin Li self.batch_size = batch_size 34*a58d3d2aSXin Li self.nb_batches = features.shape[0]//self.batch_size 35*a58d3d2aSXin Li self.features = features[:self.nb_batches*self.batch_size, :, :] 36*a58d3d2aSXin Li self.lost = lost.astype('float') 37*a58d3d2aSXin Li self.lost = self.lost[:(len(self.lost)//features.shape[1]-1)*features.shape[1]] 38*a58d3d2aSXin Li self.nb_burg_features = nb_burg_features 39*a58d3d2aSXin Li self.on_epoch_end() 40*a58d3d2aSXin Li 41*a58d3d2aSXin Li def on_epoch_end(self): 42*a58d3d2aSXin Li self.indices = np.arange(self.nb_batches*self.batch_size) 43*a58d3d2aSXin Li np.random.shuffle(self.indices) 44*a58d3d2aSXin Li offset = np.random.randint(0, high=self.features.shape[1]) 45*a58d3d2aSXin Li self.lost_offset = np.reshape(self.lost[offset:-self.features.shape[1]+offset], (-1, self.features.shape[1])) 46*a58d3d2aSXin Li self.lost_indices = np.random.randint(0, high=self.lost_offset.shape[0], size=self.nb_batches*self.batch_size) 47*a58d3d2aSXin Li 48*a58d3d2aSXin Li def __getitem__(self, index): 49*a58d3d2aSXin Li features = self.features[self.indices[index*self.batch_size:(index+1)*self.batch_size], :, :] 50*a58d3d2aSXin Li burg_lost = (np.random.rand(features.shape[0], features.shape[1]) > .1).astype('float') 51*a58d3d2aSXin Li burg_lost = np.reshape(burg_lost, (features.shape[0], features.shape[1], 1)) 52*a58d3d2aSXin Li burg_mask = np.tile(burg_lost, (1,1,self.nb_burg_features)) 53*a58d3d2aSXin Li 54*a58d3d2aSXin Li lost = self.lost_offset[self.lost_indices[index*self.batch_size:(index+1)*self.batch_size], :] 55*a58d3d2aSXin Li lost = np.reshape(lost, (features.shape[0], features.shape[1], 1)) 56*a58d3d2aSXin Li lost_mask = np.tile(lost, (1,1,features.shape[2])) 57*a58d3d2aSXin Li in_features = features*lost_mask 58*a58d3d2aSXin Li in_features[:,:,:self.nb_burg_features] = in_features[:,:,:self.nb_burg_features]*burg_mask 59*a58d3d2aSXin Li 60*a58d3d2aSXin Li #For the first frame after a loss, we don't have valid features, but the Burg estimate is valid. 61*a58d3d2aSXin Li #in_features[:,1:,self.nb_burg_features:] = in_features[:,1:,self.nb_burg_features:]*lost_mask[:,:-1,self.nb_burg_features:] 62*a58d3d2aSXin Li out_lost = np.copy(lost) 63*a58d3d2aSXin Li #out_lost[:,1:,:] = out_lost[:,1:,:]*out_lost[:,:-1,:] 64*a58d3d2aSXin Li 65*a58d3d2aSXin Li out_features = np.concatenate([features[:,:,self.nb_burg_features:], 1.-out_lost], axis=-1) 66*a58d3d2aSXin Li burg_sign = 2*burg_lost - 1 67*a58d3d2aSXin Li # last dim is 1 for received packet, 0 for lost packet, and -1 when just the Burg info is missing 68*a58d3d2aSXin Li inputs = [in_features*lost_mask, lost*burg_sign] 69*a58d3d2aSXin Li outputs = [out_features] 70*a58d3d2aSXin Li return (inputs, outputs) 71*a58d3d2aSXin Li 72*a58d3d2aSXin Li def __len__(self): 73*a58d3d2aSXin Li return self.nb_batches 74