Update app.py

app.py

@@ -1,33 +1,27 @@
 import numpy as np
 import gradio as gr
-import onnxruntime as ort
-from matplotlib import pyplot as plt
+import torch
 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]
+    #model_path = hf_hub_download(repo_id=repo_id, filename=model_name)
+    model_path = 'trained_models/cFOS_in_HC/cFOS_in_HC_ensemble_1.pt'
+    model = torch.jit.load(model_path, map_location='cpu')
+    n_channels = len(model.norm.mean)
     
-    img = img[...,:n_channels]/255
-    ort_inputs = {ort_session.get_inputs()[0].name: img.astype(np.float32)}
+    # Remove redundant channels
+    img = img[...,:n_channels]
+    inp = torch.from_numpy(img).float()
 
-    ort_outs = ort_session.run(None, ort_inputs)
+    argmax, softmax, stdeviation = model(inp)
     
-    return ort_outs[0]*255, ort_outs[2]/0.25
+    return argmax*255, stdeviation
+ 
     
 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']
+examples=[['matjesg/deepflash2_demo', 'cFOS_in_HC_ensemble.pt', 'cFOS_example.png'],
+          ['matjesg/deepflash2_demo', 'YFP_in_CTX_ensemble.pt', 'YFP_example.png']
           ]
 
 gr.Interface(inference,