from time import perf_counter
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

# import face_alignment
# import face_alignment
# import onnx
# from onnx import optimizer
import onnx
import onnxruntime
from onnxruntime import InferenceSession
# import numpy as np
from tqdm import tqdm

# os.environ['FACEALIGNMENT_USERDIR'] = 'models'
# import torch
# import cv2
# import face_alignment
from test import *

print(cv2.__version__)

img = cv2.imread(r"C:\Users\maxbe\Desktop\selfie_hd.jpg")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
face_box = [
    160,
    300,
    550,
    850,
]
# model = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, device='cpu', flip_input=False, verbose=True)
# # warmup
# model.get_landmarks_from_image(img, [face_box])
#
# # ltrb
# times = []
# for _ in tqdm(range(100)):
#     a = perf_counter()
#     landmarks = model.get_landmarks_from_image(img, [face_box])
#     b = perf_counter()
#     times.append(b - a)
# #
# net = model.face_alignment_net
#
# input_names = ["actual_input_1"]
# output_names = ["output1"]
# dummy_input = torch.randn(1, 3, 256, 256, device='cpu')
# torch.onnx.export(net.cpu().eval(), dummy_input, "FAN.onnx", verbose=True, opset_version=11, export_params=True,
#                   keep_initializers_as_inputs=True)
#
# #
# # dnn = cv2.dnn.readNetFromONNX('FAN.onnx')
# #
# onnx_graph = onnx.load('FAN.onnx')
#
# # Check the model
# onnx.checker.check_model(onnx_graph)
# print('The model is checked!')
# passes = """eliminate_deadend
#         eliminate_identity
#         eliminate_nop_dropout
#         eliminate_nop_monotone_argmax
#         eliminate_nop_pad
#         eliminate_nop_transpose
#         eliminate_unused_initializer
#         extract_constant_to_initializer
#         fuse_add_bias_into_conv
#         fuse_bn_into_conv
#         fuse_consecutive_concats
#         fuse_consecutive_log_softmax
#         fuse_consecutive_reduce_unsqueeze
#         fuse_consecutive_squeezes
#         fuse_consecutive_transposes
#         fuse_matmul_add_bias_into_gemm
#         fuse_pad_into_conv
#         fuse_transpose_into_gemm""".split('\n')
#
# passes = [x.strip() for x in passes]
# optimized_model = optimizer.optimize(
#     onnx_graph,
#     passes=passes
# )
# onnx.save(optimized_model, 'FAN_optimized.onnx')
sessionOptions = onnxruntime.SessionOptions()
# sessionOptions.log_severity_level = 0
sessionOptions.inter_op_num_threads = 8
sessionOptions.intra_op_num_threads = 1
sessionOptions.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
sessionOptions.enable_profiling = True
# sessionOptions.optimized_model_filepath = 'FAN_opt_graph.onnx'
session = InferenceSession('FAN.onnx', sess_options=sessionOptions)

# get the name of the first input of the model
input_name = session.get_inputs()[0].name

times_onnx = []
for _ in range(100):
    a = perf_counter()
    tensor, center, scale = [face[0] for face in preprocess_image(img, [face_box])]  # extracted preprocessing
    b = perf_counter()
    print("PRE: ", b - a)
    a = perf_counter()
    outputs = session.run([], {input_name: tensor})[0]
    b = perf_counter()
    print("INF: ", b - a)
    a = perf_counter()
    landmarks_onnx = postprocess_output(torch.tensor(outputs), center, scale)
    b = perf_counter()
    print("POST:", b - a)
    print()

# print(f"ONNX: {np.mean(times_onnx):.2f}s avg.")
session.end_profiling()
#
# onnx_graph = onnx.load('FAN.onnx')
#
# # Check the model
# onnx.checker.check_model(onnx_graph)
# print('The model is checked!')
# passes = """eliminate_deadend
#         eliminate_identity
#         eliminate_nop_dropout
#         eliminate_nop_monotone_argmax
#         eliminate_nop_pad
#         eliminate_nop_transpose
#         eliminate_unused_initializer
#         extract_constant_to_initializer
#         fuse_add_bias_into_conv
#         fuse_bn_into_conv
#         fuse_consecutive_concats
#         fuse_consecutive_log_softmax
#         fuse_consecutive_reduce_unsqueeze
#         fuse_consecutive_squeezes
#         fuse_consecutive_transposes
#         fuse_matmul_add_bias_into_gemm
#         fuse_pad_into_conv
#         fuse_transpose_into_gemm""".split('\n')
#
# passes = [x.strip() for x in passes]
# optimized_model = optimizer.optimize(
#     onnx_graph,
#     passes=passes
# )
# onnx.save(optimized_model, 'FAN_optimized.onnx')
#
# session = InferenceSession('FAN_optimized.onnx')
#
# # get the name of the first input of the model
# input_name = session.get_inputs()[0].name
#
# times_onnx_optimized = []
# for _ in tqdm(range(100)):
#     a = perf_counter()
#     tensor, center, scale = [face[0] for face in preprocess_image(img, [face_box])]  # extracted preprocessing
#     outputs = session.run([], {input_name: tensor})[0]
#     landmarks_onnx = postprocess_output(torch.tensor(outputs), center, scale)
#     b = perf_counter()
#     times_onnx_optimized.append(b - a)
# # print(f"PYTORCH: {np.mean(times):.2f}s avg.")
# print(f"ONNX: {np.mean(times_onnx):.2f}s avg.")
# print(f"ONNX (OPTIMIZED): {np.mean(times_onnx_optimized):.2f}s avg.")
# print('Input Name:', input_name)
# a = perf_counter()
# tensor = model.preprocess_image(img, [face_box])  # extracted preprocessing
# outputs = session.run([], {input_name: tensor})[0]
# b = perf_counter()
# print(b - a)
#
# a = perf_counter()
# tensor = model.preprocess_image(img, [face_box])[0]  # extracted preprocessing
# outputs = session.run([], {input_name: tensor})[0]
# b = perf_counter()
# print(b - a)