# NetVLAD
        x = tf.nn.l2_normalize(x, dim=-1)
        x = layers.netVLAD(x, 64)

import tensorflow as tf

import os
import cv2
from pathlib import Path
import h5py
from tqdm import tqdm
import numpy as np

import netvlad_tf.nets as nets

import tensorflow as tf
from tensorflow.python.saved_model import tag_constants
tf.contrib.resampler

model_dir = "./checkpoints/vd16_pitts30k_conv5_3_vlad_preL2_intra_white"
output_frozen_path = "./output/hf_frozen.pb"


resize_max=1024

tf.reset_default_graph()

image_batch = tf.placeholder(
    dtype=tf.float32, shape=[1, 1024, 1024, 3])
#tf.identity(image_batch, name="prInput")
print("test image_batch")
print(image_batch.name)
net_out = nets.vgg16NetvladPca(image_batch)

saver = tf.train.Saver()
sess = tf.Session()
ckpt = nets.defaultCheckpoint()
saver.restore(sess, ckpt)

path = "netvlad_tf_open-master/demo/query/query1.jpg"

img = cv2.imread(str(path))
h, w = img.shape[:2]
if max(h, w) > resize_max:
    scale = resize_max / max(h, w)
    h_new, w_new = int(round(h*scale)), int(round(w*scale))
    img = cv2.resize(img, (w_new, h_new), interpolation=cv2.INTER_LINEAR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

desc = sess.run(net_out, feed_dict={image_batch: img[None]})[0]

graph = tf.get_default_graph()

#frozen_graph = tf.graph_util.convert_variables_to_constants(
#    sess, graph.as_graph_def(),
#    output_node_names=['Placeholder', 'vgg16_netvlad_pca/l2_normalize_1']
#)
#
#frozen_graph = tf.graph_util.remove_training_nodes(frozen_graph, ['Placeholder', 'vgg16_netvlad_pca/l2_normalize_1'])
frozen_graph = tf.graph_util.convert_variables_to_constants(
    sess, graph.as_graph_def(), output_node_names=['vgg16_netvlad_pca/l2_normalize_1'])
with tf.gfile.GFile(output_frozen_path, mode='wb') as f:
    f.write(frozen_graph.SerializeToString())

summary_write = tf.summary.FileWriter("./output/log" , graph)

print("run success")

import tensorflow as tf
import numpy as np
from tensorflow.python.saved_model import tag_constants
tf.contrib.resampler
import tf2onnx

frozen_model = "./output/hf_frozen.pb"

with tf.Session() as sess:
    with tf.gfile.GFile(frozen_model, mode="rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        sess.graph.as_default()
        tf.import_graph_def(graph_def, name='')

        #onnx_graph = tf2onnx.tfonnx.process_tf_graph(sess.graph, opset = 7, input_names=["Placeholder:0"], inputs_as_nchw=["Placeholder:0"],  output_names=["vgg16_netvlad_pca/l2_normalize_1:0"])
        onnx_graph = tf2onnx.tfonnx.process_tf_graph(sess.graph, opset = 13,  input_names=["Placeholder:0"],  output_names=["vgg16_netvlad_pca/l2_normalize_1:0"])
        model_proto = onnx_graph.make_model("test")
        with open("path to onnx", "wb") as f:
            f.write(model_proto.SerializeToString())
        print("done")

import numpy as np
import onnxruntime as ort

import os
import cv2
from pathlib import Path
import h5py
from tqdm import tqdm
import numpy as np

#img = np.load("./assets/image.npz").reshape([1, 784])  
sess_ort = ort.InferenceSession("path to onnx")


path = "netvlad_tf_open-master/demo/query/query1.jpg"

resize_max=1024

img = cv2.imread(str(path))
h, w = img.shape[:2]
if max(h, w) > resize_max:
    scale = resize_max / max(h, w)
    h_new, w_new = int(round(h*scale)), int(round(w*scale))
    img = cv2.resize(img, (w_new, h_new), interpolation=cv2.INTER_LINEAR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

#img = img.astype(np.float32)/255

data = np.zeros((1, 1024, 1024, 3), np.float32) 
data[0, :, :, :] = img

#data = data.transpose(0,3,1,2) # Batch, Chanel, Height, Width

res = sess_ort.run(output_names=["vgg16_netvlad_pca/l2_normalize_1:0"], input_feed={"Placeholder:0": data})
print(res)
print("done")