text: import tensorflow as tf
  import numpy as np
  import part_detector
  import settings
  import utils
  import os
  
  from abc import abstractmethod
  from functools import lru_cache
  from scipy.stats import norm
  
  from inception_resnet_v2 import inception_resnet_v2_arg_scope, inception_resnet_v2
  
  import tensorflow.contrib.layers as layers
  
  slim = tf.contrib.slim
  
  SUMMARY_PATH = settings.LOGDIR_PATH
  
  KEY_SUMMARIES = tf.GraphKeys.SUMMARIES
  KEY_SUMMARIES_PER_JOINT = ['summary_joint_%02d' % i for i in range(16)]
  
  
  class HumanPoseNN(object):
      """
      The neural network used for pose estimation.
      """
  
      def __init__(self, log_name, heatmap_size, image_size, loss_type = 'SCE', is_training = True):
          tf.set_random_seed(0)
  
          if loss_type not in { 'MSE', 'SCE' }:
              raise NotImplementedError('Loss function should be either MSE or SCE!')
  
          self.log_name = log_name
          self.heatmap_size = heatmap_size
          self.image_size = image_size
          self.is_train = is_training
          self.loss_type = loss_type
  
          # Initialize placeholders
          self.input_tensor = tf.placeholder(
              dtype = tf.float32,
              shape = (None, image_size, image_size, 3),
              name = 'input_image')
  
          self.present_joints = tf.placeholder(
              dtype = tf.float32,
              shape = (None, 16),
              name = 'present_joints')
  
          self.inside_box_joints = tf.placeholder(
              dtype = tf.float32,
              shape = (None, 16),
              name = 'inside_box_joints')
  
          self.desired_heatmap = tf.placeholder(
              dtype = tf.float32,
              shape = (None, heatmap_size, heatmap_size, 16),
              name = 'desired_heatmap')
  
          self.desired_points = tf.placeholder(
              dtype = tf.float32,
              shape = (None, 2, 16),
              name = 'desired_points')
  
          self.network = self.pre_process(self.input_tensor)
          self.network, self.feature_tensor = self.get_network(self.network, is_training)
  
          self.sigm_network = tf.sigmoid(self.network)
          self.smoothed_sigm_network = self._get_gauss_smoothing_net(self.sigm_network, std = 0.7)
  
          self.loss_err = self._get_loss_function(loss_type)
          self.euclidean_dist = self._euclidean_dist_err()
          self.euclidean_dist_per_joint = self._euclidean_dist_per_joint_err()
  
          if is_training:
              self.global_step = tf.Variable(0, name = 'global_step', trainable = False)
  
              self.learning_rate = tf.placeholder(
                  dtype = tf.float32,
                  shape = [],
                  name = 'learning_rate')
  
              self.optimize = layers.optimize_loss(loss = self.loss_err,
                                                   global_step = self.global_step,
                                                   learning_rate = self.learning_rate,
                                                   optimizer = tf.train.RMSPropOptimizer(self.learning_rate),
                                                   clip_gradients = 2.0
                                                   )
  
          self.sess = tf.Session()
          self.sess.run(tf.global_variables_initializer())
  
          if log_name is not None:
              self._init_summaries()
  
      def _init_summaries(self):
          if self.is_train:
              logdir = os.path.join(SUMMARY_PATH, self.log_name, 'train')
  
              self.summary_writer = tf.summary.FileWriter(logdir)
              self.summary_writer_by_points = [tf.summary.FileWriter(os.path.join(logdir, 'point_%02d' % i))
                                               for i in range(16)]
  
              tf.scalar_summary('Average euclidean distance', self.euclidean_dist, collections = [KEY_SUMMARIES])
  
              for i in range(16):
                  tf.scalar_summary('Joint euclidean distance', self.euclidean_dist_per_joint[i],
                                    collections = [KEY_SUMMARIES_PER_JOINT[i]])
  
              self.create_summary_from_weights()
  
              self.ALL_SUMMARIES = tf.merge_all_summaries(KEY_SUMMARIES)
              self.SUMMARIES_PER_JOINT = [tf.merge_all_summaries(KEY_SUMMARIES_PER_JOINT[i]) for i in range(16)]
          else:
              logdir = os.path.join(SUMMARY_PATH, self.log_name, 'test')
              self.summary_writer = tf.summary.FileWriter(logdir)
  
      def _get_loss_function(self, loss_type):
          loss_dict = {
              'MSE': self._loss_mse(),
              'SCE': self._loss_cross_entropy()
          }
  
          return loss_dict[loss_type]
  
      @staticmethod
      @lru_cache()
      def _get_gauss_filter(size = 15, std = 1.0, kernel_sum = 1.0):
          samples = norm.pdf(np.linspace(-2, 2, size), 0, std)
          samples /= np.sum(samples)
          samples *= kernel_sum ** 0.5
  
          samples = np.expand_dims(samples, 0)
          weights = np.zeros(shape = (1, size, 16, 1), dtype = np.float32)
  
          for i in range(16):
              weights[:, :, i, 0] = samples
  
          return weights
  
      @staticmethod
      def _get_gauss_smoothing_net(net, size = 15, std = 1.0, kernel_sum = 1.0):
          filter_h = HumanPoseNN._get_gauss_filter(size, std, kernel_sum)
          filter_v = filter_h.swapaxes(0, 1)
  
          net = tf.nn.depthwise_conv2d(net, filter = filter_h, strides = [1, 1, 1, 1], padding = 'SAME',
                                       name = 'SmoothingHorizontal')
  
          net = tf.nn.depthwise_conv2d(net, filter = filter_v, strides = [1, 1, 1, 1], padding = 'SAME',
                                       name = 'SmoothingVertical')
  
          return net
  
      def generate_output(self, shape, presented_parts, labels, sigma):
          heatmap_dict = {
              'MSE': utils.get_gauss_heat_map(
                  shape = shape, is_present = presented_parts,
                  mean = labels, sigma = sigma),
              'SCE': utils.get_binary_heat_map(
                  shape = shape, is_present = presented_parts,
                  centers = labels, diameter = sigma)
          }
  
          return heatmap_dict[self.loss_type]
  
      def _adjust_loss(self, loss_err):
          # Shape: [batch, joints]
          loss = tf.reduce_sum(loss_err, [1, 2])
  
          # Stop error propagation of joints that are not presented
          loss = tf.multiply(loss, self.present_joints)
  
          # Compute average loss of presented joints
          num_of_visible_joints = tf.reduce_sum(self.present_joints)
          loss = tf.reduce_sum(loss) / num_of_visible_joints
  
          return loss
  
      def _loss_mse(self):
          sq = tf.squared_difference(self.sigm_network, self.desired_heatmap)
          loss = self._adjust_loss(sq)
  
          return loss
  
      def _loss_cross_entropy(self):
          ce = tf.nn.sigmoid_cross_entropy_with_logits(logits = self.network, labels = self.desired_heatmap)
          loss = self._adjust_loss(ce)
  
          return loss
  
      def _joint_highest_activations(self):
          highest_activation = tf.reduce_max(self.smoothed_sigm_network, [1, 2])
  
          return highest_activation
  
      def _joint_positions(self):
          highest_activation = tf.reduce_max(self.sigm_network, [1, 2])
          x = tf.argmax(tf.reduce_max(self.smoothed_sigm_network, 1), 1)
          y = tf.argmax(tf.reduce_max(self.smoothed_sigm_network, 2), 1)
  
          x = tf.cast(x, tf.float32)
          y = tf.cast(y, tf.float32)
          a = tf.cast(highest_activation, tf.float32)
  
          scale_coef = (self.image_size / self.heatmap_size)
          x *= scale_coef
          y *= scale_coef
  
          out = tf.stack([y, x, a])
  
          return out
  
      def _euclidean_dist_err(self):
          # Work only with joints that are presented inside frame
          l2_dist = tf.multiply(self.euclidean_distance(), self.inside_box_joints)
  
          # Compute average loss of presented joints
          num_of_visible_joints = tf.reduce_sum(self.inside_box_joints)
          l2_dist = tf.reduce_sum(l2_dist) / num_of_visible_joints
  
          return l2_dist
  
      def _euclidean_dist_per_joint_err(self):
          # Work only with joints that are presented inside frame
          l2_dist = tf.multiply(self.euclidean_distance(), self.inside_box_joints)
  
          # Average euclidean distance of presented joints
          present_joints = tf.reduce_sum(self.inside_box_joints, 0)
          err = tf.reduce_sum(l2_dist, 0) / present_joints
  
          return err
  
      def _restore(self, checkpoint_path, variables):
          saver = tf.train.Saver(variables)
          saver.restore(self.sess, checkpoint_path)
  
      def _save(self, checkpoint_path, name, variables):
          if not os.path.exists(checkpoint_path):
              os.mkdir(checkpoint_path)
  
          checkpoint_name_path = os.path.join(checkpoint_path, '%s.ckpt' % name)
  
          saver = tf.train.Saver(variables)
          saver.save(self.sess, checkpoint_name_path)
  
      def euclidean_distance(self):
          x = tf.argmax(tf.reduce_max(self.smoothed_sigm_network, 1), 1)
          y = tf.argmax(tf.reduce_max(self.smoothed_sigm_network, 2), 1)
  
          x = tf.cast(x, tf.float32)
          y = tf.cast(y, tf.float32)
  
          dy = tf.squeeze(self.desired_points[:, 0, :])
          dx = tf.squeeze(self.desired_points[:, 1, :])
  
          sx = tf.squared_difference(x, dx)
          sy = tf.squared_difference(y, dy)
  
          l2_dist = tf.sqrt(sx + sy)
  
          return l2_dist
  
      def feed_forward(self, x):
          out = self.sess.run(self.sigm_network, feed_dict = {
              self.input_tensor: x
          })
  
          return out
  
      def heat_maps(self, x):
          out = self.sess.run(self.smoothed_sigm_network, feed_dict = {
              self.input_tensor: x
          })
  
          return out
  
      def feed_forward_pure(self, x):
          out = self.sess.run(self.network, feed_dict = {
              self.input_tensor: x
          })
  
          return out
  
      def feed_forward_features(self, x):
          out = self.sess.run(self.feature_tensor, feed_dict = {
              self.input_tensor: x,
          })
  
          return out
  
      def test_euclidean_distance(self, x, points, present_joints, inside_box_joints):
          err = self.sess.run(self.euclidean_dist, feed_dict = {
              self.input_tensor: x,
              self.desired_points: points,
              self.present_joints: present_joints,
              self.inside_box_joints: inside_box_joints
          })
  
          return err
  
      def test_joint_distances(self, x, y):
          err = self.sess.run(self.euclidean_distance(), feed_dict = {
              self.input_tensor: x,
              self.desired_points: y
          })
  
          return err
  
      def test_joint_activations(self, x):
          err = self.sess.run(self._joint_highest_activations(), feed_dict = {
              self.input_tensor: x
          })
  
          return err
  
      def estimate_joints(self, x):
          out = self.sess.run(self._joint_positions(), feed_dict = {
              self.input_tensor: x
          })
  
          return out
  
      def train(self, x, heatmaps, present_joints, learning_rate, is_inside_box):
          if not self.is_train:
              raise Exception('Network is not in train mode!')
  
          self.sess.run(self.optimize, feed_dict = {
              self.input_tensor: x,
              self.desired_heatmap: heatmaps,
              self.present_joints: present_joints,
              self.learning_rate: learning_rate,
              self.inside_box_joints: is_inside_box
          })
  
      def write_test_summary(self, epoch, loss):
          loss_sum = tf.Summary()
          loss_sum.value.add(
              tag = 'Average Euclidean Distance',
              simple_value = float(loss))
          self.summary_writer.add_summary(loss_sum, epoch)
          self.summary_writer.flush()
  
      def write_summary(self, inp, desired_points, heatmaps, present_joints, learning_rate, is_inside_box,
                        write_frequency = 20, write_per_joint_frequency = 100):
          step = tf.train.global_step(self.sess, self.global_step)
  
          if step % write_frequency == 0:
              feed_dict = {
                  self.input_tensor: inp,
                  self.desired_points: desired_points,
                  self.desired_heatmap: heatmaps,
                  self.present_joints: present_joints,
                  self.learning_rate: learning_rate,
                  self.inside_box_joints: is_inside_box
              }
  
              summary, loss = self.sess.run([self.ALL_SUMMARIES, self.loss_err], feed_dict = feed_dict)
              self.summary_writer.add_summary(summary, step)
  
              if step % write_per_joint_frequency == 0:
                  summaries = self.sess.run(self.SUMMARIES_PER_JOINT, feed_dict = feed_dict)
  
                  for i in range(16):
                      self.summary_writer_by_points[i].add_summary(summaries[i], step)
  
                  for i in range(16):
                      self.summary_writer_by_points[i].flush()
  
              self.summary_writer.flush()
  
      @abstractmethod
      def pre_process(self, inp):
          pass
  
      @abstractmethod
      def get_network(self, input_tensor, is_training):
          pass
  
      @abstractmethod
      def create_summary_from_weights(self):
          pass
  
  
  class HumanPoseIRNetwork(HumanPoseNN):
      """
      The first part of our network that exposes as an extractor of spatial features. It s derived from
      Inception-Resnet-v2 architecture and modified for generating heatmaps - i.e. dense predictions of body joints.
      """
  
      FEATURES = 32
      IMAGE_SIZE = 299
      HEATMAP_SIZE = 289
      POINT_DIAMETER = 15
      SMOOTH_SIZE = 21
  
      def __init__(self, log_name = None, loss_type = 'SCE', is_training = False):
          super().__init__(log_name, self.HEATMAP_SIZE, self.IMAGE_SIZE, loss_type, is_training)
  
      def pre_process(self, inp):
          return ((inp / 255) - 0.5) * 2.0
  
      def get_network(self, input_tensor, is_training):
          # Load pre-trained inception-resnet model
          with slim.arg_scope(inception_resnet_v2_arg_scope(batch_norm_decay = 0.999, weight_decay = 0.0001)):
              net, end_points = inception_resnet_v2(input_tensor, is_training = is_training)
  
          # Adding some modification to original InceptionResnetV2 - changing scoring of AUXILIARY TOWER
          weight_decay = 0.0005
          with tf.variable_scope('NewInceptionResnetV2'):
              with tf.variable_scope('AuxiliaryScoring'):
                  with slim.arg_scope([layers.convolution2d, layers.convolution2d_transpose],
                                      weights_regularizer = slim.l2_regularizer(weight_decay),
                                      biases_regularizer = slim.l2_regularizer(weight_decay),
                                      activation_fn = None):
                      tf.summary.histogram('Last_layer/activations', net, [KEY_SUMMARIES])
  
                      # Scoring
                      net = slim.dropout(net, 0.7, is_training = is_training, scope = 'Dropout')
                      net = layers.convolution2d(net, num_outputs = self.FEATURES, kernel_size = 1, stride = 1,
                                                 scope = 'Scoring_layer')
                      feature = net
                      tf.summary.histogram('Scoring_layer/activations', net, [KEY_SUMMARIES])
  
                      # Upsampling
                      net = layers.convolution2d_transpose(net, num_outputs = 16, kernel_size = 17, stride = 17,
                                                           padding = 'VALID', scope = 'Upsampling_layer')
  
                      tf.summary.histogram('Upsampling_layer/activations', net, [KEY_SUMMARIES])
  
              # Smoothing layer - separable gaussian filters
              net = super()._get_gauss_smoothing_net(net, size = self.SMOOTH_SIZE, std = 1.0, kernel_sum = 0.2)
  
              return net, feature
  
      def restore(self, checkpoint_path, is_pre_trained_imagenet_checkpoint = False):
          all_vars = tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'InceptionResnetV2')
          if not is_pre_trained_imagenet_checkpoint:
              all_vars += tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'NewInceptionResnetV2/AuxiliaryScoring')
  
          super()._restore(checkpoint_path, all_vars)
  
      def save(self, checkpoint_path, name):
          all_vars = tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'InceptionResnetV2')
          all_vars += tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'NewInceptionResnetV2/AuxiliaryScoring')
  
          super()._save(checkpoint_path, name, all_vars)
  
      def create_summary_from_weights(self):
          with tf.variable_scope('NewInceptionResnetV2/AuxiliaryScoring', reuse = True):
              tf.summary.histogram('Scoring_layer/biases', tf.get_variable('Scoring_layer/biases'), [KEY_SUMMARIES])
              tf.summary.histogram('Upsampling_layer/biases', tf.get_variable('Upsampling_layer/biases'), [KEY_SUMMARIES])
              tf.summary.histogram('Scoring_layer/weights', tf.get_variable('Scoring_layer/weights'), [KEY_SUMMARIES])
              tf.summary.histogram('Upsampling_layer/weights', tf.get_variable('Upsampling_layer/weights'),
                                   [KEY_SUMMARIES])
  
          with tf.variable_scope('InceptionResnetV2/AuxLogits', reuse = True):
              tf.summary.histogram('Last_layer/weights', tf.get_variable('Conv2d_2a_5x5/weights'), [KEY_SUMMARIES])
              tf.summary.histogram('Last_layer/beta', tf.get_variable('Conv2d_2a_5x5/BatchNorm/beta'), [KEY_SUMMARIES])
              tf.summary.histogram('Last_layer/moving_mean', tf.get_variable('Conv2d_2a_5x5/BatchNorm/moving_mean'),
                                   [KEY_SUMMARIES])
  
  
  class PartDetector(HumanPoseNN):
      """
      Architecture of Part Detector network, as was described in https://arxiv.org/abs/1609.01743
      """
  
      IMAGE_SIZE = 256
      HEATMAP_SIZE = 256
      POINT_DIAMETER = 11
  
      def __init__(self, log_name = None, init_from_checkpoint = None, loss_type = 'SCE', is_training = False):
          if init_from_checkpoint is not None:
              part_detector.init_model_variables(init_from_checkpoint, is_training)
              self.reuse = True
          else:
              self.reuse = False
  
          super().__init__(log_name, self.HEATMAP_SIZE, self.IMAGE_SIZE, loss_type, is_training)
  
      def pre_process(self, inp):
          return inp / 255
  
      def create_summary_from_weights(self):
          pass
  
      def restore(self, checkpoint_path):
          all_vars = tf.get_collection(tf.GraphKeys.VARIABLES, scope = 'HumanPoseResnet')
          all_vars += tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'NewHumanPoseResnet/Scoring')
  
          super()._restore(checkpoint_path, all_vars)
  
      def save(self, checkpoint_path, name):
          all_vars = tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'HumanPoseResnet')
          all_vars += tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'NewHumanPoseResnet/Scoring')
  
          super()._save(checkpoint_path, name, all_vars)
  
      def get_network(self, input_tensor, is_training):
          net_end, end_points = part_detector.human_pose_resnet(input_tensor, reuse = self.reuse, training = is_training)
  
          return net_end, end_points['features']
  