# Path: mesh_tensorflow/experimental/data_aug_lib.py
# def _truncated_normal(mean, stddev):
# def _rand_noise(noise_mean, noise_dev, scale, shape):
# def projective_transform(
#     image, label, reso, image_translate_ratio, image_transform_ratio,
#     sampled_2d_slices=False):
#   def _projective_transform(data, proj_matrix, static_axis, interpolation):
# def maybe_add_noise(image, noise_shape, scale0, scale1,
#                     image_noise_probability, image_noise_ratio):
# def _gen_rand_mask(ratio_mean, ratio_stddev, scale, shape, smoothness=0):
# def maybe_gen_fake_data_based_on_real_data(
#     image, label, reso, min_fake_lesion_ratio, gen_fake_probability):
# def maybe_flip(image, label, flip_axis, flip_indicator=None):
#   def _maybe_flip(data):
# def maybe_rot180(image, label, static_axis, rot180_k=None):
#   def _maybe_rot180(data):
# def intensity_shift(
#     image, label, per_class_intensity_scale, per_class_intensity_shift):
# def image_corruption(
#     image, label, reso, image_corrupt_ratio_mean, image_corrupt_ratio_stddev):


import functools
import os
import mesh_tensorflow as mtf
import numpy as np
import tensorflow.compat.v1 as tf  # tf
from six.moves import range
from mesh_tensorflow.experimental import data_aug_lib
from tensorflow.python.platform import flags

      if dataset_str != 'train':
        # Preprocess intensity, clip to 0 ~ 1.
        # The training set is already preprocessed.
        image = tf.clip_by_value(image / 1024.0 + 0.5, 0, 1)

      image = tf.reshape(image, spatial_dims)
      label = tf.reshape(label, spatial_dims)

      if dataset_str == 'eval' and FLAGS.sampled_2d_slices:
        return _get_stacked_2d_slices(image, label)

      if FLAGS.sampled_2d_slices:
        # Take random slices of images and label
        begin_idx = tf.random_uniform(
            shape=[], minval=0,
            maxval=FLAGS.ct_resolution - FLAGS.image_c + 1, dtype=tf.int32)
        slice_begin = [0, 0, begin_idx]
        slice_size = [FLAGS.ct_resolution, FLAGS.ct_resolution, FLAGS.image_c]

        image = tf.slice(image, slice_begin, slice_size)
        label = tf.slice(label, slice_begin, slice_size)

      if dataset_str == 'train':
        for flip_axis in [0, 1, 2]:
          image, label = data_aug_lib.maybe_flip(image, label, flip_axis)
        image, label = data_aug_lib.maybe_rot180(image, label, static_axis=2)
        image = data_aug_lib.intensity_shift(
            image, label,
            FLAGS.per_class_intensity_scale, FLAGS.per_class_intensity_shift)

# Path: mesh_tensorflow/experimental/data_aug_lib.py
# def _truncated_normal(mean, stddev):
# def _rand_noise(noise_mean, noise_dev, scale, shape):
# def projective_transform(
#     image, label, reso, image_translate_ratio, image_transform_ratio,
#     sampled_2d_slices=False):
#   def _projective_transform(data, proj_matrix, static_axis, interpolation):
# def maybe_add_noise(image, noise_shape, scale0, scale1,
#                     image_noise_probability, image_noise_ratio):
# def _gen_rand_mask(ratio_mean, ratio_stddev, scale, shape, smoothness=0):
# def maybe_gen_fake_data_based_on_real_data(
#     image, label, reso, min_fake_lesion_ratio, gen_fake_probability):
# def maybe_flip(image, label, flip_axis, flip_indicator=None):
#   def _maybe_flip(data):
# def maybe_rot180(image, label, static_axis, rot180_k=None):
#   def _maybe_rot180(data):
# def intensity_shift(
#     image, label, per_class_intensity_scale, per_class_intensity_shift):
# def image_corruption(
#     image, label, reso, image_corrupt_ratio_mean, image_corrupt_ratio_stddev):


# Path: mesh_tensorflow/experimental/unet.py
import functools
import os
import mesh_tensorflow as mtf
import numpy as np
import tensorflow.compat.v1 as tf  # tf
from six.moves import range
from mesh_tensorflow.experimental import data_aug_lib
from tensorflow.python.platform import flags


      if dataset_str != 'train':
        # Preprocess intensity, clip to 0 ~ 1.
        # The training set is already preprocessed.
        image = tf.clip_by_value(image / 1024.0 + 0.5, 0, 1)

      image = tf.reshape(image, spatial_dims)
      label = tf.reshape(label, spatial_dims)

      if dataset_str == 'eval' and FLAGS.sampled_2d_slices:
        return _get_stacked_2d_slices(image, label)

      if FLAGS.sampled_2d_slices:
        # Take random slices of images and label
        begin_idx = tf.random_uniform(
            shape=[], minval=0,
            maxval=FLAGS.ct_resolution - FLAGS.image_c + 1, dtype=tf.int32)
        slice_begin = [0, 0, begin_idx]
        slice_size = [FLAGS.ct_resolution, FLAGS.ct_resolution, FLAGS.image_c]

        image = tf.slice(image, slice_begin, slice_size)
        label = tf.slice(label, slice_begin, slice_size)

      if dataset_str == 'train':
        for flip_axis in [0, 1, 2]:
          image, label = data_aug_lib.maybe_flip(image, label, flip_axis)
        image, label = data_aug_lib.maybe_rot180(image, label, static_axis=2)
        image = data_aug_lib.intensity_shift(
            image, label,
            FLAGS.per_class_intensity_scale, FLAGS.per_class_intensity_shift)

        image = data_aug_lib.image_corruption(