Update app.py

app.py

@@ -5,20 +5,16 @@ from PIL import Image
 from cellpose import models
 
 
-@st.cache()
-def load_model(model_path):
+def inference(img, model_path, **model_params):
     inf_model = models.CellposeModel(gpu=False, pretrained_model=model_path)
-    return inf_model
+    preds, flows, _ = inf_model.eval([img], **model_params)
+    return preds, flows
 
 
 if __name__ == "__main__":
 
     st.title("Sartorius Neuronal Cell Segmentation")
 
-    inf_model = load_model(
-        model_path="./cellpose_residual_on_style_on_concatenation_off_fold1_ep_649_cv_0.2834"
-    )
-
     uploaded_img = st.file_uploader(label="Upload neuronal cell image")
 
     with st.expander("View input image"):
@@ -40,14 +36,18 @@ if __name__ == "__main__":
             "resample": True,
         }
         with st.spinner("Performing segmentation. This might take a while..."):
-            preds, flows, _ = inf_model.eval([img], **model_params)
+            preds, flows = inference(
+                img=img,
+                model_path="./cellpose_residual_on_style_on_concatenation_off_fold1_ep_649_cv_0.2834",
+                **model_params
+            )
 
         fig, (ax1, ax2, ax3) = plt.subplots(1, 3)
         ax1.axis("off")
         ax2.axis("off")
         ax3.axis("off")
         ax1.set_title("Original Image")
-        ax1.imshow(img)
+        ax1.imshow(img, cmap="gray")
         ax2.set_title("Segmented image")
         ax2.imshow(preds[0])
         ax3.set_title("Image flows")