import gradio as gr
import os
import platform
from helper import CoreMLPipeline

force_tf = os.environ.get('FORCE_TF', False)
auth_key = os.environ.get('HF_TOKEN', True)

config = { "coreml_extractor_repoid":"crossprism/efficientnetv2-21k-fv-m",
           "coreml_extractor_path":"efficientnetV2M21kExtractor.mlmodel",
           "tf_extractor_repoid":"crossprism/efficientnetv2-21k-fv-m-tf",
           "tf_extractor_path":"efficientnetv2-21k-fv-m",
           "coreml_classifier_repoid":"crossprism/travel_sa_landmarks",
           "coreml_classifier_path":"LandmarksSAHead_quant8.mlpackage/Data/com.apple.CoreML/efficientnetV2M21kSALandmarksHead_quant8.mlmodel",
           "activation":"softmax"
           }
use_tf = force_tf or (platform.system() != 'Darwin')

helper = CoreMLPipeline(config, auth_key, use_tf)

def classify_image(image):
    resized = image.resize((480,480))
    return helper.classify(resized)

image = gr.Image(type='pil')
label = gr.Label(num_top_classes=3)

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