text: import json
  import logging
  import os
  import pickle
  from itertools import zip_longest
  
  import click
  import numpy as np
  from tensorflow.saved_model import signature_def_utils, builder, tag_constants, signature_constants
  from keras import backend as K
  from keras.models import Model, load_model
  from keras.layers import Input
  from keras.applications.inception_resnet_v2 import InceptionResNetV2
  from keras.applications.inception_resnet_v2 import preprocess_input
  from sklearn.neighbors import BallTree
  from PIL import Image
  
  
  logger = logging.getLogger('keras_training')
  formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(name)s - %(message)s')
  ch = logging.StreamHandler()
  ch.setFormatter(formatter)
  logger.setLevel(logging.INFO)
  logger.addHandler(ch)
  
  IMAGE_DIMENSIONS = (299, 299, 3)
  SEARCH_INDEX_FILENAME = 'search_index.pkl'
  
  
  def get_default_inception_model():
      input = Input(shape=IMAGE_DIMENSIONS)
      return InceptionResNetV2(weights='imagenet', input_tensor=input)
  
  
  def get_default_inception_model_signature():
      model = get_default_inception_model()
  
      signature = signature_def_utils.predict_signature_def(
          inputs={'input': model.input},
          outputs={
              'embedding': model.layers[-1].input,
              'softmax': model.layers[-1].output
          }
      )
      return signature
  
  
  def export_savedmodel_for_tensorflow_serving(basepath, version, signature_name, signature):
      """Exports a model signature to a SavedModel format
  
      Args:
          basepath (str): Path to store SavedModel files
          version (int): Version of the model being exported
          signature_name (str): Name of signature to request in Tensorflow Serving
          signature (signature_def): Signature to make available to Tensorflow Serving
      """
      save_path = os.path.join(basepath, str(version))
      model_builder = builder.SavedModelBuilder(save_path)
      model_builder.add_meta_graph_and_variables(
          sess=K.get_session(),
          tags=[tag_constants.SERVING],
          signature_def_map={
              signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature,
              signature_name: signature
          }
      )
      model_builder.save()
  
  
  def convert_image_to_square_rgb(image, image_dims=IMAGE_DIMENSIONS[:2]):
      """Takes an image, converts it to RGB, pad to a square image, resizes,
      and converts to numpy array.
  
      Args:
          image (PIL.Image): Image to process
  
      Returns:
          PIL.Image
      """
      if image.mode != 'RGB':
          image = image.convert('RGB')
  
      width, height = image.size
      if width != height:
          max_dimension = max(image.size)
          square_image = Image.new('RGB', (max_dimension, max_dimension), 'white')
          paste_location = ((max_dimension - width) // 2, (max_dimension - height) // 2)
          square_image.paste(image, paste_location)
          image = square_image
  
      if image.size != image_dims:
          image = image.resize(image_dims, Image.ANTIALIAS)
          
      return np.array(image)
  
  
  def preprocess_image_inception_keras(image_matrix):
      X = np.array(image_matrix, dtype=np.float32)
      return preprocess_input(X)
  
  
  def contrastive_loss(y_true, y_pred):
      """Contrastive loss from Hadsell-et-al.'06
      http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf
      """
      margin = 1
      return K.mean(y_true * K.square(y_pred) +
                    (1 - y_true) * K.square(K.maximum(margin - y_pred, 0)))
  
  
  def grouper(n, iterable, fillvalue=None):
      args = [iter(iterable)] * n
      return zip_longest(fillvalue=fillvalue, *args)
  
  
  def image_batches(rows, batch_size, image_dir=None):
      """A generator that yields a batch of images from a flat list of image
      data.
  
      Args:
          rows (List): List of (metadata, filename) pairs
          batch_size (int): The max size of each batch of images
          image_dir (Optional[str]): The root dir for where images are stored,
                                     use if filename in rows is relative
  
      Returns:
          tuple(dict, ndarray)
          """
      for batch in grouper(batch_size, rows, (None, None)):
          images = []
          payloads = []
          for metadata, filename in batch:
              if filename:
                  if image_dir:
                      filename = os.path.join(image_dir, filename)
                  image = Image.open(filename)
                  x = convert_image_to_square_rgb(image)
                  images.append(x)
                  payloads.append(metadata)
              else:
                  break
          yield (payloads, preprocess_image_inception_keras(images))
  
  
  @click.command()
  @click.option('--export', type=click.Choice(['savedmodel', 'balltree']))
  @click.option('--keras-model')
  @click.option('--labeled-data')
  @click.option('--image-dir')
  @click.option('--batch-size', default=1024)
  def cli(export, keras_model, labeled_data, image_dir, batch_size):
      if export == 'savedmodel' and keras_model:
          logger.info('Starting export of {} to SavedModel'.format(keras_model))
          model = load_model(
              keras_model,
              custom_objects={'contrastive_loss': contrastive_loss}
          )
          embedder = model.get_layer('embedding_model')
          image_tensor = embedder.inputs[0]
          embedding_tensor = embedder.outputs[0]
          signature = signature_def_utils.predict_signature_def(
              inputs={'input': image_tensor},
              outputs={
                  'embedding': embedding_tensor,
              }
          )
          export_savedmodel_for_tensorflow_serving('/app/savedmodels/inception_resnet_v2', 2, 'prediction', signature)
      elif export == 'savedmodel' and not keras_model:
          logger.info('Starting export of default SavedModel')
          signature = get_default_inception_model_signature()
          export_savedmodel_for_tensorflow_serving('/app/savedmodels/inception_resnet_v2', 1, 'prediction', signature)
      elif export == 'balltree':
          if keras_model:
              logger.info('Starting BallTree creation for {}'.format(keras_model))
              model = load_model(
                  keras_model,
                  custom_objects={'contrastive_loss': contrastive_loss}
              )
              model = model.get_layer('embedding_model')
          else:
              logger.info('Starting BallTree creation for default model')
              inception = get_default_inception_model()
              model = Model(
                  inputs=[inception.input],
                  outputs=[inception.layers[-1].input, inception.layers[-1].output],
                  name='embedding_model'
              )
  
          with open(labeled_data) as f:
              item_image_data = json.loads(f.read())
  
          rows = []
          for item_id, data in item_image_data.items():
              for image in data['images']:
                  # Edit here if you want more data in the payload for each image
                  payload = {
                      'item_id': item_id,
                      'filename': image['filename']
                  }
                  row = (payload, image['filename'])
                  rows.append(row)
  
          logger.info('Embedding {} images'.format(len(rows)))
          payloads = []
          embeddings = []
          num_batches = len(rows) // batch_size
          for i, (payload_batch, images) in enumerate(image_batches(rows, batch_size, image_dir)):
              print(' {}/{} batches done'.format(i, num_batches), end='\r')
              outputs = model.predict(images)
              if type(outputs) == list:
                  embedding, softmax = outputs
              elif type(outputs) == np.ndarray:
                  embedding = outputs
              payloads.extend(payload_batch)
              embeddings.append(embedding)
          print()
  
          logger.info('Making ball tree')
          X = np.vstack(embeddings)
          ball_tree = BallTree(X, leaf_size=40)
          index_data = {
              'metadata': payloads,
              'index': ball_tree
          }
  
          logger.info('Pickling to {}'.format(SEARCH_INDEX_FILENAME))
          with open(SEARCH_INDEX_FILENAME, 'wb') as f:
              pickle.dump(index_data, f, pickle.HIGHEST_PROTOCOL)
  
  
  if __name__ == '__main__':
      cli()
  