image_classifier/freeze_graph.py at tk · i-rtfsc/image_classifier · 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
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# -*- encoding: utf-8 -*-


import os
import shutil

from tensorflow.python.tools import freeze_graph

from base import env_utils, file_utils
from config.global_configs import ProjectConfig, TrainBaseConfig, TrainConfig, TFRecordBaseConfig, TFRecordConfig, \
    UserConfig


def freeze_session(model_dir='', frozen_out_dir='',
                   frozen_graph_filename='image_classifier',
                   output_tensor_name='Softmax',
                   gpu='0', meta_file=None):
    final_model_path = None
    for model in os.listdir(model_dir):
        if model.endswith('.pb'):
            final_model_path = model_dir

    if final_model_path is None:
        best_model_timestamp = sorted(os.listdir(model_dir))[-1]
        last_model_path = os.path.join(model_dir, best_model_timestamp)
        for model in os.listdir(last_model_path):
            if model.endswith('.pb'):
                final_model_path = last_model_path

    if final_model_path is None:
        print('error, model file was none')
        return

    file_utils.create_directory(frozen_out_dir)

    if not os.path.exists(meta_file):
        print('meta file was none, does not need copy')
    else:
        shutil.copy(meta_file, frozen_out_dir)

    output_path = os.path.join(frozen_out_dir, '{}.pb'.format(frozen_graph_filename))

    env_utils.select_gpu(gpu)

    freeze_graph.freeze_graph(input_graph=None,
                              input_saver=None,
                              input_binary=False,
                              input_checkpoint=None,
                              output_node_names=output_tensor_name,
                              restore_op_name=None,
                              filename_tensor_name=None,
                              output_graph=output_path,
                              clear_devices=False,
                              initializer_nodes=None,
                              input_saved_model_dir=final_model_path)


if __name__ == '__main__':
    project = None
    time = None
    gpu = '3'

    ProjectConfig.getDefault().update(project=project, time=time)
    UserConfig.getDefault().update()
    TFRecordConfig.getDefault().update(TFRecordBaseConfig.UPDATE_BASE)
    TFRecordConfig.getDefault().update(TFRecordBaseConfig.UPDATE_DATASET)
    TrainConfig.getDefault().update()

    freeze_session(model_dir=TrainConfig.getDefault().train_best_export_dir,
                   frozen_out_dir=TrainConfig.getDefault().model_freeze_dir,
                   frozen_graph_filename=TrainConfig.getDefault().project,
                   output_tensor_name=TrainBaseConfig.OUTPUT_TENSOR_NAME,
                   gpu=gpu,
                   meta_file=TFRecordConfig.getDefault().meta_file)