import tensorflow as tf
from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, LearningRateScheduler
from tensorflow.keras.optimizers import SGD, Adam
import augmentation
from ssd_loss import CustomLoss
from utils import bbox_utils, data_utils, io_utils, train_utils
import os


def create_mobilenet_v2():
    from tensorflow.keras.applications.mobilenet_v2 import MobileNetV2
    img_size = 224
    return MobileNetV2(include_top=False, input_shape=(img_size, img_size, 3))


def create_ssd_model():
    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
    from models.ssd_mobilenet_v2 import get_model, init_model

    load_weights = False
    backbone = 'mobilenet_v2'

    # Create SSD Model
    hyper_params = train_utils.get_hyper_params(backbone)
    hyper_params["total_labels"] = 1000
    ssd_model = get_model(hyper_params)
    ssd_custom_losses = CustomLoss(hyper_params["neg_pos_ratio"], hyper_params["loc_loss_alpha"])
    ssd_model.compile(optimizer=Adam(learning_rate=1e-3),
                    loss=[ssd_custom_losses.loc_loss_fn, ssd_custom_losses.conf_loss_fn])
    init_model(ssd_model)
    ssd_model_path = io_utils.get_model_path(backbone)
    if load_weights:
        ssd_model.load_weights(ssd_model_path)
    return ssd_model

# https://www.cnblogs.com/happyamyhope/p/11822111.html

def convert_keras_to_tflite(keras_model, tflite_model_name):
    converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
    tflite_model = converter.convert()
    with open(tflite_model_name, 'wb') as f:
        f.write(tflite_model)
    print('Convert done!')


if __name__ == '__main__':
    tf_model = create_ssd_model()
    convert_keras_to_tflite(tf_model, tflite_model_name='./tflite_model/ssd300_mobilenet_v2.tflite')