import numpy as np
import gradio as gr
import onnxruntime as ort
from matplotlib import pyplot as plt
from huggingface_hub import hf_hub_download

def create_model_for_provider(model_path, provider="CPUExecutionProvider"): 
    options = ort.SessionOptions()
    options.intra_op_num_threads = 1
    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    session = ort.InferenceSession(str(model_path), options, providers=[provider])
    session.disable_fallback()
    return session

def inference(repo_id, model_name, img):
    model = hf_hub_download(repo_id=repo_id, filename=model_name)
    ort_session = create_model_for_provider(model)
    n_channels = ort_session.get_inputs()[0].shape[-1]
    
    img = img[...,:n_channels]/255
    ort_inputs = {ort_session.get_inputs()[0].name: img.astype(np.float32)}

    ort_outs = ort_session.run(None, ort_inputs)
    
    return ort_outs[0]*255, ort_outs[2]/0.25
    
title="deepflash2"
description='deepflash2 is a deep-learning pipeline for the segmentation of ambiguous microscopic images.\n deepflash2 uses deep model ensembles to achieve more accurate and reliable results. Thus, inference time will be more than a minute in this space.'
examples=[['matjesg/deepflash2_demo', 'cFOS_ensemble.onnx', 'cFOS_example.png'],
          ['matjesg/deepflash2_demo', 'YFP_ensemble.onnx', 'YFP_example.png']
          ]

gr.Interface(inference,
             [gr.inputs.Textbox(placeholder='e.g., matjesg/cFOS_in_HC', label='repo_id'), 
              gr.inputs.Textbox(placeholder='e.g., ensemble.onnx', label='model_name'),
              gr.inputs.Image(type='numpy', label='Input image')
             ],
             [gr.outputs.Image(label='Segmentation Mask'),
              gr.outputs.Image(label='Uncertainty Map')],
             title=title,
             description=description,
             examples=examples,
             ).launch()