Example #1
0
def test_generator(test_path, target_size, as_gray=True):
    """ Image Data Generator
    Function that generates batches od data for testing from specified folder
    Reads images as grey, makes them square, scales them
    Returns images with specified pixel size
    Does preprocessing (normalization to 0-1)
    """
    for test_img in os.listdir(test_path):
        img = io.imread(os.path.join(test_path, test_img), as_gray=as_gray)
        img = img / 255.
        img = reshape_image(img, target_size)
        yield img
    test_image_names.append(_)

test_predictions = []

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')