tfdataset_from_generator.py

class ArrayDataset:
    def __init__(self, 
                 phase, 
                 array, 
                 labels, 
                 mean, std, 
                 transformations = None):
        self.phase = phase
        self.imgarr = copy.deepcopy(array)
        self.labels = copy.deepcopy(labels)
        self.mean = mean
        self.std = std
        self.transforms = transformations

    def __len__(self):
        return self.imgarr.shape[0]

    def __getitem__(self,idx):
        x = self.imgarr[idx]
        x = self.augment(x)
        y = self.labels[idx].astype(np.int64)
        return x, y
    
    def __call__(self):
        for i in range(self.__len__()):
            yield self.__getitem__(i)
            
            if i == self.__len__()-1:
                self.on_epoch_end()
    
    def normalize(self, x):
        return (x-self.mean)/self.std
            
    #shuffles the dataset at the end of each epoch
    def on_epoch_end(self):
        reidx = random.sample(population = list(range(self.__len__())),k = self.__len__())
        self.imgarr = self.imgarr[reidx]
        self.labels = self.labels[reidx]
    
    #applies randomly selected augmentations to each clip (same for each frame in the clip)
    def augment(self, x):
        if self.transforms is not None:
            x = self.transforms(x)
        x = self.normalize(x)
        return x