h5_pb.py

#-*- coding:utf-8 -*-
import os
import numpy as np
import tensorflow as tf
from keras.models import Model
from keras import backend as K
from model import get_srresnet_model
from tensorflow.python.platform import gfile
import cv2
import pdb

def h5_to_pb(h5_model, output_dir, model_name, out_prefix = "output_", log_tensorboard = False):
    K.set_learning_phase(0)
    out_nodes = []

    for i in range(len(h5_model.outputs)):
        out_nodes.append(out_prefix + str(i + 1))
        tf.identity(h5_model.outputs[i], out_prefix + str(i + 1))
    sess = K.get_session()

    from tensorflow.python.framework import graph_util,graph_io

    init_graph = sess.graph.as_graph_def()
    main_graph = graph_util.convert_variables_to_constants(sess, init_graph, out_nodes)
    graph_io.write_graph(main_graph, output_dir, name = model_name, as_text = False)

    if log_tensorboard:
        from tensorflow.python.tools import import_pb_to_tensorboard
        import_pb_to_tensorboard.import_to_tensorboard(osp.join(output_dir, model_name), output_dir)

if __name__ == '__main__':
    #hdf5 model convert to pb model
    model = get_srresnet_model()
    model_path = 'checkpoints/weights.006-8.494-30.57169-0.98235.hdf5'
    model.load_weights(model_path)
    model.summary()
    print('load model over!')
    basemodel = Model(inputs = model.input, outputs = model.output)
    h5_to_pb(basemodel, output_dir='pb_model', model_name='model_restoration.pb')
    print('pb model has saved sucessfully!')