notebooks/full_image_plots.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:100% !important; }</style>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEBUG = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run ./nb_core/root_dirs.ipynb\n",
    "setup_syspath_disentangle(DEBUG)\n",
    "%run ./nb_core/disentangle_imports.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from disentangle.analysis.plot_utils import clean_ax\n",
    "from disentangle.core.tiff_reader import load_tiff\n",
    "from disentangle.config_utils import load_config, get_configdir_from_saved_predictionfile\n",
    "from disentangle.core.data_split_type import DataSplitType\n",
    "from disentangle.scripts.evaluate import get_highsnr_data\n",
    "import ml_collections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "        # '/home/ashesh.ashesh/training/disentangle/2402/D16-M3-S0-L0/128',\n",
    "        # '/home/ashesh.ashesh/training/disentangle/2402/D16-M3-S0-L0/144',\n",
    "# 2402/D16-M3-S0-L0/144\n",
    "        # '/home/ashesh.ashesh/training/disentangle/2402/D16-M3-S0-L0/145'\n",
    "\n",
    "        # '/home/ashesh.ashesh/training/disentangle/2402/D16-M3-S0-L0/165',\n",
    "        # '/home/ashesh.ashesh/training/disentangle/2402/D16-M3-S0-L0/164',\n",
    "        # '/home/ashesh.ashesh/training/disentangle/2402/D16-M3-S0-L0/169',\n",
    "\n",
    "noise_levels = [5100]\n",
    "pred_dir = '/group/jug/ashesh/data/paper_stats/'\n",
    "\n",
    "usplit_fname = {5100: 'Test_P64_G32_M5_Sk44/pred_disentangle_2402_D16-M3-S0-L0_165.tif',\n",
    "                # 6675: 'Test_P64_G32_M5_Sk44/pred_disentangle_2402_D16-M3-S0-L0_164.tif',\n",
    "                # 6675: 'Test_P64_G32_M5_Sk44/pred_disentangle_colorfuljug_2403_D16-M3-S0-L0_1.tif',\n",
    "                4725: 'Test_P64_G32_M5_Sk44/pred_disentangle_2402_D16-M3-S0-L0_169.tif',\n",
    "                228: 'Test_PNone_G32_M10_Sk0/pred_disentangle_2403_D23-M3-S0-L0_0.tif',\n",
    "                \n",
    "                }\n",
    "\n",
    "denoiSplitNM_fname = {\n",
    "                5100: 'Test_P128_G64_M5_Sk44/pred_disentangle_2402_D16-M3-S0-L0_128.tif',\n",
    "                # 6675: 'Test_P128_G64_M5_Sk44/pred_disentangle_2402_D16-M3-S0-L0_144.tif', \n",
    "                # 6675: 'Test_P128_G64_M50_Sk44/pred_disentangle_2403_D16-M3-S0-L0_25.tif',\n",
    "                4725: 'Test_P128_G64_M5_Sk44/pred_disentangle_2402_D16-M3-S0-L0_145.tif',\n",
    "                228: 'Test_P128_G32_M10_Sk32/pred_disentangle_2402_D3-M3-S0-L0_32.tif',\n",
    "                      }\n",
    "hdn_usplit = {} #{4450: 'pred_disentangle_2402_D23-M3-S0-L0_34.tif'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _noise_model_path(nmodel_dir):\n",
    "    histfpath = None\n",
    "    gmmfpath = None\n",
    "    for fname in os.listdir(nmodel_dir):\n",
    "        if fname.startswith('HistNoiseModel'):\n",
    "            histfpath = os.path.join(nmodel_dir,fname)\n",
    "        elif fname.startswith('GMMNoiseModel'):\n",
    "            gmmfpath = os.path.join(nmodel_dir,fname)\n",
    "    return {'gmm':gmmfpath, 'hist':histfpath}\n",
    "\n",
    "def noise_model_paths(pred_file_name):\n",
    "    \"\"\"\n",
    "    denoiSplitNM_fname[noise_levels[0]]\n",
    "    \"\"\"\n",
    "    cfg = load_config(get_configdir_from_saved_predictionfile(pred_file_name))\n",
    "    nmodel1_fpath_dict = None\n",
    "    nmodel2_fpath_dict = None\n",
    "    if 'noise_model_ch1_fpath' in cfg.model and cfg.model.noise_model_ch1_fpath is not None:\n",
    "        nmodel1_fpath_dict = _noise_model_path(os.path.dirname(cfg.model.noise_model_ch1_fpath))\n",
    "    if 'noise_model_ch2_fpath' in cfg.model and cfg.model.noise_model_ch2_fpath is not None:\n",
    "        nmodel2_fpath_dict = _noise_model_path(os.path.dirname(cfg.model.noise_model_ch2_fpath))\n",
    "    return nmodel1_fpath_dict, nmodel2_fpath_dict\n",
    "\n",
    "def _get_noise_model(nmodel_fpath_dict):\n",
    "    from disentangle.nets.gmm_noise_model import GaussianMixtureNoiseModel\n",
    "    from disentangle.nets.hist_noise_model import HistNoiseModel\n",
    "    nmodel_params = np.load(nmodel_fpath_dict['gmm'])\n",
    "    gmm_model1 = GaussianMixtureNoiseModel(params=nmodel_params)\n",
    "    \n",
    "    histdata = np.load(nmodel_fpath_dict['hist'])\n",
    "    hist_model = HistNoiseModel(histdata)\n",
    "    return {'gmm':gmm_model1, 'hist':hist_model}\n",
    "\n",
    "def get_noise_models(pred_file_name):\n",
    "    nmodel1_fpath_dict, nmodel2_fpath_dict = noise_model_paths(pred_file_name)\n",
    "    nmodel1 = _get_noise_model(nmodel1_fpath_dict)\n",
    "    nmodel2 = _get_noise_model(nmodel2_fpath_dict)\n",
    "    return nmodel1, nmodel2\n",
    "\n",
    "nmodel1, nmodel2 = get_noise_models(denoiSplitNM_fname[noise_levels[0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from disentangle.analysis.plot_utils import add_subplot_axes\n",
    "\n",
    "def get_signal_from_index(signalBinIndex, n_bin, min_signal, max_signal, histBinSize):\n",
    "    querySignal_numpy = (signalBinIndex / float(n_bin) * (max_signal - min_signal) + min_signal)\n",
    "    querySignal_numpy += histBinSize / 2\n",
    "    querySignal_torch = torch.from_numpy(np.array(querySignal_numpy)).float()\n",
    "    return querySignal_torch\n",
    "\n",
    "def get_scaled_pdf(pdf, axymax, axymin, yval, factor=0.2):\n",
    "    scaled_pdf = pdf/pdf.max()\n",
    "    scaled_pdf = scaled_pdf - scaled_pdf.min()\n",
    "    scaled_pdf = scaled_pdf * (axymax - axymin)*factor + yval\n",
    "    return scaled_pdf\n",
    "\n",
    "\n",
    "# def add_signal_value(ax, signal)\n",
    "\n",
    "def plot_noise_model(signal1_index, signal2_index, histogramNoiseModel, gaussianMixtureNoiseModel, device, ax, linetxt_offset = 0.1):\n",
    "    \"\"\"Plots probability distribution P(x|s) for a certain ground truth signal.\n",
    "       Predictions from both Histogram and GMM-based Noise models are displayed for comparison.\n",
    "        Parameters\n",
    "        ----------\n",
    "        signalBinIndex: int\n",
    "            index of signal bin. Values go from 0 to number of bins (`n_bin`).\n",
    "        histogramNoiseModel: Histogram based noise model\n",
    "        gaussianMixtureNoiseModel: GaussianMixtureNoiseModel\n",
    "            Object containing trained parameters.\n",
    "        device: GPU device\n",
    "        \"\"\"\n",
    "    max_signal = histogramNoiseModel.maxv.item()\n",
    "    min_signal = histogramNoiseModel.minv.item()\n",
    "    n_bin = int(histogramNoiseModel.bins.item())\n",
    "\n",
    "    histBinSize = (max_signal - min_signal) / n_bin\n",
    "    signal1 = get_signal_from_index(signal1_index, n_bin, min_signal, max_signal, histBinSize).to(device)\n",
    "    signal2 = get_signal_from_index(signal2_index, n_bin, min_signal, max_signal, histBinSize).to(device)\n",
    "\n",
    "    queryObservations_numpy = np.arange(min_signal, max_signal, histBinSize)\n",
    "    queryObservations_numpy += histBinSize / 2\n",
    "    queryObservations = torch.from_numpy(queryObservations_numpy).float().to(device)\n",
    "    \n",
    "    gmm_pdf1 = gaussianMixtureNoiseModel.likelihood(queryObservations, signal1)\n",
    "    gmm_pdf1 = gmm_pdf1.detach().cpu().numpy()\n",
    "\n",
    "    gmm_pdf2 = gaussianMixtureNoiseModel.likelihood(queryObservations, signal2)\n",
    "    gmm_pdf2 = gmm_pdf2.detach().cpu().numpy()\n",
    "\n",
    "    # plt.figure(figsize=(12, 5))\n",
    "\n",
    "    # plt.subplot(1, 2, 1)\n",
    "    # plt.xlabel('Observation Bin')\n",
    "    # plt.ylabel('Signal Bin')\n",
    "    histogram = histogramNoiseModel.fullHist.cpu().numpy()\n",
    "    ax.imshow(histogram**0.25, cmap='gray', aspect='auto')\n",
    "    yval1 = signal1_index + 0.5\n",
    "    yval2 = signal2_index + 0.5\n",
    "    ax.axhline(y=yval1, linewidth=1, color='green', linestyle='--', alpha=0.5, label=f'{signal1.cpu().numpy():.1f}')\n",
    "    ax.axhline(y=yval2, linewidth=1, color='green', linestyle='--', alpha=0.5, label=f'{signal2.cpu().numpy():.1f}')\n",
    "\n",
    "    # plt.subplot(1, 2, 2)\n",
    "    hist_pdf1 = histogramNoiseModel.likelihood(queryObservations, signal1).cpu().numpy()\n",
    "    hist_pdf2 = histogramNoiseModel.likelihood(queryObservations, signal2).cpu().numpy()\n",
    "    ymin, ymax = ax.get_ylim()\n",
    "\n",
    "    pdf1 = get_scaled_pdf(gmm_pdf1, ymax, ymin, yval1)\n",
    "    pdf2 = get_scaled_pdf(gmm_pdf2, ymax, ymin, yval2)\n",
    "\n",
    "    step = histogram.shape[1]/pdf1.shape[0]\n",
    "    x = np.arange(0, histogram.shape[1], step=step)\n",
    "    ax.plot(x, pdf1, color='green')\n",
    "    ax.plot(x, pdf2, color='green')\n",
    "    \n",
    "    ymin, ymax = ax.get_ylim()\n",
    "    print(ymin, ymax)\n",
    "    props = dict(alpha=0)\n",
    "    fact1 = (signal1_index - ymin)/(ymax - ymin) + linetxt_offset\n",
    "    ax.text(0.77, fact1, f'{signal1.cpu().numpy():.0f}', transform=ax.transAxes, fontsize=10,\n",
    "        verticalalignment='top', bbox=props, color='green')\n",
    "    fact2 = (signal2_index - ymin)/(ymax - ymin) + linetxt_offset\n",
    "    ax.text(0.02, fact2, f'{signal2.cpu().numpy():.0f}', transform=ax.transAxes, fontsize=10,\n",
    "        verticalalignment='top', bbox=props, color='green')\n",
    "\n",
    "    # ax.legend(frameon=False, labelcolor='white', loc='upper right')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots(figsize=(6,3))\n",
    "plot_noise_model(25, 75, nmodel1['hist'], nmodel1['gmm'], 'cpu', ax, linetxt_offset=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from disentangle.utils import plotProbabilityDistribution\n",
    "# signalBinIndex=60\n",
    "# data_dict = plotProbabilityDistribution(signalBinIndex=signalBinIndex, \n",
    "#                             histogramNoiseModel=nmodel2['hist'],\n",
    "#                             gaussianMixtureNoiseModel=nmodel2['gmm'],\n",
    "#                             device='cpu')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sanity_check_config():\n",
    "    data_dicts = [usplit_fname, denoiSplitNM_fname]\n",
    "    for ith_data, ddict in enumerate(data_dicts):\n",
    "        for noise,fname in ddict.items():\n",
    "            configdir = get_configdir_from_saved_predictionfile(fname)\n",
    "            config = load_config(configdir)\n",
    "            assert 'synthetic_gaussian_scale' in config.data\n",
    "            assert config.data.synthetic_gaussian_scale == noise, f'{ith_data} {fname}: noise: {noise}, config: {config.data.synthetic_gaussian_scale}'\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sanity_check_config()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "configdir  = get_configdir_from_saved_predictionfile(denoiSplitNM_fname[noise_levels[0]])\n",
    "config = ml_collections.ConfigDict(load_config(configdir))\n",
    "highsnr_data = get_highsnr_data(config, config.datadir, DataSplitType.Test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "usplit_data = {k: load_tiff(os.path.join(pred_dir, v)) for k,v in usplit_fname.items()}\n",
    "denoiSplitNM_data = {k: load_tiff(os.path.join(pred_dir, v)) for k,v in denoiSplitNM_fname.items()}\n",
    "hdn_usplit_data = {k: load_tiff(os.path.join(pred_dir, v)) for k,v in hdn_usplit.items()}\n",
    "\n",
    "# Undoing the offset.\n",
    "for k,v in usplit_fname.items():\n",
    "     with open(os.path.join(pred_dir, v.replace('.tif', '.json')),'rb') as f:\n",
    "        offset = float(json.load(f)['offset'])\n",
    "        usplit_data[k] = usplit_data[k] + offset\n",
    "\n",
    "for k,v in denoiSplitNM_fname.items():\n",
    "    with open(os.path.join(pred_dir, v.replace('.tif', '.json')),'rb') as f:\n",
    "        offset = float(json.load(f)['offset'])\n",
    "        denoiSplitNM_data[k] = denoiSplitNM_data[k] + offset\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cropping the target to get to the same shape as the predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shape = usplit_data[noise_levels[0]].shape\n",
    "highsnr_data = highsnr_data[:, :shape[1], :shape[2]].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sanity_check_data():\n",
    "    # all shapes should be same\n",
    "    for noise_level in noise_levels:\n",
    "        shape = usplit_data[noise_level].shape\n",
    "        if noise_level in denoiSplitNM_data:\n",
    "            assert shape == denoiSplitNM_data[noise_level].shape\n",
    "        if noise_level in hdn_usplit_data:\n",
    "            assert shape == hdn_usplit_data[noise_level].shape\n",
    "        assert shape == highsnr_data.shape\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# usplit_data[noise_levels[0]] = usplit_data[noise_levels[0]][...,::-1].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sanity_check_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paper_figures_dir = '/group/jug/ashesh/data/paper_figures'\n",
    "def get_output_fpath(noise_level):\n",
    "    if 'ch1_fname' in config.data:\n",
    "        ch1str = config.data.ch1_fname.split('.')[0].replace('/','').replace('GT_all', '')\n",
    "        ch2str = config.data.ch2_fname.split('.')[0].replace('/','').replace('GT_all', '')\n",
    "    else:\n",
    "        ch1str = config.data.channel_1\n",
    "        ch2str = config.data.channel_2\n",
    "    modelid = config.workdir.strip('/').split('/')[-1]\n",
    "\n",
    "    output_filepath =os.path.join(paper_figures_dir, f'{modelid}_{noise_level}_{ch1str}_{ch2str}.png')\n",
    "    output_filepath\n",
    "    return output_filepath"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "def get_noisy_data(noise_level):\n",
    "    return highsnr_data + np.random.normal(0, noise_level, highsnr_data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.gridspec import GridSpec\n",
    "import matplotlib.pyplot as plt\n",
    "from disentangle.analysis.plot_utils import add_pixel_kde, clean_ax\n",
    "from disentangle.core.psnr import RangeInvariantPsnr\n",
    "import seaborn as sns   \n",
    "\n",
    "#### inset specific\n",
    "inset_rect=[0.05, 0.05, 0.4, 0.2]\n",
    "inset_min_labelsize=10\n",
    "color_ch_list=['goldenrod', 'cyan']\n",
    "color_pred='red'\n",
    "# insetplot_xmax_value = 30000\n",
    "# insetplot_xmin_value = -8000\n",
    "# paviaatn \n",
    "insetplot_xmax_value = 1500\n",
    "insetplot_xmin_value = -100\n",
    "\n",
    "plt_dsample = 4\n",
    "####\n",
    "data_idx = 0\n",
    "img_sz = 3\n",
    "ncol_imgs = 5\n",
    "nrow_imgs = 2\n",
    "example_spacing = 1\n",
    "grid_factor = 5\n",
    "nimgs = 1\n",
    "noise_level = noise_levels[0]\n",
    "# extra spacing for c0. It does not work. Don't know why. I think there is some integer division happening.\n",
    "c0_extra = 1\n",
    "\n",
    "noisy_data = get_noisy_data(noise_level)\n",
    "\n",
    "# for subscripts and superscripts\n",
    "params = {'mathtext.default': 'regular' }          \n",
    "plt.rcParams.update(params)\n",
    "\n",
    "def get_psnr_str(prediction, ch_idx):\n",
    "    return f'{RangeInvariantPsnr(highsnr_data[data_idx,...,ch_idx][None], prediction[data_idx,...,ch_idx][None]).item():.1f}' \n",
    "\n",
    "def add_psnr_str(ax_, psnr):\n",
    "    \"\"\"\n",
    "    Add psnr string to the axes\n",
    "    \"\"\"\n",
    "    textstr = f'PSNR\\n{psnr}'\n",
    "    props = dict(\n",
    "        boxstyle='round', \n",
    "        facecolor='gray', alpha=0.3)\n",
    "    # place a text box in upper left in axes coords\n",
    "    ax_.text(0.05, 0.95, textstr, transform=ax_.transAxes, fontsize=11,\n",
    "            verticalalignment='top', bbox=props, color='white')\n",
    "\n",
    "# extra spacing for the first and the last column.\n",
    "fig_w = ncol_imgs * img_sz + 2*c0_extra/grid_factor\n",
    "fig_h = int(img_sz * nrow_imgs + (example_spacing * (nimgs - 1)) / grid_factor )\n",
    "fig = plt.figure(figsize=(fig_w, fig_h))\n",
    "gs = GridSpec(nrows=int(grid_factor * fig_h), ncols=int(grid_factor * fig_w), hspace=0.2, wspace=0.2)\n",
    "grid_img_sz = img_sz * grid_factor\n",
    "\n",
    "# input\n",
    "ax_temp = fig.add_subplot(gs[:grid_img_sz,:grid_img_sz])\n",
    "ax_temp.imshow(np.mean(noisy_data[data_idx], axis=-1), cmap='magma')\n",
    "legend_ax = ax_temp\n",
    "\n",
    "clean_ax(ax_temp)\n",
    "\n",
    "# ax[0,0].set_title('Input')\n",
    "ax_temp = fig.add_subplot(gs[:grid_img_sz, (c0_extra+grid_img_sz):(c0_extra + grid_img_sz * 2)])\n",
    "ax_temp.imshow(noisy_data[data_idx,:,:,0], cmap='magma')\n",
    "inset_ax = add_pixel_kde(ax_temp,\n",
    "                                  inset_rect,\n",
    "                                  [noisy_data[data_idx,::plt_dsample,::plt_dsample,0],\n",
    "                                  highsnr_data[data_idx,::plt_dsample,::plt_dsample,0]],\n",
    "                                  inset_min_labelsize,\n",
    "                                  label_list=['NoisyCh1','Ch1'],\n",
    "                                  plot_kwargs_list=[{'linestyle':'--'}, {}],\n",
    "                                  color_list=[color_ch_list[0],color_ch_list[0]],\n",
    "                                  plot_xmax_value=insetplot_xmax_value,\n",
    "                                  plot_xmin_value=insetplot_xmin_value)\n",
    "inset_ax.set_xticks([])\n",
    "inset_ax.set_yticks([])\n",
    "clean_ax(ax_temp)\n",
    "\n",
    "ax_temp = fig.add_subplot(gs[grid_img_sz:grid_img_sz * 2, c0_extra+grid_img_sz:c0_extra + grid_img_sz * 2])\n",
    "ax_temp.imshow(noisy_data[data_idx,:,:,1], cmap='magma')\n",
    "inset_ax = add_pixel_kde(ax_temp,\n",
    "                                  inset_rect,\n",
    "                                  [noisy_data[data_idx,::plt_dsample,::plt_dsample,1],\n",
    "                                  highsnr_data[data_idx,::plt_dsample,::plt_dsample,1]],\n",
    "                                  inset_min_labelsize,\n",
    "                                  label_list=['NoisyCh2','Ch2'],\n",
    "                                  color_list=[color_ch_list[1],color_ch_list[1]],\n",
    "                                  plot_kwargs_list=[{'linestyle':'--'},{}],\n",
    "                                  plot_xmax_value=insetplot_xmax_value,\n",
    "                                  plot_xmin_value=insetplot_xmin_value)\n",
    "inset_ax.set_xticks([])\n",
    "inset_ax.set_yticks([])\n",
    "clean_ax(ax_temp)\n",
    "\n",
    "ax_temp = fig.add_subplot(gs[:grid_img_sz, c0_extra+grid_img_sz * 2:c0_extra+grid_img_sz * 3])\n",
    "ax_temp.imshow(usplit_data[noise_level][data_idx,...,0], cmap='magma')\n",
    "# inset_ax = add_pixel_kde(ax_temp,\n",
    "#                                   inset_rect,\n",
    "#                                   [highsnr_data[data_idx,::plt_dsample,::plt_dsample,0],\n",
    "#                                    noisy_data[data_idx,::plt_dsample,::plt_dsample,0],\n",
    "#                                   usplit_data[noise_level][data_idx,::plt_dsample,::plt_dsample,0]],\n",
    "#                                   inset_min_labelsize,\n",
    "#                                   label_list=['Ch1','input', 'Pred1'],\n",
    "#                                   color_list=[color_ch_list[0],color_ch_list[0], color_pred],\n",
    "#                                   plot_kwargs_list=[{},{'linestyle':'--'},{}],\n",
    "#                                   plot_xmax_value=insetplot_xmax_value,\n",
    "#                                   plot_xmin_value=insetplot_xmin_value)\n",
    "inset_ax = add_pixel_kde(ax_temp,\n",
    "                                  inset_rect,\n",
    "                                  [highsnr_data[data_idx,::plt_dsample,::plt_dsample,0],\n",
    "                                  usplit_data[noise_level][data_idx,::plt_dsample,::plt_dsample,0]],\n",
    "                                  inset_min_labelsize,\n",
    "                                  label_list=['Ch1', 'Pred1'],\n",
    "                                  color_list=[color_ch_list[0], color_pred],\n",
    "                                #   plot_kwargs_list=[{},{'linestyle':'--'},{}],\n",
    "                                  plot_xmax_value=insetplot_xmax_value,\n",
    "                                  plot_xmin_value=insetplot_xmin_value)\n",
    "\n",
    "# adding input to the inset.\n",
    "# sns.kdeplot(data=,\n",
    "#                     ax=inset_ax,\n",
    "#                     color=color_ch_list[0],\n",
    "#                     label='',\n",
    "#                     clip=(insetplot_xmin_value, None),\n",
    "#                     )\n",
    "\n",
    "inset_ax.set_xticks([])\n",
    "inset_ax.set_yticks([])\n",
    "add_psnr_str(ax_temp, get_psnr_str(usplit_data[noise_level], 0))\n",
    "clean_ax(ax_temp)\n",
    "\n",
    "ax_temp = fig.add_subplot(gs[grid_img_sz:grid_img_sz * 2,c0_extra+grid_img_sz * 2:c0_extra+grid_img_sz * 3])\n",
    "ax_temp.imshow(usplit_data[noise_level][data_idx,...,1], cmap='magma')\n",
    "# inset_ax = add_pixel_kde(ax_temp,\n",
    "#                                   inset_rect,\n",
    "#                                   [highsnr_data[data_idx,::plt_dsample,::plt_dsample,1],\n",
    "#                                    noisy_data[data_idx,::plt_dsample,::plt_dsample,1],\n",
    "#                                   usplit_data[noise_level][data_idx,::plt_dsample,::plt_dsample,1]],\n",
    "#                                   inset_min_labelsize,\n",
    "#                                   label_list=['Ch2','input','Pred2'],\n",
    "#                                   color_list=[color_ch_list[1],color_ch_list[1], color_pred],\n",
    "#                                   plot_kwargs_list=[{},{'linestyle':'--'},{}],\n",
    "#                                   plot_xmax_value=insetplot_xmax_value,\n",
    "#                                   plot_xmin_value=insetplot_xmin_value)\n",
    "inset_ax = add_pixel_kde(ax_temp,\n",
    "                                  inset_rect,\n",
    "                                  [highsnr_data[data_idx,::plt_dsample,::plt_dsample,1],\n",
    "                                #    noisy_data[data_idx,::plt_dsample,::plt_dsample,1],\n",
    "                                  usplit_data[noise_level][data_idx,::plt_dsample,::plt_dsample,1]],\n",
    "                                  inset_min_labelsize,\n",
    "                                  label_list=['Ch2','Pred2'],\n",
    "                                  color_list=[color_ch_list[1],color_ch_list[1], color_pred],\n",
    "                                #   plot_kwargs_list=[{},{'linestyle':'--'},{}],\n",
    "                                  plot_xmax_value=insetplot_xmax_value,\n",
    "                                  plot_xmin_value=insetplot_xmin_value)\n",
    "inset_ax.set_xticks([])\n",
    "inset_ax.set_yticks([])\n",
    "add_psnr_str(ax_temp, get_psnr_str(usplit_data[noise_level], 1))\n",
    "clean_ax(ax_temp)\n",
    "\n",
    "ax_temp = fig.add_subplot(gs[:grid_img_sz, c0_extra+grid_img_sz * 3:c0_extra+grid_img_sz * 4])\n",
    "ax_temp.imshow(denoiSplitNM_data[noise_level][data_idx,...,0], cmap='magma')\n",
    "inset_ax = add_pixel_kde(ax_temp,\n",
    "                                  inset_rect,\n",
    "                                  [highsnr_data[data_idx,::plt_dsample,::plt_dsample,0],\n",
    "                                  denoiSplitNM_data[noise_level][data_idx,::plt_dsample,::plt_dsample,0]],\n",
    "                                  inset_min_labelsize,\n",
    "                                  label_list=['Ch1','Pred1'],\n",
    "                                  color_list=[color_ch_list[0],color_pred],\n",
    "                                  plot_xmax_value=insetplot_xmax_value,\n",
    "                                  plot_xmin_value=insetplot_xmin_value)\n",
    "inset_ax.set_xticks([])\n",
    "inset_ax.set_yticks([])\n",
    "\n",
    "add_psnr_str(ax_temp, get_psnr_str(denoiSplitNM_data[noise_level], 0))\n",
    "clean_ax(ax_temp)\n",
    "ax_temp = fig.add_subplot(gs[grid_img_sz:grid_img_sz * 2, c0_extra+grid_img_sz * 3:c0_extra+grid_img_sz * 4])\n",
    "ax_temp.imshow(denoiSplitNM_data[noise_level][data_idx,...,1], cmap='magma')\n",
    "inset_ax = add_pixel_kde(ax_temp,\n",
    "                                  inset_rect,\n",
    "                                  [highsnr_data[data_idx,::plt_dsample,::plt_dsample,1],\n",
    "                                  denoiSplitNM_data[noise_level][data_idx,::plt_dsample,::plt_dsample,1]],\n",
    "                                  inset_min_labelsize,\n",
    "                                  label_list=['Ch2','Pred2'],\n",
    "                                  color_list=[color_ch_list[1],color_pred],\n",
    "                                  plot_xmax_value=insetplot_xmax_value,\n",
    "                                  plot_xmin_value=insetplot_xmin_value)\n",
    "inset_ax.set_xticks([])\n",
    "inset_ax.set_yticks([])\n",
    "\n",
    "add_psnr_str(ax_temp, get_psnr_str(denoiSplitNM_data[noise_level], 1))\n",
    "clean_ax(ax_temp)\n",
    "\n",
    "ax_temp = fig.add_subplot(gs[:grid_img_sz, 2*c0_extra+grid_img_sz * 4:2*c0_extra+grid_img_sz * 5])\n",
    "ax_temp.imshow(highsnr_data[data_idx,...,0], cmap='magma')\n",
    "legend_ch1_ax = ax_temp\n",
    "inset_ax = add_pixel_kde(ax_temp,\n",
    "                                  inset_rect,\n",
    "                                  [highsnr_data[data_idx,::plt_dsample,::plt_dsample,0]],\n",
    "                                  inset_min_labelsize,\n",
    "                                  label_list=['Ch1'],\n",
    "                                  color_list=[color_ch_list[0]],\n",
    "                                  plot_xmax_value=insetplot_xmax_value,\n",
    "                                  plot_xmin_value=insetplot_xmin_value)\n",
    "\n",
    "inset_ax.set_xticks([])\n",
    "inset_ax.set_yticks([])\n",
    "\n",
    "clean_ax(ax_temp)\n",
    "\n",
    "\n",
    "ax_temp = fig.add_subplot(gs[grid_img_sz:grid_img_sz * 2, 2*c0_extra+grid_img_sz * 4:2*c0_extra+grid_img_sz * 5])\n",
    "ax_temp.imshow(highsnr_data[data_idx,...,1], cmap='magma')\n",
    "inset_ax = add_pixel_kde(ax_temp,\n",
    "                                  inset_rect,\n",
    "                                  [highsnr_data[data_idx,::plt_dsample,::plt_dsample,1]],\n",
    "                                  inset_min_labelsize,\n",
    "                                  label_list=['Ch2'],\n",
    "                                  color_list=[color_ch_list[1]],\n",
    "                                  plot_xmax_value=insetplot_xmax_value,\n",
    "                                  plot_xmin_value=insetplot_xmin_value)\n",
    "legend_ch2_ax = ax_temp\n",
    "\n",
    "inset_ax.set_xticks([])\n",
    "inset_ax.set_yticks([])\n",
    "\n",
    "clean_ax(ax_temp)\n",
    "\n",
    "# add noise models. \n",
    "nmodel1, nmodel2 = get_noise_models(denoiSplitNM_fname[noise_level])\n",
    "\n",
    "ax_temp = fig.add_subplot(gs[grid_img_sz+1:int(grid_img_sz * 3/2) -1, 2:grid_img_sz])\n",
    "# ax_temp = fig.add_subplot(gs[grid_img_sz+grid_img_sz//4:grid_img_sz//4 + int(grid_img_sz * 3/2)+1, 1:1+grid_img_sz//2])\n",
    "clean_ax(ax_temp)\n",
    "plot_noise_model(40, 90, nmodel1['hist'], nmodel1['gmm'], 'cpu', ax_temp, linetxt_offset=0.2)\n",
    "\n",
    "ax_temp = fig.add_subplot(gs[int(grid_img_sz * 3/2)+2:2*grid_img_sz -1, 2:grid_img_sz])\n",
    "# ax_temp = fig.add_subplot(gs[grid_img_sz + grid_img_sz//4:grid_img_sz//4 + int(grid_img_sz * 3/2)+1, grid_img_sz//2+1:grid_img_sz])\n",
    "clean_ax(ax_temp)\n",
    "plot_noise_model(40, 90, nmodel2['hist'], nmodel2['gmm'], 'cpu', ax_temp, linetxt_offset=0.2)\n",
    "\n",
    "# ax_temp = fig.add_subplot(gs[grid_img_sz:int(grid_img_sz * 3/2), :grid_img_sz])\n",
    "# plot_noise_model(45, 100, nmodel1['hist'], nmodel1['gmm'], 'cpu', ax_temp)\n",
    "\n",
    "# manually setting legends\n",
    "import matplotlib.lines as mlines\n",
    "line_ch1 = mlines.Line2D([0, 1], [0, 1], color=color_ch_list[0], linestyle='-', label='$C_1$')\n",
    "line_ch2 = mlines.Line2D([0, 1], [0, 1], color=color_ch_list[1], linestyle='-', label='$C_2$')\n",
    "line_pred = mlines.Line2D([0, 1], [0, 1], color=color_pred, linestyle='-', label='Pred')\n",
    "line_noisych1 = mlines.Line2D([0, 1], [0, 1], color=color_ch_list[0], linestyle='--', label='$C^N_1$')\n",
    "line_noisych2 = mlines.Line2D([0, 1], [0, 1], color=color_ch_list[1], linestyle='--', label='$C^N_2$')\n",
    "\n",
    "legend_ch1 = legend_ch1_ax.legend(handles=[line_ch1, line_noisych1, line_pred], loc='upper right', frameon=False, labelcolor='white', \n",
    "                          prop={'size': 11})\n",
    "legend_ch2 = legend_ch2_ax.legend(handles=[line_ch2, line_noisych2, line_pred], loc='upper right', frameon=False, labelcolor='white',\n",
    "                            prop={'size': 11})\n",
    "# legend = legend_ax.legend(handles=[line_ch1, line_noisych1, line_ch2, line_noisych2, line_pred], loc='upper left', frameon=False, labelcolor='white', \n",
    "#                           prop={'size': 11})\n",
    "\n",
    "fpath = get_output_fpath(noise_level)\n",
    "# plt.savefig(fpath, dpi=200, bbox_inches='tight')\n",
    "print(f'Saved to {fpath}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(usplit_data[noise_levels[0]][0,:500,:500, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}