{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7926a818-4882-4f6a-94ce-ef769832eef5",
   "metadata": {},
   "source": [
    "# Demo about effect of the different stiching kernels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56934ef9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f37733de-593c-4c90-998d-77492fcc33f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.misc\n",
    "import numpy\n",
    "from nabu.testutils import get_data\n",
    "\n",
    "raw_data_1 = get_data('brain_phantom.npz')['data']\n",
    "raw_data_2 = numpy.random.rand(*raw_data_1.shape) *255.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa2ff82d-312e-4326-a1b4-4d3b0cf1f8a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nabu.stitching.overlap import OverlapStitchingStrategy, ZStichOverlapKernel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bae743b1-ea18-41b6-bb4b-d4966bd58445",
   "metadata": {},
   "source": [
    "### plot the two raw data side by side"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f76fe833",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(2)\n",
    "ax[0].imshow(raw_data_1)\n",
    "ax[1].imshow(raw_data_2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6ff0fff-a684-408a-bf50-8574d4956acc",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plt.title(\"raw data\")\n",
    "plt.subplot(121)\n",
    "plt.imshow(raw_data_1)\n",
    "plt.subplot(122)\n",
    "plt.imshow(raw_data_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07f8da48-4e5d-4490-bcd0-42a418c00847",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_strategy_result(overlap_strategy, img_1, img_2, with_weights=True, extra_params=None):\n",
    "    strategy = OverlapStitchingStrategy.from_value(overlap_strategy)\n",
    "    kernel = ZStichOverlapKernel(\n",
    "        frame_width=img_1.shape[1],\n",
    "        stitching_strategy=strategy,\n",
    "        overlap_size=img_1.shape[0],\n",
    "        extra_params=extra_params,\n",
    "    )\n",
    "    stitch = kernel.stitch(img_1, img_2)[0]\n",
    "    if with_weights:\n",
    "        w_1 = kernel.weights_img_1\n",
    "        w_2 = kernel.weights_img_2\n",
    "    else:\n",
    "        w_1 = w_2 = None\n",
    "    \n",
    "    fig = plt.figure(constrained_layout=True)\n",
    "    if with_weights:\n",
    "        axes = fig.subplot_mosaic(\n",
    "            [\n",
    "                [\"00\", \"01\"],\n",
    "                [\"10\", \"11\"],\n",
    "                [\"2\", \"2\"],\n",
    "            ],\n",
    "        )\n",
    "    else:\n",
    "        axes = fig.subplot_mosaic(\n",
    "            [\n",
    "                [\"2\", \"2\"],\n",
    "            ],\n",
    "        )\n",
    "    if w_1 is not None and w_2 is not None:\n",
    "        plt.title(f\"{strategy.value} strategy\")\n",
    "        # plot weights curves\n",
    "        axes[\"00\"].set_title(\"image 1 weigths\")\n",
    "        axes[\"00\"].plot(numpy.arange(len(w_1)), w_1, color=\"red\")\n",
    "        axes[\"01\"].set_title(\"image 2 weigths\")\n",
    "        axes[\"01\"].plot(numpy.arange(len(w_2)), w_2, color=\"blue\")\n",
    "\n",
    "        # plot weighted images\n",
    "        axes[\"10\"].set_title(\"image 1 with weights applied\")\n",
    "        axes[\"10\"].imshow(img_1*w_1)\n",
    "        axes[\"11\"].set_title(\"image 2 with weights applied\")\n",
    "        axes[\"11\"].imshow(img_2*w_2)\n",
    "    \n",
    "    # plot stitched images\n",
    "    axes[\"2\"].set_title(\"stitched image\")\n",
    "    axes[\"2\"].imshow(stitch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f879c010-859c-4ad1-aee8-ecdf499b188f",
   "metadata": {},
   "source": [
    "## mean strategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dfee29e-71a0-4055-9ad9-5a3c92aaeebd",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_strategy_result(\n",
    "    overlap_strategy=\"mean\",\n",
    "    img_1=raw_data_1,\n",
    "    img_2=raw_data_2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe12cb7f-890a-4565-8a43-930d3c61c232",
   "metadata": {},
   "source": [
    "## closest strategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "560038a3-ae8e-48a7-976e-3a5189eceac1",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_strategy_result(\n",
    "    overlap_strategy=\"closest\",\n",
    "    img_1=raw_data_1,\n",
    "    img_2=raw_data_2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e47ddc51-cbec-4a48-99d7-fc733a63378f",
   "metadata": {},
   "source": [
    "## linear weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2057782f-efce-4c08-ace1-7aa188d46fbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_strategy_result(\n",
    "    overlap_strategy=\"linear weights\",\n",
    "    img_1=raw_data_1,\n",
    "    img_2=raw_data_2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc48a908-2a1d-4827-98ec-e71f95444126",
   "metadata": {},
   "source": [
    "## cosinus weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da1547bb-7a40-42ff-9bfb-51a301a8a801",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_strategy_result(\n",
    "    overlap_strategy=\"cosinus weights\",\n",
    "    img_1=raw_data_1,\n",
    "    img_2=raw_data_2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1c0f1eb-8fb7-43eb-b755-5ce9a6b60e1f",
   "metadata": {},
   "source": [
    "## image minimum divergence\n",
    "\n",
    "This method will not applied a simple predetermined weights on the two images.\n",
    "\n",
    "It split the two images into two parts: high frequency and low frequency.\n",
    "\n",
    "The two low frequency parts will be stitched using a 'sinusoidal' / cosinus weights approach.\n",
    "When the two high frequency part will be stitched by taking the lower divergent pixels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "673e336c-8fda-4270-938e-3c65187a047e",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_strategy_result(\n",
    "    overlap_strategy=\"image minimum divergence\",\n",
    "    img_1=raw_data_1,\n",
    "    img_2=raw_data_2,\n",
    "    with_weights=False,\n",
    "    extra_params={\n",
    "        \"frequency_threshold\": 5,\n",
    "    },\n",
    ")\n",
    "# uncomment to compare vs 'raw' sinusoidal\n",
    "# plot_strategy_result(\n",
    "#     overlap_strategy=\"cosinus weights\",\n",
    "#     img_1=raw_data_1,\n",
    "#     img_2=raw_data_2,\n",
    "#     with_weights=False,\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52189fa6-68cc-4033-985d-ed907078c428",
   "metadata": {},
   "source": [
    "## higher signal\n",
    "\n",
    "This method will not applied predetermine weights on the two images.\n",
    "\n",
    "It will pick pixel on the image having the higher signal.\n",
    "\n",
    "A use case is that if there is some artefacts on images which creates stripes (from scintillator artefacts for example) it could be removed from this method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9918cf26-fa1b-4a16-afff-aa81be892aca",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_data = get_data(\"brain_phantom.npz\")[\"data\"]\n",
    "new_dataset_1 = raw_data_1.copy()\n",
    "new_dataset_1[40:75] = 0.0\n",
    "new_dataset_1[:, 210:245] = 0.0\n",
    "\n",
    "new_dataset_2 = raw_data_1.copy()\n",
    "new_dataset_2[:, 100:120] = 0.0\n",
    "\n",
    "fig = plt.figure()\n",
    "plt.title(\"raw data\")\n",
    "plt.subplot(121)\n",
    "plt.imshow(new_dataset_1)\n",
    "plt.subplot(122)\n",
    "plt.imshow(new_dataset_2)\n",
    "\n",
    "plot_strategy_result(\n",
    "    overlap_strategy=\"higher signal\",\n",
    "    img_1=new_dataset_1,\n",
    "    img_2=new_dataset_2,\n",
    "    with_weights=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e20aead-0364-42d6-a587-2f2b82a76302",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  },
  "nbsphinx": {
   "execute": "never"
  },
  "vscode": {
   "interpreter": {
    "hash": "27a91ab57669bd25447dc4402c0770a4a8cc30df75a096933cc563087c33cc2f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}