def _draw_inference(self, step): self.model.__class__ = efficientdet_keras.EfficientDetModel results = self.model(self.sample_image, training=False) boxes, scores, classes, valid_len = tf.nest.map_structure(np.array, results) length = valid_len[0] image = inference.visualize_image( self.sample_image[0], boxes[0][:length], classes[0].astype(np.int)[:length], scores[0][:length], label_map=self.model.config.label_map, min_score_thresh=self.min_score_thresh, max_boxes_to_draw=self.max_boxes_to_draw) with self.file_writer.as_default(): tf.summary.image('Test image', tf.expand_dims(image, axis=0), step=step) self.model.__class__ = EfficientDetNetTrain
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)