{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a9e360da",
   "metadata": {},
   "source": [
    "# Working with a sequence of images\n",
    "In this notebook, we will look at how to work with a sequence of images, and possibly change some parameters on the fly. As in all examples, we will start with setting the Keras backend."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b303c558",
   "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/pipeline/zea_sequence_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/pipeline/zea_sequence_example.ipynb)\n",
    " \n",
    "[![Hugging Face dataset](https://img.shields.io/badge/Hugging%20Face-Dataset-yellow?logo=huggingface)](https://huggingface.co/datasets/zeahub/zea-carotid-2023)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6eebf11",
   "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,
   "id": "6d6dd155",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install zea"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b41342fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"KERAS_BACKEND\"] = \"jax\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7e4af10b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: Using backend 'jax'\n"
     ]
    }
   ],
   "source": [
    "import keras\n",
    "\n",
    "import zea\n",
    "from zea.visualize import set_mpl_style\n",
    "from zea.internal.notebooks import animate_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2047ca56",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "n_frames = 15\n",
    "n_tx = 11\n",
    "n_tx_total = 127"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6de3d4f2",
   "metadata": {},
   "source": [
    "We will work with the GPU if available, and initialize using `init_device` to pick the best available device. Also, (optionally), we will set the matplotlib style for plotting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e55d6107",
   "metadata": {},
   "outputs": [],
   "source": [
    "zea.init_device(verbose=False)\n",
    "set_mpl_style()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "449d254c",
   "metadata": {},
   "source": [
    "Let's initialize a default B-mode ultrasound image formation pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "39a27a22",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = zea.Pipeline.from_default(enable_pfield=False, with_batch_dim=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c369482f",
   "metadata": {},
   "source": [
    "We will load a sequence of acquired RF data frames (carotid scan) and reconstruct a B-mode image from each frame. We will then animate the sequence of images. But first let's load the data and parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f79441a3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc3f1eb2c4e64e5fbe98738a9ebd33d9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "2_cross_bifur_right_0000_small.hdf5:   0%|          | 0.00/2.49G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mWARNING\u001b[0m This ``zea.File`` '/root/.cache/zea/huggingface/datasets/datasets--zeahub--zea-carotid-2023/snapshots/ecae16d62a54cd362afd1b69cc5d8c4a28d5404f/2_cross_bifur_right_0000_small.hdf5' was created with a legacy version of zea (<0.1.0), while you are using zea v0.1.0. It may behave in unexpected ways. Install an earlier version of zea<0.1.0 for full compatibility or re-save the file with zea v0.1.0 or later (e.g. via File.create).\n"
     ]
    }
   ],
   "source": [
    "# this can take a while to download\n",
    "file_path = \"hf://zeahub/zea-carotid-2023/2_cross_bifur_right_0000.hdf5\"  # ~25GB\n",
    "# so let's use the smaller one by default:\n",
    "file_path = \"hf://zeahub/zea-carotid-2023/2_cross_bifur_right_0000_small.hdf5\"  # ~2.5GB\n",
    "\n",
    "frames = list(range(n_frames))  # use first 15 frames for demonstration\n",
    "with zea.File(file_path, mode=\"r\") as file:\n",
    "    data = file.data.raw_data[frames]\n",
    "    parameters = file.load_parameters()\n",
    "\n",
    "parameters.set_transmits(n_tx)  # reduce number of transmits for faster processing\n",
    "parameters.zlims = (0, 0.04)  # reduce z-limits a bit for better visualizations\n",
    "\n",
    "config = zea.Config(parameters={\"dynamic_range\": (-40, 0)})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "481edac0",
   "metadata": {},
   "source": [
    "## Reconstructing a sequence of B-mode images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "eabec7f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m 1/15\u001b[0m \u001b[32m━\u001b[0m\u001b[37m━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[1m20s\u001b[0m 1s/step\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[33mDEBUG\u001b[0m [zea.Pipeline] The following input keys are not used by the pipeline: {'n_ch', 'center_frequency'}. Make sure this is intended. This warning will only be shown once.\n",
      "\u001b[1m15/15\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m5s\u001b[0m 218ms/step\n"
     ]
    }
   ],
   "source": [
    "images = []\n",
    "n_frames = data.shape[0]\n",
    "progbar = keras.utils.Progbar(n_frames, stateful_metrics=[\"frame\"])\n",
    "\n",
    "inputs = pipeline.prepare_parameters(parameters, **config.parameters)\n",
    "\n",
    "for frame_no in range(n_frames):\n",
    "    output = pipeline(data=data[frame_no, parameters.selected_transmits], **inputs)\n",
    "    image = output.pop(\"data\")\n",
    "    inputs = output\n",
    "    images.append(image)\n",
    "    progbar.update(frame_no + 1)\n",
    "\n",
    "animate_images(images, \"./bmode_sequence.gif\", parameters, interval=100, cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ed2a3f2",
   "metadata": {},
   "source": [
    "![B-mode sequence](./bmode_sequence.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "364238c4",
   "metadata": {},
   "source": [
    "## Change transmits on the fly"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9af5147a",
   "metadata": {},
   "source": [
    "We now used 11 transmits throughout for every frame. We can also sweep through the the transmits for each frame to see how it affects the image quality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "43aa9551",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m15/15\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m2s\u001b[0m 97ms/step\n"
     ]
    }
   ],
   "source": [
    "images = []\n",
    "n_frames = len(frames)\n",
    "progbar = keras.utils.Progbar(n_frames, stateful_metrics=[\"frame\"])\n",
    "\n",
    "for idx, frame_no in enumerate(frames):\n",
    "    tx_idx = int(round(idx * (n_tx_total - 1) / (n_frames - 1)))\n",
    "    parameters.set_transmits([tx_idx])\n",
    "    raw_data_frame = data[frame_no, tx_idx][None, ...]\n",
    "    inputs = pipeline.prepare_parameters(parameters)\n",
    "    output = pipeline(data=raw_data_frame, **inputs)\n",
    "    images.append(output[\"data\"])\n",
    "    progbar.update(idx + 1)\n",
    "\n",
    "animate_images(images, \"./tx_sweep.gif\", parameters, interval=100, cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1298ca72",
   "metadata": {},
   "source": [
    "![Transmit sweep](./tx_sweep.gif)"
   ]
  }
 ],
 "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": 5
}