def train_generator(batch_size,
train_path,
image_folder,
mask_folder,
target_size,
image_color_mode='grayscale',
mask_color_mode='grayscale'):
""" Image Data Generator
Function that generates batches of data (img, mask) for training
from specified folder. Returns images with specified pixel size
Does preprocessing (normalization to 0-1)
"""
# no augmentation, only rescaling
image_datagen = ImageDataGenerator(rescale=1. / 255)
mask_datagen = ImageDataGenerator(rescale=1. / 255)
image_generator = image_datagen.flow_from_directory(
train_path,
classes=[image_folder],
class_mode=None,
color_mode=image_color_mode,
target_size=target_size,
batch_size=batch_size,
seed=1)
mask_generator = mask_datagen.flow_from_directory(
train_path,
classes=[mask_folder],
class_mode=None,
color_mode=mask_color_mode,
target_size=target_size,
batch_size=batch_size,
seed=1)
train_generator = zip(image_generator, mask_generator)
for (img, mask) in train_generator:
mask = normalize_mask(mask)
yield (img, mask)
def save_results(save_path, npyfile):
""" Save Results
Function that takes predictions from U-Net model
and saves them to specified folder.
"""
for i, item in enumerate(npyfile):
img = normalize_mask(item)
img = (img * 255).astype('uint8')
# io.imsave(os.path.join(save_path,"%d_predict.png"%(i+1)),img)
cv2.imwrite(os.path.join(save_path, "%d_predict.png" % (i + 1)),
train_img)
KERNEL = np.ones((15, 15), np.uint8)
#total prediction
for index, test_image_name in enumerate(test_image_names):
sys.stdout.write("Predicting: {}/{} ...\r".format(index + 1,
len(test_image_names)))
sys.stdout.flush()
img = io.imread(os.path.join(path_to_test_images, test_image_name),
as_gray=True)
img = img / 255.
img = reshape_image(img, unet_size)
result = unet.predict_on_batch(img)
result = result[0]
new_img = normalize_mask(result)
new_img = (new_img * 255).astype('uint8')
new_img = cv2.erode(new_img, KERNEL, iterations=2)
new_img = cv2.dilate(new_img, KERNEL, iterations=3)
new_img = np.expand_dims(new_img, axis=2)
new_img = cv2.bitwise_not(new_img)
img = img[0]
img = (img * 255).astype('uint8')
img = cv2.add(img, new_img)
img[img == new_img] = 0
test_image = Image.fromarray(img).convert('RGB')
#HeatMap Processing
DILATE_KERNEL = np.ones((15, 15), np.uint8)
for img_path in os.listdir(SETTINGS_JSON['TEST_IMAGES_DIR']):
orig_img = cv2.imread(
os.path.join(SETTINGS_JSON['TEST_IMAGES_DIR'], img_path),
cv2.IMREAD_COLOR)
img = io.imread(os.path.join(SETTINGS_JSON['TEST_IMAGES_DIR'], img_path),
as_gray=True)
img = img / 255.
img = reshape_image(img, unet_size)
result = unet.predict_on_batch(img)
result = result[0]
result = normalize_mask(result)
result = (result * 255).astype('uint8')
result = cv2.dilate(result, DILATE_KERNEL, iterations=3)
result = np.expand_dims(result, axis=2)
result = cv2.bitwise_not(result)
img = img[0]
img = (img * 255).astype('uint8')
result = cv2.add(img, result)
result = cv2.bitwise_not(result)
result = cv2.cvtColor(result, cv2.COLOR_GRAY2BGR)