# -*- coding: utf-8 -*-

from __future__ import absolute_import, print_function, division

import os, sys
import tensorflow as tf
import tf_slim as slim
from tensorflow.python.tools import freeze_graph

sys.path.append('../../')
from data.io.image_preprocess import short_side_resize_for_inference_data
from libs.configs import cfgs
from libs.networks import build_whole_network

CKPT_PATH = '/home/yjr/PycharmProjects/Faster-RCNN_Tensorflow/output/trained_weights/FasterRCNN_20180517/voc_200000model.ckpt'
OUT_DIR = '../../output/Pbs'
PB_NAME = 'FasterRCNN_Res101_Pascal.pb'


def build_detection_graph():
    # 1. preprocess img
    img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3],
                              name='input_img')  # is RGB. not GBR
    raw_shape = tf.shape(img_plac)
    raw_h, raw_w = tf.to_float(raw_shape[0]), tf.to_float(raw_shape[1])

    img_batch = tf.cast(img_plac, tf.float32)
    img_batch = short_side_resize_for_inference_data(img_tensor=img_batch,
                                                     target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN,
                                                     length_limitation=cfgs.IMG_MAX_LENGTH)
    img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN)
    img_batch = tf.expand_dims(img_batch, axis=0)  # [1, None, None, 3]

    det_net = build_whole_network.DetectionNetwork(base_network_name=cfgs.NET_NAME,
                                                   is_training=False)

    detected_boxes, detection_scores, detection_category = det_net.build_whole_detection_network(
        input_img_batch=img_batch,
        gtboxes_batch=None)

    xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \
                             detected_boxes[:, 2], detected_boxes[:, 3]

    resized_shape = tf.shape(img_batch)
    resized_h, resized_w = tf.to_float(resized_shape[1]), tf.to_float(resized_shape[2])

    xmin = xmin * raw_w / resized_w
    xmax = xmax * raw_w / resized_w

    ymin = ymin * raw_h / resized_h
    ymax = ymax * raw_h / resized_h

    boxes = tf.transpose(tf.stack([xmin, ymin, xmax, ymax]))
    dets = tf.concat([tf.reshape(detection_category, [-1, 1]),
                     tf.reshape(detection_scores, [-1, 1]),
                     boxes], axis=1, name='DetResults')

    return dets


def export_frozenPB():

    tf.reset_default_graph()

    dets = build_detection_graph()

    saver = tf.train.Saver()

    with tf.Session() as sess:
        print("we have restred the weights from =====>>\n", CKPT_PATH)
        saver.restore(sess, CKPT_PATH)

        tf.train.write_graph(sess.graph_def, OUT_DIR, PB_NAME)
        freeze_graph.freeze_graph(input_graph=os.path.join(OUT_DIR, PB_NAME),
                                  input_saver='',
                                  input_binary=False,
                                  input_checkpoint=CKPT_PATH,
                                  output_node_names="DetResults",
                                  restore_op_name="save/restore_all",
                                  filename_tensor_name='save/Const:0',
                                  output_graph=os.path.join(OUT_DIR, PB_NAME.replace('.pb', '_Frozen.pb')),
                                  clear_devices=False,
                                  initializer_nodes='')

if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = ''
    export_frozenPB()

[((10, 0, 10, 25), 'sys.path.append', 'sys.path.append', ({(10, 16, 10, 24): '"""../../"""'}, {}), "('../../')", False, 'import os, sys\n'), ((22, 15, 23, 47), 'tensorflow.placeholder', 'tf.placeholder', (), '', True, 'import tensorflow as tf\n'), ((24, 16, 24, 34), 'tensorflow.shape', 'tf.shape', ({(24, 25, 24, 33): 'img_plac'}, {}), '(img_plac)', True, 'import tensorflow as tf\n'), ((27, 16, 27, 45), 'tensorflow.cast', 'tf.cast', ({(27, 24, 27, 32): 'img_plac', (27, 34, 27, 44): 'tf.float32'}, {}), '(img_plac, tf.float32)', True, 'import tensorflow as tf\n'), ((28, 16, 30, 91), 'data.io.image_preprocess.short_side_resize_for_inference_data', 'short_side_resize_for_inference_data', (), '', False, 'from data.io.image_preprocess import short_side_resize_for_inference_data\n'), ((32, 16, 32, 49), 'tensorflow.expand_dims', 'tf.expand_dims', (), '', True, 'import tensorflow as tf\n'), ((34, 14, 35, 69), 'libs.networks.build_whole_network.DetectionNetwork', 'build_whole_network.DetectionNetwork', (), '', False, 'from libs.networks import build_whole_network\n'), ((44, 20, 44, 39), 'tensorflow.shape', 'tf.shape', ({(44, 29, 44, 38): 'img_batch'}, {}), '(img_batch)', True, 'import tensorflow as tf\n'), ((63, 4, 63, 28), 'tensorflow.reset_default_graph', 'tf.reset_default_graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((67, 12, 67, 28), 'tensorflow.train.Saver', 'tf.train.Saver', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((25, 19, 25, 44), 'tensorflow.to_float', 'tf.to_float', ({(25, 31, 25, 43): 'raw_shape[0]'}, {}), '(raw_shape[0])', True, 'import tensorflow as tf\n'), ((25, 46, 25, 71), 'tensorflow.to_float', 'tf.to_float', ({(25, 58, 25, 70): 'raw_shape[1]'}, {}), '(raw_shape[1])', True, 'import tensorflow as tf\n'), ((31, 28, 31, 56), 'tensorflow.constant', 'tf.constant', ({(31, 40, 31, 55): 'cfgs.PIXEL_MEAN'}, {}), '(cfgs.PIXEL_MEAN)', True, 'import tensorflow as tf\n'), ((45, 27, 45, 56), 'tensorflow.to_float', 'tf.to_float', ({(45, 39, 45, 55): 'resized_shape[1]'}, {}), '(resized_shape[1])', True, 'import tensorflow as tf\n'), ((45, 58, 45, 87), 'tensorflow.to_float', 'tf.to_float', ({(45, 70, 45, 86): 'resized_shape[2]'}, {}), '(resized_shape[2])', True, 'import tensorflow as tf\n'), ((53, 25, 53, 59), 'tensorflow.stack', 'tf.stack', ({(53, 34, 53, 58): '[xmin, ymin, xmax, ymax]'}, {}), '([xmin, ymin, xmax, ymax])', True, 'import tensorflow as tf\n'), ((69, 9, 69, 21), 'tensorflow.Session', 'tf.Session', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((73, 8, 73, 62), 'tensorflow.train.write_graph', 'tf.train.write_graph', ({(73, 29, 73, 43): 'sess.graph_def', (73, 45, 73, 52): 'OUT_DIR', (73, 54, 73, 61): 'PB_NAME'}, {}), '(sess.graph_def, OUT_DIR, PB_NAME)', True, 'import tensorflow as tf\n'), ((54, 22, 54, 61), 'tensorflow.reshape', 'tf.reshape', ({(54, 33, 54, 51): 'detection_category', (54, 53, 54, 60): '[-1, 1]'}, {}), '(detection_category, [-1, 1])', True, 'import tensorflow as tf\n'), ((55, 21, 55, 58), 'tensorflow.reshape', 'tf.reshape', ({(55, 32, 55, 48): 'detection_scores', (55, 50, 55, 57): '[-1, 1]'}, {}), '(detection_scores, [-1, 1])', True, 'import tensorflow as tf\n'), ((74, 46, 74, 76), 'os.path.join', 'os.path.join', ({(74, 59, 74, 66): 'OUT_DIR', (74, 68, 74, 75): 'PB_NAME'}, {}), '(OUT_DIR, PB_NAME)', False, 'import os, sys\n')]