# Copyright 2020 The Flax Authors.
#
# 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.

# Lint as: python3
# Copyright 2020 The Flax Authors.
#
# 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.

"""Machine Translation example.

This script trains a Transformer on a WMT dataset.
"""

import collections
import functools
import os
import time
from absl import app
from absl import flags
from absl import logging
from flax import jax_utils
from flax import nn
from flax import optim
import bleu
import decode
import input_pipeline
import models
from flax.metrics import tensorboard
from flax.training import checkpoints
from flax.training import common_utils
import jax
from jax import random
import jax.nn
import jax.numpy as jnp
import numpy as np
import tensorflow.compat.v2 as tf

FLAGS = flags.FLAGS

flags.DEFINE_string(
    'model_dir', default=None,
    help='Directory to store model data.')

flags.DEFINE_string(
    'data_dir', default=None,
    help='Tensorflow datasets directory.')

flags.DEFINE_string(
    'vocab_path', default=None,
    help='Path to load or store sentencepiece vocab file.')

flags.DEFINE_integer(
    'vocab_size', default=32000,
    help='Vocabulary size if `vocab_path` is not given.')

flags.DEFINE_string(
    'dataset_name', default='wmt17_translate/de-en',
    help='Name of TFDS translation dataset to use.')

flags.DEFINE_string(
    'eval_dataset_name', default='wmt14_translate/de-en:test',
    help='Optional name of TFDS translation dataset to use for evaluation.')

flags.DEFINE_bool(
    'reverse_translation', default=False,
    help='Reverse the direction of translation.')

flags.DEFINE_integer(
    'batch_size', default=256,
    help='Per host batch size for training.')

flags.DEFINE_integer(
    'beam_size', default=4,
    help='Beam size for inference.')

flags.DEFINE_integer(
    'eval_frequency', default=1000,
    help='Frequency of eval during training, e.g. every 1000 steps.')

flags.DEFINE_integer(
    'num_train_steps', default=500000,
    help='Number of train steps.')

flags.DEFINE_integer(
    'num_eval_steps', default=20,
    help='Number of steps to take during evaluation.')

flags.DEFINE_float(
    'learning_rate', default=0.0625,
    help='Base learning rate.')

flags.DEFINE_integer(
    'warmup_steps', default=1000,
    help='Linear learning rate warmup.')

flags.DEFINE_float(
    'label_smoothing', default=0.1,
    help='Cross entropy loss label smoothing.')

flags.DEFINE_float(
    'weight_decay', default=0.0,
    help='Decay factor for AdamW style weight decay.')

flags.DEFINE_integer(
    'max_target_length', default=256,
    help='Maximum length cutoff for training examples.')

flags.DEFINE_integer(
    'max_eval_target_length', default=256,
    help='Maximum length cutoff for eval examples.')

flags.DEFINE_integer(
    'max_predict_length', default=256,
    help='Maximum length cutoff for predicted tokens.')

flags.DEFINE_bool(
    'share_embeddings', default=True,
    help='Inputs and targets share embedding.')

flags.DEFINE_bool(
    'logits_via_embedding', default=True,
    help='Final logit transform uses embedding matrix transpose.')

flags.DEFINE_integer(
    'num_layers', default=6,
    help='Number of transformer layers.')

flags.DEFINE_integer(
    'qkv_dim', default=1024,
    help='Size of query/key/value for attention.')

flags.DEFINE_integer(
    'emb_dim', default=1024,
    help='Size of embeddings.')

flags.DEFINE_integer(
    'mlp_dim', default=4096,
    help='Size of the MLP.')

flags.DEFINE_integer(
    'num_heads', default=16,
    help='Number of attention heads.')

flags.DEFINE_float(
    'dropout_rate', default=0.1,
    help='Dropout rate.')

flags.DEFINE_float(
    'attention_dropout_rate', default=0.1,
    help='Attention dropout rate.')

flags.DEFINE_integer(
    'random_seed', default=0,
    help='Integer for PRNG random seed.')

flags.DEFINE_bool(
    'save_checkpoints', default=True,
    help='Whether to save model checkpoints.')

flags.DEFINE_bool(
    'restore_checkpoints', default=True,
    help='Whether to restore from existing model checkpoints.')

flags.DEFINE_integer(
    'checkpoint_freq', default=10000,
    help='Save a checkpoint every these number of steps.')

flags.DEFINE_bool(
    'use_bfloat16', default=True,
    help=('Use bfloat16 mixed precision training instead of float32.'))

flags.DEFINE_string(
    'jax_backend_target', default=None,
    help=('TPU grpc target for use with cloud TPUs.'
          ' e.g. grpc://192.168.0.2:8470'))

def create_learning_rate_scheduler(
    factors='constant * linear_warmup * rsqrt_decay',
    base_learning_rate=0.5,
    warmup_steps=1000,
    decay_factor=0.5,
    steps_per_decay=20000,
    steps_per_cycle=100000):
  """Creates learning rate schedule.

  Interprets factors in the factors string which can consist of:
  * constant: interpreted as the constant value,
  * linear_warmup: interpreted as linear warmup until warmup_steps,
  * rsqrt_decay: divide by square root of max(step, warmup_steps)
  * rsqrt_normalized_decay: divide by square root of max(step/warmup_steps, 1)
  * decay_every: Every k steps decay the learning rate by decay_factor.
  * cosine_decay: Cyclic cosine decay, uses steps_per_cycle parameter.

  Args:
    factors: string, factors separated by '*' that defines the schedule.
    base_learning_rate: float, the starting constant for the lr schedule.
    warmup_steps: int, how many steps to warm up for in the warmup schedule.
    decay_factor: float, the amount to decay the learning rate by.
    steps_per_decay: int, how often to decay the learning rate.
    steps_per_cycle: int, steps per cycle when using cosine decay.

  Returns:
    a function learning_rate(step): float -> {'learning_rate': float}, the
    step-dependent lr.
  """
  factors = [n.strip() for n in factors.split('*')]

  def step_fn(step):
    """Step to learning rate function."""
    ret = 1.0
    for name in factors:
      if name == 'constant':
        ret *= base_learning_rate
      elif name == 'linear_warmup':
        ret *= jnp.minimum(1.0, step / warmup_steps)
      elif name == 'rsqrt_decay':
        ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
      elif name == 'rsqrt_normalized_decay':
        ret *= jnp.sqrt(warmup_steps)
        ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
      elif name == 'decay_every':
        ret *= (decay_factor**(step // steps_per_decay))
      elif name == 'cosine_decay':
        progress = jnp.maximum(0.0,
                               (step - warmup_steps) / float(steps_per_cycle))
        ret *= jnp.maximum(0.0,
                           0.5 * (1.0 + jnp.cos(jnp.pi * (progress % 1.0))))
      else:
        raise ValueError('Unknown factor %s.' % name)
    return jnp.asarray(ret, dtype=jnp.float32)

  return step_fn

@functools.partial(jax.jit, static_argnums=(1, 2, 3))
def create_model(key, input_shape, target_shape, model_kwargs):
  """Instantiate transformer model and associated autoregressive cache def."""
  model_def = models.Transformer.partial(**model_kwargs)
  with nn.attention.Cache().mutate() as cache_def:
    _, initial_params = model_def.init_by_shape(
        key, [(input_shape, jnp.float32), (target_shape, jnp.float32)],
        cache=cache_def)
    model = nn.Model(model_def, initial_params)
  return model, cache_def

def create_optimizer(model, learning_rate, weight_decay):
  optimizer_def = optim.Adam(
      learning_rate,
      beta1=0.9,
      beta2=0.98,
      eps=1e-9,
      weight_decay=weight_decay)
  optimizer = optimizer_def.create(model)
  return optimizer

def compute_weighted_cross_entropy(logits,
                                   targets,
                                   weights=None,
                                   label_smoothing=0.0):
  """Compute weighted cross entropy and entropy for log probs and targets.

  Args:
   logits: [batch, length, num_classes] float array.
   targets: categorical targets [batch, length] int array.
   weights: None or array of shape [batch, length].
   label_smoothing: label smoothing constant, used to determine the on and
     off values.

  Returns:
    Tuple of scalar loss and batch normalizing factor.
  """
  if logits.ndim != targets.ndim + 1:
    raise ValueError('Incorrect shapes. Got shape %s logits and %s targets' %
                     (str(logits.shape), str(targets.shape)))
  vocab_size = logits.shape[-1]
  confidence = 1.0 - label_smoothing
  low_confidence = (1.0 - confidence) / (vocab_size - 1)
  normalizing_constant = -(
      confidence * jnp.log(confidence) + (vocab_size - 1) *
      low_confidence * jnp.log(low_confidence + 1e-20))
  soft_targets = common_utils.onehot(
      targets, vocab_size, on_value=confidence, off_value=low_confidence)

  loss = -jnp.sum(soft_targets * nn.log_softmax(logits), axis=-1)
  loss = loss - normalizing_constant

  normalizing_factor = jnp.prod(targets.shape)
  if weights is not None:
    loss = loss * weights
    normalizing_factor = weights.sum()

  return loss.sum(), normalizing_factor

def compute_weighted_accuracy(logits, targets, weights=None):
  """Compute weighted accuracy for log probs and targets.

  Args:
   logits: [batch, length, num_classes] float array.
   targets: categorical targets [batch, length] int array.
   weights: None or array of shape [batch, length]

  Returns:
    Tuple of scalar loss and batch normalizing factor.
  """
  if logits.ndim != targets.ndim + 1:
    raise ValueError('Incorrect shapes. Got shape %s logits and %s targets' %
                     (str(logits.shape), str(targets.shape)))
  loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)
  normalizing_factor = jnp.prod(logits.shape[:-1])
  if weights is not None:
    loss = loss * weights
    normalizing_factor = weights.sum()

  return loss.sum(), normalizing_factor

def compute_metrics(logits, labels, weights, label_smoothing=0.0):
  """Compute summary metrics."""
  loss, weight_sum = compute_weighted_cross_entropy(logits, labels, weights,
                                                    label_smoothing)
  acc, _ = compute_weighted_accuracy(logits, labels, weights)
  metrics = {
      'loss': loss,
      'accuracy': acc,
      'denominator': weight_sum,
  }
  metrics = jax.lax.psum(metrics, axis_name='batch')
  return metrics

def train_step(optimizer,
               batch,
               learning_rate_fn,
               label_smoothing=0.0,
               use_bfloat16=False,
               dropout_rng=None):
  """Perform a single training step."""
  # X_position and X_segmentation are needed only when using 'packed examples'
  # where multiple sequences are packed into the same example with this
  # metadata.
  # if such features are not present they are ignored and the example is treated
  # like a normal, unpacked sequence example.
  train_keys = ['inputs', 'targets',
                'inputs_position', 'targets_position',
                'inputs_segmentation', 'targets_segmentation']
  (inputs, targets,
   inputs_positions, targets_positions,
   inputs_segmentation, targets_segmentation) = [
       batch.get(k, None) for k in train_keys]

  weights = jnp.where(targets > 0, 1, 0).astype(jnp.float32)

  # We handle PRNG splitting inside the top pmap to improve efficiency.
  dropout_rng, new_dropout_rng = random.split(dropout_rng)

  def loss_fn(model):
    """loss function used for training."""
    with nn.stochastic(dropout_rng):
      logits = model(
          inputs,
          targets,
          use_bfloat16=use_bfloat16,
          inputs_positions=inputs_positions,
          targets_positions=targets_positions,
          inputs_segmentation=inputs_segmentation,
          targets_segmentation=targets_segmentation,
          train=True,
          cache=None)

    loss, weight_sum = compute_weighted_cross_entropy(logits, targets, weights,
                                                      label_smoothing)
    mean_loss = loss / weight_sum
    return mean_loss, logits

  step = optimizer.state.step
  lr = learning_rate_fn(step)
  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
  (_, logits), grad = grad_fn(optimizer.target)
  grad = jax.lax.pmean(grad, 'batch')
  new_optimizer = optimizer.apply_gradient(grad, learning_rate=lr)
  metrics = compute_metrics(logits, targets, weights)
  metrics['learning_rate'] = lr

  return new_optimizer, metrics, new_dropout_rng

def eval_step(model, batch, label_smoothing=0.0, use_bfloat16=False):
  """Calculate evaluation metrics on a batch."""
  inputs, targets = batch['inputs'], batch['targets']
  weights = jnp.where(targets > 0, 1.0, 0.0)
  logits = model(inputs, targets, use_bfloat16=use_bfloat16, train=False,
                 cache=None)
  return compute_metrics(logits, targets, weights, label_smoothing)

def predict_step(inputs, model, cache, eos_id, max_decode_len,
                 use_bfloat16=False, beam_size=4):
  """Predict translation with fast decoding beam search on a batch."""
  batch_size = inputs.shape[0]

  # Prepare transformer fast-decoder call for beam search: for beam search, we
  # need to set up our decoder model to handle a batch size equal to
  # batch_size * beam_size, where each batch item's data is expanded in-place
  # rather than tiled.
  # i.e. if we denote each batch element subtensor as el[n]:
  # [el0, el1, el2] --> beamsize=2 --> [el0,el0,el1,el1,el2,el2]
  src_padding_mask = decode.flat_batch_beam_expand(
      (inputs > 0)[..., None], beam_size)
  tgt_padding_mask = decode.flat_batch_beam_expand(
      jnp.ones((batch_size, 1, 1)), beam_size)
  encoded_inputs = decode.flat_batch_beam_expand(
      model.encode(inputs, use_bfloat16=use_bfloat16,
                   train=False, cache=None), beam_size)

  def tokens_ids_to_logits(flat_ids, flat_cache):
    """Token slice to logits from decoder model."""
    # --> [batch * beam, 1, vocab]
    with flat_cache.mutate() as new_flat_cache:
      flat_logits = model.decode(encoded_inputs,
                                 src_padding_mask,
                                 flat_ids,
                                 cache=new_flat_cache,
                                 shift=False,
                                 train=False,
                                 use_bfloat16=use_bfloat16,
                                 tgt_padding_mask=tgt_padding_mask)
    # Remove singleton sequence-length dimension:
    # [batch * beam, 1, vocab] --> [batch * beam, vocab]
    flat_logits = flat_logits.squeeze(axis=1)
    return flat_logits, new_flat_cache

  # Using the above-defined single-step decoder function, run a
  # beam search over possible sequences given input encoding.
  beam_seqs, _ = decode.beam_search(
      inputs,
      cache,
      tokens_ids_to_logits,
      beam_size=beam_size,
      alpha=0.6,
      eos_id=eos_id,
      max_decode_len=max_decode_len)

  # Beam search returns [n_batch, n_beam, n_length + 1] with beam dimension
  # sorted in increasing order of log-probability.
  # Return the highest scoring beam sequence, drop first dummy 0 token.
  return beam_seqs[:, -1, 1:]

def pad_examples(x, desired_batch_size):
  """Expand batch to desired size by repeating last slice."""
  batch_pad = desired_batch_size - x.shape[0]
  return np.concatenate([x, np.tile(x[-1], (batch_pad, 1))], axis=0)

def per_host_sum_pmap(in_tree):
  """Execute psum on in_tree's leaves over one device per host."""
  host2devices = collections.defaultdict(list)
  for d in jax.devices():
    host2devices[d.host_id].append(d)
  devices = [host2devices[k][0] for k in host2devices]
  host_psum = jax.pmap(lambda x: jax.lax.psum(x, 'i'), 'i', devices=devices)
  def pre_pmap(xs):
    return jax.tree_map(lambda x: jnp.broadcast_to(x, (1,) + x.shape), xs)
  def post_pmap(xs):
    return jax.tree_map(lambda x: x[0], xs)
  return post_pmap(host_psum(pre_pmap(in_tree)))

def tohost(x):
  """Collect batches from all devices to host and flatten batch dimensions."""
  n_device, n_batch, *remaining_dims = x.shape
  return np.array(x).reshape((n_device * n_batch,) + tuple(remaining_dims))

def main(argv):
  if len(argv) > 1:
    raise app.UsageError('Too many command-line arguments.')

  if FLAGS.jax_backend_target:
    jax.config.FLAGS.jax_xla_backend = 'tpu_driver'
    jax.config.FLAGS.jax_backend_target = FLAGS.jax_backend_target

  # This seems to be necessary even when importing TF2?
  tf.enable_v2_behavior()

  # Number of local devices for this host.
  n_devices = jax.local_device_count()

  if jax.host_id() == 0:
    tf.io.gfile.makedirs(FLAGS.model_dir)
    train_summary_writer = tensorboard.SummaryWriter(
        os.path.join(FLAGS.model_dir, 'train'))
    eval_summary_writer = tensorboard.SummaryWriter(
        os.path.join(FLAGS.model_dir, 'eval'))

  if FLAGS.batch_size % n_devices:
    raise ValueError('Batch size must be divisible by the number of devices')

  vocab_path = FLAGS.vocab_path
  if vocab_path is None:
    vocab_path = os.path.join(FLAGS.model_dir, 'sentencepiece_model')
  tf.io.gfile.makedirs(os.path.split(vocab_path)[0])

  # Load Dataset
  logging.info('Initializing dataset.')
  train_ds, eval_ds, predict_ds, encoder = input_pipeline.get_wmt_datasets(
      n_devices=n_devices,
      dataset_name=FLAGS.dataset_name,
      eval_dataset_name=FLAGS.eval_dataset_name,
      shard_idx=jax.host_id(),
      shard_count=jax.host_count(),
      data_dir=FLAGS.data_dir,
      vocab_path=vocab_path,
      target_vocab_size=FLAGS.vocab_size,
      batch_size=FLAGS.batch_size,
      max_length=FLAGS.max_target_length,
      max_eval_length=FLAGS.max_eval_target_length)
  train_iter = iter(train_ds)
  vocab_size = int(encoder.vocab_size())
  eos_id = decode.EOS_ID  # Default Sentencepiece EOS token.
  def decode_tokens(toks):
    valid_toks = toks[:np.argmax(toks == eos_id) + 1].astype(np.int32)
    return encoder.detokenize(valid_toks).numpy().decode('utf-8')

  logging.info('Initializing model, optimizer, and step functions.')
  # Build Model and Optimizer
  transformer_kwargs = {
      'vocab_size': vocab_size,
      'output_vocab_size': vocab_size,
      'emb_dim': FLAGS.emb_dim,
      'num_heads': FLAGS.num_heads,
      'num_layers': FLAGS.num_layers,
      'qkv_dim': FLAGS.qkv_dim,
      'mlp_dim': FLAGS.mlp_dim,
      'max_len': max(FLAGS.max_target_length, FLAGS.max_eval_target_length),
      'share_embeddings': FLAGS.share_embeddings,
      'logits_via_embedding': FLAGS.logits_via_embedding,
  }

  start_step = 0
  rng = random.PRNGKey(FLAGS.random_seed)
  rng, init_rng = random.split(rng)
  input_shape = (FLAGS.batch_size, FLAGS.max_target_length)
  target_shape = (FLAGS.batch_size, FLAGS.max_target_length)
  model, cache_def = create_model(init_rng,
                                  input_shape,
                                  target_shape,
                                  transformer_kwargs)
  optimizer = create_optimizer(model,
                               FLAGS.learning_rate,
                               FLAGS.weight_decay)
  # We access model only from optimizer below via optimizer.target.
  del model

  if FLAGS.restore_checkpoints:
    # Restore unreplicated optimizer + model state from last checkpoint.
    optimizer = checkpoints.restore_checkpoint(FLAGS.model_dir, optimizer)
    # Grab last step.
    start_step = int(optimizer.state.step)

  # Replicate optimizer.
  optimizer = jax_utils.replicate(optimizer)

  learning_rate_fn = create_learning_rate_scheduler(
      base_learning_rate=FLAGS.learning_rate,
      warmup_steps=FLAGS.warmup_steps)

  p_train_step = jax.pmap(
      functools.partial(
          train_step,
          learning_rate_fn=learning_rate_fn,
          label_smoothing=FLAGS.label_smoothing,
          use_bfloat16=FLAGS.use_bfloat16),
      axis_name='batch')
  p_eval_step = jax.pmap(
      functools.partial(
          eval_step,
          label_smoothing=FLAGS.label_smoothing,
          use_bfloat16=FLAGS.use_bfloat16),
      axis_name='batch')
  p_pred_step = jax.pmap(
      functools.partial(predict_step, use_bfloat16=FLAGS.use_bfloat16,
                        beam_size=FLAGS.beam_size),
      axis_name='batch',
      static_broadcasted_argnums=(3, 4))  # eos token, max_length are constant

  # We init the first set of dropout PRNG keys, but update it afterwards inside
  # the main pmap'd training update for performance.
  dropout_rngs = random.split(rng, n_devices)

  logging.info('Starting training loop.')
  metrics_all = []
  t_loop_start = time.time()
  for step, batch in zip(range(start_step, FLAGS.num_train_steps), train_iter):
    # Shard data to devices and do a training step.
    batch = common_utils.shard(jax.tree_map(lambda x: x._numpy(), batch))  # pylint: disable=protected-access
    optimizer, metrics, dropout_rngs = p_train_step(
        optimizer, batch, dropout_rng=dropout_rngs)
    metrics_all.append(metrics)

    # Save a checkpoint on one host after every checkpoint_freq steps.
    if (FLAGS.save_checkpoints and step % FLAGS.checkpoint_freq == 0 and
        step > 0 and jax.host_id() == 0):
      checkpoints.save_checkpoint(FLAGS.model_dir,
                                  jax_utils.unreplicate(optimizer), step)

    # Periodic metric handling.
    if step % FLAGS.eval_frequency != 0 and step > 0:
      continue

    logging.info('Gathering training metrics.')
    # Training Metrics
    metrics_all = common_utils.get_metrics(metrics_all)
    lr = metrics_all.pop('learning_rate').mean()
    metrics_sums = jax.tree_map(jnp.sum, metrics_all)
    denominator = metrics_sums.pop('denominator')
    summary = jax.tree_map(lambda x: x / denominator, metrics_sums)  # pylint: disable=cell-var-from-loop
    summary['learning_rate'] = lr
    steps_per_eval = FLAGS.eval_frequency if step != 0 else 1
    steps_per_sec = steps_per_eval / (time.time() - t_loop_start)
    t_loop_start = time.time()
    if jax.host_id() == 0:
      train_summary_writer.scalar('steps per second', steps_per_sec, step)
      for key, val in summary.items():
        train_summary_writer.scalar(key, val, step)
      train_summary_writer.flush()
    metrics_all = []
    logging.info('train in step: %d, loss: %.4f', step, summary['loss'])

    # Eval Metrics
    logging.info('Gathering evaluation metrics.')
    t_eval_start = time.time()
    eval_metrics = []
    eval_iter = iter(eval_ds)
    for _, eval_batch in zip(range(FLAGS.num_eval_steps), eval_iter):
      eval_batch = jax.tree_map(lambda x: x._numpy(), eval_batch)  # pylint: disable=protected-access
      eval_batch = common_utils.shard(eval_batch)
      metrics = p_eval_step(optimizer.target, eval_batch)
      eval_metrics.append(metrics)
    eval_metrics = common_utils.get_metrics(eval_metrics)
    eval_metrics_sums = jax.tree_map(jnp.sum, eval_metrics)
    eval_denominator = eval_metrics_sums.pop('denominator')
    eval_summary = jax.tree_map(
        lambda x: x / eval_denominator,  # pylint: disable=cell-var-from-loop
        eval_metrics_sums)
    if jax.host_id() == 0:
      for key, val in eval_summary.items():
        eval_summary_writer.scalar(key, val, step)
      eval_summary_writer.flush()
    logging.info('eval in step: %d, loss: %.4f', step, eval_summary['loss'])
    logging.info('eval time: %.4f s step %d', time.time()-t_eval_start, step)

    # Translation and BLEU Score.
    logging.info('Translating evaluation dataset.')
    t_inference_start = time.time()
    predict_iter = iter(predict_ds)
    sources, references, predictions = [], [], []
    for _, pred_batch in enumerate(predict_iter):
      pred_batch = jax.tree_map(lambda x: x._numpy(), pred_batch)  # pylint: disable=protected-access
      # Handle final odd-sized batch by padding instead of dropping it.
      cur_pred_batch_size = pred_batch['inputs'].shape[0]
      if cur_pred_batch_size % n_devices:
        padded_size = int(
            np.ceil(cur_pred_batch_size / n_devices) * n_devices)
        pred_batch = jax.tree_map(
            lambda x: pad_examples(x, padded_size), pred_batch)  # pylint: disable=cell-var-from-loop
      pred_batch = common_utils.shard(pred_batch)
      per_device_batchsize = pred_batch['inputs'].shape[1]
      cache_dtype = jnp.bfloat16 if FLAGS.use_bfloat16 else jnp.float32
      cache = jax_utils.replicate(
          cache_def.initialize_cache((per_device_batchsize,
                                      FLAGS.max_predict_length),
                                     dtype=cache_dtype))
      predicted = p_pred_step(pred_batch['inputs'],
                              optimizer.target,
                              cache,
                              eos_id,
                              FLAGS.max_predict_length)
      predicted = tohost(predicted)
      inputs = tohost(pred_batch['inputs'])
      targets = tohost(pred_batch['targets'])
      # Iterate through non-padding examples of batch.
      for i, s in enumerate(predicted[:cur_pred_batch_size]):
        sources.append(decode_tokens(inputs[i]))
        references.append(decode_tokens(targets[i]))
        predictions.append(decode_tokens(s))
    logging.info('Translation: %d predictions %d references %d sources.',
                 len(predictions), len(references), len(sources))
    logging.info('Translation time: %.4f s step %d.',
                 time.time() - t_inference_start, step)

    # Calculate BLEU score for translated eval corpus against reference.
    bleu_matches = bleu.bleu_partial(references, predictions)
    all_bleu_matches = per_host_sum_pmap(bleu_matches)
    bleu_score = bleu.complete_bleu(*all_bleu_matches)
    # Save translation samples for tensorboard.
    exemplars = ''
    for n in np.random.choice(np.arange(len(predictions)), 8):
      exemplars += f'{sources[n]}\n\n{references[n]}\n\n{predictions[n]}\n\n'
    if jax.host_id() == 0:
      eval_summary_writer.scalar('bleu', bleu_score, step)
      eval_summary_writer.text('samples', exemplars, step)
      eval_summary_writer.flush()
    logging.info('Translation BLEU Score %.4f', bleu_score)

if __name__ == '__main__':
  app.run(main)


import keras
'''lib loading error prevention'''
import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'
import numpy as np
from keras.utils import to_categorical
from keras.models import load_model
import scipy as sp

"""========================"""
"""tensorflow configuration"""
""""======================="""

import tensorflow as tf
from keras import backend as K
num_cores = 48

num_CPU = 1
num_GPU = 1

config = tf.ConfigProto(intra_op_parallelism_threads=num_cores,
                        inter_op_parallelism_threads=num_cores,
                        allow_soft_placement=True,
                        device_count={'CPU': num_CPU, 'GPU': num_GPU})
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.9)
config.gpu_options.allow_growth = True

session = tf.Session(config=config)
K.set_session(session)
init = tf.global_variables_initializer()
session.run(init)

from sig_tool import normalize
from plain_model import PlainModel, hyper_params, gen_ckpt_prefix, gen_pp_data_dir_name
from plain_data_make import PlainPreprocessor, load_kfold_names, preload
from icbeb_tool import  load_icbeb2019_label
import pickle
import re
import numpy as np
import matplotlib.pyplot as plt

from scipy import io
from sig_tool import  med_smooth

'''global variables'''
fold = 5
fs = 500 # sampling frequency of the data
preprocessor = PlainPreprocessor(hyper_params)

def calc(seg):
    seg = med_smooth(seg, 10)
    qrs_interval = []
    qrs = []

    qrs_start = -1
    qrs_stop = -1
    for idx in range(len(seg)):
        if seg[idx] == 1:
            if qrs_start == -1:
                # new qrs segmentation
                qrs_start = idx
            else:
                continue
        else:
            if qrs_start >= 0:
                qrs_stop = idx
                qrs_interval.append((qrs_start, qrs_stop))

                qrs_start = -1
                qrs_stop = -1
            else:
                continue
    idx = 0
    while idx < len(qrs_interval):
        # searching for
        interval = qrs_interval[idx]
        central = (interval[1] + interval[0]) // 2
        idx += 1

        # if interval[1]-interval[0] < 20:
        #     # qrs glitch
        #     # searching for next qrs
        #     idx_next = idx + 1
        #     if idx_next == len(qrs_interval):
        #         central = (interval[1]+interval[0])//2
        #         idx += 1
        #     else:
        #         while idx_next < len(qrs_interval):
        #             interval_next = qrs_interval[idx_next]
        #             if interval_next[1]-interval_next[0] < 20:
        #                 if interval_next[1]-interval[0] < 160:
        #                     idx_next += 1
        #                     if idx_next == len(qrs_interval):
        #                         central = (interval_next[1]+interval[0])//2
        #                         idx = idx_next
        #                         break
        #                 else:
        #                     central = (qrs_interval[idx_next-1][1] + interval[0])//2
        #                     idx = idx_next
        #                     break
        #             else:
        #                 central = (interval_next[1] + interval[0])//2
        #                 idx = idx_next + 1
        #                 break
        # else:
        #     central = (interval[1]+interval[0])//2
        #     idx += 1

        '''calibrate central with -20 ms if fir applied'''
        qrs.append(central)

    fs_ = 500
    r_hr = np.array([loc for loc in qrs if ((loc > 5.5 * fs_) and (loc < 5000 - 0.5 * fs_))])
    # rr = []
    # r_prev = -1
    # for r in qrs:
    #     if r < 5.5 * 500 or r > 5000 - 0.5 * 500:
    #         continue
    #
    #     if r_prev < 0:
    #         r_prev = r
    #     else:
    #         rr.append(r - r_prev)
    #         r_prev = r

    if len(r_hr) == 0:
        hr = 60
    else:
        # hr = 60 / (np.mean(rr) / 500)
        hr = np.round( 60 * fs_ / np.mean(np.diff(r_hr)))

    # plt.plot(ECG[0])
    # plt.plot(seg)
    # rwaves = np.zeros(len(seg))
    # for r in qrs:
    #     rwaves[r] = 1
    # plt.plot(rwaves)
    # plt.show()

    # fs_ = 500
    # sub_qrs_end = np.where((np.array(qrs) <= 9.575 * fs_))[0]
    # sub_qrs_start = np.where((np.array(qrs) >= 0.425 * fs_))[0]
    # if len(sub_qrs_start) >= 1 and len(sub_qrs_end) >= 1:
    #     qrs_inrange = qrs[sub_qrs_start[0]:sub_qrs_end[-1]]
    # elif len(sub_qrs_start) >= 1:
    #     qrs_inrange = qrs[sub_qrs_start[0]:]
    # elif len(sub_qrs_end) >= 1:
    #     qrs_inrange = qrs[:sub_qrs_end[-1]]
    # else:
    #     qrs_inrange = qrs[:]

    return hr, qrs

def score(r_ref, hr_ref, r_ans, hr_ans, fs_, thr_):
    HR_score = 0
    record_flags = np.ones(len(r_ref))
    for i in range(len(r_ref)):
        FN = 0
        FP = 0
        TP = 0
        hr_der = abs(int(hr_ans[i]) - int(hr_ref[i]))
        if hr_der <= 0.02 * hr_ref[i]:
            HR_score = HR_score + 1
        elif hr_der <= 0.05 * hr_ref[i]:
            HR_score = HR_score + 0.75
        elif hr_der <= 0.1 * hr_ref[i]:
            HR_score = HR_score + 0.5
        elif hr_der <= 0.2 * hr_ref[i]:
            HR_score = HR_score + 0.25

        ref_count = 0
        for j in range(len(r_ref[i])):
            # if r_ref[i][j] < 0.5 * fs_ or r_ref[i][j] > 9.5 * fs_:
            #     continue
            # ref_count += 1

            loc = np.where(np.abs(r_ans[i] - r_ref[i][j]) <= thr_ * fs_)[0]
            if j == 0:
                err = np.where((r_ans[i] >= 0.5 * fs_) & (r_ans[i] <= r_ref[i][j] - thr_ * fs_))[0]
            elif j == len(r_ref[i]) - 1:
                err = np.where((r_ans[i] >= r_ref[i][j] + thr_ * fs_) & (r_ans[i] <= 9.5 * fs_))[0]
            else:
                err = np.where((r_ans[i] >= r_ref[i][j] + thr_ * fs_) & (r_ans[i] <= r_ref[i][j + 1] - thr_ * fs_))[0]

            FP = FP + len(err)
            if len(loc) >= 1:
                TP += 1
                FP = FP + len(loc) - 1
            elif len(loc) == 0:
                FN += 1

        if FN + FP > 1:
            record_flags[i] = 0
        elif FN == 1 and FP == 0:
            record_flags[i] = 0.3
        elif FN == 0 and FP == 1:
            record_flags[i] = 0.7

    rec_acc = round(np.sum(record_flags) / len(r_ref), 4)
    hr_acc = round(HR_score / len(r_ref), 4)

    print( 'QRS_acc: {}'.format(rec_acc))
    print('HR_acc: {}'.format(hr_acc))
    print('Scoring complete.')

    return record_flags, rec_acc, hr_acc

if __name__ == '__main__':
    train_paired, val_paired = pickle.load(open('shuffle_names.dat', 'rb'))

    models = []
    # model = load_model('models/rematch_ckpt_plain_merge_40_96086_0_012_0.0262_0.0317_0.9919_0.9882.h5')
    # model = load_model('models/rematch_ckpt_plain_merge_40_31093_2_006_0.1236_0.1534_0.9930_0.9779.h5')
    # model = load_model('models/rematch_ckpt_plain_merge_40_55824_0_015_0.0262_0.0234_0.9946_0.9952.h5')
    # model = load_model('models/rematch_ckpt_plain_merge_40_38633_1_053_0.0560_0.0956_0.9823_0.9647.h5')
    # model = load_model('models/rematch_ckpt_plain_40_50937_0_043_0.0603_0.0837_0.9705.h5')

    # model = load_model('models/rematch_ckpt_plain_rev2_40_98498_1_581_0.0652_0.0754_0.9738_0.9712.h5')
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev2_40_28799_0_055_0.0439_0.0861_0.9847_0.9712.h5')
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev2_40_90826_0_063_0.0573_0.1105_0.9821_0.9652.h5')
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev3_40_32153_1_466_0.0743_0.0687_0.9706_0.9720.h5')
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev3_40_54196_0_147_0.0746_0.0738_0.9700_0.9710.h5')
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev3_40_27053_0_409_0.0679_0.0659_0.9729_0.9729.h5')
    # models.append(model)

    '''entry v2'''
    # model = load_model('models/rematch_ckpt_plain_rev4_40_19031_0_031_0.0587_0.0711_0.9763_0.9722.h5')
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_40_92098_0_047_0.0654_0.0733_0.9740_0.9720.h5')
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_40_11085_0_037_0.0669_0.0644_0.9729_0.9741.h5')
    # models.append(model)

    '''test entry'''
    # # Unet++ models
    # # model = load_model('models/rematch_ckpt_plain_rev4_40_90350_1_040_0.0583_0.0603_0.9771_0.9757.h5')
    # # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_40_42101_0_038_0.0582_0.0607_0.9773_0.9769.h5')
    # models.append(model)
    # # model = load_model('models/rematch_ckpt_plain_rev4_40_52454_2_030_0.0667_0.0619_0.9747_0.9765.h5')
    # # models.append(model)
    #
    # # stacked lstm with attention
    # model = load_model('models/rematch_ckpt_plain_rev4_40_77525_0_036_0.0606_0.0622_0.9754_0.9751.h5')
    # models.append(model)
    #
    # # U-net conv
    # # model = load_model('models/rematch_ckpt_plain_rev4_40_81909_0_055_0.0734_0.0633_0.9710_0.9751.h5')
    # # models.append(model)
    # # model = load_model('models/rematch_ckpt_plain_rev4_50_11185_0_047_0.0909_0.0733_0.9648_0.9717.h5')
    # # models.append(model)
    #
    # # U-net LSTM
    # model = load_model('models/rematch_ckpt_plain_rev4_30_91712_0_061_0.0610_0.0576_0.9751_0.9770.h5')
    # models.append(model)

    '''v.2.3 finetune'''
    #
    # model = load_model('models/rematch_ckpt_plain_rev4_30_96625_0_001_0.0510_0.0587_0.9789_0.9771.h5') # U-net LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_52643_0_001_0.0562_0.0622_0.9772_0.9746.h5') # stacked LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_59847_0_001_0.0532_0.0606_0.9792_0.9768.h5') # U-net++ LSTM
    # models.append(model)

    '''v.2.4 finetune'''

    # model = load_model('models/rematch_ckpt_plain_rev4_30_69622_0_002_0.0498_0.0590_0.9793_0.9770.h5') # U-net LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_91612_0_004_0.0511_0.0613_0.9792_0.9761.h5') # stacked LSTM
    # model = load_model('models/rematch_ckpt_plain_rev4_30_52643_0_001_0.0562_0.0622_0.9772_0.9746.h5') # stacked LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_17321_2_002_0.0517_0.0608_0.9799_0.9768.h5') # U-net++ LSTM
    # models.append(model)

    '''v.2.5 with sig0'''
    # model = load_model('models/rematch_ckpt_plain_rev4_30_sig0_88967_0_055_0.0613_0.0627_0.9750_0.9741.h5') # HR-net LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_sig0_66164_1_053_0.0634_0.0619_0.9742_0.9755.h5') # U-net LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_sig0_80957_0_194_0.0514_0.0851_0.9849_0.9729.h5') # U-net++ LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_sig0_35157_0_001_0.0534_0.0613_0.9781_0.9744.h5') # stacked LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_sig0_30705_0_001_0.0564_0.0603_0.9770_0.9763.h5') #U-net++ LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_sig0_87538_0_001_0.0587_0.0618_0.9776_0.9768.h5') #U-net LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_sig0_23347_0_001_0.0518_0.0633_0.9787_0.9743.h5') #stacked LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_sig0_63087_0_001_0.0498_0.0639_0.9795_0.9749.h5') #HRnet LSTM
    # models.append(model)

    '''v.3 test'''
    # model = load_model('models/rematch_ckpt_plain_rev4_30_sig0_37092_0_081_0.0609_0.0626_0.9751_0.9750.h5') #HRnet LSTM
    # models.append(model)
    # model = load_model('models/ckpt_ex_rev5_lstmunet_9014_0_30_082_0.0314_0.0828_0.9869_0.9748. h5') #HRnet LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_unetconvlstm_sig71698_0_189_0.2243_0.2186_0.9737_0.9761.h5') #HRnet LSTM
    # models.append(model)

    '''v.3 kfold bagging lstm unet'''
    model = load_model('models/tune_fold10_30_netlstm_sig1_89383_4_002_0.0553_0.0452_0.9776_0.9811.h5')
    models.append(model)
    model = load_model('models/tune_fold10_30_netlstm_sig1_99576_3_007_0.0579_0.0484_0.9765_0.9801.h5')
    models.append(model)
    model = load_model('models/tune_fold10_30_netlstm_sig1_59925_7_003_0.0641_0.0546_0.9744_0.9781.h5')
    models.append(model)
    # model = load_model('models/tune_fold10_30_netlstm_sig1_24285_9_004_0.0651_0.0704_0.9740_0.9715.h5')
    # models.append(model)
    model = load_model('models/tune_fold10_30_netlstm_sig1_64082_1_007_0.0628_0.0658_0.9746_0.9732.h5')
    models.append(model)

    data_dir = gen_pp_data_dir_name()

    batch_x = []
    batch_x0 = []
    batch_y = []

    cnt = 0
    ref_hr = []
    ref_r = []
    # val = val_paired[:]
    # val = val_paired[:1900]
    val = val_paired[600:800]
    for name, offset in val:
        train_sig, train_label, pre_train_sig, pre_train_label = pickle.load(open(os.path.join(data_dir, name+'.dat'), 'rb'))

        RPOS_PATH = 'dat/icbeb2019/ref/'
        ref_path = os.path.join(RPOS_PATH, 'R_'+ str.split(name,'_')[1])
        r = io.loadmat(ref_path)['R_peak'].flatten()

        # mitdb label and sig extraction
        # label = train_label[offset:offset+hyper_params['crop_len']]
        # pre_train_label = pre_train_label[offset:offset+hyper_params['crop_len']]
        #
        # r = np.array([i for i in np.where(label == 1)])
        # pre_train_sig = pre_train_sig[offset:offset+hyper_params['crop_len']]
        # plt.plot(pre_train_sig[:,1])
        # plt.plot(pre_train_label)
        # plt.show()

        r = r[(r >= 0.5*500) & (r <= 9.5*500)]

        ref_r.append(r)
        r_hr = np.array([loc for loc in r if ((loc > 5.5 * 500) and (loc < 5000 - 0.5 * 500))])
        hr = round( 60 * 500 / np.mean(np.diff(r_hr)))
        ref_hr.append(hr)

        sig = pre_train_sig
        # sig = normalize(sig)
        if np.isnan(sig).any():
            continue

        # plt.plot(train_sig)
        # plt.plot(sig[:,0])
        # plt.plot(sig[:,1])
        # plt.plot(pre_train_label)
        # plt.legend(['raw', 'sig', 'diff', 'label'])
        # plt.show()

        # batch_x.append(np.transpose(sig))
        batch_x.append(pre_train_sig[:,1])
        batch_x0.append(pre_train_sig[:,0])
        batch_y.append(pre_train_label)

    batch_x = np.reshape(batch_x, newshape=(len(batch_x), hyper_params['crop_len'], 1))
    batch_x0 = np.reshape(batch_x0, newshape=(len(batch_x), hyper_params['crop_len'], 1))
    segs = []
    for model in models:
        seg = model.predict(np.array(batch_x))
        segs.append(seg)
    # seg = models[0].predict(np.array(batch_x))
    # segs.append(seg)
    # seg = models[1].predict(np.array(batch_x0))
    # segs.append(seg)
    # segs.append(seg[1])
    # seg = models[1].predict(np.array(batch_x))
    # segs.append(seg)
    # seg = models[3].predict(np.array(batch_x0))
    # segs.append(seg)
    # seg = models[4].predict(np.array(batch_x0))
    # segs.append(seg)
    # seg = models[5].predict(np.array(batch_x0))
    # segs.append(seg)

    # store the data and seg into mat
    sp.io.savemat('smooth_demo.mat', {'x':batch_x, 'y':batch_y, 'seg':segs, 'ref_r':ref_r})

    FS = 500
    THR = 0.075

    # score for individual networks
    ss_esti = []
    for seg in segs:
        ss = np.argmax(seg, axis=2)

        hr_ans = []
        r_ans = []
        for s in ss:
            hr, qrs = calc(s)
            hr_ans.append(hr)
            r_ans.append(np.array(qrs))

        record_flags, _, _ = score(ref_r, ref_hr, r_ans, hr_ans, FS, THR)

        ss_esti.append(ss)

    # score for ensemble networks

    # linear ensemble
    seg = np.average(segs, axis=0)
    seg = np.argmax(seg, axis=2)

    #voting
    # ss_sum = np.sum(ss_esti, axis=0)
    # seg = []
    # for ss in ss_sum:
    #     seg.append([0 if s <= 2 else 1 for s in ss])

    hr_ans = []
    r_ans = []
    for s in seg:
        hr, qrs = calc(s)
        hr_ans.append(hr)
        r_ans.append(np.array(qrs))

    record_flags, _, _ = score(ref_r, ref_hr, r_ans, hr_ans, FS, THR)

    for idx in range(len(val)):
        if record_flags[idx] >= 0.9:
            continue
        r_seq = np.zeros(5000)
        for r in r_ans[idx]:
            r_seq[r] = 1
        r_seq_ref = np.zeros(5000)
        for r in ref_r[idx]:
            r_seq_ref[r] = 1

        plt.figure(figsize=(16,16))
        plt.title(str(val_paired[idx][0])+' '+ str(record_flags[idx]))
        plt.plot(batch_x[idx]+1)
        plt.plot(seg[idx])
        plt.plot(r_seq)
        plt.plot(batch_y[idx]-1)
        plt.plot(r_seq_ref-1)
        plt.legend(['sig', 'seg', 'label'])
        plt.show()


# -*- encoding:utf-8 -*-
# @Time    : 2019/9/15 20:43
# @Author  : gfjiang
# @Site    :
# @File    : f1_score.py
# @Software: PyCharm
import numpy as np
from collections import defaultdict

from .mean_ap import get_cls_results, tpfp_imagenet, tpfp_default
from cvtools.cocotools.coco import COCO
from cvtools.utils.boxes import x1y1wh_to_x1y1x2y2

def get_coco_anns(file, imgs=None):
    gt_bboxes = []
    gt_labels = []
    coco = COCO(file)
    if not imgs:
        imgs = [img_info['file_name']
                for img_info in coco.dataset['images']]
    imgToAnns = defaultdict(list)
    for ann in coco.dataset['annotations']:
        imgToAnns[coco.imgs[ann['image_id']]['file_name']].append(ann)
    anns = [imgToAnns[img] if img in imgToAnns.keys() else []
            for img in imgs]
    for ann in anns:
        bboxes = np.array([bbox['bbox'] for bbox in ann])
        if len(bboxes) > 0:
            bboxes = x1y1wh_to_x1y1x2y2(bboxes)
        labels = np.array([bbox['category_id'] for bbox in ann])
        gt_bboxes.append(bboxes)
        gt_labels.append(labels)
    return gt_bboxes, gt_labels

def f1_score(det_results,
             gt_bboxes,
             gt_labels,
             gt_ignore=None,
             scale_ranges=None,
             iou_thr=0.5,
             dataset=None,
             print_summary=True):
    """Evaluate F1 score of a dataset.

    Args:
        det_results (list): a list of list, [[cls1_det, cls2_det, ...], ...]
        gt_bboxes (list): ground truth bboxes of each image, a list of K*4
            array.
        gt_labels (list): ground truth labels of each image, a list of K array
        gt_ignore (list): gt ignore indicators of each image, a list of K array
        scale_ranges (list, optional): [(min1, max1), (min2, max2), ...]
        iou_thr (float): IoU threshold
        dataset (None or str or list): dataset name or dataset classes, there
            are minor differences in metrics for different datsets, e.g.
            "voc07", "imagenet_det", etc.
        print_summary (bool): whether to print the mAP summary

    Returns:
        tuple: (mAP, [dict, dict, ...])
    """
    assert len(det_results) == len(gt_bboxes) == len(gt_labels)
    if gt_ignore is not None:
        assert len(gt_ignore) == len(gt_labels)
        for i in range(len(gt_ignore)):
            assert len(gt_labels[i]) == len(gt_ignore[i])
    area_ranges = ([(rg[0]**2, rg[1]**2) for rg in scale_ranges]
                   if scale_ranges is not None else None)
    num_scales = len(scale_ranges) if scale_ranges is not None else 1
    eval_results = []
    cls_scores = []
    num_classes = len(det_results[0])  # positive class num
    gt_labels = [
        label if label.ndim == 1 else label[:, 0] for label in gt_labels
    ]
    for i in range(num_classes):
        # get gt and det bboxes of this class
        cls_dets, cls_gts, cls_gt_ignore = get_cls_results(
            det_results, gt_bboxes, gt_labels, gt_ignore, i)
        # calculate tp and fp for each image
        tpfp_func = (
            tpfp_imagenet if dataset in ['det', 'vid'] else tpfp_default)
        tpfp = [
            tpfp_func(cls_dets[j], cls_gts[j], cls_gt_ignore[j], iou_thr,
                      area_ranges) for j in range(len(cls_dets))
        ]
        tp, fp = tuple(zip(*tpfp))
        # calculate gt number of each scale, gts ignored or beyond scale
        # are not counted
        num_gts = np.zeros(num_scales, dtype=int)
        for j, bbox in enumerate(cls_gts):
            if area_ranges is None:
                num_gts[0] += np.sum(np.logical_not(cls_gt_ignore[j]))
            else:
                gt_areas = (bbox[:, 2] - bbox[:, 0] + 1) * (
                    bbox[:, 3] - bbox[:, 1] + 1)
                for k, (min_area, max_area) in enumerate(area_ranges):
                    num_gts[k] += np.sum(
                        np.logical_not(cls_gt_ignore[j]) &
                        (gt_areas >= min_area) & (gt_areas < max_area))
        # sort all det bboxes by score, also sort tp and fp
        cls_dets = np.vstack(cls_dets)
        num_dets = cls_dets.shape[0]
        sort_inds = np.argsort(-cls_dets[:, -1])
        cls_scores.append(cls_dets[sort_inds][:, 4])
        tp = np.hstack(tp)[:, sort_inds]
        fp = np.hstack(fp)[:, sort_inds]
        # calculate recall and precision with tp and fp
        tp = np.cumsum(tp, axis=1)
        fp = np.cumsum(fp, axis=1)
        eps = np.finfo(np.float32).eps
        recalls = tp / np.maximum(num_gts[:, np.newaxis], eps)
        precisions = tp / np.maximum((tp + fp), eps)
        f1 = 2 * precisions * recalls / (precisions + recalls)
        eval_results.append(f1)
    max_f1_ids = [f1s.argmax(axis=1)[0] for f1s in eval_results]
    best_f1 = [f1s[:, i] for i, f1s in zip(max_f1_ids, eval_results)]
    best_th = [scores[i] for i, scores in zip(max_f1_ids, cls_scores)]
    print('best f1: {}, best score threshold: {}'.format(best_f1, best_th))
    return eval_results, cls_scores


##
# coding: utf-8
##
# # v01
# - created 15-11-17
# - final version for the cont integration uc102 (short version)
# - processes a single image locally stored
# - stores locally an output image and json with the same filename as the image
# - currently using FasterRCNN trained on COCO model
# - place in /usr/src/listener/
# - there are internal parameters

# In[ ]:

import numpy as np
import os
import six.moves.urllib as urllib
import tensorflow as tf
import pandas
from collections import defaultdict
from PIL import Image
import cv2
from tqdm import tqdm
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as vis_util
import json
from keras.models import load_model
from keras.applications.imagenet_utils import preprocess_input
from keras.preprocessing import image
from keras import backend as K
import gc

# Path to frozen Keras model
PATH_TO_CRISIS_MODEL = './model/vgg_places'

# Path to json for decoding predictions
CRISIS_INDEX_PATH = './model/crisis_index.json'

def decode_predictions(preds):
    class_index = json.load(open(CRISIS_INDEX_PATH))
    sorted_preds = np.argsort(-preds[0])
    results = []
    for idx in sorted_preds:
        results += [[class_index[str(idx)], preds[0][idx]]]
    return results

def preprocess_image(img_np):
    img = image.array_to_img(img_np)
    img = img.resize(size=(224,224))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    #x = np.transpose(x, (0, 1, 2, 3))
    #x = x[::-1, :, :, :]
    x = preprocess_input(x, mode='caffe')
    return x

# Path to frozen detection graph. This is the actual model that is used for the object detection.
PATH_TO_CKPT = './model/frozen_inference_graph.pb'

# List of the strings that is used to add correct label for each box.
PATH_TO_LABELS = './model/label_map.pbtxt'
NUM_CLASSES = 90

detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.GraphDef()
    with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')

#config = tf.ConfigProto()
#config.gpu_options.allow_growth=True

label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)

print('Analyzer imported')

#PARAMETERS
confidence_drop = 0.6 #bellow this confidence score boxes will be dropped
upper_size_thres = 0.9 #above this size (percentage of screen) boxes will be dropped

def box_trans(box, height, width): #transform from (ymin, xmin, ymax, xmax)[norm] to (left, top, width, height)[non-norm]
    box = [box[0] * height, box[1] * width, box[2] * height, box[3] * width]
    box = list(map(int, box))
    box = [box[1], box[0], abs(box[3]-box[1]), abs(box[2]-box[0])]
    return box

def analyze(img_np, file_name, timestamp):

    crisis_model = load_model(PATH_TO_CRISIS_MODEL)
    preds = crisis_model.predict(preprocess_image(img_np))
    crisis_prediction = decode_predictions(preds)
    K.clear_session()

    #with detection_graph.as_default():
    sess =  tf.Session(graph=detection_graph)

    #LOAD IMAGE
    frame_np = img_np
    write_on = frame_np.copy()
    height = frame_np.shape[0]
    width = frame_np.shape[1]

    # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
    frame_np_expanded = np.expand_dims(frame_np, axis=0)
    image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
    # Each box represents a part of the image where a particular object was detected.
    boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
    # Each score represent how level of confidence for each of the objects.
    # Score is shown on the result image, together with the class label.
    scores = detection_graph.get_tensor_by_name('detection_scores:0')
    classes = detection_graph.get_tensor_by_name('detection_classes:0')
    num_detections = detection_graph.get_tensor_by_name('num_detections:0')
    (boxes, scores, classes, num_detections) = sess.run(
        [boxes, scores, classes, num_detections],
        feed_dict={image_tensor: frame_np_expanded})
    to_del=()
    #delete unwanted boxes (too large, non car labels, etc)
    for i in range(boxes.shape[1]):
        if((np.abs(boxes[0][i][0]-boxes[0][i][2])*np.abs(boxes[0][i][1]-boxes[0][i][3]) >= upper_size_thres) or
            (scores[0][i] < confidence_drop) or
            (classes[0][i] not in [1, 2, 3, 4, 6, 8, 9])):
            to_del = to_del + (i,)
            num_detections[0] = num_detections[0] - 1
    boxes = np.delete(boxes, to_del, axis=1)
    scores = np.delete(scores, to_del, axis=1)
    classes = np.delete(classes, to_del, axis=1)
    boxes_t=[] #initiallize transformed bozes storage
    for box in np.squeeze(boxes, axis=0): # transform detected boxes
        box = box_trans(box, height, width)
        boxes_t = boxes_t + [box]

    #visualize detection
    vis_util.visualize_boxes_and_labels_on_image_array(
        write_on,
        np.squeeze(boxes, axis=0),
        np.squeeze(classes, axis=0).astype(np.int32),
        np.squeeze(scores, axis=0),
        category_index,
        use_normalized_coordinates=True,
        max_boxes_to_draw=500,
        line_thickness=1,
        min_score_thresh=confidence_drop)
    to_write = Image.fromarray(write_on)
    to_write.save('./output/'+file_name+'_output.jpg')
    target_dict = []
    for idx, box in enumerate(np.squeeze(boxes, axis=0)):
        risk = np.random.random_sample()
        box = box_trans(box, height, width)
        target_dict += [{"left":box[0],
                        "top":box[1],
                        "width":box[2],
                        "height":box[3],
                        "risk":risk,
                        "type":category_index[int(np.squeeze(classes, axis=0)[idx])]['name'],
                        "confidence":float("{0:.3f}".format(np.squeeze(scores, axis=0)[idx]))}]
    son = {'image':{"name":file_name, "width":width, "height":height, "timestamp":timestamp, "crisis_type":crisis_prediction[0][0], "crisis_level":"severe", "target":target_dict}}
    with open('./output/'+file_name+'_output.json', 'w') as outfile:
        json.dump(son, outfile, sort_keys=False, indent=1)
    outfile.close()
    sess.close()


# -*- coding: utf-8 -*-

import os
import csv
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC, LinearSVC
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.externals import joblib
from ModeAnalysisRestructured import mode_analysis
from scipy.sparse import csr_matrix

def normalise_dictionary(mode_dictionary):
    if (mode_dictionary['Abnormal replacing'] > 0):
        mode_dictionary['Abnormal replacing'] = 1
    if (mode_dictionary['Abnormal string concatenation 1'] > 0):
        mode_dictionary['Abnormal string concatenation 1'] = 1
    if (mode_dictionary['Abnormal string concatenation 2'] > 0):
        mode_dictionary['Abnormal string concatenation 2'] = 1
    if (mode_dictionary['Abnormal true/false statement 1'] > 0):
        mode_dictionary['Abnormal true/false statement 1'] = 1
    if (mode_dictionary['Array replacing number'] > 0):
        mode_dictionary['Array replacing number'] = 1
    if (mode_dictionary['Assign keyword to variable'] > 0):
        mode_dictionary['Assign keyword to variable'] = 1
    if (mode_dictionary['Exceedingly long heximal string'] > 0):
        mode_dictionary['Exceedingly long heximal string'] = 1
    if (mode_dictionary['Exceedingly long \\x escape string'] > 0):
        mode_dictionary['Exceedingly long \\x escape string'] = 1
    if (mode_dictionary['Expression replacing number'] > 0):
        mode_dictionary['Experssion replacing number'] = 1
    if (mode_dictionary['Function replacing assignment'] > 1):
        mode_dictionary['Function replacing assignment'] = 1
    if (mode_dictionary['Number encoded script'] < 4):
        mode_dictionary['Number encoded script'] = mode_dictionary['Number encoded script'] / float(3)
    else:
        mode_dictionary['Number encoded script'] = 1
    # if (mode_dictionary['Random Variable Name'] > 0):
    #     mode_dictionary['Random Variable Name'] = 1
    if (mode_dictionary['Exceedingly long heximal string'] > 0):
        mode_dictionary['Exceedingly long heximal string'] = 1
    if (mode_dictionary['\\x escaped string count'] < 25):
        mode_dictionary['\\x escaped string count'] = mode_dictionary['\\x escaped string count'] / float(25)
    else:
        mode_dictionary['\\x escaped string count'] = 1
    if (mode_dictionary['Too much whitespace in string'] > 0):
        mode_dictionary['Too much whitespace in string'] = 1
    if (mode_dictionary['XOR encoding'] > 0):
        mode_dictionary['XOR encoding'] = 1
    if (mode_dictionary['Exceedingly long array'] < 4):
        mode_dictionary['Exceedingly long array'] = mode_dictionary['Exceedingly long array'] / float(3)
    else:
        mode_dictionary['Exceedingly long array'] = 1
    if (mode_dictionary['Abnormal % escape string'] > 0):
        mode_dictionary['Abnormal % escape string'] = 1
    if (mode_dictionary['Abnormal % escape string'] > 0):
        mode_dictionary['Abnormal % escape string'] = 1
    if (mode_dictionary['Abnormal function call'] < 10):
        mode_dictionary['Abnormal function call'] = mode_dictionary['Abnormal function call'] / float(10)
    else:
        mode_dictionary['Abnormal function call'] = 1
    if (mode_dictionary['Abnormal if logic'] > 0):
        mode_dictionary['Abnormal if logic'] = 1
    if (mode_dictionary['Abnormal String.fromCharCode'] > 0):
        mode_dictionary['Abnormal String.fromCharCode'] = 1
    if (mode_dictionary['Keyword in string'] < 5):
        mode_dictionary['Keyword in string'] = mode_dictionary['Keyword in string'] / float(4)
    else:
        mode_dictionary['Keyword in string'] = 1
    # if (mode_dictionary['Keyword in string'] > 0):
    #     mode_dictionary['Keyword in string'] = 1
    return mode_dictionary

def dictionary_to_array(mode_dictionary):
    array_vector = []
    array_vector.append(mode_dictionary['Minimisation'])
    array_vector.append(mode_dictionary['Exceedingly long mapping'])
    array_vector.append(mode_dictionary['Keyword Concatenation'])
    array_vector.append(mode_dictionary['Exceedingly long heximal string'])
    array_vector.append(mode_dictionary['Exceedingly long \\x escape string'])
    array_vector.append(mode_dictionary['Number encoded script'])
    array_vector.append(mode_dictionary['Abnormal % escape string'])
    array_vector.append(mode_dictionary['Continuous number string'])
    array_vector.append(mode_dictionary['\\x escaped string count'])
    array_vector.append(mode_dictionary['Expression replacing number'])
    array_vector.append(mode_dictionary['Too much \\x escape characters in string'])
    array_vector.append(mode_dictionary['Exceedingly long array'])
    array_vector.append(mode_dictionary['Function replacing assignment'])
    array_vector.append(mode_dictionary['Abnormal string concatenation 1'])
    array_vector.append(mode_dictionary['Abnormal string concatenation 2'])
    array_vector.append(mode_dictionary['Abnormal string concatenation 3'])
    array_vector.append(mode_dictionary['Random Variable Name'])
    array_vector.append(mode_dictionary['Too much single variable'])
    array_vector.append(mode_dictionary['XOR encoding'])
    array_vector.append(mode_dictionary['Variable name too long'])
    array_vector.append(mode_dictionary['Too much whitespace in string'])
    array_vector.append(mode_dictionary['Character seperated programme'])
    array_vector.append(mode_dictionary['Abnormal function call'])
    array_vector.append(mode_dictionary['Abnormal if logic'])
    array_vector.append(mode_dictionary['Abnormal true/false statement 1'])
    array_vector.append(mode_dictionary['Abnormal true/false statement 2'])
    array_vector.append(mode_dictionary['Array replacing number'])
    array_vector.append(mode_dictionary['Assign keyword to variable'])
    array_vector.append(mode_dictionary['Abnormal replacing'])
    array_vector.append(mode_dictionary['Abnormal String.fromCharCode'])
    array_vector.append(mode_dictionary['Keyword in string'])
    array_vector.append(mode_dictionary['Base 64 encoding'])
    array_vector.append(mode_dictionary['Mysterious mode 1'])
    array_vector.append(mode_dictionary['Too much special characters in string'])
    array_vector.append(mode_dictionary['Special character string'])
    return array_vector

def get_data(abnormal_path, normal_path, keyword_path):
    raw_data = []
    abnormal_result_array = []
    l = []
    for root, dirs, files in os.walk(abnormal_path):
        l.extend(files)
    for f in l:
        vector = dictionary_to_array(normalise_dictionary(mode_analysis(abnormal_path + '\\' + f, keyword_path)))
        vector.append(1)
        abnormal_result_array.append(vector)
    normal_result_array = []
    for root, dirs, files in os.walk(normal_path):
        for file in files:
            programme_path = normal_path + '\\' + file
            try:
                vector = dictionary_to_array(normalise_dictionary(mode_analysis(programme_path, keyword_path)))
                vector.append(0)
                normal_result_array.append(vector)
            except:
               pass
    raw_data.extend(abnormal_result_array)
    raw_data.extend(normal_result_array)
    return raw_data

def test_model(raw_data):
    # Build DataFrame
    entry_length = len(raw_data[0])
    df = pd.DataFrame(np.array(raw_data))
    X = df.drop((entry_length - 1), axis=1)
    y = df[entry_length-1]
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.10)
    X_train_converted = csr_matrix(X_train)
    X_test_converted = csr_matrix(X_test)
    # for array_vector in X_test:
    #     X_test_converted.append(csr_matrix(array_vector))
    classifier = SVC(kernel='linear')
    # classifier = LinearSVC()
    classifier.fit(X_train_converted, y_train)
    y_pred = classifier.predict(X_test_converted)
    # print classifier.support_
    print classifier.coef_
    # print classifier.decision_function(X_test)
    print confusion_matrix(y_test, y_pred)
    print classification_report(y_test, y_pred)

def get_classifier(raw_data):
    # Build DataFrame
    entry_length = len(raw_data[0])
    df = pd.DataFrame(np.array(raw_data))
    X = df.drop((entry_length - 1), axis=1)
    y = df[entry_length-1]
    classifier = SVC(kernel='linear')
    classifier.fit(X, y)
    joblib.dump(classifier, 'prediction_obfuscation.pkl')
    return classifier

def predict_obfuscation(programme_path, keyword_path):
    classifier = joblib.load('prediction_obfuscation.pkl')
    vector = dictionary_to_array(normalise_dictionary(mode_analysis(programme_path, keyword_path)))
    vector_list = []
    vector_list.append(vector)
    prediction = classifier.predict(vector_list)
    if (prediction[0] == 1):
        return True
    else:
        return False

def get_model():
    abnormal_path = 'E:\\encrypted_obfuscated_Javascript_programme_analysis\\Model Testing\\Abnormal Train'
    normal_path = 'E:\\encrypted_obfuscated_Javascript_programme_analysis\\Model Testing\\Normal Train'
    keyword_path = 'JavaScriptKeywords.txt'
    get_classifier(get_data(abnormal_path, normal_path, keyword_path))

# test_model(get_data('E:\\encrypted_obfuscated_JavaScript_programme_analysis\\Virus', 'E:\\encrypted_obfuscated_JavaScript_programme_analysis\\NormalProgrammes', 'JavaScriptKeywords.txt'))
# get_model()


# -*- coding: utf-8 -*-
import numpy as np
import multiprocessing
from scipy.fft import fftshift as sfftshift, ifftshift as sifftshift
from scipy.fft import fft as sfft, rfft as srfft, ifft as sifft, irfft as sirfft
from pathlib import Path
from . import backend as S
from ..configs import IS_PARALLEL

try:
    from torch.fft import (fft as tfft, rfft as trfft,
                           ifft as tifft, irfft as tirfft,
                           fftshift as tfftshift, ifftshift as tifftshift)
except ImportError:
    pass

try:
    import pyfftw
    pyfftw.interfaces.cache.enable()
    pyfftw.interfaces.cache.set_keepalive_time(600)
except ImportError:
    pyfftw = None

UTILS_DIR = Path(__file__).parent

__all__ = [
    'fft',
    'rfft',
    'ifft',
    'irfft',
    'fftshift',
    'ifftshift',
    'FFT',
    'FFT_GLOBAL',
]

#############################################################################

class FFT():
    """Global class for ssqueezepy FFT methods.

    Will use GPU via PyTorch if environment flag `'SSQ_GPU'` is set to `'1'`.
    Will use `scipy.fft` or `pyfftw` depending on `patience` argument (and
    whether `pyfftw` is installed).
    Both will use `threads` CPUs to accelerate computing.

    In a nutshell, if you plan on re-running FFT on input of same shape and dtype,
    prefer `patience=1`, which introduces a lengthy first-time overhead but may
    compute significantly faster afterwards.

    # Arguments (`fft`, `rfft`, `ifft`, `irfft`):
        x: np.ndarray
            1D or 2D.

        axis: int
            FFT axis. One of `0, 1, -1`.

        patience: int / tuple[int, int]
            If int:
                0: will use `scipy.fft`
                1: `pyfftw` with flag `'FFTW_PATIENT'`
                2: `pyfftw` with flag `'FFTW_EXHAUSTIVE'`
            Else, if tuple, second element specifies `planning_timelimit`
            passed to `pyfftw.FFTW` (so tuple requires `patience[0] != 0`).

            Set `planning_timelimit = None` to allow planning to finish,
            but beware; `patience = 1` can take hours for large inputs, and `2`
            even longer.

        astensor: bool (default False)
            If computing on GPU, whether to return as `torch.Tensor` (if False,
            will move to CPU and convert to `numpy.ndarray`).

        n: int / None
            Only for `irfft`; length of original input. If None, will default to
            `2*(x.shape[axis] - 1)`.

    __________________________________________________________________________
    # Arguments (`__init__`):
        planning_timelimit: int
            Default.

        wisdom_dir: str
            Where to save wisdom to or load from. Empty string means
            `ssqueezepy/utils/`.

        threads: int
            Number of CPU threads to use. -1 = maximum.

        patience: int
            Default `patience`.

        cache_fft_objects: bool (default False)
            If True, `pyfftw` objects generated throughout session are stored in
            `FFT._input_history`, and retrieved if all of below match:
                `(x.shape, x.dtype, real, patience, n)`
            where `patience` includes `planning_timelimit` as a tuple.
            Default False since loading from wisdom is very fast anyway.

        verbose: bool (default True)
            Controls whether a message is printed upon `patience >= 1`.
    __________________________________________________________________________
    **Wisdom**

    `pyfftw` uses "wisdom", basically storing and reusing generated FFT plans
    if input attributes match:
        (`x.shape`, `x.dtype`, `axis`, `flags`, `planning_timelimit`)
    `flags` and `planning_timelimit` are set via `patience`.

    With each `pyfftw` use, `save_wisdom()` is called, writing to `wisdom32` and
    `wisdom64` bytes files in `ssqueezepy/utils`. Each time ssqueezepy runs in a
    new session, `load_wisdom()` is called to load these values, so wisdom is
    only expansive.
    """
    def __init__(self, planning_timelimit=120, wisdom_dir=UTILS_DIR, threads=None,
                 patience=0, cache_fft_objects=False, verbose=1):
        self.planning_timelimit = planning_timelimit
        self.wisdom_dir = wisdom_dir
        self._user_threads = threads
        self._patience = patience  # default patience
        self._process_patience(patience)  # error if !=0 and pyfftw not installed
        self.cache_fft_objects = cache_fft_objects
        self.verbose = verbose

        if pyfftw is not None:
            pyfftw.config.NUM_THREADS = self.threads

            self._wisdom32_path = str(Path(self.wisdom_dir, 'wisdom32'))
            self._wisdom64_path = str(Path(self.wisdom_dir, 'wisdom64'))
            self._wisdom32, self._wisdom64 = b'', b''
            self._input_history = {}
            self.load_wisdom()

    @property
    def threads(self):
        """Set dynamically if `threads` wasn't passed in __init__."""
        if self._user_threads is None:
            return (multiprocessing.cpu_count() if IS_PARALLEL() else 1)
        return self._user_threads

    @property
    def patience(self):
        """Setter will also set `planning_timelimit` if setting to tuple."""
        return self._patience

    @patience.setter
    def patience(self, value):
        self._validate_patience(value)
        if isinstance(value, tuple):
            self._patience, self.planning_timelimit = value
        else:
            self._patience = value

    #### Main methods #########################################################
    def fft(self, x, axis=-1, patience=None, astensor=False):
        """See `help(ssqueezepy.utils.FFT)`."""
        out = self._maybe_gpu('fft', x, dim=axis, astensor=astensor)
        if out is not None:
            return out

        patience = self._process_patience(patience)
        if patience == 0:
            return sfft(x, axis=axis, workers=self.threads)

        fft_object = self._get_save_fill(x, axis, patience, real=False)
        return fft_object()

    def rfft(self, x, axis=-1, patience=None, astensor=False):
        """See `help(ssqueezepy.utils.FFT)`."""
        out = self._maybe_gpu('rfft', x, dim=axis, astensor=astensor)
        if out is not None:
            return out

        patience = self._process_patience(patience)
        if patience == 0:
            return srfft(x, axis=axis, workers=self.threads)

        fft_object = self._get_save_fill(x, axis, patience, real=True)
        return fft_object()

    def ifft(self, x, axis=-1, patience=None, astensor=False):
        """See `help(ssqueezepy.utils.FFT)`."""
        out = self._maybe_gpu('ifft', x, dim=axis, astensor=astensor)
        if out is not None:
            return out

        patience = self._process_patience(patience)
        if patience == 0:
            return sifft(x, axis=axis, workers=self.threads)

        fft_object = self._get_save_fill(x, axis, patience, real=False,
                                         inverse=True)
        return fft_object()

    def irfft(self, x, axis=-1, patience=None, astensor=False, n=None):
        """See `help(ssqueezepy.utils.FFT)`."""
        out = self._maybe_gpu('irfft', x, dim=axis, astensor=astensor, n=n)
        if out is not None:
            return out

        patience = self._process_patience(patience)
        if patience == 0:
            return sirfft(x, axis=axis, workers=self.threads, n=n)

        fft_object = self._get_save_fill(x, axis, patience, real=True,
                                         inverse=True, n=n)
        return fft_object()

    def fftshift(self, x, axes=-1, astensor=False):
        out = self._maybe_gpu('fftshift', x, dim=axes, astensor=astensor)
        if out is not None:
            return out
        return sfftshift(x, axes=axes)

    def ifftshift(self, x, axes=-1, astensor=False):
        out = self._maybe_gpu('ifftshift', x, dim=axes, astensor=astensor)
        if out is not None:
            return out
        return sifftshift(x, axes=axes)

    def _maybe_gpu(self, name, x, astensor=False, **kw):
        if S.is_tensor(x):
            fn = {'fft': tfft, 'ifft': tifft,
                  'rfft': trfft, 'irfft': tirfft,
                  'fftshift': tfftshift, 'ifftshift': tifftshift}[name]
            out = fn(S.asarray(x), **kw)
            return out if astensor else out.cpu().numpy()
        return None

    #### FFT makers ###########################################################
    def _get_save_fill(self, x, axis, patience, real, inverse=False, n=None):
        fft_object = self.get_fft_object(x, axis, patience, real, inverse, n)
        self.save_wisdom()
        fft_object.input_array[:] = x
        return fft_object

    def get_fft_object(self, x, axis, patience=1, real=False, inverse=False,
                       n=None):
        combo = (x.shape, x.dtype, axis, real, n)
        if self.cache_fft_objects and combo in self._input_history:
            fft_object = self._input_history[combo]
        else:
            fft_object = self._get_fft_object(x, axis, patience, real, inverse, n)
            if self.cache_fft_objects:
                self._input_history[combo] = fft_object
        return fft_object

    def _get_fft_object(self, x, axis, patience, real, inverse, n):
        (shapes, dtypes, flags, planning_timelimit, direction
         ) = self._process_input(x, axis, patience, real, inverse, n)
        shape_in, shape_out = shapes
        dtype_in, dtype_out = dtypes

        a = pyfftw.empty_aligned(shape_in,  dtype=dtype_in)
        b = pyfftw.empty_aligned(shape_out, dtype=dtype_out)
        fft_object = pyfftw.FFTW(a, b, axes=(axis,), flags=flags,
                                 planning_timelimit=planning_timelimit,
                                 direction=direction, threads=self.threads)
        return fft_object

    def _process_input(self, x, axis, patience, real, inverse, n):
        self._validate_input(x, axis, real, patience, inverse)

        # patience, planning time, forward / inverse
        if isinstance(patience, tuple):
            patience, planning_timelimit = patience
        else:
            planning_timelimit = self.planning_timelimit
        flags = ['FFTW_PATIENT'] if patience == 1 else ['FFTW_EXHAUSTIVE']
        direction = 'FFTW_BACKWARD' if inverse else 'FFTW_FORWARD'

        # shapes
        shape_in = x.shape
        shape_out = self._get_output_shape(x, axis, real, inverse, n)

        # dtypes
        double = x.dtype in (np.float64, np.cfloat)
        cdtype = 'complex128' if double else 'complex64'
        rdtype = 'float64'    if double else 'float32'
        dtype_in  = rdtype if (real and not inverse) else cdtype
        dtype_out = rdtype if (real and inverse)     else cdtype

        # notify user of procedure
        if self.verbose:
            if planning_timelimit is None:
                adjective = "very long" if patience == 2 else "long"
                print("Planning optimal FFT algorithm; this may "
                      "take %s..." % adjective)
            else:
                print("Planning optimal FFT algorithm; this will take up to "
                      "%s secs" % planning_timelimit)

        return ((shape_in, shape_out), (dtype_in, dtype_out), flags,
                planning_timelimit, direction)

    def _get_output_shape(self, x, axis, real=False, inverse=False, n=None):
        if not inverse:
            n_fft = x.shape[axis]
            fft_out_len = (n_fft//2 + 1) if real else n_fft
        else:
            if real:
                n_fft = n if (n is not None) else 2*(x.shape[axis] - 1)
            else:
                n_fft = x.shape[axis]
            fft_out_len = n_fft

        if x.ndim != 1:
            shape = list(x.shape)
            shape[axis] = fft_out_len
            shape = tuple(shape)
        else:
            shape = (fft_out_len,)
        return shape

    #### Misc #################################################################
    def load_wisdom(self):
        for name in ('wisdom32', 'wisdom64'):
            path = getattr(self, f"_{name}_path")
            if Path(path).is_file():
                with open(path, 'rb') as f:
                    setattr(self, f"_{name}", f.read())
        pyfftw.import_wisdom((self._wisdom64, self._wisdom32, b''))

    def save_wisdom(self):
        """Will overwrite."""
        self._wisdom64, self._wisdom32, _ = pyfftw.export_wisdom()
        for name in ('wisdom32', 'wisdom64'):
            path = getattr(self, f"_{name}_path")
            with open(path, 'wb') as f:
                f.write(getattr(self, f"_{name}"))

    def _validate_input(self, x, axis, real, patience, inverse):
        """Assert is single/double precision and is 1D/2D."""
        supported = ('float32', 'float64', 'complex64', 'complex128')
        dtype = str(x.dtype)
        if dtype not in supported:
            raise TypeError("unsupported `x.dtype`: %s " % dtype
                            + "(must be one of: %s)" % ', '.join(supported))
        if (real and not inverse) and dtype.startswith('complex'):
            raise TypeError("`x` cannot be complex for `rfft`")

        if axis not in (0, 1, -1):
            raise ValueError("unsupported `axis`: %s " % axis
                             + "; must be 0, 1, or -1")

        self._validate_patience(patience)

    def _validate_patience(self, patience):
        if not isinstance(patience, (int, tuple)):
            raise TypeError("`patience` must be int or tuple "
                            "(got %s)" % type(patience))
        elif isinstance(patience, int):
            from .common import assert_is_one_of
            assert_is_one_of(patience, 'patience', (0, 1, 2))

    def _process_patience(self, patience):
        patience = patience if (patience is not None) else self.patience
        if pyfftw is None and patience != 0:
            raise ValueError("`patience != 0` requires `pyfftw` installed.")
        return patience

FFT_GLOBAL = FFT()

fft   = FFT_GLOBAL.fft
rfft  = FFT_GLOBAL.rfft
ifft  = FFT_GLOBAL.ifft
irfft = FFT_GLOBAL.irfft
fftshift  = FFT_GLOBAL.fftshift
ifftshift = FFT_GLOBAL.ifftshift


# ----------------------------------------------------------------------------
# Copyright 2015 Nervana Systems Inc.
# 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.
# ----------------------------------------------------------------------------
"""
Defines PASCAL_VOC datatset handling.
"""
import numpy as np
from neon.data.dataloader_transformers import DataLoaderTransformer
from neon.util.persist import get_data_cache_or_nothing

class ObjectLocalization(DataLoaderTransformer):
    """
    Transforms the dataloader inputs for provisioning to the Faster-RCNN model.
    Provided dataloader must be of the (image, localization) type.

    To support the Faster-RCNN model, this transformer:
    1. allocates buffers to start the targets for the classification layers. These
       buffers are populated during fprop by part of the network.
    2. stores metadata for each image that are populated with each minibatch.
    3. provides functions for the Faster-RCNN network to access the above target
       and metdata buffers.
    4. Repacks the buffers provided by the dataloader for the tree structure of the
       Faster-RCNN model.

    Args:
        dataloader (DataLoader): dataloader to transform
        frcn_rois_per_img (int): how many rois to sample to train frcnn
    """

    def __init__(self, dataloader, frcn_rois_per_img, *args, **kwargs):
        super(ObjectLocalization, self).__init__(dataloader, None, *args, **kwargs)

        # obtain settings from the dataloader config
        self.CLASSES = self.dataloader.config['localization']['class_names']
        self.conv_scale = self.dataloader.config['localization']['scaling_factor']
        self.conv_height = int(np.floor(self.dataloader.config['image']['height'] *
                                        self.conv_scale))
        self.conv_width = int(np.floor(self.dataloader.config['image']['width'] *
                                       self.conv_scale))
        self.rpn_rois_per_img = self.dataloader.config['localization']['rois_per_image']

        self.num_classes = len(self.CLASSES)
        self.frcn_rois_per_img = frcn_rois_per_img

        # create placeholders for data needed by ProposalLayer
        self.im_shape = None  # image shape
        self.gt_boxes = None  # ground truth boxes
        self.num_gt_boxes = None  # number of gt boxes
        self.gt_classes = None  # class label for each gt box
        self.im_scale = None  # image scale
        self.difficult = None  # whether the gt boxes are marked difficult

        # allocate buffers
        self.allocate()

    def allocate(self):

        # Data needed for training the RPN are provided by the DataLoader

        # For training, Fast-RCNN needs:
        # 1. class labels
        # 2. bbox targets
        #
        # The above are computed during fprop by the ProposalLayer,
        # so here we create the buffers to pass that to layer.
        self.dev_y_frcn_labels = self.be.zeros(
            (self.num_classes, self.frcn_rois_per_img), dtype=np.int32)
        self.dev_y_frcn_labels_mask = self.be.zeros(
            (self.num_classes, self.frcn_rois_per_img), dtype=np.int32)
        self.dev_y_frcn_bbtargets = self.be.zeros(
            (self.num_classes * 4, self.frcn_rois_per_img), dtype=np.float32)
        self.dev_y_frcn_bbmask = self.be.zeros(
            (self.num_classes * 4, self.frcn_rois_per_img), dtype=np.float32)

    def get_target_buffers(self):

        return ((self.dev_y_frcn_labels, self.dev_y_frcn_labels_mask),
                (self.dev_y_frcn_bbtargets, self.dev_y_frcn_bbmask))

    def get_metadata_buffers(self):

        return (self.im_shape, self.im_scale, self.gt_boxes, self.gt_classes,
                self.num_gt_boxes, self.difficult)

    def transform(self, t):
        # unpack buffers
        (img, bbtargets, bbtargets_mask, labels, labels_mask, self.im_shape, gt_boxes,
            self.num_gt_boxes, self.gt_classes, self.im_scale, self.difficult) = t

        self.gt_boxes = gt_boxes.reshape((-1, 4))

        # repack for faster-rcnn model
        X = img
        Y = ((labels, labels_mask),
             (bbtargets, bbtargets_mask),
             ((self.dev_y_frcn_labels, self.dev_y_frcn_labels_mask),
              (self.dev_y_frcn_bbtargets, self.dev_y_frcn_bbmask)))

        return (X, Y)

def PASCALVOC(manifest_file, manifest_root, rois_per_img=256,
              height=1000, width=1000, inference=False):
    """
    Returns the aeon dataloader configuration for PASCAL VOC dataset.
    """

    CLASSES = ('__background__',
               'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair',
               'cow', 'diningtable', 'dog', 'horse',
               'motorbike', 'person', 'pottedplant',
               'sheep', 'sofa', 'train', 'tvmonitor')

    # if inference, we do not shuffle or flip
    do_transforms = not inference

    image_decode_cfg = dict(height=height, width=width, flip_enable=do_transforms,
                            crop_enable=False, fixed_aspect_ratio=True)
    localization_cfg = dict(rois_per_image=rois_per_img,
                            class_names=CLASSES, scaling_factor=1. / 16)

    return dict(
        manifest_filename=manifest_file,
        manifest_root=manifest_root,
        type='image,localization',
        image=image_decode_cfg,
        cache_directory=get_data_cache_or_nothing(subdir='pascalvoc_cache'),
        shuffle_every_epoch=do_transforms,
        shuffle_manifest=do_transforms,
        minibatch_size=1,
        macrobatch_size=100,
        localization=localization_cfg,
    )

def KITTI(manifest_file, manifest_root, rois_per_img=256,
          height=375, width=1242, inference=False):
    """
    Returns the aeon dataloader configuration for KITTI dataset.
    """

    CLASSES = ('__background__',
               'Car', 'Van', 'Truck',
               'Pedestrian', 'Person_sitting', 'Cyclist', 'Tram',
               'Misc', 'DontCare')

    # if inference, we do not shuffle or flip
    do_transforms = not inference

    image_decode_cfg = dict(height=height, width=width, flip_enable=do_transforms,
                            crop_enable=False, fixed_aspect_ratio=True)
    localization_cfg = dict(rois_per_image=rois_per_img,
                            class_names=CLASSES, scaling_factor=1. / 16)

    return dict(
        manifest_filename=manifest_file,
        manifest_root=manifest_root,
        type='image,localization',
        image=image_decode_cfg,
        cache_directory=get_data_cache_or_nothing(subdir='kitti_cache'),
        shuffle_every_epoch=do_transforms,
        shuffle_manifest=do_transforms,
        minibatch_size=1,
        macrobatch_size=100,
        localization=localization_cfg,
    )


import cv2
import numpy as np
import face_alignment
from skimage import io
import torch
import torch.nn.functional as F
import json
import os
from sklearn.neighbors import NearestNeighbors
from pathlib import Path
import argparse

def euler2rot(euler_angle):
    batch_size = euler_angle.shape[0]
    theta = euler_angle[:, 0].reshape(-1, 1, 1)
    phi = euler_angle[:, 1].reshape(-1, 1, 1)
    psi = euler_angle[:, 2].reshape(-1, 1, 1)
    one = torch.ones((batch_size, 1, 1), dtype=torch.float32,
                     device=euler_angle.device)
    zero = torch.zeros((batch_size, 1, 1), dtype=torch.float32,
                       device=euler_angle.device)
    rot_x = torch.cat((
        torch.cat((one, zero, zero), 1),
        torch.cat((zero, theta.cos(), theta.sin()), 1),
        torch.cat((zero, -theta.sin(), theta.cos()), 1),
    ), 2)
    rot_y = torch.cat((
        torch.cat((phi.cos(), zero, -phi.sin()), 1),
        torch.cat((zero, one, zero), 1),
        torch.cat((phi.sin(), zero, phi.cos()), 1),
    ), 2)
    rot_z = torch.cat((
        torch.cat((psi.cos(), -psi.sin(), zero), 1),
        torch.cat((psi.sin(), psi.cos(), zero), 1),
        torch.cat((zero, zero, one), 1)
    ), 2)
    return torch.bmm(rot_x, torch.bmm(rot_y, rot_z))

parser = argparse.ArgumentParser()
parser.add_argument('--id', type=str,
                    default='obama', help='identity of target person')
parser.add_argument('--step', type=int,
                    default=0, help='step for running')

args = parser.parse_args()
id = args.id
vid_file = os.path.join('dataset', 'vids', id+'.mp4')
if not os.path.isfile(vid_file):
    print('no video')
    exit()

id_dir = os.path.join('dataset', id)
Path(id_dir).mkdir(parents=True, exist_ok=True)
ori_imgs_dir = os.path.join('dataset', id, 'ori_imgs')
Path(ori_imgs_dir).mkdir(parents=True, exist_ok=True)
parsing_dir = os.path.join(id_dir, 'parsing')
Path(parsing_dir).mkdir(parents=True, exist_ok=True)
head_imgs_dir = os.path.join('dataset', id, 'head_imgs')
Path(head_imgs_dir).mkdir(parents=True, exist_ok=True)
com_imgs_dir = os.path.join('dataset', id, 'com_imgs')
Path(com_imgs_dir).mkdir(parents=True, exist_ok=True)

running_step = args.step

# # Step 0: extract wav & deepspeech feature, better run in terminal to parallel with
# below commands since this may take a few minutes
if running_step == 0:
    print('--- Step0: extract deepspeech feature ---')
    wav_file = os.path.join(id_dir, 'aud.wav')
    extract_wav_cmd = 'ffmpeg -i ' + vid_file + ' -f wav -ar 16000 ' + wav_file
    os.system(extract_wav_cmd)
    extract_ds_cmd = 'python data_util/deepspeech_features/extract_ds_features.py --input=' + id_dir
    os.system(extract_ds_cmd)
    exit()

# Step 1: extract images
if running_step == 1:
    print('--- Step1: extract images from vids ---')
    cap = cv2.VideoCapture(vid_file)
    frame_num = 0
    while(True):
        _, frame = cap.read()
        if frame is None:
            break
        cv2.imwrite(os.path.join(ori_imgs_dir, str(frame_num) + '.jpg'), frame)
        frame_num = frame_num + 1
    cap.release()
    exit()

# Step 2: detect lands
if running_step == 2:
    print('--- Step 2: detect landmarks ---')
    fa = face_alignment.FaceAlignment(
        face_alignment.LandmarksType._2D, flip_input=False)
    for image_path in os.listdir(ori_imgs_dir):
            if image_path.endswith('.jpg'):
                input = io.imread(os.path.join(ori_imgs_dir, image_path))[:, :, :3]
                preds = fa.get_landmarks(input)
                if len(preds) > 0:
                    lands = preds[0].reshape(-1, 2)[:,:2]
                    np.savetxt(os.path.join(ori_imgs_dir, image_path[:-3] + 'lms'), lands, '%f')

max_frame_num = 100000
valid_img_ids = []
for i in range(max_frame_num):
    if os.path.isfile(os.path.join(ori_imgs_dir, str(i) + '.lms')):
        valid_img_ids.append(i)
valid_img_num = len(valid_img_ids)
tmp_img = cv2.imread(os.path.join(ori_imgs_dir, str(valid_img_ids[0])+'.jpg'))
h, w = tmp_img.shape[0], tmp_img.shape[1]

# Step 3: face parsing
if running_step == 3:
    print('--- Step 3: face parsing ---')
    face_parsing_cmd = 'python data_util/face_parsing/test.py --respath=dataset/' + \
        id + '/parsing --imgpath=dataset/' + id + '/ori_imgs'
    os.system(face_parsing_cmd)

# Step 4: extract bc image
if running_step == 4:
    print('--- Step 4: extract background image ---')
    sel_ids = np.array(valid_img_ids)[np.arange(0, valid_img_num, 20)]
    all_xys = np.mgrid[0:h, 0:w].reshape(2, -1).transpose()
    distss = []
    for i in sel_ids:
        parse_img = cv2.imread(os.path.join(id_dir, 'parsing', str(i) + '.png'))
        bg = (parse_img[..., 0] == 255) & (
            parse_img[..., 1] == 255) & (parse_img[..., 2] == 255)
        fg_xys = np.stack(np.nonzero(~bg)).transpose(1, 0)
        nbrs = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(fg_xys)
        dists, _ = nbrs.kneighbors(all_xys)
        distss.append(dists)
    distss = np.stack(distss)
    print(distss.shape)
    max_dist = np.max(distss, 0)
    max_id = np.argmax(distss, 0)
    bc_pixs = max_dist > 5
    bc_pixs_id = np.nonzero(bc_pixs)
    bc_ids = max_id[bc_pixs]
    imgs = []
    num_pixs = distss.shape[1]
    for i in sel_ids:
        img = cv2.imread(os.path.join(ori_imgs_dir, str(i) + '.jpg'))
        imgs.append(img)
    imgs = np.stack(imgs).reshape(-1, num_pixs, 3)
    bc_img = np.zeros((h*w, 3), dtype=np.uint8)
    bc_img[bc_pixs_id, :] = imgs[bc_ids, bc_pixs_id, :]
    bc_img = bc_img.reshape(h, w, 3)
    max_dist = max_dist.reshape(h, w)
    bc_pixs = max_dist > 5
    bg_xys = np.stack(np.nonzero(~bc_pixs)).transpose()
    fg_xys = np.stack(np.nonzero(bc_pixs)).transpose()
    nbrs = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(fg_xys)
    distances, indices = nbrs.kneighbors(bg_xys)
    bg_fg_xys = fg_xys[indices[:, 0]]
    print(fg_xys.shape)
    print(np.max(bg_fg_xys), np.min(bg_fg_xys))
    bc_img[bg_xys[:, 0], bg_xys[:, 1],
        :] = bc_img[bg_fg_xys[:, 0], bg_fg_xys[:, 1], :]
    cv2.imwrite(os.path.join(id_dir, 'bc.jpg'), bc_img)

# Step 5: save training images
if running_step == 5:
    print('--- Step 5: save training images ---')
    bc_img = cv2.imread(os.path.join(id_dir, 'bc.jpg'))
    for i in valid_img_ids:
        parsing_img = cv2.imread(os.path.join(parsing_dir, str(i) + '.png'))
        head_part = (parsing_img[:, :, 0] == 255) & (
            parsing_img[:, :, 1] == 0) & (parsing_img[:, :, 2] == 0)
        bc_part = (parsing_img[:, :, 0] == 255) & (
            parsing_img[:, :, 1] == 255) & (parsing_img[:, :, 2] == 255)
        img = cv2.imread(os.path.join(ori_imgs_dir, str(i) + '.jpg'))
        img[bc_part] = bc_img[bc_part]
        cv2.imwrite(os.path.join(com_imgs_dir, str(i) + '.jpg'), img)
        img[~head_part] = bc_img[~head_part]
        cv2.imwrite(os.path.join(head_imgs_dir, str(i) + '.jpg'), img)

# Step 6: estimate head pose
if running_step == 6:
    print('--- Estimate Head Pose ---')
    est_pose_cmd = 'python data_util/face_tracking/face_tracker.py --idname=' + \
        id + ' --img_h=' + str(h) + ' --img_w=' + str(w) + \
        ' --frame_num=' + str(max_frame_num)
    os.system(est_pose_cmd)
    exit()

# Step 7: save transform param & write config file
if running_step == 7:
    print('--- Step 7: Save Transform Param ---')
    params_dict = torch.load(os.path.join(id_dir, 'track_params.pt'))
    focal_len = params_dict['focal']
    euler_angle = params_dict['euler']
    trans = params_dict['trans'] / 10.0
    valid_num = euler_angle.shape[0]
    train_val_split = int(valid_num*10/11)
    train_ids = torch.arange(0, train_val_split)
    val_ids = torch.arange(train_val_split, valid_num)
    rot = euler2rot(euler_angle)
    rot_inv = rot.permute(0, 2, 1)
    trans_inv = -torch.bmm(rot_inv, trans.unsqueeze(2))
    pose = torch.eye(4, dtype=torch.float32)
    save_ids = ['train', 'val']
    train_val_ids = [train_ids, val_ids]
    mean_z = -float(torch.mean(trans[:, 2]).item())
    for i in range(2):
        transform_dict = dict()
        transform_dict['focal_len'] = float(focal_len[0])
        transform_dict['cx'] = float(w/2.0)
        transform_dict['cy'] = float(h/2.0)
        transform_dict['frames'] = []
        ids = train_val_ids[i]
        save_id = save_ids[i]
        for i in ids:
            i = i.item()
            frame_dict = dict()
            frame_dict['img_id'] = int(valid_img_ids[i])
            frame_dict['aud_id'] = int(valid_img_ids[i])
            pose[:3, :3] = rot_inv[i]
            pose[:3, 3] = trans_inv[i, :, 0]
            frame_dict['transform_matrix'] = pose.numpy().tolist()
            lms = np.loadtxt(os.path.join(
                ori_imgs_dir, str(valid_img_ids[i]) + '.lms'))
            min_x, max_x = np.min(lms, 0)[0], np.max(lms, 0)[0]
            cx = int((min_x+max_x)/2.0)
            cy = int(lms[27, 1])
            h_w = int((max_x-cx)*1.5)
            h_h = int((lms[8, 1]-cy)*1.15)
            rect_x = cx - h_w
            rect_y = cy - h_h
            if rect_x < 0:
                rect_x = 0
            if rect_y < 0:
                rect_y = 0
            rect_w = min(w-1-rect_x, 2*h_w)
            rect_h = min(h-1-rect_y, 2*h_h)
            rect = np.array((rect_x, rect_y, rect_w, rect_h), dtype=np.int32)
            frame_dict['face_rect'] = rect.tolist()
            transform_dict['frames'].append(frame_dict)
        with open(os.path.join(id_dir, 'transforms_' + save_id + '.json'), 'w') as fp:
            json.dump(transform_dict, fp, indent=2, separators=(',', ': '))

    dir_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
    testskip = int(val_ids.shape[0]/7)

    HeadNeRF_config_file = os.path.join(id_dir, 'HeadNeRF_config.txt')
    with open(HeadNeRF_config_file, 'w') as file:
        file.write('expname = ' + id + '_head\n')
        file.write('datadir = ' + os.path.join(dir_path, 'dataset', id) + '\n')
        file.write('basedir = ' + os.path.join(dir_path,
                                            'dataset', id, 'logs') + '\n')
        file.write('near = ' + str(mean_z-0.2) + '\n')
        file.write('far = ' + str(mean_z+0.4) + '\n')
        file.write('testskip = ' + str(testskip) + '\n')
    Path(os.path.join(dir_path, 'dataset', id, 'logs', id + '_head')
        ).mkdir(parents=True, exist_ok=True)

    ComNeRF_config_file = os.path.join(id_dir, 'TorsoNeRF_config.txt')
    with open(ComNeRF_config_file, 'w') as file:
        file.write('expname = ' + id + '_com\n')
        file.write('datadir = ' + os.path.join(dir_path, 'dataset', id) + '\n')
        file.write('basedir = ' + os.path.join(dir_path,
                                            'dataset', id, 'logs') + '\n')
        file.write('near = ' + str(mean_z-0.2) + '\n')
        file.write('far = ' + str(mean_z+0.4) + '\n')
        file.write('testskip = ' + str(testskip) + '\n')
    Path(os.path.join(dir_path, 'dataset', id, 'logs', id + '_com')
        ).mkdir(parents=True, exist_ok=True)

    ComNeRFTest_config_file = os.path.join(id_dir, 'TorsoNeRFTest_config.txt')
    with open(ComNeRFTest_config_file, 'w') as file:
        file.write('expname = ' + id + '_com\n')
        file.write('datadir = ' + os.path.join(dir_path, 'dataset', id) + '\n')
        file.write('basedir = ' + os.path.join(dir_path,
                                            'dataset', id, 'logs') + '\n')
        file.write('near = ' + str(mean_z-0.2) + '\n')
        file.write('far = ' + str(mean_z+0.4) + '\n')
        file.write('with_test = ' + str(1) + '\n')
        file.write('test_pose_file = transforms_val.json' + '\n')

    print(id + ' data processed done!')