{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Carotid artery segmentation\n",
    "This section demonstrates how to use the `CarotidSegmenter` model from the `zea` library to segment the carotid artery in ultrasound images. The model predicts a mask highlighting the carotid region for a given input image.\n",
    "\n",
    "For more details, see the [original paper](https://doi.org/10.1016/j.cmpb.2022.107037):  \n",
    "- Luuk van Knippenberg, \"Unsupervised domain adaptation method for segmenting cross-sectional CCA images\", Computers in Biology and Medicine, 2022."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tue-bmd/zea/blob/main/docs/source/notebooks/models/carotid_segmentation_example.ipynb)\n",
    " \n",
    "[![View on GitHub](https://img.shields.io/badge/GitHub-View%20Source-blue?logo=github)](https://github.com/tue-bmd/zea/blob/main/docs/source/notebooks/models/carotid_segmentation_example.ipynb)\n",
    " \n",
    "[![Hugging Face model](https://img.shields.io/badge/Hugging%20Face-Model-yellow?logo=huggingface)](https://huggingface.co/zeahub/carotid-segmenter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "‼️ **Important:** This notebook is optimized for **GPU/TPU**. Code execution on a **CPU** may be very slow.\n",
    "\n",
    "If you are running in Colab, please enable a hardware accelerator via:\n",
    "\n",
    "**Runtime → Change runtime type → Hardware accelerator → GPU/TPU** 🚀."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install zea"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: Using backend 'tensorflow'\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "# NOTE: should be `tensorflow` or `jax` for EchoNetDynamic\n",
    "os.environ[\"KERAS_BACKEND\"] = \"tensorflow\"\n",
    "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from keras import ops\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import requests\n",
    "from io import BytesIO\n",
    "\n",
    "from zea import init_device, log\n",
    "from zea.visualize import plot_shape_from_mask\n",
    "from zea.visualize import set_mpl_style\n",
    "\n",
    "init_device(verbose=False)\n",
    "set_mpl_style()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: Available built-in zea presets for CarotidSegmenter: ['carotid-segmenter']\n"
     ]
    }
   ],
   "source": [
    "from zea.models.carotid_segmenter import CarotidSegmenter\n",
    "\n",
    "presets = list(CarotidSegmenter.presets.keys())\n",
    "log.info(f\"Available built-in zea presets for CarotidSegmenter: {presets}\")\n",
    "\n",
    "model = CarotidSegmenter.from_preset(\"carotid-segmenter\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_imgs = 1\n",
    "url = \"https://raw.githubusercontent.com/tue-bmd/zea/main/docs/source/notebooks/assets/carotid.png\"\n",
    "\n",
    "response = requests.get(url)\n",
    "img = Image.open(BytesIO(response.content)).convert(\"L\")\n",
    "img_np = np.asarray(img).astype(np.float32) / 255.0\n",
    "\n",
    "img_np = img_np[None, ..., None]  # add batch and channel dimensions\n",
    "batch = ops.convert_to_tensor(img_np)\n",
    "\n",
    "masks = model(batch)\n",
    "masks = ops.squeeze(masks, axis=-1)\n",
    "masks_clipped = ops.where(masks > 0.5, 1, 0)\n",
    "masks_clipped = ops.convert_to_numpy(masks_clipped)\n",
    "\n",
    "# stack batch twice to get 2 rows\n",
    "batch_stacked = ops.concatenate([batch, batch])\n",
    "\n",
    "# Plot the original image and its mask side by side, with each row as an example\n",
    "fig, axes = plt.subplots(n_imgs, 2, figsize=(6, 3 * n_imgs))\n",
    "axes = np.atleast_2d(axes)\n",
    "\n",
    "for ax_img, ax_mask, img_arr, mask_arr in zip(axes[:, 0], axes[:, 1], batch[..., 0], masks_clipped):\n",
    "    ax_img.imshow(img_arr, cmap=\"gray\", vmin=0, vmax=1)\n",
    "    ax_img.set_title(\"Input Image\")\n",
    "    ax_img.axis(\"off\")\n",
    "    ax_mask.imshow(img_arr, cmap=\"gray\", vmin=0, vmax=1)\n",
    "    plot_shape_from_mask(ax_mask, mask_arr, color=\"red\", alpha=0.5)\n",
    "    ax_mask.set_title(\"Predicted Mask\")\n",
    "    ax_mask.axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"carotid_output.png\")\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Carotid Example Output](./carotid_output.png)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}