# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model defination for the Mask-RCNN Model.

Defines model_fn of Mask-RCNN for TF Estimator. The model_fn includes Mask-RCNN
model architecture, loss function, learning rate schedule, and evaluation
procedure.

"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
import re
from absl import logging
import numpy as np
import six
import tensorflow.compat.v1 as tf
import tensorflow.compat.v2 as tf2

import anchors
import fpn
import heads
import learning_rates
import losses
import postprocess_ops
import resnet
import roi_ops
import spatial_transform_ops
import training_ops
import sys
sys.path.append('tpu/models/official/mnasnet')
import mnasnet_models

def create_optimizer(learning_rate, params):
  """Creates optimized based on the specified flags."""
  if params['optimizer'] == 'momentum':
    optimizer = tf.train.MomentumOptimizer(
        learning_rate, momentum=params['momentum'])
  elif params['optimizer'] == 'adam':
    optimizer = tf.train.AdamOptimizer(learning_rate)
  elif params['optimizer'] == 'adadelta':
    optimizer = tf.train.AdadeltaOptimizer(learning_rate)
  elif params['optimizer'] == 'adagrad':
    optimizer = tf.train.AdagradOptimizer(learning_rate)
  elif params['optimizer'] == 'rmsprop':
    optimizer = tf.train.RMSPropOptimizer(
        learning_rate, momentum=params['momentum'])
  elif params['optimizer'] == 'lars':
    try:
      from tensorflow.contrib.opt import LARSOptimizer  # pylint: disable=g-import-not-at-top

      optimizer = LARSOptimizer(
          learning_rate,
          momentum=params['momentum'],
          weight_decay=params['lars_weight_decay'],
          skip_list=['batch_normalization', 'bias'])
    except ImportError as e:
      logging.exception('LARSOptimizer is currently not supported '
                        'in TensorFlow 2.x.')
      raise e

  else:
    raise ValueError('Unsupported optimizer type %s.' % params['optimizer'])
  return optimizer

def remove_variables(variables, prefix):
  """Removes low-level variables from the input.

  Removing low-level parameters (e.g., initial convolution layer) from training
  usually leads to higher training speed and slightly better testing accuracy.
  The intuition is that the low-level architecture (e.g., ResNet-50) is able to
  capture low-level features such as edges; therefore, it does not need to be
  fine-tuned for the detection task.

  Args:
    variables: all the variables in training
    prefix: prefix of backbone

  Returns:
    var_list: a list containing variables for training

  """
  # Freeze at conv2 based on reference model.
  # Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/core/config.py#L194  # pylint: disable=line-too-long
  remove_list = []
  remove_list.append(prefix + 'conv2d/')
  remove_list.append(prefix + 'batch_normalization/')
  for i in range(1, 11):
    remove_list.append(prefix + 'conv2d_{}/'.format(i))
    remove_list.append(prefix + 'batch_normalization_{}/'.format(i))

  def _is_kept(variable):
    for rm_str in remove_list:
      if rm_str in variable.name:
        return False
    return True

  var_list = [v for v in variables if _is_kept(v)]
  return var_list

def compute_model_statistics(batch_size, is_training=True):
  """Compute number of parameters and FLOPS."""
  num_trainable_params = np.sum(
      [np.prod(var.get_shape().as_list()) for var in tf.trainable_variables()])
  logging.info('number of trainable params: %d', num_trainable_params)

  options = tf.profiler.ProfileOptionBuilder.float_operation()
  options['output'] = 'none'
  flops = tf.profiler.profile(
      tf.get_default_graph(), options=options).total_float_ops
  flops_per_image = flops / batch_size
  if is_training:
    logging.info(
        'number of FLOPS per image: %f in training', flops_per_image)
  else:
    logging.info(
        'number of FLOPS per image: %f in eval', flops_per_image)

def build_model_graph(features, labels, is_training, params):
  """Builds the forward model graph."""
  use_batched_nms = (not params['use_tpu'] and params['use_batched_nms'])
  is_gpu_inference = (not is_training and use_batched_nms)
  model_outputs = {}

  if is_training and params['transpose_input']:
    if (params['backbone'].startswith('resnet') and
        params['conv0_space_to_depth_block_size'] > 0):
      features['images'] = tf.transpose(features['images'], [2, 0, 1, 3])
    else:
      features['images'] = tf.transpose(features['images'], [3, 0, 1, 2])

  batch_size, image_height, image_width, _ = (
      features['images'].get_shape().as_list())

  conv0_space_to_depth_block_size = 0
  if (is_training and
      (params['backbone'].startswith('resnet') and
       params['conv0_space_to_depth_block_size'] > 0)):
    conv0_space_to_depth_block_size = params['conv0_space_to_depth_block_size']
    image_height *= conv0_space_to_depth_block_size
    image_width *= conv0_space_to_depth_block_size

  if 'source_ids' not in features:
    features['source_ids'] = -1 * tf.ones([batch_size], dtype=tf.float32)

  all_anchors = anchors.Anchors(params['min_level'], params['max_level'],
                                params['num_scales'], params['aspect_ratios'],
                                params['anchor_scale'],
                                (image_height, image_width))

  if 'resnet' in params['backbone']:
    with tf.variable_scope(params['backbone']):
      resnet_fn = resnet.resnet_v1(
          params['backbone'],
          conv0_kernel_size=params['conv0_kernel_size'],
          conv0_space_to_depth_block_size=conv0_space_to_depth_block_size,
          num_batch_norm_group=params['num_batch_norm_group'])
      backbone_feats = resnet_fn(
          features['images'],
          (params['is_training_bn'] and is_training))
  elif 'mnasnet' in params['backbone']:
    with tf.variable_scope(params['backbone']):
      _, endpoints = mnasnet_models.build_mnasnet_base(
          features['images'],
          params['backbone'],
          training=(params['is_training_bn'] and is_training),
          override_params={'use_keras': False})

      backbone_feats = {
          2: endpoints['reduction_2'],
          3: endpoints['reduction_3'],
          4: endpoints['reduction_4'],
          5: endpoints['reduction_5'],
      }
  else:
    raise ValueError('Not a valid backbone option: %s' % params['backbone'])

  fpn_feats = fpn.fpn(
      backbone_feats, params['min_level'], params['max_level'])
  model_outputs.update({
      'fpn_features': fpn_feats,
  })

  rpn_score_outputs, rpn_box_outputs = heads.rpn_head(
      fpn_feats,
      params['min_level'], params['max_level'],
      len(params['aspect_ratios'] * params['num_scales']))

  if is_training:
    rpn_pre_nms_topn = params['rpn_pre_nms_topn']
    rpn_post_nms_topn = params['rpn_post_nms_topn']
  else:
    rpn_pre_nms_topn = params['test_rpn_pre_nms_topn']
    rpn_post_nms_topn = params['test_rpn_post_nms_topn']

  rpn_box_scores, rpn_box_rois = roi_ops.multilevel_propose_rois(
      rpn_score_outputs,
      rpn_box_outputs,
      all_anchors,
      features['image_info'],
      rpn_pre_nms_topn,
      rpn_post_nms_topn,
      params['rpn_nms_threshold'],
      params['rpn_min_size'],
      bbox_reg_weights=None,
      use_batched_nms=use_batched_nms)
  rpn_box_rois = tf.to_float(rpn_box_rois)
  if is_training:
    rpn_box_rois = tf.stop_gradient(rpn_box_rois)
    rpn_box_scores = tf.stop_gradient(rpn_box_scores)

  if is_training:
    # Sampling
    box_targets, class_targets, rpn_box_rois, proposal_to_label_map = (
        training_ops.proposal_label_op(
            rpn_box_rois,
            labels['gt_boxes'],
            labels['gt_classes'],
            batch_size_per_im=params['batch_size_per_im'],
            fg_fraction=params['fg_fraction'],
            fg_thresh=params['fg_thresh'],
            bg_thresh_hi=params['bg_thresh_hi'],
            bg_thresh_lo=params['bg_thresh_lo']))

  # Performs multi-level RoIAlign.
  box_roi_features = spatial_transform_ops.multilevel_crop_and_resize(
      fpn_feats, rpn_box_rois, output_size=7, is_gpu_inference=is_gpu_inference)

  class_outputs, box_outputs, _ = heads.box_head(
      box_roi_features, num_classes=params['num_classes'],
      mlp_head_dim=params['fast_rcnn_mlp_head_dim'])

  if not is_training:
    if is_gpu_inference:
      generate_detections_fn = postprocess_ops.generate_detections_gpu
    else:
      generate_detections_fn = postprocess_ops.generate_detections_tpu
    detections = generate_detections_fn(
        class_outputs,
        box_outputs,
        rpn_box_rois,
        features['image_info'],
        params['test_rpn_post_nms_topn'],
        params['test_detections_per_image'],
        params['test_nms'],
        params['bbox_reg_weights'])

    model_outputs.update({
        'num_detections': detections[0],
        'detection_boxes': detections[1],
        'detection_classes': detections[2],
        'detection_scores': detections[3],
    })
  else:
    encoded_box_targets = training_ops.encode_box_targets(
        rpn_box_rois, box_targets, class_targets, params['bbox_reg_weights'])
    model_outputs.update({
        'rpn_score_outputs': rpn_score_outputs,
        'rpn_box_outputs': rpn_box_outputs,
        'class_outputs': class_outputs,
        'box_outputs': box_outputs,
        'class_targets': class_targets,
        'box_targets': encoded_box_targets,
        'box_rois': rpn_box_rois,
    })

  # Faster-RCNN mode.
  if not params['include_mask']:
    return model_outputs

  # Mask sampling
  if not is_training:
    selected_box_rois = model_outputs['detection_boxes']
    class_indices = model_outputs['detection_classes']
    # If using GPU for inference, delay the cast until when Gather ops show up
    # since GPU inference supports float point better.
    # TODO(laigd): revisit this when newer versions of GPU libraries is
    # released.
    if not is_gpu_inference:
      class_indices = tf.to_int32(class_indices)
  else:
    (selected_class_targets, selected_box_targets, selected_box_rois,
     proposal_to_label_map) = (
         training_ops.select_fg_for_masks(
             class_targets, box_targets, rpn_box_rois,
             proposal_to_label_map,
             max_num_fg=int(
                 params['batch_size_per_im'] * params['fg_fraction'])))
    class_indices = tf.to_int32(selected_class_targets)

  mask_roi_features = spatial_transform_ops.multilevel_crop_and_resize(
      fpn_feats,
      selected_box_rois,
      output_size=14,
      is_gpu_inference=is_gpu_inference)
  mask_outputs = heads.mask_head(
      mask_roi_features,
      class_indices,
      num_classes=params['num_classes'],
      mrcnn_resolution=params['mrcnn_resolution'],
      is_gpu_inference=is_gpu_inference)

  # Print #parameters and #FLOPs in model.
  compute_model_statistics(batch_size, is_training=is_training)

  if is_training:
    mask_targets = training_ops.get_mask_targets(
        selected_box_rois, proposal_to_label_map, selected_box_targets,
        labels['cropped_gt_masks'], params['mrcnn_resolution'])
    model_outputs.update({
        'mask_outputs': mask_outputs,
        'mask_targets': mask_targets,
        'selected_class_targets': selected_class_targets,
    })
  else:
    model_outputs.update({
        'detection_masks': tf.nn.sigmoid(mask_outputs),
    })

  return model_outputs

def _build_assigment_map(optimizer, prefix=None, skip_variables_regex=None):
  """Generate assigment map for loading checkpoints."""
  optimizer_vars = set([var.name for var in optimizer.variables()])
  all_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=prefix)
  if not prefix:
    prefix = ''
  assignment_map = {}
  for var in all_vars:
    if var.name not in optimizer_vars:
      var_name = var.name
      # Trim the index of the variable.
      if ':' in var_name:
        var_name = var_name[:var_name.rindex(':')]
      if skip_variables_regex and re.match(skip_variables_regex,
                                           var_name[len(prefix):]):
        continue
      assignment_map[var_name[len(prefix):]] = var
  return assignment_map

def _model_fn(features, labels, mode, params, variable_filter_fn=None):
  """Model defination for the Mask-RCNN model based on ResNet.

  Args:
    features: the input image tensor and auxiliary information, such as
      `image_info` and `source_ids`. The image tensor has a shape of
      [batch_size, height, width, 3]. The height and width are fixed and equal.
    labels: the input labels in a dictionary. The labels include score targets
      and box targets which are dense label maps. The labels are generated from
      get_input_fn function in data/dataloader.py
    mode: the mode of TPUEstimator including TRAIN, EVAL, and PREDICT.
    params: the dictionary defines hyperparameters of model. The default
      settings are in default_hparams function in this file.
    variable_filter_fn: the filter function that takes trainable_variables and
      returns the variable list after applying the filter rule.

  Returns:
    tpu_spec: the TPUEstimatorSpec to run training, evaluation, or prediction.
  """
  if mode == tf.estimator.ModeKeys.PREDICT:
    if params['include_groundtruth_in_features'] and ('labels' in features):
      # In include groundtruth for eval.
      labels = features['labels']
    else:
      labels = None
    if 'features' in features:
      features = features['features']
      # Otherwise, it is in export mode, the features is past in directly.

  if params['precision'] == 'bfloat16':
    with tf.tpu.bfloat16_scope():
      model_outputs = build_model_graph(
          features, labels,
          (mode == tf.estimator.ModeKeys.TRAIN or
           mode == tf.estimator.ModeKeys.EVAL),
          params)
      model_outputs.update({
          'source_id': features['source_ids'],
          'image_info': features['image_info'],
      })
      def cast_outputs_to_float(d):
        for k, v in sorted(six.iteritems(d)):
          if isinstance(v, dict):
            cast_outputs_to_float(v)
          else:
            d[k] = tf.cast(v, tf.float32)
      cast_outputs_to_float(model_outputs)
  else:
    model_outputs = build_model_graph(
        features, labels,
        (mode == tf.estimator.ModeKeys.TRAIN or
         mode == tf.estimator.ModeKeys.EVAL),
        params)
    model_outputs.update({
        'source_id': features['source_ids'],
        'image_info': features['image_info'],
    })

  # First check if it is in PREDICT or EVAL mode to fill out predictions.
  # Predictions are used during the eval step to generate metrics.
  predictions = {}
  if (mode == tf.estimator.ModeKeys.PREDICT or
      mode == tf.estimator.ModeKeys.EVAL):
    if 'orig_images' in features:
      model_outputs['orig_images'] = features['orig_images']
    if labels and params['include_groundtruth_in_features']:
      # Labels can only be embedded in predictions. The predition cannot output
      # dictionary as a value.
      predictions.update(labels)
    model_outputs.pop('fpn_features', None)
    predictions.update(model_outputs)
    # If we are doing PREDICT, we can return here.
    if mode == tf.estimator.ModeKeys.PREDICT:
      if params['use_tpu']:
        return tf.estimator.tpu.TPUEstimatorSpec(mode=mode,
                                                 predictions=predictions)
      return tf.estimator.EstimatorSpec(mode=mode,
                                        predictions=predictions)

  # Set up training loss and learning rate.
  global_step = tf.train.get_or_create_global_step()
  if params['learning_rate_type'] == 'step':
    learning_rate = learning_rates.step_learning_rate_with_linear_warmup(
        global_step,
        params['init_learning_rate'],
        params['warmup_learning_rate'],
        params['warmup_steps'],
        params['learning_rate_levels'],
        params['learning_rate_steps'])
  elif params['learning_rate_type'] == 'cosine':
    learning_rate = learning_rates.cosine_learning_rate_with_linear_warmup(
        global_step,
        params['init_learning_rate'],
        params['warmup_learning_rate'],
        params['warmup_steps'],
        params['total_steps'])
  else:
    raise ValueError('Unsupported learning rate type: `{}`!'
                     .format(params['learning_rate_type']))
  # score_loss and box_loss are for logging. only total_loss is optimized.
  total_rpn_loss, rpn_score_loss, rpn_box_loss = losses.rpn_loss(
      model_outputs['rpn_score_outputs'], model_outputs['rpn_box_outputs'],
      labels, params)

  (total_fast_rcnn_loss, fast_rcnn_class_loss,
   fast_rcnn_box_loss) = losses.fast_rcnn_loss(
       model_outputs['class_outputs'], model_outputs['box_outputs'],
       model_outputs['class_targets'], model_outputs['box_targets'], params)
  # Only training has the mask loss. Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/model_builder.py  # pylint: disable=line-too-long
  if mode == tf.estimator.ModeKeys.TRAIN and params['include_mask']:
    mask_loss = losses.mask_rcnn_loss(
        model_outputs['mask_outputs'], model_outputs['mask_targets'],
        model_outputs['selected_class_targets'], params)
  else:
    mask_loss = 0.
  if variable_filter_fn and ('resnet' in params['backbone']):
    var_list = variable_filter_fn(tf.trainable_variables(),
                                  params['backbone'] + '/')
  else:
    var_list = tf.trainable_variables()
  l2_regularization_loss = params['l2_weight_decay'] * tf.add_n([
      tf.nn.l2_loss(v)
      for v in var_list
      if 'batch_normalization' not in v.name and 'bias' not in v.name
  ])
  total_loss = (total_rpn_loss + total_fast_rcnn_loss + mask_loss +
                l2_regularization_loss)

  if mode == tf.estimator.ModeKeys.EVAL:
    # Predictions can only contain a dict of tensors, not a dict of dict of
    # tensors. These outputs are not used for eval purposes.
    del predictions['rpn_score_outputs']
    del predictions['rpn_box_outputs']
    if params['use_tpu']:
      # For now, just return loss against the eval dataset.
      # In the future, calculate COCO metrics against the eval dataset.
      return tf.estimator.tpu.TPUEstimatorSpec(mode=mode,
                                               predictions=predictions,
                                               eval_metrics=None,
                                               loss=total_loss)
    return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions,
                                      eval_metrics=None, loss=total_loss)

  host_call = None
  if mode == tf.estimator.ModeKeys.TRAIN:
    optimizer = create_optimizer(learning_rate, params)
    if params['use_tpu']:
      optimizer = tf.tpu.CrossShardOptimizer(optimizer)

    scaffold_fn = None
    if params['warm_start_path']:

      def warm_start_scaffold_fn():
        logging.info(
            'model_fn warm start from: %s,', params['warm_start_path'])
        assignment_map = _build_assigment_map(
            optimizer,
            prefix=None,
            skip_variables_regex=params['skip_checkpoint_variables'])
        tf.train.init_from_checkpoint(params['warm_start_path'], assignment_map)
        return tf.train.Scaffold()

      scaffold_fn = warm_start_scaffold_fn

    elif params['checkpoint']:

      def backbone_scaffold_fn():
        """Loads pretrained model through scaffold function."""
        # Exclude all variable of optimizer.
        vars_to_load = _build_assigment_map(
            optimizer,
            prefix=params['backbone'] + '/',
            skip_variables_regex=params['skip_checkpoint_variables'])
        tf.train.init_from_checkpoint(params['checkpoint'], vars_to_load)
        if not vars_to_load:
          raise ValueError('Variables to load is empty.')
        return tf.train.Scaffold()

      scaffold_fn = backbone_scaffold_fn

    # Batch norm requires update_ops to be added as a train_op dependency.
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    grads_and_vars = optimizer.compute_gradients(total_loss, var_list)
    if params['global_gradient_clip_ratio'] > 0:
      # Clips the gradients for training stability.
      # Refer: https://arxiv.org/abs/1211.5063
      with tf.name_scope('clipping'):
        old_grads, variables = zip(*grads_and_vars)
        num_weights = sum(
            g.shape.num_elements() for g in old_grads if g is not None)
        clip_norm = params['global_gradient_clip_ratio'] * math.sqrt(
            num_weights)
        logging.info(
            'Global clip norm set to %g for %d variables with %d elements.',
            clip_norm, sum(1 for g in old_grads if g is not None),
            num_weights)
        gradients, _ = tf.clip_by_global_norm(old_grads, clip_norm)
    else:
      gradients, variables = zip(*grads_and_vars)
    grads_and_vars = []
    # Special treatment for biases (beta is named as bias in reference model)
    # Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/optimizer.py#L113  # pylint: disable=line-too-long
    for grad, var in zip(gradients, variables):
      if grad is not None and ('beta' in var.name or 'bias' in var.name):
        grad = 2.0 * grad
      grads_and_vars.append((grad, var))

    with tf.control_dependencies(update_ops):
      train_op = optimizer.apply_gradients(
          grads_and_vars, global_step=global_step)

    if params['use_host_call']:
      def host_call_fn(global_step, total_loss, total_rpn_loss, rpn_score_loss,
                       rpn_box_loss, total_fast_rcnn_loss, fast_rcnn_class_loss,
                       fast_rcnn_box_loss, mask_loss, l2_regularization_loss,
                       learning_rate):
        """Training host call. Creates scalar summaries for training metrics.

        This function is executed on the CPU and should not directly reference
        any Tensors in the rest of the `model_fn`. To pass Tensors from the
        model to the `metric_fn`, provide as part of the `host_call`. See
        https://www.tensorflow.org/api_docs/python/tf/estimator/tpu/TPUEstimatorSpec
        for more information.

        Arguments should match the list of `Tensor` objects passed as the second
        element in the tuple passed to `host_call`.

        Args:
          global_step: `Tensor with shape `[batch, ]` for the global_step.
          total_loss: `Tensor` with shape `[batch, ]` for the training loss.
          total_rpn_loss: `Tensor` with shape `[batch, ]` for the training RPN
            loss.
          rpn_score_loss: `Tensor` with shape `[batch, ]` for the training RPN
            score loss.
          rpn_box_loss: `Tensor` with shape `[batch, ]` for the training RPN
            box loss.
          total_fast_rcnn_loss: `Tensor` with shape `[batch, ]` for the
            training Mask-RCNN loss.
          fast_rcnn_class_loss: `Tensor` with shape `[batch, ]` for the
            training Mask-RCNN class loss.
          fast_rcnn_box_loss: `Tensor` with shape `[batch, ]` for the
            training Mask-RCNN box loss.
          mask_loss: `Tensor` with shape `[batch, ]` for the training Mask-RCNN
            mask loss.
          l2_regularization_loss: `Tensor` with shape `[batch, ]` for the
            regularization loss.
          learning_rate: `Tensor` with shape `[batch, ]` for the learning_rate.

        Returns:
          List of summary ops to run on the CPU host.
        """
        # Outfeed supports int32 but global_step is expected to be int64.
        global_step = tf.reduce_mean(global_step)
        # Host call fns are executed FLAGS.iterations_per_loop times after one
        # TPU loop is finished, setting max_queue value to the same as number of
        # iterations will make the summary writer only flush the data to storage
        # once per loop.
        with (tf2.summary.create_file_writer(
            params['model_dir'],
            max_queue=params['iterations_per_loop']).as_default()):
          with tf2.summary.record_if(True):
            tf2.summary.scalar(
                'total_loss', tf.reduce_mean(total_loss), step=global_step)
            tf2.summary.scalar(
                'total_rpn_loss', tf.reduce_mean(total_rpn_loss),
                step=global_step)
            tf2.summary.scalar(
                'rpn_score_loss', tf.reduce_mean(rpn_score_loss),
                step=global_step)
            tf2.summary.scalar(
                'rpn_box_loss', tf.reduce_mean(rpn_box_loss), step=global_step)
            tf2.summary.scalar(
                'total_fast_rcnn_loss', tf.reduce_mean(total_fast_rcnn_loss),
                step=global_step)
            tf2.summary.scalar(
                'fast_rcnn_class_loss', tf.reduce_mean(fast_rcnn_class_loss),
                step=global_step)
            tf2.summary.scalar(
                'fast_rcnn_box_loss', tf.reduce_mean(fast_rcnn_box_loss),
                step=global_step)
            if params['include_mask']:
              tf2.summary.scalar(
                  'mask_loss', tf.reduce_mean(mask_loss), step=global_step)
            tf2.summary.scalar(
                'l2_regularization_loss',
                tf.reduce_mean(l2_regularization_loss),
                step=global_step)
            tf2.summary.scalar(
                'learning_rate', tf.reduce_mean(learning_rate),
                step=global_step)

            return tf.summary.all_v2_summary_ops()

      # To log the loss, current learning rate, and epoch for Tensorboard, the
      # summary op needs to be run on the host CPU via host_call. host_call
      # expects [batch_size, ...] Tensors, thus reshape to introduce a batch
      # dimension. These Tensors are implicitly concatenated to
      # [params['batch_size']].
      global_step_t = tf.reshape(global_step, [1])
      total_loss_t = tf.reshape(total_loss, [1])
      total_rpn_loss_t = tf.reshape(total_rpn_loss, [1])
      rpn_score_loss_t = tf.reshape(rpn_score_loss, [1])
      rpn_box_loss_t = tf.reshape(rpn_box_loss, [1])
      total_fast_rcnn_loss_t = tf.reshape(total_fast_rcnn_loss, [1])
      fast_rcnn_class_loss_t = tf.reshape(fast_rcnn_class_loss, [1])
      fast_rcnn_box_loss_t = tf.reshape(fast_rcnn_box_loss, [1])
      mask_loss_t = tf.reshape(mask_loss, [1])
      l2_regularization_loss = tf.reshape(l2_regularization_loss, [1])
      learning_rate_t = tf.reshape(learning_rate, [1])
      host_call = (host_call_fn,
                   [global_step_t, total_loss_t, total_rpn_loss_t,
                    rpn_score_loss_t, rpn_box_loss_t, total_fast_rcnn_loss_t,
                    fast_rcnn_class_loss_t, fast_rcnn_box_loss_t,
                    mask_loss_t, l2_regularization_loss, learning_rate_t])
  else:
    train_op = None
    scaffold_fn = None

  if params['use_tpu']:
    return tf.estimator.tpu.TPUEstimatorSpec(
        mode=mode,
        loss=total_loss,
        train_op=train_op,
        host_call=host_call,
        scaffold_fn=scaffold_fn)
  return tf.estimator.EstimatorSpec(
      mode=mode, loss=total_loss, train_op=train_op)

def mask_rcnn_model_fn(features, labels, mode, params):
  """Mask-RCNN model."""
  with tf.variable_scope('', reuse=tf.AUTO_REUSE):
    return _model_fn(
        features,
        labels,
        mode,
        params,
        variable_filter_fn=remove_variables)


import cv2
import numpy as np
import pyautogui
import random
import time
from imutils.object_detection import non_max_suppression
from PIL import Image
import pytesseract
import os

pytesseract.pytesseract.tesseract_cmd = "/usr/local/Cellar/tesseract/4.0.0/bin/tesseract"

def region_grabber(region):
    """
    grabs a region (topx, topy, bottomx, bottomy)
    to the tuple (topx, topy, width, height)
    input : a tuple containing the 4 coordinates of the region to capture
    output : a PIL image of the area selected.
    """
    x1 = region[0]
    y1 = region[1]
    width = region[2]-x1
    height = region[3]-y1

    return pyautogui.screenshot(region=(x1,y1,width,height))

def image_search_area(image, x1, y1, x2, y2, precision=0.8, im=None):
    """
    Searchs for an image within an area
    input :
    image : path to the image file (see opencv imread for supported types)
    x1 : top left x value
    y1 : top left y value
    x2 : bottom right x value
    y2 : bottom right y value
    precision : the higher, the lesser tolerant and fewer false positives are found default is 0.8
    im : a PIL image, usefull if you intend to search the same unchanging region for several elements
    returns :
    the top left corner coordinates of the element if found as an array [x,y] or [-1,-1] if not
    """
    if im is None :
        im = region_grabber(region=(x1, y1, x2, y2))
        #im.save('testarea.png') usefull for debugging purposes, this will save the captured region as "testarea.png"

    img_rgb = np.array(im)
    img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2GRAY)
    template = cv2.imread(image, 0)

    res = cv2.matchTemplate(img_gray, template, cv2.TM_CCOEFF_NORMED)
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
    if max_val < precision:
        return [-1, -1]
    return max_loc

def click_image(image, pos, action, timestamp, offset=5):
    """
    click on the center of an image with a bit of random.
    eg, if an image is 100*100 with an offset of 5 it may click at 52,50 the first time and then 55,53 etc
    Usefull to avoid anti-bot monitoring while staying precise.
    this function doesn't search for the image, it's only ment for easy clicking on the images.
    input :
    image : path to the image file (see opencv imread for supported types)
    pos : array containing the position of the top left corner of the image [x,y]
    action : button of the mouse to activate : "left" "right" "middle", see pyautogui.click documentation for more info
    time : time taken for the mouse to move from where it was to the new position
    """
    img = cv2.imread(image)
    height, width, channels = img.shape
    pyautogui.moveTo(pos[0] + r(width / 2, offset), pos[1] + r(height / 2,offset),
                     timestamp)
    pyautogui.click(button=action)

def image_search(image, precision=0.8):
    """
    Searchs for an image on the screen
    input :
    image : path to the image file (see opencv imread for supported types)
    precision : the higher, the lesser tolerant and fewer false positives are found default is 0.8
    im : a PIL image, usefull if you intend to search the same unchanging region for several elements
    returns :
    the top left corner coordinates of the element if found as an array [x,y] or [-1,-1] if not
    """
    im = pyautogui.screenshot()
    #im.save('testarea.png') usefull for debugging purposes, this will save the captured region as "testarea.png"
    img_rgb = np.array(im)
    img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2GRAY)
    template = cv2.imread(image, 0)
    template.shape[::-1]

    res = cv2.matchTemplate(img_gray, template, cv2.TM_CCOEFF_NORMED)
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
    if max_val < precision:
        return [-1,-1]
    return max_loc

def imagesearch_loop(image, timesample, precision=0.8):
    """
    Searchs for an image on screen continuously until it's found.
    input :
    image : path to the image file (see opencv imread for supported types)
    time : Waiting time after failing to find the image
    precision : the higher, the lesser tolerant and fewer false positives are found default is 0.8
    returns :
    the top left corner coordinates of the element if found as an array [x,y]
    """
    pos = image_search(image, precision)
    while pos[0] == -1:
        print(image+" not found, waiting")
        time.sleep(timesample)
        pos = image_search(image, precision)
    return pos

def imagesearch_numLoop(image, timesample, maxSamples, precision=0.8):
    """
    Searchs for an image on screen continuously until it's found or max number of samples reached.
    input :
    image : path to the image file (see opencv imread for supported types)
    time : Waiting time after failing to find the image
    maxSamples: maximum number of samples before function times out.
    precision : the higher, the lesser tolerant and fewer false positives are found default is 0.8
    returns :
    the top left corner coordinates of the element if found as an array [x,y]
    """
    pos = image_search(image, precision)
    count = 0
    while pos[0] == -1:
        print(image+" not found, waiting")
        time.sleep(timesample)
        pos = image_search(image, precision)
        count = count + 1
        if count>maxSamples:
            break
    return pos

def imagesearch_region_loop(image, timesample, x1, y1, x2, y2, precision=0.8):
    """
    Searchs for an image on a region of the screen continuously until it's found.
    input :
    image : path to the image file (see opencv imread for supported types)
    time : Waiting time after failing to find the image
    x1 : top left x value
    y1 : top left y value
    x2 : bottom right x value
    y2 : bottom right y value
    precision : the higher, the lesser tolerant and fewer false positives are found default is 0.8
    returns :
    the top left corner coordinates of the element as an array [x,y]
    """
    pos = image_search_area(image, x1,y1,x2,y2, precision)

    while pos[0] == -1:
        time.sleep(timesample)
        pos = image_search_area(image, x1, y1, x2, y2, precision)
    return pos

def imagesearch_count(image, precision=0.9):
    """
    Searches for an image on the screen and counts the number of occurrences.
    input :
    image : path to the target image file (see opencv imread for supported types)
    precision : the higher, the lesser tolerant and fewer false positives are found default is 0.9
    returns :
    the number of times a given image appears on the screen.
    optionally an output image with all the occurances boxed with a red outline.
    """
    img_rgb = pyautogui.screenshot()
    img_rgb = np.array(img_rgb)
    img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2GRAY)
    template = cv2.imread(image, 0)
    w, h = template.shape[::-1]
    res = cv2.matchTemplate(img_gray, template, cv2.TM_CCOEFF_NORMED)
    loc = np.where(res >= precision)
    count = 0
    for pt in zip(*loc[::-1]):  # Swap columns and rows
        #cv2.rectangle(img_rgb, pt, (pt[0] + w, pt[1] + h), (0, 0, 255), 2) // Uncomment to draw boxes around found occurances
        count = count + 1
    #cv2.imwrite('result.png', img_rgb) // Uncomment to write output image with boxes drawn around occurances
    return count

def r(num, rand):
    return num + rand*random.random()

def save_textarea(image, east="/Users/luizeduardocartolano/Dropbox/DUDU/python-test/automate_test/frozen_east_text_detection.pb",
                   min_confidence=0.35, width=640, height=640):
    # load the input image and grab the image dimensions
    image = cv2.imread(image)
    orig = image.copy()
    (origH, origW) = image.shape[:2]

    # set the new width and height and then determine the ratio in change
    # for both the width and height
    (newW, newH) = (width, height)
    rW = origW / float(newW)
    rH = origH / float(newH)

    # resize the image and grab the new image dimensions
    image = cv2.resize(image, (newW, newH))
    (H, W) = image.shape[:2]

    # define the two output layer names for the EAST detector model that
    # we are interested -- the first is the output probabilities and the
    # second can be used to derive the bounding box coordinates of text
    layerNames = [
        "feature_fusion/Conv_7/Sigmoid",
        "feature_fusion/concat_3"]

    # load the pre-trained EAST text detector
    print("[INFO] loading EAST text detector...")
    net = cv2.dnn.readNet(east)

    # construct a blob from the image and then perform a forward pass of
    # the model to obtain the two output layer sets
    blob = cv2.dnn.blobFromImage(image, 1.0, (W, H),
        (123.68, 116.78, 103.94), swapRB=True, crop=False)
    net.setInput(blob)
    (scores, geometry) = net.forward(layerNames)

    # decode the predictions, then  apply non-maxima suppression to
    # suppress weak, overlapping bounding boxes
    (rects, confidences) = decode_predictions(scores, geometry, min_confidence)
    boxes = non_max_suppression(np.array(rects), probs=confidences)

    # initialize the list of results
    results = []

    # loop over the bounding boxes
    for (startX, startY, endX, endY) in boxes:
        # scale the bounding box coordinates based on the respective
        # ratios
        startX = int(startX * rW)
        startY = int(startY * rH)
        endX = int(endX * rW)
        endY = int(endY * rH)

        temp = orig[startY:startY+(endY-startY), startX:startX+(endX-startX)]
        cv2.imshow("Text Detection", temp)
        # cv2.imshow("Text Detection", orig)
        cv2.waitKey(0)

def text_detection(image, east="/Users/luizeduardocartolano/Dropbox/DUDU/python-test/automate_test/frozen_east_text_detection.pb",
                   min_confidence=0.35, width=640, height=640):
    # load the input image and grab the image dimensions
    image = cv2.imread(image)
    orig = image.copy()
    (origH, origW) = image.shape[:2]

    # set the new width and height and then determine the ratio in change
    # for both the width and height
    (newW, newH) = (width, height)
    rW = origW / float(newW)
    rH = origH / float(newH)

    # resize the image and grab the new image dimensions
    image = cv2.resize(image, (newW, newH))
    (H, W) = image.shape[:2]

    # define the two output layer names for the EAST detector model that
    # we are interested -- the first is the output probabilities and the
    # second can be used to derive the bounding box coordinates of text
    layerNames = [
        "feature_fusion/Conv_7/Sigmoid",
        "feature_fusion/concat_3"]

    # load the pre-trained EAST text detector
    print("[INFO] loading EAST text detector...")
    net = cv2.dnn.readNet(east)

    # construct a blob from the image and then perform a forward pass of
    # the model to obtain the two output layer sets
    blob = cv2.dnn.blobFromImage(image, 1.0, (W, H),
        (123.68, 116.78, 103.94), swapRB=True, crop=False)
    net.setInput(blob)
    (scores, geometry) = net.forward(layerNames)

    # decode the predictions, then  apply non-maxima suppression to
    # suppress weak, overlapping bounding boxes
    (rects, confidences) = decode_predictions(scores, geometry, min_confidence)
    boxes = non_max_suppression(np.array(rects), probs=confidences)

    # initialize the list of results
    results = []

    # loop over the bounding boxes
    for (startX, startY, endX, endY) in boxes:
        # scale the bounding box coordinates based on the respective
        # ratios
        startX = int(startX * rW)
        startY = int(startY * rH)
        endX = int(endX * rW)
        endY = int(endY * rH)

        # extract the actual padded ROI
        cv2.rectangle(orig, (startX, startY), (endX, endY), (0, 0, 255), 2)

        # show image
        cv2.imshow("Text Detection", orig)
        cv2.waitKey(0)

def decode_predictions(scores, geometry, min_confidence):
    # grab the number of rows and columns from the scores volume, then
    # initialize our set of bounding box rectangles and corresponding
    # confidence scores
    (numRows, numCols) = scores.shape[2:4]
    rects = []
    confidences = []

    # loop over the number of rows
    for y in range(0, numRows):
        # extract the scores (probabilities), followed by the
        # geometrical data used to derive potential bounding box
        # coordinates that surround text
        scoresData = scores[0, 0, y]
        xData0 = geometry[0, 0, y]
        xData1 = geometry[0, 1, y]
        xData2 = geometry[0, 2, y]
        xData3 = geometry[0, 3, y]
        anglesData = geometry[0, 4, y]

        # loop over the number of columns
        for x in range(0, numCols):
            # if our score does not have sufficient probability,
            # ignore it
            if scoresData[x] < min_confidence:
                continue

            # compute the offset factor as our resulting feature
            # maps will be 4x smaller than the input image
            (offsetX, offsetY) = (x * 4.0, y * 4.0)

            # extract the rotation angle for the prediction and
            # then compute the sin and cosine
            angle = anglesData[x]
            cos = np.cos(angle)
            sin = np.sin(angle)

            # use the geometry volume to derive the width and height
            # of the bounding box
            h = xData0[x] + xData2[x]
            w = xData1[x] + xData3[x]

            # compute both the starting and ending (x, y)-coordinates
            # for the text prediction bounding box
            endX = int(offsetX + (cos * xData1[x]) + (sin * xData2[x]))
            endY = int(offsetY - (sin * xData1[x]) + (cos * xData2[x]))
            startX = int(endX - w)
            startY = int(endY - h)

            # add the bounding box coordinates and probability score
            # to our respective lists
            rects.append((startX, startY, endX, endY))
            confidences.append(scoresData[x])

    # return a tuple of the bounding boxes and associated confidences
    return (rects, confidences)

def extract_text(image, east="/Users/luizeduardocartolano/Dropbox/DUDU/python-test/automate_test/frozen_east_text_detection.pb", min_confidence=0.35, width=640, height=640, padding=0):
    # load the input image and grab the image dimensions
    image = cv2.imread(image)
    orig = image.copy()
    (origH, origW) = image.shape[:2]

    # set the new width and height and then determine the ratio in change
    # for both the width and height
    (newW, newH) = (width, height)
    rW = origW / float(newW)
    rH = origH / float(newH)

    # resize the image and grab the new image dimensions
    image = cv2.resize(image, (newW, newH))
    (H, W) = image.shape[:2]

    # define the two output layer names for the EAST detector model that
    # we are interested -- the first is the output probabilities and the
    # second can be used to derive the bounding box coordinates of text
    layerNames = [
        "feature_fusion/Conv_7/Sigmoid",
        "feature_fusion/concat_3"]

    # load the pre-trained EAST text detector
    print("[INFO] loading EAST text detector...")
    net = cv2.dnn.readNet(east)

    # construct a blob from the image and then perform a forward pass of
    # the model to obtain the two output layer sets
    blob = cv2.dnn.blobFromImage(image, 1.0, (W, H),
        (123.68, 116.78, 103.94), swapRB=True, crop=False)
    net.setInput(blob)
    (scores, geometry) = net.forward(layerNames)

    # decode the predictions, then  apply non-maxima suppression to
    # suppress weak, overlapping bounding boxes
    (rects, confidences) = decode_predictions(scores, geometry, min_confidence)
    boxes = non_max_suppression(np.array(rects), probs=confidences)

    # initialize the list of results
    results = []

    # loop over the bounding boxes
    for (startX, startY, endX, endY) in boxes:
        # scale the bounding box coordinates based on the respective
        # ratios
        startX = int(startX * rW)
        startY = int(startY * rH)
        endX = int(endX * rW)
        endY = int(endY * rH)

        # in order to obtain a better OCR of the text we can potentially
        # apply a bit of padding surrounding the bounding box -- here we
        # are computing the deltas in both the x and y directions
        dX = int((endX - startX) * padding)
        dY = int((endY - startY) * padding)

        # apply padding to each side of the bounding box, respectively
        startX = max(0, startX - dX)
        startY = max(0, startY - dY)
        endX = min(origW, endX + (dX * 2))
        endY = min(origH, endY + (dY * 2))

        # extract the actual padded ROI
        roi = orig[startY:endY, startX:endX]

        # in order to apply Tesseract v4 to OCR text we must supply
        # (1) a language, (2) an OEM flag of 4, indicating that the we
        # wish to use the LSTM neural net model for OCR, and finally
        # (3) an OEM value, in this case, 7 which implies that we are
        # treating the ROI as a single line of text
        config = ("-l eng --oem 1 --psm 7")
        text = pytesseract.image_to_string(roi, config=config)

        # add the bounding box coordinates and OCR'd text to the list
        # of results
        results.append(((startX, startY, endX, endY), text))

    # sort the results bounding box coordinates from top to bottom
    results = sorted(results, key=lambda r:r[0][1])

    # loop over the results
    for ((startX, startY, endX, endY), text) in results:
        # display the text OCR'd by Tesseract
        print("OCR TEXT")
        print("========")
        print("{}\n".format(text))

        # strip out non-ASCII text so we can draw the text on the image
        # using OpenCV, then draw the text and a bounding box surrounding
        # the text region of the input image
        text = "".join([c if ord(c) < 128 else "" for c in text]).strip()
        output = orig.copy()
        cv2.rectangle(output, (startX, startY), (endX, endY),
            (0, 0, 255), 2)
        cv2.putText(output, text, (startX, startY - 20),
            cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 3)

        # show the output image
        cv2.imshow("Text Detection", output)
        cv2.waitKey(0)


from Utilitis import overlay_dicom_map, load_nii, save_nii, resize_mask_array
import numpy as np
from scipy.optimize import curve_fit
from abc import ABC, abstractmethod
import csv, os
import glob
from multiprocessing import Pool, cpu_count
import scipy.io as sio
import pandas as pd

class AbstractFitting(ABC):
    """
    Abstracted class that provides the essential functions for the various fit classes.

    Args:
        dim (int): Dimension of the images (2 or 3D)

        folder (str): path to the dicom images, here it is to be noted that in the derived classes
            the function fit is implemented, depending on the respective implementation the path must be
            passed correctly

        fit_funktion (*function): Pointer to the respective fit function

        bounds (tuple([len(n)], [len(n)]): Bounds values of the fit function bounds = ([x_min, y_min, z_min],
            [x_max, y_max, z_max]).
            Note: The bounds are handed in according to the scipy.optimize.curve_fit convention.
    """

    def __init__(self, dim, folder, fit_function, bounds=None, config_fit=None):
        self.dim = dim
        self.folder = folder
        self.bounds = bounds
        self.x = None
        self.mask = None
        self.fit_map = None
        self.r_squares = None
        self.array = None
        self.fit_function = fit_function
        self.sort = False
        self.config_fit = config_fit
        self.num_cls = 0

    @abstractmethod
    def load(self):
        pass

    def load_mask(self, counts = None):
        """
        Automatic import of the mask, it is assumed that the mask is in the same path as the commit path
            (see init function: folder). The file must be named as mask.nii.gz!
        Args:
            None

        Returns:
            None
        """
        nii_files = glob.glob(self.folder + '/mask.nii.gz')

        if len(nii_files) == 1:
            mask, affine, header = load_nii(nii_files[0])
        else:
            nii_files = glob.glob(self.folder + '/*.nii.gz')
            if len(nii_files) == 1:
                if 'map.nii.gz' in nii_files[0]:
                    nii_files = nii_files[1:]
        mask, affine, header = load_nii(nii_files[0])
        if counts is not None:
            for i in range(100):
                if i not in counts:
                    mask[mask == i] = 0
        self.set_mask(mask, affine, header)

    def set_mask(self, mask, affine, header):
        """
        set mask manual

        Args:
            mask (np.array): Integer Array, hierbei entspricht jede Zahl einer segmentiereten Klasse
            affine (np.array): An array that tells you the position of the image array data in a reference space.
            header (): image metadata (data about the data) describing the image, usually in the form of an image header

        Returns:
            None
        """
        self.mask = {
            'mask': mask,
            'affine': affine,
            'header': header
        }
        self.num_cls = mask.max()

    def get_map(self):
        return self.fit_map

    def run(self, multiprocessing=False, x=None):
        """
        Starts the essential functions
        """
        self.fit(multiprocessing, x)
        self.overlay_map()
        return self.save_results()

    def fit(self, multiprocessing=False, x=None):
        """
        Calculates the fit_map (array in the image dimension where the fitted times are stored).

        Args:
            x (np.array): times of the different acquisitions. If x = None, it is assumed that the times could be
                read out from the dicom header when importing the image.

        Returns:
            fit_map (np.array): array with the fitted times (T1, T2, T2star, .....)

        Note: It is important that the data and the array have already been read in beforehand. The quality of the fits
            (R^2) is saved as a class variable, but not returned.
        """
        x = self.x if x is None else x
        mask = self.mask['mask']
        if x is None or mask is None:
            return
        assert self.array.shape[0] == len(x), 'The passed times: {} do not match the dimension ({}) of the ' \
                                              'loaded Dicom files!!'.format(x, self.array.shape[0])
        mask, self.array = resize_mask_array(mask, self.array)
        self.mask['mask'] = mask
        fit_map = np.zeros(mask.shape)
        r_squares = np.zeros(mask.shape)

        if multiprocessing:
            with Pool(cpu_count() - 2) as pool:
                idxs, map, r_square = zip(*pool.map(fit_slice_process,
                                                    [(fit_map[:, :, i], r_squares[:, :, i], self.array[:, :, :, i],
                                                      mask[:, :, i],
                                                      x, self.fit_function, self.bounds, i, self.config_fit[:, :, i] if
                                                      self.config_fit is not None else None) for i in
                                                     range(self.array.shape[-1])]))
                for i in range(len(idxs)):
                    fit_map[:, :, idxs[i]], r_squares[:, :, idxs[i]] = map[idxs[i]], r_square[idxs[i]]
        else:
            for i in range(self.array.shape[-1]):
                config_fit = None if self.config_fit is None else self.config_fit[:, :, i]
                fit_map[:, :, i], r_squares[:, :, i] = fit_slice(self.array[:, :, :, i], mask[:, :, i],
                                                                 x, self.fit_function, self.bounds, config_fit)

        self.fit_map = fit_map
        self.r_squares = r_squares
        return fit_map

    def overlay_map(self, fit_map=None):
        """
        Overlays and saves the calculated map

        Args:
            fit_map (np.array): map to overlay, if fit_map = None, then it is assumed that the map was
                previously calculated and therefore stored in self.fit_map
        Outputs:
            map_dyn_{03d}.format(i): Overlaid images saved as .png
        """
        fit_map = self.fit_map if fit_map is None else fit_map
        mask = self.mask['mask']
        clim = np.nanmax(fit_map)
        # deleted all old *.pngs
        for file in glob.glob(self.folder + os.sep + '*.png'):
            os.remove(file)
        for i in range(fit_map.shape[-1]):
            file = self.folder + r'\map_dyn_{:03d}.png'.format(i)
            file_map = self.folder + r'\mask_dyn_{:03d}.png'.format(i)
            try:
                os.remove(file)
            except FileNotFoundError:
                pass
            if np.nanmax(fit_map[:, :, i]) > 0:
                overlay_dicom_map(self.array[-1, :, :, i], fit_map[:, :, i], [0, clim], file)
                overlay_dicom_map(self.array[-1, :, :, i], np.copy(mask[:, :, i]), [0, self.num_cls], file_map)

    def save_results(self):
        """
        Saves the calculated results.

        Output:
            Map.nii.gz : Calculated results as nii array, can be visualized e.g. with ITK-Snap or overlaid on the images
                in the following adapted with other functions.
            _results.csv: CSV file in which for each class separately mean, standard deviation, minimum, maximum and
                number of pixels is calculated.
        """
        save_nii(self.fit_map, self.mask['affine'], self.mask['header'], self.folder + '/Map.nii.gz')
        save_as_mat(self.folder + '\Results.mat', mask=self.mask['mask'], map=self.fit_map)
        results = {}
        for i in range(1, int(self.mask['mask'].max()) + 1):
            m = self.mask['mask'].copy()
            m = np.where(m == i, 1, 0)

            fit_map = np.multiply(self.fit_map, m)
            k = fit_map.copy()
            k[k > 0] = 1
            fit_map = np.where(fit_map != 0.0, fit_map, np.nan)

            r_squares = np.multiply(self.r_squares, m)
            r_squares = np.where(r_squares != 0, r_squares, np.nan)

            results[str(i)] = ['%.2f' % np.nanmean(fit_map), '%.2f' % np.nanstd(fit_map),
                               '%.2f' % np.nanmin(fit_map), '%.2f' % np.nanmax(fit_map),
                               '%.2f' % np.nansum(k) + '/' + '%.2f' % np.sum(m),
                               '%.2f' % np.nanmean(r_squares)]
        with open(self.folder + '_results.csv', mode='w', newline='') as csv_file:
            writer = csv.writer(csv_file, delimiter=';')
            writer.writerow(['mask_index', 'mean', 'std', 'min', 'max', 'Pixels', 'Mean R^2'])
            for key, value in results.items():
                value = [v.replace('.', ',') for v in value]
                writer.writerow([key] + value)
        return results

    def load_result(self):
        df = pd.read_csv(self.folder + '_results.csv', delimiter=';')
        return df.to_dict()

def save_as_mat(file, mask, map):
    sio.savemat(file, {'mask': mask, 'map': map})

def fit_slice_process(data):
    data = list(data)
    data[0], data[1] = fit_slice(data[2], data[3], data[4], data[5], data[6], config_fit=data[8], min_r_squared=0.75)
    return data[7], data[0], data[1]

def fit_slice(d_slice, mask, x, fit, bounds, config_fit = None, min_r_squared=0):
    if bounds is not None:
        bounds_ = ([bounds[0][0], bounds[0][1], bounds[0][2]],
                     [bounds[1][0], bounds[1][1], bounds[1][2]])
    else:
        bounds_ = None
    """
    Fits one slice

    Args:
        d_slice (np.array): dicom array [times, rows, columns].
        mask (np.array): [Rows, Columns].
        x (list): list with the different time points ([time_1, time_2, time_3, time_4, ...])
        fit (*function): Pointer to the respective fit function
        bounds (tuple([len(n)], [len(n)]): Bounds values of the fit function bounds = ([x_min, y_min, z_min],
            [x_max, y_max, z_max]).
            Note: The bounds are handed in according to the scipy.optimize.curve_fit convention.
        min_r_squared (float): Grenzwert über dem R^2 liegen muss, damit der Pixel eingeschlossen wurde.

    Returns:
        fit_map (np.array): array with the fitted times (T1, T2, T2star, .....)
        r_squares (np.array): array with the calculated r_squares
    """
    fit_map = np.full((d_slice.shape[1], d_slice.shape[2]), np.nan)
    r_squares = fit_map.copy()
    if mask.max() == 0:
        return fit_map, r_squares

    args = np.argwhere(mask != 0)
    for row, column in args:
        y = d_slice[:, row, column]
        try:
            y = y / y.max()
        except ValueError:
            continue
        try:
            if config_fit is not None:
                fit_ = fit((config_fit[row, column]))
            else:
                fit_ = fit
            if bounds is not None:
                param, param_cov = curve_fit(fit_, x, y, bounds=bounds_, xtol=0.01, maxfev=1200)
            else:
                param, param_cov = curve_fit(fit_, x, y, xtol=0.1)
        except RuntimeError:
            continue
        except ValueError:
            continue
        residuals = y - fit_(np.array(x), param[0], param[1], param[2])
        ss_res = np.sum(residuals ** 2)
        ss_tot = np.sum((y - np.mean(y)) ** 2)
        r_squared = 1 - (ss_res / ss_tot)
        if r_squared < min_r_squared:
            continue
        if bounds is not None:
            if param[1] <= bounds[0][1] or param[1] >= bounds[1][1]:
                continue
        fit_map[row, column] = param[1]
        r_squares[row, column] = r_squared
    return fit_map, r_squares


# IMPORTS
from tkinter.filedialog import askopenfilename
import json
import importlib
import threading
import time
from tabulate import tabulate
import numpy as np
import matplotlib.pyplot as plt
import os
import shutil
import webbrowser
from wikipedia import wikipedia
import colorama
from termcolor import cprint
from win10toast import ToastNotifier

# CLASS CONTANING GLOBAL VARIABLES
class GlobalVariables:
    totalGoals = 0
    totalAssists = 0
    statObjectsList = []
    completePlayerStats = []
    tour = type
    dictFilename = ""
    txtFilename = ""
    pathDelimiter = "../"
    tkinterFilePath = ""
    maxPercentageGoalInvolvement = 0
    mostValuablePlayer = ""
    mostValuablePlayerWikiInfo = ""
    AG_PlayerDataDict = {}
    HM_PlayerDataDict = {}
    NS_PlayerDataDict = {}
    TZ_PlayerDataDict = {}
    script_helper = type
    Search_Number = 0

# CLASS FOR CALCULATING PLAYER STATS
class PlayerStats():
    def __init__(self, name, goals, assists):
        self.name = name
        self.goals = goals
        self.assists = assists
        self.percentage_goal_involvement = np.around(((self.goals*3 + self.assists*2) / (
            GlobalVariables.totalGoals*3 + GlobalVariables.totalAssists*2) * 100), decimals=2)
        self.single_player_complete_stat = []
        self.single_player_complete_stat.extend(
            [self.name, self.goals, self.assists, self.percentage_goal_involvement])
        GlobalVariables.completePlayerStats.append(
            self.single_player_complete_stat)
        if self.percentage_goal_involvement > GlobalVariables.maxPercentageGoalInvolvement:
            GlobalVariables.maxPercentageGoalInvolvement = self.percentage_goal_involvement
            GlobalVariables.mostValuablePlayer = self.name

# FINDING LAST MODIFIED FILE
def get_new_file(filesPath):
    filename = ""
    cont = 'n'
    list_of_files = os.listdir(filesPath)
    if len(list_of_files) == 0:
        print("\nUnable to find any file")
    else:
        modified_time_of_new_file = np.format_float_scientific(
            os.path.getmtime(filesPath+list_of_files[0]))
        filename = list_of_files[0]
        for i in list_of_files:
            mtime = np.format_float_scientific(os.path.getmtime(filesPath+i))
            if mtime > modified_time_of_new_file:
                modified_time_of_new_file = mtime
                filename = i
        if filename.endswith(".txt"):
            print("\nNew file found: ", filename)
            cont = input("Continue? (y/n): ").lower().strip()
        else:
            print("\nUnable to find any file")
            cont = 'n'
    if cont == 'y' or cont == '':
        return filename
    elif cont == 'n':
        manual_file_select = input(
            "Enter filename manually (eg. jan1) or Enter to skip: ").lower().strip()
        manual_file_select = manual_file_select+".txt"
        if manual_file_select not in list_of_files:
            print("Cannot find", manual_file_select+". Opening filepicker")
            abs_txt_path = askopenfilename(
                filetypes=[('Text Document', '*.txt')])
            tfilename = str(abs_txt_path).split('/')[-1]
            GlobalVariables.tkinterFilePath = abs_txt_path
            return tfilename
        else:
            return manual_file_select
    else:
        raise Exception("Invalid input. Exiting")

# MAKE A COPY OF FILE IF THE FILE IS CHOSEN WITH tkinter
def write_to_file(tkinter_chosen_file_path):
    if tkinter_chosen_file_path != GlobalVariables.pathDelimiter+"files/"+GlobalVariables.txtFilename:
        shutil.copyfile(tkinter_chosen_file_path,
                        GlobalVariables.pathDelimiter+"files/"+GlobalVariables.txtFilename)

# TO GET THE RESPECTIVE WORKING DICTIONARY FROM script_helper.py
def get_working_dict_from_helper():
    GlobalVariables.script_helper = importlib.import_module("script_helper")
    GlobalVariables.tour, GlobalVariables.dictFilename = GlobalVariables.script_helper.get_working_dict(
        GlobalVariables.txtFilename)

# CALCULATING TOTAL GOALS
def sum_of_goals(every_player_stats):
    goal_list = []
    goal_list = np.array(goal_list)
    for i in every_player_stats:
        goal_list = np.append(goal_list, i[1])
    return (int(np.sum(goal_list)))

# CALCULATING TOTAL ASSISTS
def sum_of_assists(every_player_stats):
    assist_list = []
    assist_list = np.array(assist_list)
    for i in every_player_stats:
        assist_list = np.append(assist_list, i[2])
    return (int(np.sum(assist_list)))

# DOING WIKIPEDIA SEARCH ON MVP
def wiki_search():
    while GlobalVariables.Search_Number < 3:
        search_term = GlobalVariables.mostValuablePlayer
        category = "football"
        try:
            keyword = wikipedia.search(search_term)[
                GlobalVariables.Search_Number]
            WikiInfo = str(wikipedia.summary(
                keyword, sentences=2, auto_suggest=False))
            if WikiInfo.lower().find(category) != -1:
                GlobalVariables.mostValuablePlayerWikiInfo = WikiInfo
                return
            else:
                GlobalVariables.Search_Number += 1
                wiki_search()
        except:
            GlobalVariables.mostValuablePlayerWikiInfo = "None"

# SAVING EVERY PLAYER'S TOTAL STATS
def record_player_stats(list_element, player_dict, alph):
    if player_dict != {}:
        if list_element[0] in player_dict.keys():
            goal_updated = player_dict[list_element[0]][0] + list_element[1]
            assist_updated = player_dict[list_element[0]][1] + list_element[2]
            count = player_dict[list_element[0]][2][2] + 1
            avg_goals = np.around(goal_updated / count, decimals=2)
            avg_assists = np.around(assist_updated / count, decimals=2)
            val_updated = {list_element[0]: [
                goal_updated, assist_updated, [avg_goals, avg_assists, count]]}
            if alph == "AG":
                GlobalVariables.AG_PlayerDataDict.update(val_updated)
            elif alph == "HM":
                GlobalVariables.HM_PlayerDataDict.update(val_updated)
            elif alph == "NS":
                GlobalVariables.NS_PlayerDataDict.update(val_updated)
            else:
                GlobalVariables.TZ_PlayerDataDict.update(val_updated)
        else:
            remaining_lists_dict = {list_element[0]: [
                list_element[1], list_element[2], [list_element[1], list_element[2], 1]]}
            if alph == "AG":
                GlobalVariables.AG_PlayerDataDict.update(remaining_lists_dict)
            elif alph == "HM":
                GlobalVariables.HM_PlayerDataDict.update(remaining_lists_dict)
            elif alph == "NS":
                GlobalVariables.NS_PlayerDataDict.update(remaining_lists_dict)
            else:
                GlobalVariables.TZ_PlayerDataDict.update(remaining_lists_dict)
    if player_dict == {}:
        remaining_lists_dict = {list_element[0]: [
            list_element[1], list_element[2], [list_element[1], list_element[2], 1]]}
        if alph == "AG":
            GlobalVariables.AG_PlayerDataDict.update(remaining_lists_dict)
        elif alph == "HM":
            GlobalVariables.HM_PlayerDataDict.update(remaining_lists_dict)
        elif alph == "NS":
            GlobalVariables.NS_PlayerDataDict.update(remaining_lists_dict)
        else:
            GlobalVariables.TZ_PlayerDataDict.update(remaining_lists_dict)

# TABULATE PRETTY PRINTING
def tabular_display(table, headers):
    print("\n")
    print(tabulate(table, headers, tablefmt="pretty"))

# MAIN
def main():

    # DISPLAYING TITLE
    display_statements_list = ["PES 21 myClub Tour"]
    width = len(display_statements_list[0])
    colorama.init()
    print('+-' + '-' * width + '-+')
    for s in display_statements_list:
        cprint('| {0:^{1}} |'.format(s, width), color='green')
        print('+-' + '-'*(width) + '-+')

    # FINDING CORRECT PATH pathDelimiter FOR BAT EXECUTION
    try:
        os.listdir(GlobalVariables.pathDelimiter+"files")
    except FileNotFoundError:
        GlobalVariables.pathDelimiter = "./"

    # FINDING THE LATEST MODIFIED FILE
    GlobalVariables.txtFilename = get_new_file(
        GlobalVariables.pathDelimiter+"files/")
    if GlobalVariables.txtFilename == '':
        raise FileNotFoundError("No file chosen")

    # GETTING THE RELEVANT TOUR DICTIONARY
    get_working_dict_from_helper_thread = threading.Thread(
        target=get_working_dict_from_helper, daemon=False)
    get_working_dict_from_helper_thread.start()

    # READING FROM THE TXT FILE
    every_player_stats = []
    single_player_stat = []
    try:
        if GlobalVariables.tkinterFilePath == "":
            openfilepath = GlobalVariables.pathDelimiter+"files/"+GlobalVariables.txtFilename
        else:
            openfilepath = GlobalVariables.tkinterFilePath
            write_to_file(GlobalVariables.tkinterFilePath)
        with open(file=openfilepath, mode="r", encoding="utf-8") as working_file:
            for line in working_file:
                if line.startswith("Player Name"):
                    continue
                else:
                    line = line.strip().split("\t")
                    param1 = str(line[0])
                    param2 = int(line[1])
                    param3 = int(line[2])
                    single_player_stat.extend([param1, param2, param3])
                    every_player_stats.append(list(single_player_stat))
                    single_player_stat.clear()
    except FileNotFoundError:
        raise FileNotFoundError("There is no such file in the directory\n")

    # FINDING TOTAL GOALS & ASSIST NUMBERS
    GlobalVariables.totalGoals = sum_of_goals(every_player_stats)
    GlobalVariables.totalAssists = sum_of_assists(every_player_stats)

    # STORING ALL OBJECTS
    for item in every_player_stats:
        GlobalVariables.statObjectsList.append(
            PlayerStats(item[0], item[1], item[2]))

    # WIKIPEDIA SEARCH THREAD
    wiki_search_thread = threading.Thread(target=wiki_search, daemon=False)
    print(GlobalVariables.mostValuablePlayerWikiInfo)
    wiki_search_thread.start()

    # SAVING EVERY PLAYERS TOTAL STATS
    analysis_mode = input(
        "Turn on player analysis mode? (y/n): ").lower().strip()
    if analysis_mode == 'y' or analysis_mode == '':
        try:
            with open(GlobalVariables.pathDelimiter+"playerdata/player_record_ag.json", mode="r", encoding="utf8") as input_json:
                try:
                    player_dict_ag = json.load(input_json)
                except:
                    player_dict_ag = {}
            with open(GlobalVariables.pathDelimiter+"playerdata/player_record_hm.json", mode="r", encoding="utf8") as input_json:
                try:
                    player_dict_hm = json.load(input_json)
                except:
                    player_dict_hm = {}
            with open(GlobalVariables.pathDelimiter+"playerdata/player_record_ns.json", mode="r", encoding="utf8") as input_json:
                try:
                    player_dict_ns = json.load(input_json)
                except:
                    player_dict_ns = {}
            with open(GlobalVariables.pathDelimiter+"playerdata/player_record_tz.json", mode="r", encoding="utf8") as input_json:
                try:
                    player_dict_tz = json.load(input_json)
                except:
                    player_dict_tz = {}
        except FileNotFoundError:
            with open(GlobalVariables.pathDelimiter+"playerdata/player_record_ag.json", mode="w+", encoding="utf8"):
                player_dict_ag = {}
            with open(GlobalVariables.pathDelimiter+"playerdata/player_record_hm.json", mode="w+", encoding="utf8"):
                player_dict_hm = {}
            with open(GlobalVariables.pathDelimiter+"playerdata/player_record_ns.json", mode="w+", encoding="utf8"):
                player_dict_ns = {}
            with open(GlobalVariables.pathDelimiter+"playerdata/player_record_tz.json", mode="w+", encoding="utf8"):
                player_dict_tz = {}
        alphabets_P1 = list(map(chr, range(65, 72)))
        alphabets_P2 = list(map(chr, range(72, 78)))
        alphabets_P3 = list(map(chr, range(78, 84)))
        alphabets_P4 = list(map(chr, range(84, 91)))
        for i in GlobalVariables.completePlayerStats:
            player_name_first_char = i[0][0]
            if player_name_first_char in alphabets_P1:
                record_player_stats(i, player_dict_ag, "AG")
            elif player_name_first_char in alphabets_P2:
                record_player_stats(i, player_dict_hm, "HM")
            elif player_name_first_char in alphabets_P3:
                record_player_stats(i, player_dict_ns, "NS")
            elif player_name_first_char in alphabets_P4:
                record_player_stats(i, player_dict_tz, "TZ")
            else:
                print("Unicode error:", player_name_first_char,
                      "in", i[0]+". Skipping")
                pass
        player_dict_ag.update(GlobalVariables.AG_PlayerDataDict)
        player_dict_hm.update(GlobalVariables.HM_PlayerDataDict)
        player_dict_ns  .update(GlobalVariables.NS_PlayerDataDict)
        player_dict_tz.update(GlobalVariables.TZ_PlayerDataDict)
        GlobalVariables.AG_PlayerDataDict.clear()
        GlobalVariables.HM_PlayerDataDict.clear()
        GlobalVariables.NS_PlayerDataDict.clear()
        GlobalVariables.TZ_PlayerDataDict.clear()
        with open(GlobalVariables.pathDelimiter+"playerdata/player_record_ag.json", mode="w+", encoding="utf8") as output_json:
            json.dump(player_dict_ag, output_json, sort_keys=True, indent=4)
        with open(GlobalVariables.pathDelimiter+"playerdata/player_record_hm.json", mode="w+", encoding="utf8") as output_json:
            json.dump(player_dict_hm, output_json, sort_keys=True, indent=4)
        with open(GlobalVariables.pathDelimiter+"playerdata/player_record_ns.json", mode="w+", encoding="utf8") as output_json:
            json.dump(player_dict_ns, output_json, sort_keys=True, indent=4)
        with open(GlobalVariables.pathDelimiter+"playerdata/player_record_tz.json", mode="w+", encoding="utf8") as output_json:
            json.dump(player_dict_tz, output_json, sort_keys=True, indent=4)

    # SORTING TABLE
    table = GlobalVariables.completePlayerStats
    sorting_method = input(
        "\nSpecify sorting method (1:Default | 2:Name | 3:Goals | 4:Assist | 5:Percentage involvement): ").lower().strip()
    if sorting_method == "1" or sorting_method == '':
        pass
    elif sorting_method == '2':
        table.sort(key=lambda x: x[0], reverse=False)
    elif sorting_method == '3':
        table.sort(key=lambda x: x[1], reverse=True)
    elif sorting_method == '4':
        table.sort(key=lambda x: x[2], reverse=True)
    elif sorting_method == '5':
        table.sort(key=lambda x: x[3], reverse=True)
    else:
        print("Input not valid. Default sorting selected")
        pass

    # DISPLAYING THE COMPLETE PLAYER STATS IN A PRETTY TABULAR FORMAT ALONG WITH TOTAL GOALS & ASSISTS
    tabular_display(table, ["Name", "Goals", "Assits", "% Involvement"])
    print("Total Goals:", GlobalVariables.totalGoals,
          "\tTotal Assists:", GlobalVariables.totalAssists)

    # DISPLAYING MVP NAME AND SHOWING WINDOWS 10 TOAST
    print("\nMVP: ", end='')
    try:
        cprint(GlobalVariables.mostValuablePlayer, color='green')
    except:
        print(GlobalVariables.mostValuablePlayer)
    try:
        n = ToastNotifier()
        n.show_toast("PES 21 myCLUB Tour", "Most valuable player is " +
                     GlobalVariables.mostValuablePlayer, icon_path=GlobalVariables.pathDelimiter+"internal/icons/logo.ico", threaded=True)
    except:
        pass

    # DISPLAYING MVP'S WIKIPEDIA INFO
    disp_mvp_info = input("\nDo you want to know about " +
                          GlobalVariables.mostValuablePlayer+"? (y/n): ").lower().strip()
    if disp_mvp_info == 'y' or disp_mvp_info == '':
        wiki_search_thread.join()
        if GlobalVariables.mostValuablePlayerWikiInfo == "" or GlobalVariables.mostValuablePlayerWikiInfo == "None":
            print("\n[Sorry. Wikipedia search was not successful]\n")
        else:
            print("\n"+str(GlobalVariables.mostValuablePlayerWikiInfo)+"\n")
        srch_about_mvp = input("Want to know more (y/n): ").lower().strip()
        if srch_about_mvp == 'y':
            try:
                url = "https://www.google.com.tr/search?q={}".format(
                    GlobalVariables.mostValuablePlayer)
                print("Searching about",
                      GlobalVariables.mostValuablePlayer+". Opening web-browser")
                time.sleep(0.3)
                webbrowser.open_new_tab(url)
            except:
                raise Exception("\nError in opening the webbrowser")

    # UPDATING THE WORKING DICTIONARY WITH NEW FILE'S TOTAL GOALS
    get_working_dict_from_helper_thread.join()
    year_suffix = GlobalVariables.script_helper.get_year_suffix(GlobalVariables.txtFilename)
    raw_file_name = GlobalVariables.txtFilename.split(".txt")[0]
    updated_key = year_suffix + '-' + \
        (GlobalVariables.txtFilename[0:3] + "-" + raw_file_name[3:]).title()
    GlobalVariables.tour.update({updated_key: GlobalVariables.totalGoals})
    with open(GlobalVariables.pathDelimiter+"internal/"+GlobalVariables.dictFilename+".json", mode="w+", encoding="utf8") as dict_json:
        json.dump(GlobalVariables.tour, dict_json, indent=4)

    # SHOWING THE GRAPH
    graph_show = input("\nSave graphical data? (y/n): ").lower().strip()
    if graph_show == 'y' or graph_show == '':
        key_lists = []
        value_lists = []
        for key, value in GlobalVariables.tour.items():
            key_lists.append(key)
            value_lists.append(value)
        x_pos = np.arange(len(key_lists))
        y_pos = value_lists
        plt.rcdefaults()
        plt.bar(x_pos, y_pos, align='center', alpha=0.5)
        plt.xticks(x_pos, key_lists)
        plt_title = GlobalVariables.script_helper.get_year(GlobalVariables.txtFilename) + " " + GlobalVariables.script_helper.getQuarter(GlobalVariables.txtFilename)
        plt.title(plt_title)
        plt.ylabel('Number of goals')
        plt.xlabel('Tour event')
        for i, v in enumerate(y_pos):
            plt.text(x=i, y=v+1, s=str(v))
        plt_graph_filepath = GlobalVariables.pathDelimiter + "statistics/" + plt_title + ".jpg"
        mng = plt.get_current_fig_manager()
        mng.window.state("zoomed")
        plt.savefig(plt_graph_filepath, format='JPEG')
        plt.show()
        print("\nGraph saved at", plt_graph_filepath)

    print("\nDONE!")

if __name__ == "__main__":
    main()

# Code developed by
# https://github.com/gokulmanohar


import os
import time
from PIL import Image
import numpy as np
from glob import glob
import cv2
from torch._C import dtype

import torchvision
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets

from peleenet import PeleeNet
from eval import AverageMeter
# /mnt/data/ILSVRC2012/ILSVRC2012_val/val/n01440764/ILSVRC2012_val_00000293.JPEG 0
# /mnt/data/ILSVRC2012/ILSVRC2012_val/val/n01531178/ILSVRC2012_val_00004243.JPEG 11
# /mnt/data/ILSVRC2012/ILSVRC2012_val/val/n02097474/ILSVRC2012_val_00000802.JPEG 200
# /mnt/data/ILSVRC2012/ILSVRC2012_val/val/n02669723/ILSVRC2012_val_00000851.JPEG 400
# /mnt/data/ILSVRC2012/ILSVRC2012_val/val/n03532672/ILSVRC2012_val_00000575.JPEG 600
# /mnt/data/ILSVRC2012/ILSVRC2012_val/val/n04243546/ILSVRC2012_val_00000857.JPEG 800

filename = '/mnt/data/ILSVRC2012/ILSVRC2012_val/val/n02669723/ILSVRC2012_val_00000851.JPEG'
max_batch_size = 32
onnx_model_path = 'peleeNetBatch32.onnx'
onnx_batch = 32
# onnx_model_path = 'resnet50.onnx'

calibDataPath   = "./data/cache/"
cacheFile       = calibDataPath + "calib.cache"
iGpu            = 0
calibCount      = 10
batchSize       = max_batch_size
inputSize       = (3,224,224)

TRT_LOGGER = trt.Logger()  # This logger is required to build an engine
# IInt8EntropyCalibrator2/IInt8EntropyCalibrator/IInt8MinMaxCalibrator/IInt8LegacyCalibrator
class MyCalibrator(trt.IInt8LegacyCalibrator):
    def __init__(self, calibCount, inputShape, dataPath, cacheFile):
        trt.IInt8LegacyCalibrator.__init__(self)                                              # 基类默认构造函数
        self.calibCount     = calibCount
        self.shape          = inputShape
        self.calibDataSet   = self.loadData(dataPath)                                      # 需要自己实现一个读数据的函数
        self.cacheFile      = cacheFile
        self.calibData      = np.zeros(self.shape, dtype=np.float32)
        self.dIn            = cuda.mem_alloc(trt.volume(self.shape) * trt.float32.itemsize)     # 准备好校正用的设备内存
        self.oneBatch       = self.batchGenerator()

    def batchGenerator(self):                                                                   # calibrator 的核心，一个提供数据的生成器
        for i in range(self.calibCount):
            print("> calibration ", i)
            idx = np.random.choice(self.calibDataSet.shape[0], self.shape[0], replace=False)
            self.calibData = self.calibDataSet[idx,:]  # 随机选取数据
            yield np.ascontiguousarray(self.calibData, dtype=np.float32)                        # 调整数据格式后抛出

    def get_batch_size(self):                                                                   # TensorRT 会调用，不能改函数名
        return self.shape[0]

    def get_batch(self, names):                                                                 # TensorRT 会调用，不能改函数名，老版本 TensorRT 的输入参数个数可能不一样
        try:
            data = next(self.oneBatch)                                                          # 生成下一组校正数据，拷贝到设备并返回设备地址，否则退出
            cuda.memcpy_htod(self.dIn, data)
            return [int(self.dIn)]
        except StopIteration:
            return None

    def read_calibration_cache(self):                                                           # TensorRT 会调用，不能改函数名
        if os.path.exists(self.cacheFile):
            print( "cahce file: %s" %(self.cacheFile) )
            f = open(self.cacheFile, "rb")
            cache = f.read()
            f.close()
            return cache

    def write_calibration_cache(self, cache):                                                   # TensorRT 会调用，不能改函数名
        print( "cahce file: %s" %(self.cacheFile) )
        f = open(self.cacheFile, "wb")
        f.write(cache)
        f.close()

    def loadData(dataPath,imgformate = ".JPEG"):
        images_list = glob("./data/val/" + "*.JPEG")
        data_np = np.zeros([len(images_list), 3, 224, 224])
        for i,filename in enumerate(images_list):
            data_np[i,:] = get_img_np_nchw(filename)

        return data_np

def getBatchData(filelists):
    batch = max_batch_size
    out_batch = np.zeros([batch,inputSize[0],inputSize[1],inputSize[2]],dtype=np.float32)
    for i in range(len(filelists)//max_batch_size):
        for j, filename in enumerate(filelists[max_batch_size*i:max_batch_size*(i+1)]):
            out_batch[j,:] = get_img_np_nchw(filename).astype(dtype=np.float32)
        yield np.ascontiguousarray(out_batch)

def getBatchLabel(filelists):
    batch = max_batch_size
    out_batch = np.zeros([batch],dtype=np.int)
    for i in range(len(filelists)//max_batch_size):
        for j, lab in enumerate(filelists[max_batch_size*i:max_batch_size*(i+1)]):
            out_batch[j] = lab
        yield np.ascontiguousarray(out_batch)

def get_img_np_nchw(filename):
    image = cv2.imread(filename)
    image_cv = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    off = 32
    image_cv = cv2.resize(image_cv, (224+off, 224+off))
    image_cv = image_cv[16:240,16:240,:]
    miu = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    img_np = np.array(image_cv, dtype=float) / 255.
    r = (img_np[:, :, 0] - miu[0]) / std[0]
    g = (img_np[:, :, 1] - miu[1]) / std[1]
    b = (img_np[:, :, 2] - miu[2]) / std[2]
    img_np_t = np.array([r, g, b])
    img_np_nchw = np.expand_dims(img_np_t, axis=0)
    return img_np_nchw

class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        """Within this context, host_mom means the cpu memory and device means the GPU memory
        """
        self.host = host_mem
        self.device = device_mem

    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        return self.__str__()

def allocate_buffers(engine):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()
    for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # Allocate host and device buffers
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        # Append the device buffer to device bindings.
        bindings.append(int(device_mem))
        # Append to the appropriate list.
        if engine.binding_is_input(binding):
            inputs.append(HostDeviceMem(host_mem, device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem, device_mem))
    return inputs, outputs, bindings, stream

def get_engine(calib, max_batch_size=1, onnx_file_path="", engine_file_path="", \
               fp16_mode=False, int8_mode=False, save_engine=True):
    """Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it."""

    def build_engine(max_batch_size, save_engine):
        """Takes an ONNX file and creates a TensorRT engine to run inference with"""
        explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        with trt.Builder(TRT_LOGGER) as builder, \
                builder.create_network(explicit_batch) as network, \
                trt.OnnxParser(network, TRT_LOGGER) as parser:

            builder.max_workspace_size = 1 << 30  # Your workspace size

            builder.max_batch_size = max_batch_size
            # pdb.set_trace()
            builder.fp16_mode = fp16_mode
            builder.int8_mode = int8_mode
            builder.int8_calibrator = calib

            # Parse model file
            if not os.path.exists(onnx_file_path):
                quit('ONNX file {} not found'.format(onnx_file_path))

            print('Loading ONNX file from path {}...'.format(onnx_file_path))
            with open(onnx_file_path, 'rb') as model:
                print('Beginning ONNX file parsing')
                parser.parse(model.read())
            # network.mark_output(network.get_layer(network.num_layers-1).get_output(0))
            # last_layer = network.get_layer(network.num_layers - 1)
            # network.mark_output(last_layer.get_output(0))

            print('Completed parsing of ONNX file')
            print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))

            engine = builder.build_cuda_engine(network)
            print("Completed creating Engine")

            if save_engine:
                with open(engine_file_path, "wb") as f:
                    f.write(engine.serialize())
            return engine

    if os.path.exists(engine_file_path):
        # If a serialized engine exists, load it instead of building a new one.
        print("Reading engine from file {}".format(engine_file_path))
        with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())
    else:
        return build_engine(max_batch_size, save_engine)

def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
    # Transfer data from CPU to the GPU.
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    # Run inference.
    context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
    # Transfer predictions back from the GPU.
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # Synchronize the stream
    stream.synchronize()
    # Return only the host outputs.
    return [out.host for out in outputs]

def postprocess_the_outputs(h_outputs, shape_of_output):
    h_outputs = h_outputs.reshape(*shape_of_output)
    return h_outputs

#data loader config
test_data_path = "/mnt/data/ILSVRC2012/ILSVRC2012_val/"
input_dim = 224
workers = 1

def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res

def main():
    valdir = os.path.join(test_data_path, 'val')
    # valdir = args.data
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    val_dataset = datasets.ImageFolder(
        valdir,
        transforms.Compose([
            transforms.Resize(input_dim+32),
            transforms.CenterCrop(input_dim),
            transforms.ToTensor(),
            normalize,
        ]))

    # val_loader = torch.utils.data.DataLoader(val_dataset,
    #     batch_size=max_batch_size, shuffle=False,
    #     num_workers=workers, pin_memory=True)

    num_classes = len(val_dataset.classes)
    print('Total classes: ',num_classes)

    #---------------------------------------
    calib = MyCalibrator(calibCount, (batchSize,) + inputSize, calibDataPath, cacheFile)

    # These two modes are dependent on hardwares
    fp16_mode = False
    int8_mode = True
    trt_engine_path = './model_fp16_{}_int8_{}_maxbatch{}_{}.trt'.format(fp16_mode, int8_mode, max_batch_size,onnx_model_path)
    # Build an engine
    engine = get_engine(calib, max_batch_size, onnx_model_path, trt_engine_path, fp16_mode, int8_mode)
    context = engine.create_execution_context()
    # Allocate buffers for input and output
    inputs, outputs, bindings, stream = allocate_buffers(engine) # input, output: host # bindings

    img_np_nchw = get_img_np_nchw(filename)
    img_np_nchw = img_np_nchw.astype(dtype=np.float32)

    shape_of_output = (max_batch_size*onnx_batch, 1000)
    # shape_of_output = (max_batch_size, 1000)

    # # Do inference
    # # Load data to the buffer
    inputs[0].host = img_np_nchw.reshape(-1)

    # inputs[1].host = ... for multiple input
    t1 = time.time()
    trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream) # numpy data
    t2 = time.time()
    feat = trt_outputs[0].reshape(*shape_of_output)
    # feat = postprocess_the_outputs(trt_outputs[0], shape_of_output)
    print(feat.argmax())
    print('TensorRT ok')

    top1 = AverageMeter()
    top5 = AverageMeter()
    batch_time = AverageMeter()

    # use dataLoader gen input
    # for i, (input, target) in enumerate(val_loader):
    imglists = [filename[0] for filename in val_dataset.imgs]
    imglabels = [filename[1] for filename in val_dataset.imgs]
    batch_data_gen = getBatchData(imglists)
    batch_lab_gen = getBatchLabel(imglabels)
    i=0
    for batch_data in batch_data_gen:
        batch_lab = next(batch_lab_gen)
        inputs[0].host = batch_data.reshape(-1)
        end = time.time()
        trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream,batch_size=batchSize)
        feat = postprocess_the_outputs(trt_outputs[0], shape_of_output)

        tensor_output = torch.from_numpy(feat[0:batchSize,:])
        # np.savetxt('b32_res.txt',feat[0:32,:])
        target = torch.from_numpy(batch_lab)

        prec1, prec5 = accuracy(tensor_output.data, target, topk=(1, 5))
        top1.update(prec1[0], (len(imglists)//max_batch_size)*max_batch_size)
        top5.update(prec5[0], (len(imglists)//max_batch_size)*max_batch_size)
        batch_time.update(time.time() - end)
        end = time.time()

        i = i+1
        if i % 10 == 0:
            print('Test: [{0}]\t'
                  'Time {batch_time.val:.5f} ({batch_time.avg:.5f})\t'
                #   'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                   i, batch_time=batch_time,
                #    loss=losses,
                   top1=top1, top5=top5))

    print("Inference time with the TensorRT engine: {}".format(t2-t1))
    print('All completed!')

if __name__ == '__main__':
    main()


from __future__ import print_function
from __future__ import absolute_import
from __future__ import division

import numpy as np
import random
import os
import rdkit
from rdkit import Chem
import csv
from gln.mods.mol_gnn.mol_utils import SmartsMols, SmilesMols
from multiprocessing import Process, Queue
import time

from gln.data_process.data_info import DataInfo, load_train_reactions
from gln.common.reactor import Reactor

class DataSample(object):
    def __init__(self, prod, center, template, label=None, neg_centers=None, neg_tpls=None,
                 reaction=None, neg_reactions=None):
        self.prod = prod
        self.center = center
        self.template = template
        self.label = label
        self.neg_centers = neg_centers
        self.neg_tpls = neg_tpls
        self.reaction = reaction
        self.neg_reactions = neg_reactions

def _rand_sample_except(candidates, exclude, k=None):
    assert len(candidates)
    if k is None:
        if len(candidates) == 1:
            assert exclude is None or candidates[0] == exclude
            return candidates[0]
        else:
            while True:
                c = np.random.choice(candidates)
                if exclude is None or c != exclude:
                    break
            return c
    else:
        if k <= 0 or len(candidates) <= k:
            return [c for c in candidates if exclude is None or c != exclude]
        cand_indices = np.random.permutation(len(candidates))[:k]
        selected = []
        for i in cand_indices:
            c = candidates[i]
            if exclude is None or c != exclude:
                selected.append(c)
            if k <= 0:
                continue
            if len(selected) >= k:
                break
        return selected

def worker_softmax(worker_id, seed, args):
    np.random.seed(seed)
    random.seed(seed)
    num_epochs = 0
    part_id = 0
    train_reactions = load_train_reactions(args)
    while True:
        if num_epochs % args.epochs_per_part == 0:
            DataInfo.load_cooked_part('train', part_id)
            tot_num = len(train_reactions)
            part_size = tot_num // args.num_parts + 1
            indices = range(part_id * part_size, min((part_id + 1) * part_size, tot_num))
            indices = list(indices)
            part_id = (part_id + 1) % args.num_parts
        random.shuffle(indices)
        for sample_idx in indices:
            rxn_type, rxn_smiles = train_reactions[sample_idx]

            if sample_idx in DataInfo.train_pos_maps:
                pos_tpls, weights = DataInfo.train_pos_maps[sample_idx]
                pos_tpl_idx = pos_tpls[np.argmax(np.random.multinomial(1, weights))]
                rxn_type, rxn_template = DataInfo.unique_templates[pos_tpl_idx]
            else:
                continue

            reactants, _, prod = rxn_smiles.split('>')
            cano_prod = DataInfo.smiles_cano_map[prod]
            sm_prod, _, _ = rxn_template.split('>')
            cano_sm_prod = DataInfo.smarts_cano_map[sm_prod]

            # negative samples of prod centers
            assert (rxn_type, cano_prod) in DataInfo.prod_center_maps
            prod_center_cand_idx = DataInfo.prod_center_maps[(rxn_type, cano_prod)]

            neg_center_idxes = _rand_sample_except(prod_center_cand_idx, DataInfo.prod_smarts_idx[cano_sm_prod], args.neg_num)
            neg_centers = [DataInfo.prod_cano_smarts[c] for c in neg_center_idxes]

            # negative samples of templates
            assert cano_sm_prod in DataInfo.unique_tpl_of_prod_center
            assert rxn_type in DataInfo.unique_tpl_of_prod_center[cano_sm_prod]
            neg_tpl_idxes = _rand_sample_except(DataInfo.unique_tpl_of_prod_center[cano_sm_prod][rxn_type], pos_tpl_idx, args.neg_num)
            tpl_cand_idx = []
            for c in neg_centers:
                tpl_cand_idx += DataInfo.unique_tpl_of_prod_center[c][rxn_type]
            if len(tpl_cand_idx):
                neg_tpl_idxes += _rand_sample_except(tpl_cand_idx, pos_tpl_idx, args.neg_num)
            neg_tpls = [DataInfo.unique_templates[i][1] for i in neg_tpl_idxes]

            sample = DataSample(prod=cano_prod, center=cano_sm_prod, template=rxn_template,
                                neg_centers=neg_centers, neg_tpls=neg_tpls)

            if args.retro_during_train:
                sample.reaction = DataInfo.get_cano_smiles(reactants) + '>>' + cano_prod
                sample.neg_reactions = []
                if len(DataInfo.neg_reactions_all[sample_idx]):
                    neg_reacts = DataInfo.neg_reactions_all[sample_idx]
                    if len(neg_reacts):
                        neg_reactants = _rand_sample_except(neg_reacts, None, args.neg_num)
                        sample.neg_reactions = [DataInfo.neg_reacts_list[r] + '>>' + cano_prod for r in neg_reactants]
            if len(sample.neg_tpls) or len(sample.neg_reactions):
                yield (worker_id, sample)
        num_epochs += 1

def worker_process(worker_func, worker_id, seed, data_q, *args):
    worker_gen = worker_func(worker_id, seed, *args)
    for t in worker_gen:
        data_q.put(t)

def data_gen(num_workers, worker_func, worker_args, max_qsize=16384, max_gen=-1, timeout=60):
    cnt = 0
    data_q = Queue(max_qsize)

    if num_workers == 0:  # single process generator
        worker_gen = worker_func(-1, np.random.randint(10000), *worker_args)
        while True:
            worker_id, data_sample = next(worker_gen)
            yield data_sample
            cnt += 1
            if max_gen > 0 and cnt >= max_gen:
                break
        return

    worker_procs = [Process(target=worker_process, args=[worker_func, i, np.random.randint(10000), data_q] + worker_args) for i in range(num_workers)]
    for p in worker_procs:
        p.start()
    last_update = [time.time()] * num_workers
    while True:
        if data_q.empty():
            time.sleep(0.1)
        if not data_q.full():
            for i in range(num_workers):
                if time.time() - last_update[i] > timeout:
                    print('worker', i, 'is dead')
                    worker_procs[i].terminate()
                    while worker_procs[i].is_alive():  # busy waiting for the stop of the process
                        time.sleep(0.01)
                    worker_procs[i] = Process(target=worker_process, args=[worker_func, i, np.random.randint(10000), data_q] + worker_args)
                    print('worker', i, 'restarts')
                    worker_procs[i].start()
                    last_update[i] = time.time()
        try:
            sample = data_q.get_nowait()
        except:
            continue
        cnt += 1
        worker_id, data_sample = sample
        last_update[worker_id] = time.time()
        yield data_sample
        if max_gen > 0 and cnt >= max_gen:
            break

    print('stopping')
    for p in worker_procs:
        p.terminate()
    for p in worker_procs:
        p.join()