def export_model()

in 07_training/serverlessml/flowers/classifier/model.py [0:0]

• 25 lines of code
• 1 McCabe index (conditional complexity)

def export_model(model, outdir, IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS):
    def create_preproc_image_of_right_size(filename):
        return create_preproc_image(filename, IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS)

    @tf.function(input_signature=[tf.TensorSpec([None,], dtype=tf.string)])
    def predict_flower_type(filenames):
        input_images = tf.map_fn(
            create_preproc_image_of_right_size,
            filenames,
            fn_output_signature=tf.float32
        )
        batch_pred = model(input_images) # same as model.predict()
        top_prob = tf.math.reduce_max(batch_pred, axis=[1])
        pred_label_index = tf.math.argmax(batch_pred, axis=1)
        pred_label = tf.gather(tf.convert_to_tensor(CLASS_NAMES), pred_label_index)
        return {
            'probability': top_prob,
            'flower_type_int': pred_label_index,
            'flower_type_str': pred_label
        }
    
    outpath = os.path.join(outdir, 'flowers_model')
    cleanup_dir(outpath)
    model.save(outpath,
          signatures={
              'serving_default': predict_flower_type
          })