# Path: dltk/core/metrics.py
# def dice(predictions, labels, num_classes):
# def abs_vol_difference(predictions, labels, num_classes):
# def crossentropy(predictions, labels, logits=True):
# 
# Path: dltk/utils.py
# def sliding_window_segmentation_inference(session,
#                                           ops_list,
#                                           sample_dict,
#                                           batch_size=1,
#                                           striding=None):
#     """
#     Utility function to perform sliding window inference for segmentation
# 
#     Args:
#         session (tf.Session): TensorFlow session to run ops with
# 
#         ops_list (array_like): Operators to fetch assemble with sliding window
# 
#         sample_dict (dict): Dictionary with tf.Placeholder keys mapping the
#         placeholders to their respective input
# 
#         batch_size (int, optional): Number of sliding windows to batch for
#             calculation
# 
#         striding (array_like): Striding of the sliding window. Can be used to
#             adjust overlap etc.
# 
#     Returns:
#         list: List of np.arrays corresponding to the assembled outputs of
#             ops_list
#     """
# 
#     # TODO: asserts
#     assert batch_size > 0, 'Batch size has to be 1 or bigger'
# 
#     pl_shape = list(sample_dict.keys())[0].get_shape().as_list()
# 
#     pl_bshape = pl_shape[1:-1]
# 
#     inp_shape = list(list(sample_dict.values())[0].shape)
#     inp_bshape = inp_shape[1:-1]
# 
#     out_dummies = [np.zeros(
#         [inp_shape[0], ] + inp_bshape + [op.get_shape().as_list()[-1]]
#         if len(op.get_shape().as_list()) == len(inp_shape) else []) for op in ops_list]
# 
#     out_dummy_counter = [np.zeros_like(o) for o in out_dummies]
# 
#     op_shape = list(ops_list[0].get_shape().as_list())
#     op_bshape = op_shape[1:-1]
# 
#     out_diff = np.array(pl_bshape) - np.array(op_bshape)
# 
#     padding = [[0, 0]] + [[diff // 2, diff - diff // 2] for diff
#                           in out_diff] + [[0, 0]]
# 
#     padded_dict = {k: np.pad(v, padding, mode='constant') for k, v
#                    in sample_dict.items()}
# 
#     f_bshape = list(padded_dict.values())[0].shape[1:-1]
# 
#     if not striding:
#         striding = (list(np.maximum(1, np.array(op_bshape) // 2))
#                     if all(out_diff == 0) else op_bshape)
# 
#     sw = SlidingWindow(f_bshape, pl_bshape, striding=striding)
#     out_sw = SlidingWindow(inp_bshape, op_bshape, striding=striding)
# 
#     if batch_size > 1:
#         slicers = []
#         out_slicers = []
# 
#     done = False
#     while True:
#         try:
#             slicer = next(sw)
#             out_slicer = next(out_sw)
#         except StopIteration:
#             done = True
# 
#         if batch_size == 1:
#             sw_dict = {k: v[slicer] for k, v in padded_dict.items()}
#             op_parts = session.run(ops_list, feed_dict=sw_dict)
# 
#             for idx in range(len(op_parts)):
#                 out_dummies[idx][out_slicer] += op_parts[idx]
#                 out_dummy_counter[idx][out_slicer] += 1
#         else:
#             slicers.append(slicer)
#             out_slicers.append(out_slicer)
#             if len(slicers) == batch_size or done:
#                 slices_dict = {k: np.concatenate(
#                     [v[slicer] for slicer in slicers], 0) for k, v in padded_dict.items()}
# 
#                 all_op_parts = session.run(ops_list, feed_dict=slices_dict)
# 
#                 zipped_parts = zip(*[np.array_split(part, len(slicers)) for
#                                      part in all_op_parts])
# 
#                 for out_slicer, op_parts in zip(out_slicers, zipped_parts):
#                     for idx in range(len(op_parts)):
#                         out_dummies[idx][out_slicer] += op_parts[idx]
#                         out_dummy_counter[idx][out_slicer] += 1
# 
#                 slicers = []
#                 out_slicers = []
# 
#         if done:
#             break
# 
#     return [o / c for o, c in zip(out_dummies, out_dummy_counter)]


import argparse
import os
import time
import numpy as np
import pandas as pd
import tensorflow as tf
import SimpleITK as sitk
from tensorflow.contrib import predictor
from dltk.core import metrics as metrics
from dltk.utils import sliding_window_segmentation_inference
from reader import read_fn
    print('Loading from {}'.format(export_dir))
    my_predictor = predictor.from_saved_model(export_dir)

    # Fetch the output probability op of the trained network
    y_prob = my_predictor._fetch_tensors['y_prob']
    num_classes = y_prob.get_shape().as_list()[-1]

    # Iterate through the files, predict on the full volumes and compute a Dice
    # coefficient
    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 = np.expand_dims(output['features']['x'], axis=0)
        lbl = np.expand_dims(output['labels']['y'], axis=0)

        # Do a sliding window inference with our DLTK wrapper
        pred = sliding_window_segmentation_inference(
            session=my_predictor.session,
            ops_list=[y_prob],
            sample_dict={my_predictor._feed_tensors['x']: img},
            batch_size=32)[0]

        # Calculate the prediction from the probabilities
        pred = np.argmax(pred, -1)

        # Calculate the Dice coefficient

# Path: dltk/core/metrics.py
# def dice(predictions, labels, num_classes):
# def abs_vol_difference(predictions, labels, num_classes):
# def crossentropy(predictions, labels, logits=True):
# 
# Path: dltk/utils.py
# def sliding_window_segmentation_inference(session,
#                                           ops_list,
#                                           sample_dict,
#                                           batch_size=1,
#                                           striding=None):
#     """
#     Utility function to perform sliding window inference for segmentation
# 
#     Args:
#         session (tf.Session): TensorFlow session to run ops with
# 
#         ops_list (array_like): Operators to fetch assemble with sliding window
# 
#         sample_dict (dict): Dictionary with tf.Placeholder keys mapping the
#         placeholders to their respective input
# 
#         batch_size (int, optional): Number of sliding windows to batch for
#             calculation
# 
#         striding (array_like): Striding of the sliding window. Can be used to
#             adjust overlap etc.
# 
#     Returns:
#         list: List of np.arrays corresponding to the assembled outputs of
#             ops_list
#     """
# 
#     # TODO: asserts
#     assert batch_size > 0, 'Batch size has to be 1 or bigger'
# 
#     pl_shape = list(sample_dict.keys())[0].get_shape().as_list()
# 
#     pl_bshape = pl_shape[1:-1]
# 
#     inp_shape = list(list(sample_dict.values())[0].shape)
#     inp_bshape = inp_shape[1:-1]
# 
#     out_dummies = [np.zeros(
#         [inp_shape[0], ] + inp_bshape + [op.get_shape().as_list()[-1]]
#         if len(op.get_shape().as_list()) == len(inp_shape) else []) for op in ops_list]
# 
#     out_dummy_counter = [np.zeros_like(o) for o in out_dummies]
# 
#     op_shape = list(ops_list[0].get_shape().as_list())
#     op_bshape = op_shape[1:-1]
# 
#     out_diff = np.array(pl_bshape) - np.array(op_bshape)
# 
#     padding = [[0, 0]] + [[diff // 2, diff - diff // 2] for diff
#                           in out_diff] + [[0, 0]]
# 
#     padded_dict = {k: np.pad(v, padding, mode='constant') for k, v
#                    in sample_dict.items()}
# 
#     f_bshape = list(padded_dict.values())[0].shape[1:-1]
# 
#     if not striding:
#         striding = (list(np.maximum(1, np.array(op_bshape) // 2))
#                     if all(out_diff == 0) else op_bshape)
# 
#     sw = SlidingWindow(f_bshape, pl_bshape, striding=striding)
#     out_sw = SlidingWindow(inp_bshape, op_bshape, striding=striding)
# 
#     if batch_size > 1:
#         slicers = []
#         out_slicers = []
# 
#     done = False
#     while True:
#         try:
#             slicer = next(sw)
#             out_slicer = next(out_sw)
#         except StopIteration:
#             done = True
# 
#         if batch_size == 1:
#             sw_dict = {k: v[slicer] for k, v in padded_dict.items()}
#             op_parts = session.run(ops_list, feed_dict=sw_dict)
# 
#             for idx in range(len(op_parts)):
#                 out_dummies[idx][out_slicer] += op_parts[idx]
#                 out_dummy_counter[idx][out_slicer] += 1
#         else:
#             slicers.append(slicer)
#             out_slicers.append(out_slicer)
#             if len(slicers) == batch_size or done:
#                 slices_dict = {k: np.concatenate(
#                     [v[slicer] for slicer in slicers], 0) for k, v in padded_dict.items()}
# 
#                 all_op_parts = session.run(ops_list, feed_dict=slices_dict)
# 
#                 zipped_parts = zip(*[np.array_split(part, len(slicers)) for
#                                      part in all_op_parts])
# 
#                 for out_slicer, op_parts in zip(out_slicers, zipped_parts):
#                     for idx in range(len(op_parts)):
#                         out_dummies[idx][out_slicer] += op_parts[idx]
#                         out_dummy_counter[idx][out_slicer] += 1
# 
#                 slicers = []
#                 out_slicers = []
# 
#         if done:
#             break
# 
#     return [o / c for o, c in zip(out_dummies, out_dummy_counter)]


# Path: examples/applications/MRBrainS13_tissue_segmentation/deploy.py
import argparse
import os
import time
import numpy as np
import pandas as pd
import tensorflow as tf
import SimpleITK as sitk
from tensorflow.contrib import predictor
from dltk.core import metrics as metrics
from dltk.utils import sliding_window_segmentation_inference
from reader import read_fn

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

    # Fetch the output probability op of the trained network
    y_prob = my_predictor._fetch_tensors['y_prob']
    num_classes = y_prob.get_shape().as_list()[-1]

    # Iterate through the files, predict on the full volumes and compute a Dice
    # coefficient
    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 = np.expand_dims(output['features']['x'], axis=0)
        lbl = np.expand_dims(output['labels']['y'], axis=0)

        # Do a sliding window inference with our DLTK wrapper
        pred = sliding_window_segmentation_inference(
            session=my_predictor.session,
            ops_list=[y_prob],
            sample_dict={my_predictor._feed_tensors['x']: img},
            batch_size=32)[0]

        # Calculate the prediction from the probabilities
        pred = np.argmax(pred, -1)

        # Calculate the Dice coefficient

        dsc = np.nanmean(metrics.dice(pred, lbl, num_classes)[1:])