Update app.py

app.py

@@ -1,11 +1,10 @@
 import numpy as np
 import gradio as gr
+import imageio.v2 as imageio
 import onnxruntime as ort
 from matplotlib import pyplot as plt
 from huggingface_hub import hf_hub_download
 
-model = hf_hub_download(repo_id="matjesg/cFOS_in_HC", filename="ensemble.onnx")
-
 def create_model_for_provider(model_path, provider="CPUExecutionProvider"): 
     options = ort.SessionOptions()
     options.intra_op_num_threads = 1
@@ -14,26 +13,29 @@ def create_model_for_provider(model_path, provider="CPUExecutionProvider"):
     session.disable_fallback()
     return session
 
-ort_session = create_model_for_provider(model)
-
-def inference(img):
-    
-    img = img[...,:1]/255
+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
-
-  
+    return ort_outs[0]*255, ort_outs[2]/0.25
+    
 title="deepflash2"
-description="deepflash2 is a deep-learning pipeline for segmentation of ambiguous microscopic images."
-examples=[['cFOS_example.png']]
+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/cFOS_in_HC', 'ensemble.onnx', '0001_cFOS.png']]
 
 gr.Interface(inference,
-             gr.inputs.Image(type="numpy"),
-             gr.outputs.Image(),
+             [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