example/model_usage.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bioimage Model Zoo Core  Example notebook\n",
    "\n",
    "This notebook shows how to interact with the `bioimageio.core` programmatically to explore, load, use, and export content from the [BioImage Model Zoo](https://bioimage.io).\n",
    "\n",
    "\n",
    "quick links:\n",
    "- [Create an input sample for a given model](#create_sample_for_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Activate human readable output error messages and load dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0.1. Install necessary dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import matplotlib\n",
    "    import torch\n",
    "\n",
    "    import bioimageio.core\n",
    "except ImportError:\n",
    "    %pip install bioimageio.core==0.6.7 torch==2.3.1 matplotlib==3.9.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0.2.Enable pretty_validation_errors\n",
    "\n",
    "This function displays validation errors in a human readable format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioimageio.spec.pretty_validation_errors import (\n",
    "    enable_pretty_validation_errors_in_ipynb,\n",
    ")\n",
    "\n",
    "enable_pretty_validation_errors_in_ipynb()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0.3. Load general dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load general dependencies\n",
    "from pprint import pprint\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from imageio.v2 import imread\n",
    "\n",
    "from bioimageio.spec.utils import download\n",
    "\n",
    "\n",
    "# Function to display input and prediction output images\n",
    "def show_images(sample_tensor, prediction_tensor):\n",
    "    input_array = sample_tensor.members[\"input0\"].data\n",
    "\n",
    "    # Check for the number of channels to enable display\n",
    "    input_array = np.squeeze(input_array)\n",
    "    if len(input_array.shape) > 2:\n",
    "        input_array = input_array[0]\n",
    "\n",
    "    output_array = prediction_tensor.members[\"output0\"].data\n",
    "\n",
    "    # Check for the number of channels to enable display\n",
    "    output_array = np.squeeze(output_array)\n",
    "    if len(output_array.shape) > 2:\n",
    "        output_array = output_array[0]\n",
    "\n",
    "    plt.figure()\n",
    "    ax1 = plt.subplot(1, 2, 1)\n",
    "    ax1.set_title(\"Input\")\n",
    "    ax1.axis(\"off\")\n",
    "    plt.imshow(input_array)\n",
    "    ax2 = plt.subplot(1, 2, 2)\n",
    "    ax2.set_title(\"Prediction\")\n",
    "    ax2.axis(\"off\")\n",
    "    plt.imshow(output_array)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Load a model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Inspect available models in the Bioimage Model Zoo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Go to https://bioimage.io to browser available models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Load model from the BioImage Model Zoo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`bioimage.io` resources may be identified via their bioimage.io __ID__, e.g. \"affable-shark\" or the [__DOI__](https://doi.org/) of their [__Zenodo__](https://zenodo.org/) backup.\n",
    "\n",
    "Both of these options may be version specific (\"affable-shark/1\" or a version specific [__Zenodo__](https://zenodo.org/) backup [__DOI__](https://doi.org/)).\n",
    "\n",
    "Alternatively, any rdf.yaml source, single file or in a .zip, may be loaded by providing its __local path__ or __URL__."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BMZ_MODEL_ID = \"\"  # \"affable-shark\"\n",
    "BMZ_MODEL_DOI = \"\"  # \"10.5281/zenodo.6287342\"\n",
    "BMZ_MODEL_URL = \"https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/affable-shark/draft/files/rdf.yaml\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`load_description` is a function of the `bioimageio.spec` package, but as it is a sub-package of `bioimageio.core` it can also be called from it by `bioimageio.core.load_description`.\n",
    "\n",
    "To learn more about the functionalities of the `bioimageio.spec` package, see the [bioimageio.spec package example notebook](https://github.com/bioimage-io/spec-bioimage-io/blob/main/example/load_model_and_create_your_own.ipynb), also available as a [Google Colab](https://colab.research.google.com/github/bioimage-io/spec-bioimage-io/blob/main/example/load_model_and_create_your_own.ipynb) notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioimageio.core import load_description\n",
    "\n",
    "# Load the model description\n",
    "# ------------------------------------------------------------------------------\n",
    "if BMZ_MODEL_ID != \"\":\n",
    "    model = load_description(BMZ_MODEL_ID)\n",
    "    print(\n",
    "        f\"\\nThe model '{model.name}' with ID '{BMZ_MODEL_ID}' has been correctly loaded.\"\n",
    "    )\n",
    "elif BMZ_MODEL_DOI != \"\":\n",
    "    model = load_description(BMZ_MODEL_DOI)\n",
    "    print(\n",
    "        f\"\\nThe model '{model.name}' with DOI '{BMZ_MODEL_DOI}' has been correctly loaded.\"\n",
    "    )\n",
    "elif BMZ_MODEL_URL != \"\":\n",
    "    model = load_description(BMZ_MODEL_URL)\n",
    "    print(\n",
    "        f\"\\nThe model '{model.name}' with URL '{BMZ_MODEL_URL}' has been correctly loaded.\"\n",
    "    )\n",
    "else:\n",
    "    print(\"\\nPlease specify a model ID, DOI or URL\")\n",
    "\n",
    "if \"draft\" in BMZ_MODEL_ID or \"draft\" in BMZ_MODEL_DOI or \"draft\" in BMZ_MODEL_URL:\n",
    "    print(\n",
    "        f\"\\nThis is the DRAFT version of '{model.name}'. \\nDraft versions have not been reviewed by the Bioimage Model Zoo Team and may contain harmful code. Run with caution.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Inspect the model metadata\n",
    "\n",
    "Let's inspect all the model metadata. For a step-by-step inspection refer to [bioimageio.spec package example notebook](https://github.com/bioimage-io/spec-bioimage-io/blob/main/example/load_model_and_create_your_own.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"\\n Covers of the model '{model.name}' are: \")\n",
    "for cover in model.covers:\n",
    "    cover_data = imread(download(cover).path)\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(cover_data)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Test the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `bioimageio.core.test_model` function can be used to fully test the model.\n",
    "This is done by running the predicition on the test input(s) and checking that they agree with the test output(s) provided in the model documentation.\n",
    "\n",
    "This test should be run before using the model to ensure that it works properly.\n",
    "\n",
    "----\n",
    "\n",
    "`bioimageio.core.test_model` returns a validation dictionary with 'status'='passed'/'failed' and other detailed information that can be inspected by calling `.display()` on it.\n",
    "\n",
    "The validation summary will indicate:\n",
    "- the versions of the `bioimageio.spec` and `bioimageio.core` libraries used to run the validation\n",
    "- the status of several validation steps\n",
    "    - ✔️: Success\n",
    "    - 🔍: information about the validation context\n",
    "    - ⚠: Warning\n",
    "    - ❌: Error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioimageio.core import test_model\n",
    "\n",
    "test_summary = test_model(model)\n",
    "test_summary.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Running a prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`bioimageio.core` implements the functionality to run a prediction with models described in the `bioimage.io` format.\n",
    "\n",
    "This includes functions to run predictions on `numpy.ndarray`/`xarray.DataArray` as input and convenience functions to run predictions for images stored on disc.\n",
    "\n",
    "### 3.1. Load the test image and convert into a tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioimageio.spec.model import v0_5\n",
    "from bioimageio.spec.utils import load_array\n",
    "\n",
    "assert isinstance(model, v0_5.ModelDescr)\n",
    "input_image = load_array(model.inputs[0].test_tensor)\n",
    "print(f\"array shape: {input_image.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a `Tensor` (light wrapper around an `xarray.DataArray`) from the test input image. \n",
    "\n",
    "`bioimageio.core.Tensors/xarray.DataArrays` are like numpy arrays, but they have annotated axes.\n",
    "\n",
    "The axes are used to validate that the axes of the input image match the axes expected by the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioimageio.core import Tensor\n",
    "\n",
    "test_input_tensor = Tensor.from_numpy(input_image, dims=model.inputs[0].axes)\n",
    "\n",
    "# print the axis annotations ('dims') and the shape of the input array\n",
    "print(f\"tensor shape: {test_input_tensor.tagged_shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A collection of tensors is called a `Sample`.\n",
    "\n",
    "In the case of the `affable-shark` model it only has one input, but for models with multiple inputs a `Sample` includes a tensor for each input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioimageio.core import Sample\n",
    "\n",
    "sample = Sample(members={\"raw\": test_input_tensor}, stat=None, id=\"sample-from-numpy\")\n",
    "\n",
    "sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`bioimageio.core` provides the helper function `create_sample_for_model` to automatically create the `Sample` for the given model.\n",
    "<a id='create_sample_for_model'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioimageio.core.digest_spec import create_sample_for_model\n",
    "from bioimageio.spec.utils import download\n",
    "\n",
    "input_paths = {ipt.id: download(ipt.test_tensor).path for ipt in model.inputs}\n",
    "print(f\"input paths: {input_paths}\")\n",
    "assert isinstance(model, v0_5.ModelDescr)\n",
    "sample = create_sample_for_model(\n",
    "    model=model, inputs=input_paths, sample_id=\"my_demo_sample\"\n",
    ")\n",
    "\n",
    "sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is also  a helper function `get_test_inputs` to directly import the test input sample for a given model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioimageio.core.digest_spec import get_test_inputs\n",
    "\n",
    "test_sample = get_test_inputs(model)\n",
    "\n",
    "test_sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2. Create a prediciton pipeline\n",
    "\n",
    "The `prediction_pipeline` function is used to run a prediction with a given model.\n",
    "\n",
    "It applies the __pre-processing__, if indicated in the model rdf.yaml, runs __inference__ with the model and applies the __post-processing__, again if specified in the model rdf.yaml.\n",
    "\n",
    "The `devices` argument can be used to specify which device(s), CPU, a single GPU, or multiple GPUs (not implemented yet), to use for inference with the model.\n",
    "\n",
    "The default is `devices=None`, this will use a __GPU__ if available, otherwise it uses the __CPU__.\n",
    "\n",
    "\n",
    "The `weight_format` argument can be used to specify which of the model's available weight formats to use.\n",
    "\n",
    "The deafult is `weight_format=None`, this will use the weight format with highest priority (as defined by bioimageio.core).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioimageio.core import create_prediction_pipeline\n",
    "\n",
    "devices = None\n",
    "weight_format = None\n",
    "\n",
    "prediction_pipeline = create_prediction_pipeline(\n",
    "    model, devices=devices, weight_format=weight_format\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the new prediction pipeline to run a prediction for the previously loaded test image.\n",
    "\n",
    "The prediction pipeline returns a `Sample` object, which will be displayed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction: Sample = prediction_pipeline.predict_sample_without_blocking(sample)\n",
    "\n",
    "# show the prediction result\n",
    "show_images(sample, prediction)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3. Prediction without a PredicitionPipeline\n",
    "\n",
    "`bioimageio.core` has two convenience functions `predict` and `predict_many` which allow the prediction of images without creating a `PredictionPipeline`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioimageio.core import predict  # , predict_many\n",
    "\n",
    "# predict_many(model=model, inputs=[sample])\n",
    "\n",
    "prediction: Sample = predict(model=model, inputs=sample)\n",
    "\n",
    "# show the prediction result\n",
    "show_images(sample, prediction)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3. Recover input and output tensors as numpy arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example code shows how to recover the image information from the input and output tensors as numpy arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np_input_list = []\n",
    "np_output_list = []\n",
    "\n",
    "# iterate over the number of tensors inside the input sample\n",
    "for ipt in range(len(sample.members.keys())):\n",
    "    input_array = sample.members[f\"input{ipt}\"].data\n",
    "\n",
    "    # Check for the number of channels to enable display\n",
    "    input_array = np.squeeze(input_array)\n",
    "    if len(input_array.shape) > 2:\n",
    "        input_array = input_array[0]\n",
    "\n",
    "    np_input_list.append(input_array)\n",
    "\n",
    "\n",
    "# iterate over the number of tensors inside the output prediction\n",
    "for out in range(len(prediction.members.keys())):\n",
    "    output_array = prediction.members[f\"output{ipt}\"].data\n",
    "\n",
    "    # Check for the number of channels to enable display\n",
    "    output_array = np.squeeze(output_array)\n",
    "    if len(output_array.shape) > 2:\n",
    "        output_array = output_array[0]\n",
    "\n",
    "    np_output_list.append(output_array)\n",
    "\n",
    "plt.imshow(np_input_list[0])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bioimageio-core",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}