bioengine_efficientsam_demo.py

import micropip
await micropip.install(["pyotritonclient", "kaibu-utils", "pyodide-http", "requests"])

import pyodide_http
pyodide_http.patch_all()  # Patch all libraries

from kaibu_utils import fetch_image
import matplotlib.pyplot as plt
import numpy as np
from pyotritonclient import execute


image = await fetch_image("https://raw.githubusercontent.com/yformer/EfficientSAM/main/figs/examples/dogs.jpg")
input_image = image.transpose(2, 0, 1)[None].astype(np.float32) / 255.0
input_points = np.array([[[[580, 350], [650, 350]]]], dtype=np.float32)
# batch_size, num_queries, num_points
input_labels = np.array([[[1, 1]]], dtype=np.float32)
orig_im_size = np.array(input_image.shape[2:], dtype=np.int64)


async def run():
    ret = await execute(
        [input_image],
        server_url='http://127.0.0.1:9520/triton', # <-------------change the server url here
        model_name="efficientsam-encoder",
    )
    embeddings = ret["image_embeddings"]
    print(embeddings.shape)
    ret = await execute(
        [
            input_labels,
            orig_im_size,
            input_points,
            embeddings,
        ],
        server_url='http://127.0.0.1:9520/triton', # <-------------change the server url here
        model_name="efficientsam-decoder",
    )
    mask = ret['output_masks']
    print(mask.shape)
    _, axes = plt.subplots(1, 2)
    axes[0].imshow(image)
    axes[0].scatter(input_points[0, 0, :, 0], input_points[0, 0, :, 1], marker='*', c='r', s=100)
    axes[1].imshow(mask[0, 0, 0]>0)
    plt.show()

await run()