tensorflow-yolov3/freeze_yolo_tf.py at master · joiskash/tensorflow-yolov3 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
# -*- coding: utf-8 -*-
"""Freeze tf-yolo model to use for inference. Change ckpt_file variable and model_export_path as needed

Example:

        $ python freeze_yolo_tf.py

"""

import tensorflow as tf
from core.yolov3 import YOLOV3
import tempfile
import os, shutil


ckpt_file = "C:/Users/Kashyap/bkp/source/repos/rdt-reader/Model_KH_EXP/yolov3rot_test_loss=0.7907.ckpt-438" # Model to freeze
output_node_names = ["input/input_data", "pred_sbbox/concat_2", "pred_mbbox/concat_2", "pred_lbbox/concat_2"]
model_export_path = "C:/Users/Kashyap/bkp/source/repos/rdt-reader/tensorflow-yolov3-models/models/Flu_audere" # export directory
version = 3
export_path = os.path.join(model_export_path, str(version))


if __name__ == "__main__":

    with tf.name_scope('input'):
        input_data = tf.placeholder(dtype=tf.float32, name='input_data',shape=(1,512,512,3))

    model = YOLOV3(input_data, trainable=False)

    with tf.Session() as sess:


        saver = tf.train.Saver()
        saver.restore(sess, ckpt_file)

        trainable = tf.placeholder(tf.bool, name='trainable')

        pred_sbbox, pred_mbbox, pred_lbbox = model.pred_sbbox, model.pred_mbbox, model.pred_lbbox

        builder = tf.saved_model.builder.SavedModelBuilder(export_path)
        classification_inputs = tf.saved_model.utils.build_tensor_info(input_data)
        classification_train = tf.saved_model.utils.build_tensor_info(trainable)

        classification_pred_sbbox = tf.saved_model.utils.build_tensor_info(pred_sbbox)
        classification_pred_mbbox = tf.saved_model.utils.build_tensor_info(pred_mbbox)
        classification_pred_lbbox = tf.saved_model.utils.build_tensor_info(pred_lbbox)


        classification_signature = (
            tf.saved_model.signature_def_utils.build_signature_def(
                inputs={"input":classification_inputs},

                outputs={

                    "pred_sbbox":
                        classification_pred_sbbox,

                    "pred_mbbox":

                        classification_pred_mbbox,

                    "pred_lbbox":

                        classification_pred_lbbox


                },

                method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME))

        builder.add_meta_graph_and_variables(

            sess, [tf.saved_model.tag_constants.SERVING],

            signature_def_map={

                tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:

                    classification_signature,

            },

            main_op=tf.tables_initializer(),

            strip_default_attrs=True,clear_devices=True)


        builder.save()