{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1692bbb0",
   "metadata": {},
   "source": [
    "# Speckle2Self: ultrasound speckle reduction\n",
    "\n",
    "[![arXiv](https://img.shields.io/badge/arXiv-Paper-b31b1b.svg)](https://arxiv.org/abs/2507.06828)   [![GitHub](https://img.shields.io/badge/GitHub-Code-black?logo=github)](https://github.com/noseefood/speckle2self)\n",
    "\n",
    "This notebook demonstrates how to perform **self-supervised ultrasound speckle reduction** using [Speckle2Self](https://arxiv.org/abs/2507.06828) within the [zea](https://github.com/tue-bmd/zea) framework. We apply the model to an **in-vivo carotid artery** scan from [zeahub/zea-carotid-2023](https://huggingface.co/datasets/zeahub/zea-carotid-2023), matching the domain on which the `speckle2self-invivo` weights were trained.\n",
    "\n",
    "For more information on the method, see the [original paper](https://arxiv.org/abs/2507.06828):\n",
    "- Lin, Huiping, et al. \"Speckle2Self: Learning Self-Supervised Despeckling with Attention Mechanism for SAR Images.\" Remote Sensing 17.23 (2025): 3840."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95796f53",
   "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/speckle2self_despeckling_example.ipynb)   [![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/speckle2self_despeckling_example.ipynb)\n",
    "[![Hugging Face model](https://img.shields.io/badge/Hugging%20Face-Model-yellow?logo=huggingface)](https://huggingface.co/zeahub/speckle2self-invivo)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9f98e76",
   "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": "3192d444",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install zea"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f254ad42",
   "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",
    "os.environ[\"KERAS_BACKEND\"] = \"tensorflow\"\n",
    "\n",
    "import numpy as np\n",
    "import keras\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import zea\n",
    "from zea.ops import (\n",
    "    Beamform,\n",
    "    EnvelopeDetect,\n",
    "    Pipeline,\n",
    "    Cast,\n",
    "    Demodulate,\n",
    "    Downsample,\n",
    "    Normalize,\n",
    "    LogCompress,\n",
    ")\n",
    "from zea.visualize import set_mpl_style"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d394b567",
   "metadata": {},
   "outputs": [],
   "source": [
    "zea.init_device(verbose=False)\n",
    "set_mpl_style()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "65f38c13",
   "metadata": {},
   "outputs": [],
   "source": [
    "carotid_path = \"hf://zeahub/picmus/in_vivo/carotid_cross/carotid_cross_expe_dataset_rf/carotid_cross_expe_dataset_rf.hdf5\"\n",
    "frame_idx = 0\n",
    "n_tx = 11  # transmits per frame (fewer -> faster; more -> better quality)\n",
    "dynamic_range = (-40, 0)  # dB"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c543bf7",
   "metadata": {},
   "source": [
    "## Load carotid data\n",
    "\n",
    "We load a single frame of an in-vivo carotid scan from [zeahub/zea-carotid-2023](https://huggingface.co/datasets/zeahub/zea-carotid-2023) on Hugging Face. Note that this is raw ultrasound RF channel data. Therefore we will first need to beamform and envelope-detect the data before feeding it to the Speckle2Self model. We will do this using the `zea.Pipeline` framework, which provides efficient implementations of these operations. For more details see the [pipeline example notebook](../pipeline/zea_pipeline_example.ipynb) or the [API reference](https://zea.readthedocs.io/en/latest/pipeline.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9204a2ea",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5dca0d0a0936453bbf7e8088a23814f2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "in_vivo/carotid_cross/carotid_cross_expe(…):   0%|          | 0.00/58.6M [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: Running compute_pfield and caching the result to /root/.cache/zea/cached_funcs/compute_pfield_3b656eda410eab8a0941a76c49516a15.pkl.\n"
     ]
    }
   ],
   "source": [
    "with zea.File(carotid_path, revision=\"v0.1.0\") as f:\n",
    "    data = f.data.raw_data[frame_idx]\n",
    "    parameters = f.load_parameters()\n",
    "\n",
    "parameters.set_transmits(n_tx)\n",
    "parameters.zlims = (0, 0.04)\n",
    "parameters.grid_size_x = 500\n",
    "parameters.grid_size_z = 800\n",
    "\n",
    "pipeline = Pipeline(\n",
    "    operations=[\n",
    "        Cast(dtype=\"float32\"),\n",
    "        Demodulate(),\n",
    "        Downsample(factor=2),\n",
    "        Beamform(\n",
    "            beamformer=\"delay_and_sum\",\n",
    "            enable_pfield=True,\n",
    "            num_patches=100,\n",
    "        ),\n",
    "        EnvelopeDetect(),\n",
    "        Normalize(),\n",
    "        LogCompress(),\n",
    "    ],\n",
    "    with_batch_dim=False,\n",
    "    jit_options=\"pipeline\",\n",
    ")\n",
    "\n",
    "inputs = pipeline.prepare_parameters(parameters, dynamic_range=dynamic_range)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "11293b68",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Beamform + envelope-detect\n",
    "data = data[parameters.selected_transmits]\n",
    "out = pipeline(data=data, **inputs)\n",
    "envelope = out[pipeline.output_key]\n",
    "envelope = keras.ops.convert_to_numpy(envelope)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb165880",
   "metadata": {},
   "source": [
    "## Load Speckle2Self model\n",
    "\n",
    "`Speckle2Self` can be loaded directly from Hugging Face:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5a8ae8b7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a69084d3862496ebc0f360e93b8b12d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/409 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "283fba8dcbeb4a23bb24bc3f73fcece7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.weights.h5:   0%|          | 0.00/15.9M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from zea.models.speckle2self import Speckle2Self\n",
    "\n",
    "model = Speckle2Self.from_preset(\"hf://zeahub/speckle2self-invivo\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02d99bd7",
   "metadata": {},
   "source": [
    "## Preprocess and run inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b8967837",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run despeckling (Speckle2Self) model inference\n",
    "despeckled = model(envelope[None, ..., None])\n",
    "despeckled = keras.ops.convert_to_numpy(despeckled)\n",
    "despeckled = despeckled.squeeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c600d1a6",
   "metadata": {},
   "source": [
    "## Visualize results\n",
    "\n",
    "For a simple comparison, we display one frame before and after Speckle2Self despeckling.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dc3b0d26",
   "metadata": {},
   "outputs": [],
   "source": [
    "xlims_mm = [v * 1e3 for v in parameters.xlims]\n",
    "zlims_mm = [v * 1e3 for v in parameters.zlims]\n",
    "extent = [xlims_mm[0], xlims_mm[1], zlims_mm[1], zlims_mm[0]]\n",
    "\n",
    "# gamma correction for better visualization\n",
    "despeckled_viz = np.power(despeckled, 1.3)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(8, 4), squeeze=False)\n",
    "\n",
    "axes[0, 0].imshow(envelope, cmap=\"gray\", extent=extent)\n",
    "axes[0, 0].set_title(\"B-Mode\")\n",
    "axes[0, 0].set_xlabel(\"X (mm)\")\n",
    "axes[0, 0].set_ylabel(\"Z (mm)\")\n",
    "\n",
    "axes[0, 1].imshow(despeckled_viz, cmap=\"gray\", extent=extent)\n",
    "axes[0, 1].set_title(\"Despeckled B-Mode (Speckle2Self)\")\n",
    "axes[0, 1].set_xlabel(\"X (mm)\")\n",
    "axes[0, 1].set_ylabel(\"Z (mm)\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"speckle2self_output.png\", bbox_inches=\"tight\", dpi=100)\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3bb233c",
   "metadata": {},
   "source": [
    "**Speckle2Self result on in-vivo carotid data**\n",
    "\n",
    "![Image](./speckle2self_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": 5
}