test for augmentation extended

The test now also shows the masks

test_augmentation.py

 import pathlib
 import numpy as np
 import matplotlib.pylab as plt
+
+from helpers.AnnotationLocationLoader import AnnotationLocationLoader
 from helpers.DataLoader import DataLoader
 
 '''
@@ -17,19 +19,23 @@ for element in class_index_map:
 
 
 train_dir = pathlib.Path('input/training_data_mini')
+annotation_path = './input/mouth_annotations/'
 batch_size = 5
 img_width = 500
 img_height = 500
 channels = 3
 batch_count_dataset1 = 9
 
+annot_loader = AnnotationLocationLoader(images_base_folder=train_dir, mouth_annotations_folder=annotation_path)
+
 train_loader = DataLoader(data_path=train_dir,
                           batch_size=batch_size,
                           should_size_dataset1=batch_count_dataset1,
                           img_width=img_width,
                           img_height=img_height,
                           channels=channels,
-                          augment=True
+                          augment=True,
+                          annotation=annot_loader
                           )
 
 train_dataset = train_loader.load_dataset()
@@ -38,7 +44,8 @@ def plot_images(dataset, n_images, samples_per_image):
     output = np.zeros((img_height * n_images, img_width * samples_per_image, 3))
 
     row = 0
-    for (images, _) in dataset.repeat(samples_per_image).batch(n_images):
+    for (img_and_mask, _) in dataset.repeat(samples_per_image).batch(n_images):
+        images = img_and_mask['img']
         output[:, row*img_height:(row+1)*img_height] = np.vstack(images.numpy())
         row += 1
 
@@ -48,16 +55,20 @@ def plot_images(dataset, n_images, samples_per_image):
 
 
 count = 0
-for batch_img, batch_label in train_dataset.as_numpy_iterator():
+for batch_img_and_mask, batch_label in train_dataset.as_numpy_iterator():
     for i in range(len(batch_label)):
-        img = batch_img[i]
+        img = batch_img_and_mask['img'][i]#['img']
+        mask = batch_img_and_mask['mouth'][i]#['mouth']
         label = batch_label[i]
         count += 1
         print(img.shape)
         label = index_class_map[np.argmax(label)]
-        plt.title(label)
+        plt.title(label + str(count))
         plt.imshow(img)
         plt.show()
+        plt.title(label + str(count))
+        plt.imshow(mask)
+        plt.show()
     print("new batch")
 print("{} image(s) processed".format(count))