在 tf.data.Dataset.from_generator() 上应用增强

如何解决在 tf.data.Dataset.from_generator() 上应用增强

我有这个增强代码：

class CustomAugment(object):
   def __call__(self,sample):        
       sample = self._random_apply(tf.image.flip_left_right,sample,p=0.5)
       sample = self._random_apply(self._color_jitter,p=0.8)
       sample = self._random_apply(self._color_drop,p=0.2)

       return sample

   def _color_jitter(self,x,s=1):
       
       x = tf.image.random_brightness(x,max_delta=0.8*s)
       x = tf.image.random_contrast(x,lower=1-0.8*s,upper=1+0.8*s)
       x = tf.image.random_saturation(x,upper=1+0.8*s)
       x = tf.image.random_hue(x,max_delta=0.2*s)
       x = tf.clip_by_value(x,1)
       return x
   
   def _color_drop(self,x):
       x = tf.image.rgb_to_grayscale(x)
       x = tf.tile(x,[1,1,3])
       return x
   
   def _random_apply(self,func,p):
       return tf.cond(
         tf.less(tf.random.uniform([],minval=0,maxval=1,dtype=tf.float32),tf.cast(p,tf.float32)),lambda: func(x),lambda: x)

这就是我导入图像数据集的方式：

train_ds = tf.data.Dataset.from_generator(path)

我想在我的 train_ds 上应用此增强功能，请问，我该如何继续？

解决方法

首先，您应该使用 tf.keras.sequence 的子类创建一个自定义生成器，然后您可以实现 __getitem__ 和 __len__ 方法。

class CustomGenerator(tf.keras.utils.Sequence):

    def __init__(self,df,X_col,y_col,batch_size,input_size=(width,height,channels),shuffle=True):
    
        self.df = df.copy()
        self.X_col = X_col
        self.y_col = y_col
        self.batch_size = batch_size
        self.input_size = input_size
    
        self.n = len(self.df)
        self.n_name = df[y_col['label']].nunique()
    
    def on_epoch_end(self):
        pass    
    
    def __getitem__(self,index):
        batches = self.df[index * self.batch_size:(index + 1) * 
                          self.batch_size]
        X,y = self.__get_data(batches)        
        return X,y
    
    def __len__(self):
        return self.n // self.batch_size
    
    def __get_output(self,label,num_classes):
        return tf.keras.utils.to_categorical(label,num_classes=num_classes)    
    
    def __get_input(self,path,target_size):
        # Load Image using PIL
        img = Image.open(self.base_path + path)
        img = np.array(img)
        
        # Your Augmentation
        img = CustomAugment(img)
        return img /255


    def __get_data(self,batches):
        # Generates data containing batch_size samples

        img_path_batch = batches[self.X_col['img']]
        label_batch = batches[self.y_col['label']]

        X_batch = np.asarray([self.__get_input(x,self.input_size)
                              for x in img_path_batch])
        y_batch = np.asarray([self.__get_output(y)
                              for y in label_batch])

        return X_batch,y_batch

如您所见，您将在 __get_input 方法中扩充您的样本。

要使用这个类：

traingen = CustomDataGen(df,base_path=IMGS_DIR,X_col={'img':'img'},y_col={'label': 'label'},max_label_len=11,batch_size=16,input_size=IMAGE_SIZE)

注意：如果您需要在 tf.data 上使用生成器，您应该像这样使用它：

train_dataset = tf.data.Dataset.from_generator(lambda: traingen,output_types = (tf.float32,tf.int32),output_shapes = ([None,width,channels],[None,num_classes]))