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 login
import os
import PIL

#login()

read_key = os.environ.get('HF_TOKEN', True)
extractor_path = hf_hub_download(repo_id="crossprism/efficientnetv221k-M", filename="efficientnetV2M21kExtractor.mlmodel", 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 = ct.models.MLModel(extractor_path)
print(f"Loading classifier...{classifier_path}")
classifier = ct.models.MLModel(classifier_path)

def classify_image(image):
    image = image.resize((480,480))
    features = extractor.predict({"image":image})
    print(features)
    features = features["Identity"]
    isDing = classifier.predict({"features":features[0]})
    print(isDing)
    isDing = isDing["Identity"]
    return {'ding': isDing["ding"]}

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()