Commit 94796478 authored by sjjsmuel's avatar sjjsmuel

test for augmentation extended

The test now also shows the masks
parent ccf2b7da
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))
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment