run_bioengine_model.py

import micropip
await micropip.install(['imjoy-rpc', 'kaibu-utils', 'pyodide-http', 'requests', 'numpy'])
# Patch requests
import pyodide_http
pyodide_http.patch_all()  # Patch all libraries

import numpy as np
import matplotlib.pyplot as plt
from imjoy_rpc.hypha import connect_to_server
from kaibu_utils import fetch_image

SERVER_URL = "https://ai.imjoy.io"
image = await fetch_image('https://zenodo.org/api/records/6647674/files/sample_input_0.tif/content')


async def run_model(model_id, input_image, kwargs, server_url=SERVER_URL):
    server = await connect_to_server(
        {"server_url": server_url}
    )
    triton = await server.get_service("triton-client")
    if model_id == "cellpose":
        # Run cellpose model
        ret = await triton.execute(
            inputs=[input_image, kwargs],
            model_name="cellpose-python",
            decode_json=True,
        )
        return ret["mask"]
    else:
        # Run inference
        ret = await triton.execute(
            inputs=[{"inputs": [input_image], "model_id": model_id}],
            model_name="bioengine-model-runner",
            serialization="imjoy",
        )
        result = ret["result"]
        mask = result['outputs'][0]
        return mask

mask = await run_model("cellpose", image[None, ...].astype(np.float32), {"diameter": 10})
# display the output
fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12, 4))
ax1.imshow(image)
ax1.set_title('input image')
ax2.imshow(mask[0])
ax2.set_title('cellpose predicted mask')

mask = await run_model("affable-shark", image[None, None, ...], {})

ax3.imshow(mask[0][0])
ax3.set_title('affable-shark predicted mask')

plt.show()