def export_saved_model(self, saved_model_dir, batch_size=None, pre_mode='infer', post_mode='global'): """Saves the model to Tensorflow SavedModel. Args: saved_model_dir: Folder path for saved model. batch_size: Batch size to be saved in saved_model. pre_mode: Pre-processing Mode in ExportModel, must be {None, 'infer'}. post_mode: Post-processing Mode in ExportModel, must be {None, 'global', 'per_class'}. """ # Create EfficientDetModel with latest checkpoint. config = self.config tf.keras.backend.clear_session() model = efficientdet_keras.EfficientDetModel(config=config) model.build((batch_size, *config.image_size, 3)) if config.model_dir: util_keras.restore_ckpt(model, config.model_dir, config['moving_average_decay'], skip_mismatch=False) else: # EfficientDetModel is random initialized without restoring the # checkpoint. This is mainly used in object_detector_test and shouldn't be # used if we want to export trained model. tf.compat.v1.logging.warn('Need to restore the checkpoint for ' 'EfficientDet.') # Gets tf.TensorSpec. if pre_mode is None: # Input is the preprocessed image that's already resized to a certain # input shape. input_spec = tf.TensorSpec( shape=[batch_size, *config.image_size, 3], dtype=tf.float32, name='images') else: # Input is that raw image that can be in any input shape, input_spec = tf.TensorSpec(shape=[batch_size, None, None, 3], dtype=tf.uint8, name='images') export_model = inference.ExportModel(model, pre_mode=pre_mode, post_mode=post_mode) tf.saved_model.save( export_model, saved_model_dir, signatures=export_model.__call__.get_concrete_function(input_spec))
def build(self, params_override=None): """Build model and restore checkpoints.""" params = copy.deepcopy(self.params) if params_override: params.update(params_override) config = hparams_config.get_efficientdet_config(self.model_name) config.override(params) if self.only_network: self.model = efficientdet_keras.EfficientDetNet(config=config) else: self.model = efficientdet_keras.EfficientDetModel(config=config) image_size = utils.parse_image_size(params['image_size']) self.model.build((self.batch_size, *image_size, 3)) util_keras.restore_ckpt(self.model, self.ckpt_path, self.params['moving_average_decay'], skip_mismatch=False)
def main(_): # pylint: disable=line-too-long # Prepare images and checkpoints: please run these commands in shell. # !mkdir tmp # !wget https://user-images.githubusercontent.com/11736571/77320690-099af300-6d37-11ea-9d86-24f14dc2d540.png -O tmp/img.png # !wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientdet/coco/efficientdet-d0.tar.gz -O tmp/efficientdet-d0.tar.gz # !tar zxf tmp/efficientdet-d0.tar.gz -C tmp imgs = [np.array(Image.open(FLAGS.image_path))] # Create model config. config = hparams_config.get_efficientdet_config(FLAGS.model_name) config.is_training_bn = False config.image_size = '1920x1280' config.nms_configs.score_thresh = 0.4 config.nms_configs.max_output_size = 100 config.override(FLAGS.hparams) # Use 'mixed_float16' if running on GPUs. policy = tf.keras.mixed_precision.Policy('float32') tf.keras.mixed_precision.set_global_policy(policy) tf.config.run_functions_eagerly(FLAGS.debug) # Create and run the model. model = efficientdet_keras.EfficientDetModel(config=config) model.build((None, None, None, 3)) model.load_weights(tf.train.latest_checkpoint(FLAGS.model_dir)) model.summary() class ExportModel(tf.Module): def __init__(self, model): super().__init__() self.model = model @tf.function def f(self, imgs): return self.model(imgs, training=False, post_mode='global') imgs = tf.convert_to_tensor(imgs, dtype=tf.uint8) export_model = ExportModel(model) if FLAGS.saved_model_dir: tf.saved_model.save( export_model, FLAGS.saved_model_dir, signatures=export_model.f.get_concrete_function( tf.TensorSpec(shape=(None, None, None, 3), dtype=tf.uint8))) export_model = tf.saved_model.load(FLAGS.saved_model_dir) boxes, scores, classes, valid_len = export_model.f(imgs) # Visualize results. for i, img in enumerate(imgs): length = valid_len[i] img = inference.visualize_image( img, boxes[i].numpy()[:length], classes[i].numpy().astype(np.int)[:length], scores[i].numpy()[:length], label_map=config.label_map, min_score_thresh=config.nms_configs.score_thresh, max_boxes_to_draw=config.nms_configs.max_output_size) output_image_path = os.path.join(FLAGS.output_dir, str(i) + '.jpg') Image.fromarray(img).save(output_image_path) print('writing annotated image to %s' % output_image_path)