import tensorflow as tf
import gradio as gr
import coremltools as ct
import numpy as np
import requests
import huggingface_hub as hf
from huggingface_hub import hf_hub_download
from huggingface_hub import snapshot_download
from huggingface_hub import login
import os
import PIL

#login()


read_key = os.environ.get('HF_TOKEN', True)
#extractor2_path = hf_hub_download(repo_id="crossprism/efficientnetv221k-M", filename="efficientnetV2M21kExtractor.mlmodel", use_auth_token  = read_key)
extractor_path = snapshot_download(repo_id="crossprism/efficientnetv2-21k-fv-m", use_auth_token  = read_key)
classifier_path = hf_hub_download(repo_id="crossprism/tesla_sentry_dings", filename="tesla_sentry_door_ding.mlpackage/Data/com.apple.CoreML/tesla_door_dings.mlmodel", use_auth_token = read_key)


print(f"Loading extractor...{extractor_path}")
extractor = tf.saved_model.load(extractor_path+"/efficientnetv2-21k-fv-m")

#extractor2 = ct.models.MLModel(extractor2_path)
print(f"Loading classifier...{classifier_path}")
classifier = ct.models.MLModel(classifier_path)
spec = classifier.get_spec()
labels = spec.neuralNetworkClassifier.stringClassLabels.vector
image = PIL.Image.open('test.jpg')

def makeKerasModel(labels, coreml_classifier):
    input = tf.keras.Input(shape = (1280))
    x = tf.keras.layers.Dense(len(labels), activation = "sigmoid")(input)
    model = tf.keras.Model(input,x, trainable = False)
    weights = np.array(coreml_classifier.layers[0].innerProduct.weights.floatValue)
    weights = weights.reshape((len(labels),1280))
    #weights = weights.reshape((1280,len(labels)))
    weights = weights.T
    bias = np.array(coreml_classifier.layers[0].innerProduct.bias.floatValue)    
    model.set_weights([weights,bias])    
    return model

#Only MacOS can run inference on CoreML models. Convert it to tensorflow to match the tf feature extractor
tf_classifier = makeKerasModel(labels, spec.neuralNetworkClassifier)

def classify_image(image):
    resized = image.resize((480,480))
    image = tf.image.convert_image_dtype(resized, tf.float32)
    image = tf.reshape(image, [1,480,480,3])
    features = extractor.signatures['serving_default'](image)
    #features2 = extractor2.predict({"image":resized})
    #print(features)
    #print(features2)
    #features2 = features2["Identity"]
    #isDing = classifier.predict({"features":features2[0]})
    #isDing = isDing["Identity"]
    input = {"input_1":features["output_1"]}
    p = tf_classifier.predict(input)
    #print(p)
    isDing = {}
    for i,label in enumerate(labels):
        isDing[label] = p[i]        
    #print(isDing)
    return {'ding': isDing["ding"]}
#classify_image(image)
image = gr.Image(type='pil')
label = gr.Label(num_top_classes=3)

gr.Interface(fn=classify_image, inputs=image, outputs=label, examples = [["test.jpg"]]).launch()