import streamlit as st
import numpy as np
from PIL import Image
from cellpose import models, io, plot

def inference(image, model_path, **model_params):
    img = image

    model_inference = models.CellposeModel(gpu=False, pretrained_model=model_path)
    preds, flows, _ = model_inference.eval([img], **model_params)
    return preds, flows


if __name__ == "__main__":

    st.title("Sartorius Cell Segmentation")

    uploaded_img = st.file_uploader(label="Upload neuronal cell image")
    if uploaded_img is not None:
        st.image(uploaded_img)
    segment = st.button("Perform segmentation")
    if uploaded_img is not None and segment:
        img = Image.open(uploaded_img)
        img = np.array(img)

        model_params = {
            "diameter": 19.0,
            "channels": [0, 0],
            "augment": True,
            "resample": True,
        }
        with st.spinner("Performing segmentation. This might take a while..."):
            preds, flows = inference(
                image=img,
                model_path="./cellpose_residual_on_style_on_concatenation_off_fold1_ep_649_cv_0.2834",
                **model_params
            )

        st.image(preds[0])