diff --git a/train.py b/train.py
index 7e3e390..130c7f4 100644
--- a/train.py
+++ b/train.py
@@ -380,7 +380,10 @@ def run(_config, n_classes, n_epochs, input_height,
         dir_flow_train_labels = os.path.join(dir_train, 'labels')
         
         classes = os.listdir(dir_flow_train_labels)
-        num_rows =len(classes)
+        if augmentation:
+            num_rows = len(classes)*(len(thetha) + 1)
+        else:
+            num_rows = len(classes)
         #ls_test = os.listdir(dir_flow_train_labels)
 
         #f1score_tot = [0]
@@ -390,7 +393,7 @@ def run(_config, n_classes, n_epochs, input_height,
         model.compile(loss="binary_crossentropy",
                             optimizer = opt_adam,metrics=['accuracy'])
         for i in range(n_epochs):
-            history = model.fit(generate_arrays_from_folder_reading_order(dir_flow_train_labels, dir_flow_train_imgs, n_batch, input_height, input_width, n_classes), steps_per_epoch=num_rows / n_batch, verbose=1)
+            history = model.fit(generate_arrays_from_folder_reading_order(dir_flow_train_labels, dir_flow_train_imgs, n_batch, input_height, input_width, n_classes, thetha, augmentation), steps_per_epoch=num_rows / n_batch, verbose=1)
             model.save( os.path.join(dir_output,'model_'+str(i+indexer_start) ))
             
             with open(os.path.join(os.path.join(dir_output,'model_'+str(i)),"config.json"), "w") as fp:
diff --git a/utils.py b/utils.py
index d7ddb99..50c21af 100644
--- a/utils.py
+++ b/utils.py
@@ -363,6 +363,11 @@ def rotation_not_90_func(img, label, thetha):
     return rotate_max_area(img, rotated, rotated_label, thetha)
 
 
+def rotation_not_90_func_single_image(img, thetha):
+    rotated = imutils.rotate(img, thetha)
+    return rotate_max_area(img, rotated, thetha)
+
+
 def color_images(seg, n_classes):
     ann_u = range(n_classes)
     if len(np.shape(seg)) == 3:
@@ -410,7 +415,7 @@ def IoU(Yi, y_predi):
     #print("Mean IoU: {:4.3f}".format(mIoU))
     return mIoU
 
-def generate_arrays_from_folder_reading_order(classes_file_dir, modal_dir, batchsize, height, width, n_classes):
+def generate_arrays_from_folder_reading_order(classes_file_dir, modal_dir, batchsize, height, width, n_classes, thetha, augmentation=False):
     all_labels_files = os.listdir(classes_file_dir)
     ret_x= np.zeros((batchsize, height, width, 3))#.astype(np.int16)
     ret_y= np.zeros((batchsize, n_classes)).astype(np.int16)
@@ -433,6 +438,22 @@ def generate_arrays_from_folder_reading_order(classes_file_dir, modal_dir, batch
                 ret_x= np.zeros((batchsize, height, width, 3))#.astype(np.int16)
                 ret_y= np.zeros((batchsize, n_classes)).astype(np.int16)
                 batchcount = 0
+                
+            if augmentation:
+                for thetha_i in thetha:
+                    img_rot = rotation_not_90_func_single_image(img, thetha_i)
+                    
+                    ret_x[batchcount, :,:,0] = img_rot[:,:,0]/3.0
+                    ret_x[batchcount, :,:,2] = img_rot[:,:,2]/3.0
+                    ret_x[batchcount, :,:,1] = img_rot[:,:,1]/5.0
+
+                    ret_y[batchcount, :] =  label_class
+                    batchcount+=1
+                    if batchcount>=batchsize:
+                        yield (ret_x, ret_y)
+                        ret_x= np.zeros((batchsize, height, width, 3))#.astype(np.int16)
+                        ret_y= np.zeros((batchsize, n_classes)).astype(np.int16)
+                        batchcount = 0
 
 def data_gen(img_folder, mask_folder, batch_size, input_height, input_width, n_classes, task='segmentation'):
     c = 0