diff --git a/bioimageio_colab/onnx_runtime.py b/bioimageio_colab/onnx_runtime.py
new file mode 100644
index 0000000..12aaca7
--- /dev/null
+++ b/bioimageio_colab/onnx_runtime.py
@@ -0,0 +1,232 @@
+# Extra pip installations:
+# pip install opencv-python
+# pip install onnxruntime
+
+#%% Imports
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import onnxruntime as ort
+import os
+import urllib.request
+from hypha_rpc.sync import connect_to_server
+from tifffile import imread
+from kaibu_utils import mask_to_features
+
+
+#%% Read example image
+image = imread("example_image.tif")
+if len(image.shape) == 3 and image.shape[0] == 3:
+    image = np.transpose(image, [1, 2, 0])
+assert image.shape == (512, 512, 3)
+
+# Save the image to png
+cv2.imwrite("example_image.png", cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
+
+# Reverse the color channels to match matplotlib's default color mapping
+image = image[..., ::-1]
+
+
+# %% Connect to SAM service and compute image embedding
+if not os.path.exists("example_image_embeddings.bin"):
+    server_url = "https://hypha.aicell.io"
+    workspace_name = "bioimageio-colab"
+    service_id = "microsam"
+    model_name = "vit_b_lm"
+
+    client = connect_to_server({"server_url": server_url, "method_timeout": 5})
+    sid = f"{workspace_name}/{service_id}"
+    svc = client.get_service(sid, {"mode": "random"})
+    features = svc.compute_embedding(model_name=model_name, image=image)
+
+    print(features.shape)
+
+    # Save the embeddings to a binary file (.bin)
+    features.tofile("example_image_embeddings.bin")
+else:
+    # Load the embeddings from the binary file
+    features = np.fromfile("example_image_embeddings.bin", dtype=np.float32)
+
+    # Reshape the embeddings to the expected shape
+    features = features.reshape(1, 256, 64, 64)
+
+
+# %% Load the SAM decoder model
+model_url = "https://uk1s3.embassy.ebi.ac.uk/public-datasets/sam-vit_b_lm-decoder/1/model.onnx"
+local_model_path = "model.onnx"
+
+if not os.path.exists(local_model_path):
+    urllib.request.urlretrieve(model_url, local_model_path)
+
+model = ort.InferenceSession(local_model_path)
+
+# %%
+def model_data(clicks, image_embedding, model_scale):
+    # Initialize variables
+    point_coords = None
+    point_labels = None
+
+    # Check for input click prompts
+    if clicks:
+        n = len(clicks)
+
+        # Initialize arrays for (n+1) points (including padding point)
+        point_coords = np.zeros((n + 1, 2), dtype=np.float32)
+        point_labels = np.zeros((n + 1,), dtype=np.float32)
+
+        # Add clicks and scale coordinates
+        for i, click in enumerate(clicks):
+            point_coords[i, 0] = click["x"] * model_scale["samScale"]
+            point_coords[i, 1] = click["y"] * model_scale["samScale"]
+            point_labels[i] = click["clickType"]
+
+        # Add the padding point
+        point_coords[n, 0] = 0.0
+        point_coords[n, 1] = 0.0
+        point_labels[n] = -1.0
+
+        # Convert to tensors
+        point_coords_tensor = ort.OrtValue.ortvalue_from_numpy(
+            np.expand_dims(point_coords, axis=0)
+        )
+        point_labels_tensor = ort.OrtValue.ortvalue_from_numpy(
+            np.expand_dims(point_labels, axis=0)
+        )
+    else:
+        return None
+
+    # Image size tensor
+    image_size_tensor = ort.OrtValue.ortvalue_from_numpy(
+        np.array([model_scale["height"], model_scale["width"]], dtype=np.float32)
+    )
+
+    # Mask input and has_mask_input
+    mask_input = ort.OrtValue.ortvalue_from_numpy(
+        np.zeros((1, 1, 256, 256), dtype=np.float32)
+    )
+    has_mask_input = ort.OrtValue.ortvalue_from_numpy(np.array([0], dtype=np.float32))
+
+    return {
+        "image_embeddings": image_embedding,
+        "point_coords": point_coords_tensor,
+        "point_labels": point_labels_tensor,
+        "mask_input": mask_input,
+        "has_mask_input": has_mask_input,
+        "orig_im_size": image_size_tensor,
+    }
+
+# Define inputs
+point_coords = [(80, 80)]
+clicks = [
+    {"x": coord[0], "y": coord[1], "clickType": 1} for coord in point_coords
+]
+img_height, img_width = image.shape[:2]
+model_scale = {
+    "height": img_height,  # original image height
+    "width": img_width,  # original image width
+    "samScale": 1024 / max(img_height, img_width),  # scaling factor
+}
+
+# Generate feeds
+feeds = model_data(clicks, features, model_scale)
+
+#%% Run the model
+output_names = ["masks"]
+masks = model.run(output_names, feeds)
+
+# %% Plot the mask
+mask = masks[0].squeeze()
+
+# Display the mask using matplotlib
+plt.imshow(mask, cmap='viridis')
+plt.colorbar(label="Normalized Intensity")
+plt.title("Visualized Mask")
+plt.axis("off")
+plt.show()
+
+# %%
+
+# Convert the mask to a binary mask
+binary_mask = (mask > 0).astype(np.uint8)
+
+# Find contours
+contours, hierarchy = cv2.findContours(
+    binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+)
+
+# Create an empty image to draw contours
+contour_image = np.zeros_like(binary_mask)
+
+# Draw contours on the empty image
+cv2.drawContours(contour_image, contours, -1, (255), thickness=1)  # Thickness can be adjusted
+
+# Visualize the contours
+plt.figure(figsize=(10, 5))
+plt.subplot(1, 2, 1)
+plt.title("Binary Mask")
+plt.imshow(binary_mask, cmap="gray")
+plt.axis("off")
+
+plt.subplot(1, 2, 2)
+plt.title("Contours")
+plt.imshow(contour_image, cmap="gray")
+plt.axis("off")
+plt.show()
+
+# %%
+image_overlay = image.copy()
+
+# Draw contours on the RGB image
+cv2.drawContours(image_overlay, contours, -1, (255, 0, 255), thickness=2)  # Green contours
+
+# Plot the overlay
+plt.figure(figsize=(8, 8))
+plt.title("Overlay of Contours on Original Image")
+plt.imshow(cv2.cvtColor(image_overlay, cv2.COLOR_BGR2RGB))  # Convert BGR to RGB for plotting
+plt.axis("off")
+plt.show()
+
+# %%
+def contours_to_geojson(contours):
+    """
+    Convert OpenCV contours into a GeoJSON structure.
+    
+    Parameters:
+        contours (list): List of OpenCV contours.
+        
+    Returns:
+        dict: GeoJSON object.
+    """
+    geojson = {
+        "type": "FeatureCollection",
+        "features": []
+    }
+    
+    for contour in contours:
+        # Convert contour to a list of [x, y] points
+        points = contour[:, 0, :].tolist()
+        
+        # Ensure the polygon is closed (first point == last point)
+        if points[0] != points[-1]:
+            points.append(points[0])
+        
+        # Create a GeoJSON feature
+        feature = {
+            "type": "Feature",
+            "geometry": {
+                "type": "Polygon",
+                "coordinates": [points]
+            },
+            "properties": {}  # Add any desired properties here
+        }
+        geojson["features"].append(feature)
+    
+    return geojson
+
+# Example usage
+geojson_data = contours_to_geojson(contours)
+
+
+# %%
+features = mask_to_features(binary_mask)
+# %%
diff --git a/plugins/bioimageio-colab-annotator.imjoy.html b/plugins/bioimageio-colab-annotator.imjoy.html
index 281fa6c..e4147a9 100644
--- a/plugins/bioimageio-colab-annotator.imjoy.html
+++ b/plugins/bioimageio-colab-annotator.imjoy.html
@@ -14,190 +14,513 @@
   "env": "",
   "permissions": [],
   "requirements": [
-    "https://cdn.jsdelivr.net/npm/hypha-rpc@latest/dist/hypha-rpc-websocket.min.js"
+    "https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.47/dist/hypha-rpc-websocket.min.js",
+    "https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js",
+    "https://docs.opencv.org/4.5.0/opencv.js"
   ],
   "dependencies": []
 }
 </config>
 
 <script lang="javascript">
-class BioImageIOColabAnnotator {
-    constructor() {
-        this.image = null; // Current image
-        this.mask = null; // Current mask
-        this.filename = null; // Filename of the current image
-        this.imageLayer = null; // Layer displaying the image
-        this.annotationLayer = null; // Layer displaying the annotations
-        this.edgeColor = "magenta"; // Default edge color for annotations
-        this.modelName = "vit_b_lm"; // Model name for the embeddings
+    
+const loadSamDecoder = async (modelID) => {
+  const modelUrls = {
+    // "sam_vit_b": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",  // TODO: Replace with ONNX decoder model
+    "sam_vit_b_lm": "https://uk1s3.embassy.ebi.ac.uk/public-datasets/sam-vit_b_lm-decoder/1/model.onnx",
+    // "sam_vit_b_em_organelles": "https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/noisy-ox/1/files/vit_b.pt",  // TODO: Replace with ONNX decoder model
+  };
+
+  const modelUrl = modelUrls[modelID];
+  console.log('Starting to created ONNX model from', modelUrl);
+  const modelPromise = ort.InferenceSession.create(modelUrl)
+    .then((model) => {
+      console.log('ONNX model created', model);
+      return model;
+    })
+    .catch((error) => {
+      console.error('Error creating ONNX model:', error);
+      throw error; // Propagate the error
+    });
+
+  return modelPromise;
+};
+
+const getServices = async (ctx) => {
+  // Extract configuration settings and check if they are valid
+  const config = ctx.config || {};
+  config.serverUrl = config.serverUrl || "https://hypha.aicell.io";
+  config.token = config.token || null;
+  config.imageProviderId = config.imageProviderId || null;
+  config.samServiceId = config.samServiceId || "bioimageio-colab/microsam";
+
+  // Connect to the Hypha server
+  console.log(`Connecting to server ${config.serverUrl}...`);
+  await api.showMessage(`Connecting to server ${config.serverUrl}...`);
+  const server = await hyphaWebsocketClient.connectToServer({
+    server_url: config.serverUrl,
+    token: config.token,
+  });
+
+  // Get the current workspace and user ID
+  const currentWorkspace = server.config.user.scope.current_workspace;
+  const userID = server.config.user.id;
+  console.log(`Connected to workspace ${currentWorkspace} as user ${userID}.`);
+
+  let dataService = null;
+  let samService = null;
+
+  if (config.imageProviderId) {
+    // Get the image provider service from the server
+    try {
+      dataService = await server.getService(config.imageProviderId);
+      console.log(`Received image provider service with ID: ${config.imageProviderId}`);
+    } catch (e) {
+      console.error(e);
+      await api.alert(
+        `The image provider cannot be reached (ID: ${config.imageProviderId}). Please check if the service is running.`
+      );
+    }
+  } else {
+    const msg = "No annotation service ID provided in the configuration. Using example image.";
+    console.log(msg);
+    await api.alert(msg);
+  }
+
+  // Get the SAM service from the server
+  try {
+    samService = await server.getService(config.samServiceId, { mode: "first" });
+    console.log(`Received SAM service with ID: ${config.samServiceId}`);
+  } catch (e) {
+    samService = null;
+    console.error(e);
+    await api.alert(
+      `The SAM service is currently not reachable (ID: ${config.samServiceId}). Please wait a few minutes and reload the page to try again.`,
+      { duration: 6000 }
+    );
+  }
+
+  // Return the services
+  return [ dataService, samService ];
+};
+
+const setImageLayer = async (viewer, dataService) => { 
+    let image;
+    let filename;
+    if (dataService) {
+      // Fetch a random image from the data provider if available
+      console.log("Loading random image from the data provider...");
+      [ image, filename ] = await dataService.get_random_image();
+      
+    } else {
+      // Load an example image if no data provider is available
+      console.log("Loading example image...");
+      image = "https://raw.githubusercontent.com/bioimage-io/bioimageio-colab/main/data/example_image.png";
+      filename = "example_image";
     }
+    console.log("Image loaded:", filename);
+    let imageLayer = await viewer.view_image(image, { name: filename });
+    console.log("Image displayed:", filename);
+
+    return imageLayer;
+};
+
+const getImageFromLayer = async (imageLayer) => {
+  const vtkImage = await imageLayer.get_image();
+  const pointData = await vtkImage.getPointData()
+  const scalars = await pointData.getScalars();
+  
+  // Pixel data
+  let pixelArray = await scalars.getData();
+  const originalDtype = pixelArray.constructor;
 
-    async setup() {
-        // No setup actions required for now
+  // Data type
+  let dType = originalDtype.name;
+  dType = dType.replace('Array', '').toLowerCase();
+  
+  // Dimensions
+  const dimensions = await vtkImage.getDimensions();
+  const width = dimensions[0];
+  const height = dimensions[1];
+  const nChannels = await scalars.getNumberOfComponents();
+
+  if (nChannels === 4) {
+    // Convert RGBA to RGB
+    console.log('The image is in RGBA format. Converting to RGB...');
+    const rgbArray = new originalDtype(width * height * 3);
+    for (let i = 0; i < width * height; i++) {
+      rgbArray[3 * i] = pixelArray[4 * i];  // R
+      rgbArray[3 * i + 1] = pixelArray[4 * i + 1];  // G
+      rgbArray[3 * i + 2] = pixelArray[4 * i + 2];  // B
     }
+    pixelArray = rgbArray;
+  }
+  
+  // Create the image object
+  const image = {
+    _rtype: "ndarray",
+    _rvalue: pixelArray,
+    _rshape: [width, height, 3],
+    _rdtype: dType,
+  };
+  console.log('===== image =====>', image)
+  return image;
+};
 
-    async run(ctx) {
-        // Extract configuration settings
-        const config = ctx.config || {};
-        const serverUrl = config.server_url || "https://hypha.aicell.io";
-        const workspace = config.workspace;
-        const token = config.token;
-        const annotationServiceId = config.annotation_service_id;
-        if(!annotationServiceId){
-            await api.alert("Please provide the annotation service ID in the configuration.");
-            return;
-        }
-        const samServiceId = "bioimageio-colab/microsam";
-        
-        // Create and display the viewer window
-        const viewer = await api.createWindow({src: "https://kaibu.org/#/app", fullscreen: true});
-        await api.showMessage(`Connecting to server ${serverUrl}...`);
-        // Login before connecting and then use userid
-        // TODO: Add login functionality
-
-        // Connect to the Hypha server
-        const server = await hyphaWebsocketClient.connectToServer({
-            server_url: serverUrl,
-            token: token,
-            workspace: workspace,
-        });
-
-        // Get the user ID
-        const userID = server.config.user.id;
-        await api.showMessage(`Connected to workspace ${server.config.user.scope.current_workspace} as user ${userID}.`);
-
-        // Get the bioimageio-colab service from the server
-        let dataProvider;
-        try {
-            dataProvider = await server.getService(annotationServiceId);
-        } catch (e) {
-            await api.alert(`Failed to get the bioimageio-colab annotation service (id=${annotationServiceId}). (Error: ${e})`);
-            return;
-        }
-
-        // Get the SAM service from the server
-        let samService;
-        try {
-            samService = await server.getService(samServiceId, {"mode": "random"});
-        } catch (e) {
-            samService = null;
-            await api.showMessage(`Failed to get the bioimageio-colab SAM service (id=${samServiceId}). Please try again later.`);
-        }
-
-        // Flag to check if the image embedding is already calculated
-        let embeddingIsCalculated = false; 
-
-        // Function to get a new image and set up the viewer
-        const getImage = async () => {
-            if (this.image !== null) {
-                // Remove existing layers if there is any image loaded
-                await viewer.remove_layer({id: this.imageLayer.id});
-                await viewer.remove_layer({id: this.annotationLayer.id});
-            }
+const createTensorFromUint8Array = (data, shape) => {
+  const dst = new ArrayBuffer(data.byteLength);
+  new Uint8Array(dst).set(new Uint8Array(data));
+  return new ort.Tensor("float32", new Float32Array(dst), shape);
+};
 
-            // Fetch a random image from the service
-            [this.image, this.filename] = await dataProvider.get_random_image();
-            this.imageLayer = await viewer.view_image(this.image, {name: "image"});
+const computeEmbedding = async (samService, image, modelID) => {
+  // TODO: Load precomputed embedding for sample images
 
-            // Clear any previous image embeddings from the SAM service
-            if (samService) {
-                embeddingIsCalculated = false;
-                await samService.clear_cache();
-            }
+  // Compute the embedding for the image
+  if (!samService) {
+      const msg = "No SAM service available. Embedding computation was skipped.";
+      console.log(msg);
+      await api.showMessage(msg);
+      return;
+  }
+
+  console.log(`Computing embedding for image with model ${modelID}...`);
+  const embeddingPromise = samService.compute_embedding(image, modelID)
+    .then(embeddingResult => {
+      console.log("===== embeddingResult =====>", embeddingResult)
+
+      // Set the model scale based on the image dimensions
+      // TODO: check how to use embeddingResult.input_size (directly calculate modelScale)
+      const longSideLength = Math.max(...embeddingResult["input_size"]);
+      const w = image._rshape[0];
+      const h = image._rshape[1];
+      const samScale = longSideLength / Math.max(h, w);
+      const modelScale = { height: h, width: w, samScale };
+      console.log("===== modelScale =====>", modelScale)
+
+      // Create the embedding tensor
+      const embeddingTensor = createTensorFromUint8Array(
+        embeddingResult.features._rvalue,
+        embeddingResult.features._rshape
+      )
+      console.log("===== embeddingTensor =====>", embeddingTensor)
+
+      return { embeddingTensor, modelScale };
+    })
+    .catch(error => {
+      // Catch any errors during the embedding calculation or tensor preparation
+      console.error('An error occurred while preparing the embedding:', error);
+      throw error; // Propagate the error to be handled later
+    });
+
+  return embeddingPromise;
+};
+
+const prepareModelData = ({ embeddingResult, coordinates }) => {
+  let pointCoords;
+  let pointLabels;
+  let pointCoordsTensor;
+  let pointLabelsTensor;
+
+  const embeddingTensor = embeddingResult.embeddingTensor;
+  const modelScale = embeddingResult.modelScale;
+
+  // Prepare input for the model
+  const clicks = [ {
+      x: coordinates[0], 
+      y: coordinates[1], 
+      clickType: 1
+  } ];
+
+  console.log('===== clicks =====>', clicks)
+
+  // Check there are input click prompts
+  if (clicks) {
+    let n = clicks.length;
+
+    // If there is no box input, a single padding point with 
+    // label -1 and coordinates (0.0, 0.0) should be concatenated
+    // so initialize the array to support (n + 1) points.
+    pointCoords = new Float32Array(2 * (n + 1));
+    pointLabels = new Float32Array(n + 1);
+
+    // Add clicks and scale to what SAM expects
+    for (let i = 0; i < n; i++) {
+      pointCoords[2 * i] = clicks[i].x * modelScale.samScale;
+      pointCoords[2 * i + 1] = clicks[i].y * modelScale.samScale;
+      pointLabels[i] = clicks[i].clickType;
+    }
+
+    // Add in the extra point/label when only clicks and no box
+    // The extra point is at (0, 0) with label -1
+    pointCoords[2 * n] = 0.0;
+    pointCoords[2 * n + 1] = 0.0;
+    pointLabels[n] = -1.0;
+
+    // Create the tensor
+    pointCoordsTensor = new ort.Tensor("float32", pointCoords, [1, n + 1, 2]);
+    pointLabelsTensor = new ort.Tensor("float32", pointLabels, [1, n + 1]);
+  }
+  const imageSizeTensor = new ort.Tensor("float32", [
+    modelScale.height,
+    modelScale.width,
+  ]);
+
+  if (pointCoordsTensor === undefined || pointLabelsTensor === undefined)
+    return;
+
+  // There is no previous mask, so default to an empty tensor
+  const maskInput = new ort.Tensor(
+    "float32",
+    new Float32Array(256 * 256),
+    [1, 1, 256, 256]
+  );
+  // There is no previous mask, so default to 0
+  const hasMaskInput = new ort.Tensor("float32", [0]);
+
+  const feeds = {
+    image_embeddings: embeddingTensor,
+    point_coords: pointCoordsTensor,
+    point_labels: pointLabelsTensor,
+    orig_im_size: imageSizeTensor,
+    mask_input: maskInput,
+    has_mask_input: hasMaskInput,
+  };
+  console.log('===== feeds =====>', feeds)
+  return feeds;
+};
 
-            // Add the annotation functionality to the interface
-            this.annotationLayer = await viewer.add_shapes([], {
-                shape_type: "polygon",
-                draw_edge_color: this.edgeColor,
-                name: "annotation",
-                _rintf: true,
-                // Callback for adding a new feature (annotation point)
-                add_feature_callback: async (shape) => {
-                    if (shape.geometry.type === "Point") {
-                        if (samService) {
-                            // The point coordinates need to be reversed to match the coordinate convention of SAM
-                            const pointCoords = [shape.geometry.coordinates.reverse()];
-                            const pointLabels = pointCoords.map(() => 1); // All points have a label of 1
-                            let features = [];
-
-                            // Compute embeddings if not already computed for the image
-                            try {
-                                if (!embeddingIsCalculated) {
-                                    await api.showMessage("Computing embedding and segmenting image...");
-                                } else {
-                                    await api.showMessage("Segmenting...");
-                                }
-                                features = await samService.segment(this.modelName, this.image, pointCoords, pointLabels);
-                                embeddingIsCalculated = true;
-                            } catch (e) {
-                                await api.showMessage(`Failed to compute the image embedding. (Error: ${e})`);
-                                return;
-                            }
-
-                            // Add the segmented features as polygons to the annotation layer
-                            for (let coords of features) {
-                                const polygon = {
-                                    type: "Feature",
-                                    coordinates: coords,
-                                    geometry: {
-                                        type: "Polygon",
-                                        coordinates: [coords],
-                                    },
-                                    properties: {
-                                        edge_color: this.edgeColor,
-                                        edge_width: 2,
-                                        size: 7,
-                                    },
-                                };
-                                this.annotationLayer.add_feature(polygon);
-                            }
-                        } else {
-                            await api.showMessage("Warning: SAM service not available. Please try again later.");
-                        }
-                    }
-                }
+const processMaskToGeoJSON = (masks) => {
+  // Dimensions of the mask (batch, channels, width, height)
+  const [b, c, width, height] = masks.dims;
+
+  // 1. Apply threshold to create binary mask
+  const binaryMask = new Uint8Array(width * height);
+  for (let i = 0; i < masks.data.length; i++) {
+    binaryMask[i] = masks.data[i] > 0.0 ? 255 : 0;
+  }
+
+  // Convert binaryMask to an OpenCV.js Mat
+  const binaryMat = new cv.Mat(height, width, cv.CV_8UC1);
+  binaryMat.data.set(binaryMask);
+
+  // 2. Use OpenCV.js to find contours
+  const contours = new cv.MatVector();
+  const hierarchy = new cv.Mat();
+  cv.findContours(
+    binaryMat,
+    contours,
+    hierarchy,
+    cv.RETR_EXTERNAL, // Retrieve only external contours
+    cv.CHAIN_APPROX_SIMPLE // Compress horizontal, vertical, and diagonal segments
+  );
+
+  // 3. Process contours into GeoJSON-compatible features
+  const features = [];
+
+  if (contours.size() > 0) {
+    // Pick the largest contour as the main object
+    let largestContourIndex = 0;
+    let largestContourSize = 0;
+    for (let i = 0; i < contours.size(); i++) {
+      const c = contours.get(i);
+      if (c.rows > largestContourSize) {
+        largestContourSize = c.rows;
+        largestContourIndex = i;
+      }
+    }
+
+    const largestContour = contours.get(largestContourIndex);
+    const pts = [];
+
+    for (let i = 0; i < largestContour.rows; i++) {
+      const x = largestContour.intPtr(i)[0]; // x coordinate
+      const y = largestContour.intPtr(i)[1]; // y coordinate
+      pts.push([x, y]);
+    }
+
+    // Close the polygon if not closed
+    if (
+      pts.length > 0 &&
+      (pts[0][0] !== pts[pts.length - 1][0] || pts[0][1] !== pts[pts.length - 1][1])
+    ) {
+      pts.push(pts[0]);
+    }
+
+    // Add the polygon to the features array
+    features.push(pts);
+  }
+
+  console.log('===== contours =====>', features);
+
+  // Clean up memory
+  contours.delete();
+  hierarchy.delete();
+  binaryMat.delete();
+
+  return features;
+};
+
+const addMaskToLayer = async (annotationLayer, masks, edgeColor) => {
+  // Process the masks into GeoJSON-compatible features
+  const features = processMaskToGeoJSON(masks);
+
+  // Add the segmented features as polygons to the annotation layer
+  for (let coords of features) {
+      const polygon = {
+          type: "Feature",
+          coordinates: coords,
+          geometry: {
+              type: "Polygon",
+              coordinates: [coords],
+          },
+          properties: {
+              edge_color: edgeColor,
+              edge_width: 2,
+              size: 7,
+          },
+      };
+      annotationLayer.add_feature(polygon);
+  }
+
+  return annotationLayer;
+};
+
+const setAnnotationLayer = async (viewer, edgeColor, samService, embeddingPromise, modelPromise) => {
+  // Add the annotation functionality to the interface
+  let annotationLayer
+  annotationLayer = await viewer.add_shapes([], {
+      shape_type: "polygon",
+      draw_edge_color: edgeColor,
+      name: "annotation",
+      _rintf: true,
+      // Callback for adding a new feature (annotation point)
+      add_feature_callback: async (shape) => {
+          if (shape.geometry.type === "Point") {
+            if (!samService) {
+              const msg = "No SAM service available. Embedding computation was skipped.";
+              console.log(msg);
+              await api.showMessage(msg);
+              return;
+            }
+            // Segment the image and add the mask to the annotation layer
+            const embeddingResult = await embeddingPromise;
+            const feeds = prepareModelData({
+                embeddingResult: embeddingResult,
+                coordinates: shape.geometry.coordinates,
             });
-        };
-
-        // Function to save the annotation
-        const saveAnnotation = async () => {
-            if (!this.annotationLayer) return;
-            const annotation = await this.annotationLayer.get_features();
-            if (annotation.features.length > 0) {
-                await dataProvider.save_annotation(this.filename, annotation, [this.image._rshape[0], this.image._rshape[1]]);
-                await api.showMessage("Annotation saved.");
-            } else {
-                await api.showMessage("No annotation provided. Saving was skipped.");
+            if (feeds === undefined) {
+              console.log("No input data available for model.");
+              return;
             }
-        };
-
-        // Function to load the next image
-        const nextImage = async () => {
-            await saveAnnotation();
-            await getImage();
-        };
-
-        // Add a control widget with a button to load the next image
-        await viewer.add_widget({
-            _rintf: true,
-            name: "Control",
-            type: "control",
-            elements: [
-                {
-                    type: "button",
-                    label: "Save Annotation",
-                    callback: nextImage,
-                }
-            ],
-        });
-
-        // Load the initial image
-        await getImage();
-        await api.showMessage("Ready to annotate!");
+            const model = await modelPromise; 
+            const results = await model.run(feeds);
+            console.log("===== results =====>", results);
+            annotationLayer = await addMaskToLayer(annotationLayer, results["masks"], edgeColor);
+          }
+      }
+  });
+  return annotationLayer;
+};
+
+const saveAnnotation = async (dataService, imageLayer, annotationLayer) => {
+    // Do not save if no data service or annotation layer is available
+    
+    if (!dataService) {
+      const msg = "No data service provided. Saving was skipped.";
+      console.log(msg);
+      await api.showMessage(msg);
+      return;
+    }
+    // Get the annotation features from the layer
+    if (!annotationLayer) {
+      const msg = "No annotation provided. Saving was skipped.";
+      console.log(msg);
+      await api.showMessage(msg);
+      return;
     }
+    const annotation = await annotationLayer.get_features();
+    if (annotation.features.length > 0) {
+        const image = await imageLayer.get_image();
+        const dimensions = await image.getDimensions();
+        await dataService.save_annotation(filename, annotation, [dimensions[0], dimensions[1]]);
+        const msg = "Annotation saved.";
+        console.log(msg);
+        await api.showMessage(msg);
+    } else {
+      const msg = "No annotation provided. Saving was skipped.";
+      console.log(msg);
+      await api.showMessage(msg);
+    }
+};
+
+
+// Define the BioImageIOColabAnnotator class
+class BioImageIOColabAnnotator {
+  constructor() {
+    this.imageLayer = null; // Layer displaying the image
+    this.annotationLayer = null; // Layer displaying the annotations
+    this.edgeColor = "magenta"; // Default edge color for annotations
+    this.modelID = "sam_vit_b_lm"; // Model name for the embedding
+    this.modelPromise = null; // Promise for loading the model
+  }
+
+  async setup() {
+    // No setup actions required for now
+  }
+
+  async run(ctx) {
+    // Create and display the viewer window
+    const viewer = await api.createWindow({src: "https://kaibu.org/#/app", fullscreen: true});
+
+    // Get the services
+    const [ dataService, samService ] = await getServices(ctx);
+
+    const setModel = async () => {
+      // Load the decoder for the selected model
+      this.modelPromise = loadSamDecoder(this.modelID);
+    };
+
+    // Function to load an image and display it in the viewer with the annotation layer
+    const setImageAnnotation = async () => {
+      // First remove existing layers from the viewer
+      await viewer.clear_layers();
+      // Then load new image and annotation layers
+      this.imageLayer = await setImageLayer(viewer, dataService);
+      const image = await getImageFromLayer(imageLayer);
+      const embeddingPromise = computeEmbedding(samService, image, this.modelID);
+      this.annotationLayer = await setAnnotationLayer(viewer, this.edgeColor, samService, embeddingPromise, this.modelPromise);
+    };
+    
+    // Function to load the next image
+    const nextImage = async () => {
+      // Save the current annotation if available
+      await saveAnnotation(dataService, this.imageLayer, this.annotationLayer);
+      // Load the next image
+      await setImageAnnotation();
+    };
+    
+    // Add a control widget with a button to load the next image
+    await viewer.add_widget({
+        _rintf: true,
+        name: "Control",
+        type: "control",
+        elements: [
+            {
+                type: "button",
+                label: "Save Annotation",
+                callback: nextImage,
+            }
+        ],
+    });
+
+    // Start loading the model and the first image
+    await setModel();
+    await setImageAnnotation();
+    await api.showMessage("Ready to annotate!");
+  }
 }
 
 // Export the annotator class
 api.export(new BioImageIOColabAnnotator());
 </script>
-
diff --git a/pyproject.toml b/pyproject.toml
index 4960ed9..2faa3d1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,7 @@ requires = ["setuptools", "wheel"]
 
 [project]
 name = "bioimageio-colab"
-version = "0.2.0"
+version = "0.2.9"
 readme = "README.md"
 description = "Collaborative image annotation and model training with human in the loop."
 dependencies = [