tensorflow–将已有模型的输入改为 base64 images

背景

在计算机视觉方向，平时训练得到的网络一般都是图像作为输入，但是如果要部署到服务端事情就会变得不同，一般服务端接收到的图像是经过base64编码之后的数据，那么这就不能直接喂给你的模型，此时有两种解决方案：一、在服务段进行处理，先将base64转为图像再喂给你的模型，这个过程你有权限或者你的服务端同事会写代码都是可以的；二、假如你没有权限，碰巧你的服务端同事也不会写代码，那你只能把含辛茹苦训好的模型改造一番，本文就是在面对这种情况的时候总结出来的。

方法

1. 重新训练 – 大可不必

2. 添加前置处理即可
   tf
   代码中类别标签的映射部分如果不需要，可以删掉，这样代码会简洁很多，前处理结束之后直接导出模型即可。

   def preprocess_image(image_buffer):
     """Preprocess JPEG encoded bytes to 3D float Tensor."""
   
     # Decode the string as an RGB JPEG.
     # Note that the resulting image contains an unknown height and width
     # that is set dynamically by decode_jpeg. In other words, the height
     # and width of image is unknown at compile-time.
     # tf.io.decode_jpeg if tf.__version__ >= 2
     image = tf.image.decode_jpeg(image_buffer, channels=3)
     # After this point, all image pixels reside in [0,1)
     # until the very end, when they're rescaled to (-1, 1).  The various
     # adjust_* ops all require this range for dtype float.
     image = tf.image.convert_image_dtype(image, dtype=tf.float32)
     # Crop the central region of the image with an area containing 87.5% of
     # the original image.
     image = tf.image.central_crop(image, central_fraction=0.875)
     # Resize the image to the original height and width.
     image = tf.expand_dims(image, 0)
     image = tf.image.resize_bilinear(
         image, [FLAGS.image_size, FLAGS.image_size], align_corners=False)
     image = tf.squeeze(image, [0])
     # Finally, rescale to [-1,1] instead of [0, 1)
     image = tf.subtract(image, 0.5)
     image = tf.multiply(image, 2.0)
     return image
   

   NUM_CLASSES = 1000
   NUM_TOP_CLASSES = 5
   
   WORKING_DIR = os.path.dirname(os.path.realpath(__file__))
   SYNSET_FILE = os.path.join(WORKING_DIR, 'imagenet_lsvrc_2015_synsets.txt')
   METADATA_FILE = os.path.join(WORKING_DIR, 'imagenet_metadata.txt')
   
   
   def export():
     # Create index->synset mapping
     synsets = []
     with open(SYNSET_FILE) as f:
       synsets = f.read().splitlines()
     # Create synset->metadata mapping
     texts = {}
     with open(METADATA_FILE) as f:
       for line in f.read().splitlines():
         parts = line.split('\t')
         assert len(parts) == 2
         texts[parts[0]] = parts[1]
   
     with tf.Graph().as_default():
       # Build inference model.
       # Please refer to Tensorflow inception model for details.
   
       # Input transformation.
       serialized_tf_example = tf.placeholder(tf.string, name='tf_example')
       feature_configs = {
           'image/encoded': tf.FixedLenFeature(
               shape=[], dtype=tf.string),
       }
       tf_example = tf.parse_example(serialized_tf_example, feature_configs)
       jpegs = tf_example['image/encoded']
       images = tf.map_fn(preprocess_image, jpegs, dtype=tf.float32)
   
       # Run inference.
       logits, _ = inception_model.inference(images, NUM_CLASSES + 1)
   
       # Transform output to topK result.
       values, indices = tf.nn.top_k(logits, NUM_TOP_CLASSES)
   
       # Create a constant string Tensor where the i'th element is
       # the human readable class description for the i'th index.
       # Note that the 0th index is an unused background class
       # (see inception model definition code).
       class_descriptions = ['unused background']
       for s in synsets:
         class_descriptions.append(texts[s])
       class_tensor = tf.constant(class_descriptions)
   
       table = tf.contrib.lookup.index_to_string_table_from_tensor(class_tensor)
       classes = table.lookup(tf.to_int64(indices))
   
       # Restore variables from training checkpoint.
       variable_averages = tf.train.ExponentialMovingAverage(
           inception_model.MOVING_AVERAGE_DECAY)
       variables_to_restore = variable_averages.variables_to_restore()
       saver = tf.train.Saver(variables_to_restore)
       with tf.Session() as sess:
         # Restore variables from training checkpoints.
         ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_dir)
         if ckpt and ckpt.model_checkpoint_path:
           saver.restore(sess, ckpt.model_checkpoint_path)
           # Assuming model_checkpoint_path looks something like:
           #   /my-favorite-path/imagenet_train/model.ckpt-0,
           # extract global_step from it.
           global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
           print('Successfully loaded model from %s at step=%s.' %
                 (ckpt.model_checkpoint_path, global_step))
         else:
           print('No checkpoint file found at %s' % FLAGS.checkpoint_dir)
           return
   
         # Export inference model.
         output_path = os.path.join(
             tf.compat.as_bytes(FLAGS.output_dir),
             tf.compat.as_bytes(str(FLAGS.model_version)))
         print('Exporting trained model to', output_path)
         builder = tf.saved_model.builder.SavedModelBuilder(output_path)
   
         # Build the signature_def_map.
         classify_inputs_tensor_info = tf.saved_model.utils.build_tensor_info(
             serialized_tf_example)
         classes_output_tensor_info = tf.saved_model.utils.build_tensor_info(
             classes)
         scores_output_tensor_info = tf.saved_model.utils.build_tensor_info(values)
   
         classification_signature = (
             tf.saved_model.signature_def_utils.build_signature_def(
                 inputs={
                     tf.saved_model.signature_constants.CLASSIFY_INPUTS:
                         classify_inputs_tensor_info
                 },
                 outputs={
                     tf.saved_model.signature_constants.CLASSIFY_OUTPUT_CLASSES:
                         classes_output_tensor_info,
                     tf.saved_model.signature_constants.CLASSIFY_OUTPUT_SCORES:
                         scores_output_tensor_info
                 },
                 method_name=tf.saved_model.signature_constants.
                 CLASSIFY_METHOD_NAME))
   
         predict_inputs_tensor_info = tf.saved_model.utils.build_tensor_info(jpegs)
         prediction_signature = (
             tf.saved_model.signature_def_utils.build_signature_def(
                 inputs={'images': predict_inputs_tensor_info},
                 outputs={
                     'classes': classes_output_tensor_info,
                     'scores': scores_output_tensor_info
                 },
                 method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME
             ))
   
         legacy_init_op = tf.group(
             tf.tables_initializer(), name='legacy_init_op')
         builder.add_meta_graph_and_variables(
             sess, [tf.saved_model.tag_constants.SERVING],
             signature_def_map={
                 'predict_images':
                     prediction_signature,
                 tf.saved_model.signature_constants.
                 DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                     classification_signature,
             },
             legacy_init_op=legacy_init_op)
   
         builder.save()
         print('Successfully exported model to %s' % FLAGS.output_dir)
   
   

   tf.keras
   其实也是一样的，不过看着更简单。

   import tensorflow as tf
   	
   model_path = './checkpoint/model.h5'
   my_model = tf.keras.models.load_model(model_path)
   
   def preprocess_and_decode(img_str, new_shape=[224,224]):
       img = tf.io.decode_base64(img_str)
       img = tf.image.decode_jpeg(img, channels=3)
       img = tf.image.resize_images(img, new_shape, method=tf.image.ResizeMethod.BILINEAR, align_corners=False)
       # if you need to squeeze your input range to [0,1] or [-1,1] do it here
       return img
   InputLayer = tf.keras.Input(shape = (1,),dtype="string")
   OutputLayer = tf.keras.layers.Lambda(lambda img : tf.map_fn(lambda im : preprocess_and_decode(im[0]), img, dtype="float32"))(InputLayer)
   base64_model = tf.keras.Model(InputLayer,OutputLayer) 
   
   
   base64_input = base64_model.input
   final_output = my_model(base64_model.output)
   new_model = tf.keras.Model(base64_input,final_output)
   
   new_model.save('./base64_model/base64.h5') # # new_model.save('./base64_model/base64',save_format='tf')
   

   最后将模型导出为savedmodel格式重新上线。

结论

• 主要是为了保持结构上的完整。
• 但其实对预测结果的后处理（标签映射等）也是非常有必要的，掌控你能掌控的。
• 记一个待办：这里的 tf.keras 部分缺少后处理，添加后就可以利用 new_model.save('./base64_model/base64',save_format='tf') 一步到位直接保存成 savedmodel 格式进行部署了。
  (其实我在上线的时候用了已有的 h5->pb, pb->savedmodel 的代码，在转savedmode的过程中做了与第一个例子类似的后处理，