# Path: dltk/io/augmentation.py
# def extract_random_example_array(image_list,
#                                  example_size=[1, 64, 64],
#                                  n_examples=1):
#     """Randomly extract training examples from image (and a corresponding label).
#         Returns an image example array and the corresponding label array.
# 
#     Args:
#         image_list (np.ndarray or list or tuple): image(s) to extract random
#             patches from
#         example_size (list or tuple): shape of the patches to extract
#         n_examples (int): number of patches to extract in total
# 
#     Returns:
#         np.ndarray, np.ndarray: class-balanced patches extracted from full
#         images with the shape [batch, example_size..., image_channels]
#     """
# 
#     assert n_examples > 0
# 
#     was_singular = False
#     if isinstance(image_list, np.ndarray):
#         image_list = [image_list]
#         was_singular = True
# 
#     assert all([i_s >= e_s for i_s, e_s in zip(image_list[0].shape, example_size)]), \
#         'Image must be bigger than example shape'
#     assert (image_list[0].ndim - 1 == len(example_size) or image_list[0].ndim == len(example_size)), \
#         'Example size doesnt fit image size'
# 
#     for i in image_list:
#         if len(image_list) > 1:
#             assert (i.ndim - 1 == image_list[0].ndim or i.ndim == image_list[0].ndim or i.ndim + 1 == image_list[0].ndim),\
#                 'Example size doesnt fit image size'
# 
#             assert all([i0_s == i_s for i0_s, i_s in zip(image_list[0].shape, i.shape)]), \
#                 'Image shapes must match'
# 
#     rank = len(example_size)
# 
#     # Extract random examples from image and label
#     valid_loc_range = [image_list[0].shape[i] - example_size[i] for i in range(rank)]
# 
#     rnd_loc = [np.random.randint(valid_loc_range[dim], size=n_examples)
#                if valid_loc_range[dim] > 0
#                else np.zeros(n_examples, dtype=int) for dim in range(rank)]
# 
#     examples = [[]] * len(image_list)
#     for i in range(n_examples):
#         slicer = [slice(rnd_loc[dim][i], rnd_loc[dim][i] + example_size[dim])
#                   for dim in range(rank)]
# 
#         for j in range(len(image_list)):
#             ex_image = image_list[j][slicer][np.newaxis]
#             # Concatenate and return the examples
#             examples[j] = np.concatenate((examples[j], ex_image), axis=0) \
#                 if (len(examples[j]) != 0) else ex_image
# 
#     if was_singular:
#         return examples[0]
#     return examples


import argparse
import os
import time
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.contrib import predictor
from dltk.io.augmentation import extract_random_example_array
from reader import read_fn
    # From the model_path, parse the latest saved model and restore a
    # predictor from it
    export_dir = [os.path.join(args.model_path, o) for o in sorted(
        os.listdir(args.model_path)) if os.path.isdir(
        os.path.join(args.model_path, o)) and o.isdigit()][-1]

    print('Loading from {}'.format(export_dir))
    my_predictor = predictor.from_saved_model(export_dir)

    # Iterate through the files, predict on the full volumes and compute a Dice
    # coefficient
    mae = []
    for output in read_fn(file_references=file_names,
                          mode=tf.estimator.ModeKeys.EVAL,
                          params=READER_PARAMS):
        t0 = time.time()

        # Parse the read function output and add a dummy batch dimension as
        # required
        img = output['features']['x']
        lbl = output['labels']['y']
        test_id = output['img_id']

        # We know, that the training input shape of [64, 96, 96] will work with
        # our model strides, so we collect several crops of the test image and
        # average the predictions. Alternatively, we could pad or crop the input
        # to any shape that is compatible with the resolution scales of the
        # model:

        num_crop_predictions = 4

# Path: dltk/io/augmentation.py
# def extract_random_example_array(image_list,
#                                  example_size=[1, 64, 64],
#                                  n_examples=1):
#     """Randomly extract training examples from image (and a corresponding label).
#         Returns an image example array and the corresponding label array.
# 
#     Args:
#         image_list (np.ndarray or list or tuple): image(s) to extract random
#             patches from
#         example_size (list or tuple): shape of the patches to extract
#         n_examples (int): number of patches to extract in total
# 
#     Returns:
#         np.ndarray, np.ndarray: class-balanced patches extracted from full
#         images with the shape [batch, example_size..., image_channels]
#     """
# 
#     assert n_examples > 0
# 
#     was_singular = False
#     if isinstance(image_list, np.ndarray):
#         image_list = [image_list]
#         was_singular = True
# 
#     assert all([i_s >= e_s for i_s, e_s in zip(image_list[0].shape, example_size)]), \
#         'Image must be bigger than example shape'
#     assert (image_list[0].ndim - 1 == len(example_size) or image_list[0].ndim == len(example_size)), \
#         'Example size doesnt fit image size'
# 
#     for i in image_list:
#         if len(image_list) > 1:
#             assert (i.ndim - 1 == image_list[0].ndim or i.ndim == image_list[0].ndim or i.ndim + 1 == image_list[0].ndim),\
#                 'Example size doesnt fit image size'
# 
#             assert all([i0_s == i_s for i0_s, i_s in zip(image_list[0].shape, i.shape)]), \
#                 'Image shapes must match'
# 
#     rank = len(example_size)
# 
#     # Extract random examples from image and label
#     valid_loc_range = [image_list[0].shape[i] - example_size[i] for i in range(rank)]
# 
#     rnd_loc = [np.random.randint(valid_loc_range[dim], size=n_examples)
#                if valid_loc_range[dim] > 0
#                else np.zeros(n_examples, dtype=int) for dim in range(rank)]
# 
#     examples = [[]] * len(image_list)
#     for i in range(n_examples):
#         slicer = [slice(rnd_loc[dim][i], rnd_loc[dim][i] + example_size[dim])
#                   for dim in range(rank)]
# 
#         for j in range(len(image_list)):
#             ex_image = image_list[j][slicer][np.newaxis]
#             # Concatenate and return the examples
#             examples[j] = np.concatenate((examples[j], ex_image), axis=0) \
#                 if (len(examples[j]) != 0) else ex_image
# 
#     if was_singular:
#         return examples[0]
#     return examples


# Path: examples/applications/IXI_HH_age_regression_resnet/deploy.py
import argparse
import os
import time
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.contrib import predictor
from dltk.io.augmentation import extract_random_example_array
from reader import read_fn

    # From the model_path, parse the latest saved model and restore a
    # predictor from it
    export_dir = [os.path.join(args.model_path, o) for o in sorted(
        os.listdir(args.model_path)) if os.path.isdir(
        os.path.join(args.model_path, o)) and o.isdigit()][-1]

    print('Loading from {}'.format(export_dir))
    my_predictor = predictor.from_saved_model(export_dir)

    # Iterate through the files, predict on the full volumes and compute a Dice
    # coefficient
    mae = []
    for output in read_fn(file_references=file_names,
                          mode=tf.estimator.ModeKeys.EVAL,
                          params=READER_PARAMS):
        t0 = time.time()

        # Parse the read function output and add a dummy batch dimension as
        # required
        img = output['features']['x']
        lbl = output['labels']['y']
        test_id = output['img_id']

        # We know, that the training input shape of [64, 96, 96] will work with
        # our model strides, so we collect several crops of the test image and
        # average the predictions. Alternatively, we could pad or crop the input
        # to any shape that is compatible with the resolution scales of the
        # model:

        num_crop_predictions = 4

        crop_batch = extract_random_example_array(