example_med_unet_cluster.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example U-Net organ clustering"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports + model initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-26 10:31:30.315745: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-05-26 10:31:30.413464: I tensorflow/core/util/util.cc:169] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2023-05-26 10:31:30.437905: E tensorflow/stream_executor/cuda/cuda_blas.cc:2981] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "/etinfo/users2/qlanglois/.local/lib/python3.9/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.23.4\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n",
      "/etinfo/users2/qlanglois/.local/lib/python3.9/site-packages/tensorflow_addons/utils/tfa_eol_msg.py:23: UserWarning: \n",
      "\n",
      "TensorFlow Addons (TFA) has ended development and introduction of new features.\n",
      "TFA has entered a minimal maintenance and release mode until a planned end of life in May 2024.\n",
      "Please modify downstream libraries to take dependencies from other repositories in our TensorFlow community (e.g. Keras, Keras-CV, and Keras-NLP). \n",
      "\n",
      "For more information see: https://github.com/tensorflow/addons/issues/2807 \n",
      "\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensorflow version : 2.10.0\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import importlib\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "\n",
    "from sklearn.utils import shuffle as sklearn_shuffle\n",
    "\n",
    "from loggers import set_level\n",
    "from utils import plot, plot_multiple, set_display_options\n",
    "from datasets import get_dataset, prepare_dataset, test_dataset_time, train_test_split\n",
    "from models.detection import med_unet_clusterer\n",
    "from models import get_pretrained\n",
    "from utils.med_utils import *\n",
    "from models.model_utils import is_model_name\n",
    "\n",
    "\n",
    "set_display_options()\n",
    "\n",
    "tf.config.set_visible_devices(tf.config.list_physical_devices('GPU')[:1], 'GPU')\n",
    "\n",
    "input_size = (None, None, 1)\n",
    "model_name = 'cluster_unet_cosine_v2'\n",
    "\n",
    "print(\"Tensorflow version : {}\".format(tf.__version__))\n",
    "#print('# GPU(s) : {}'.format(len(tf.config.list_logical_devices('GPU'))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Model creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializing model with kwargs : {'model': {'architecture_name': 'totalsegmentator', 'input_shape': (None, None, None, 1), 'output_dim': 32, 'final_activation': None, 'normalize_output': True, 'norm_type': 'instance'}}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-26 10:31:10.692884: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Please cite the following paper when using nnUNet:\n",
      "\n",
      "Isensee, F., Jaeger, P.F., Kohl, S.A.A. et al. \"nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation.\" Nat Methods (2020). https://doi.org/10.1038/s41592-020-01008-z\n",
      "\n",
      "\n",
      "If you have questions or suggestions, feel free to open an issue at https://github.com/MIC-DKFZ/nnUNet\n",
      "\n",
      "nnUNet_raw_data_base is not defined and nnU-Net can only be used on data for which preprocessed files are already present on your system. nnU-Net cannot be used for experiment planning and preprocessing like this. If this is not intended, please read nnunet/paths.md for information on how to set this up properly.\n",
      "nnUNet_preprocessed is not defined and nnU-Net can not be used for preprocessing or training. If this is not intended, please read nnunet/pathy.md for information on how to set this up.\n",
      "RESULTS_FOLDER is not defined and nnU-Net cannot be used for training or inference. If this is not intended behavior, please read nnunet/paths.md for information on how to set this up\n",
      "folds is None so we will automatically look for output folders (not using 'all'!)\n",
      "found the following folds:  ['/etinfo/users2/qlanglois/.totalsegmentator/nnunet/results/nnUNet/3d_fullres/Task256_TotalSegmentator_3mm_1139subj/nnUNetTrainerV2_ep8000_nomirror__nnUNetPlansv2.1/fold_0']\n",
      "using the following model files:  ['/etinfo/users2/qlanglois/.totalsegmentator/nnunet/results/nnUNet/3d_fullres/Task256_TotalSegmentator_3mm_1139subj/nnUNetTrainerV2_ep8000_nomirror__nnUNetPlansv2.1/fold_0/model_final_checkpoint.model']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/etinfo/users2/qlanglois/phd_segmentation/models/weights_converter.py:101: UserWarning: Unable to determine the root based on candidates : {'conv_blocks_localization': 40, 'conv_blocks_context': 48, 'tu': 5, 'seg_outputs': 5}\n",
      "  warnings.warn('Unable to determine the root based on candidates : {}'.format(parts))\n",
      "/etinfo/users2/qlanglois/phd_segmentation/models/weights_converter.py:101: UserWarning: Unable to determine the root based on candidates : {'conv_blocks_context': 48, 'tu': 5, 'conv_blocks_localization': 40, 'seg_outputs': 1}\n",
      "  warnings.warn('Unable to determine the root based on candidates : {}'.format(parts))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weights transfered successfully !\n",
      "Initializing submodel : `model` !\n",
      "Submodel model saved in pretrained_models/cluster_unet_euclidian_v2/saving/model.json !\n",
      "Model cluster_unet_euclidian_v2 initialized successfully !\n",
      "\n",
      "========== cluster_unet_euclidian_v2 ==========\n",
      "Sub model model\n",
      "- Inputs \t: (None, None, None, None, 1)\n",
      "- Outputs \t: (None, None, None, None, 32)\n",
      "- Number of layers \t: 123\n",
      "- Number of parameters \t: 30.477 Millions\n",
      "- Model not compiled\n",
      "\n",
      "Transfer-learning from : totalsegmentator\n",
      "Already trained on 0 epochs (0 steps)\n",
      "\n",
      "- Image size : (None, None, 1)\n",
      "- Normalization style : mean\n",
      "- Voxel dims : (1.5, 1.5, 1.5)\n",
      "- # frames : variable\n",
      "- Embedding dim   : 32\n",
      "- Distance metric : euclidian\n",
      "\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    tf.config.set_visible_devices([], 'GPU')\n",
    "except:\n",
    "    print('Unable to modify the visible devices')\n",
    "\n",
    "importlib.reload(med_unet_clusterer)\n",
    "\n",
    "if 'highres' in model_name:\n",
    "    voxel_dims = (0.5, 0.5, 1.5)\n",
    "elif 'lowres' in model_name:\n",
    "    voxel_dims = (3., 3., 3.)\n",
    "else:\n",
    "    voxel_dims = (1.5, 1.5, 1.5)\n",
    "\n",
    "config = {\n",
    "    'input_size' : input_size,\n",
    "    'voxel_dims' : voxel_dims,\n",
    "    'n_frames'   : None,\n",
    "    'pad_value'  : 0,\n",
    "\n",
    "    'normalize'       : False if 'cosine' in model_name else True,\n",
    "    'embedding_dim'   : 32,\n",
    "    'distance_metric' : model_name.split('_')[2],\n",
    "    \n",
    "    'image_normalization' : 'mean',\n",
    "    \n",
    "    'norm_type' : 'instance'\n",
    "}\n",
    "\n",
    "if 'scratch' in model_name:\n",
    "    config.update({\n",
    "        # Architecture config\n",
    "        'n_stages'   : 4,\n",
    "        'n_conv_per_stage'    : 1,\n",
    "        'up_n_conv_per_stage' : lambda i: min(i, 1),\n",
    "        'filters'     : list(np.array([16, 32, 64, 128])),\n",
    "        'bnorm'       : 'never',\n",
    "        'activation'  : 'leaky',\n",
    "        'drop_rate'   : lambda i: 0. if i == 0 else 0.25,\n",
    "\n",
    "        'n_middle_stages' : 2,\n",
    "        'n_middle_conv'   : 2,\n",
    "        'middle_filters'  : 64,\n",
    "        'middle_bnorm'    : 'never',\n",
    "\n",
    "        'concat_mode'     : lambda i: 'concat' if i > 0 else None,\n",
    "    })\n",
    "else:\n",
    "    config['pretrained_name'] = 'totalsegmentator'\n",
    "\n",
    "model = med_unet_clusterer.MedUNetClusterer(\n",
    "    nom = model_name, ** config\n",
    ")\n",
    "\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = med_unet_clusterer.MedUNetClusterer.from_pretrained(\n",
    "    nom = model_name, pretrained_name = model_name[:-2]\n",
    ")\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "========== Summary of model ==========\n",
      "\n",
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " input_image (InputLayer)       [(None, None, None,  0           []                               \n",
      "                                 None, 1)]                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d (ZeroPadding3D)  (None, None, None,   0          ['input_image[0][0]']            \n",
      "                                None, 1)                                                          \n",
      "                                                                                                  \n",
      " conv_blocks_context/0/blocks/0  (None, None, None,   896        ['zero_padding3d[0][0]']         \n",
      " /conv (Conv3D)                 None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/0/blocks/0  (None, None, None,   64         ['conv_blocks_context/0/blocks/0/\n",
      " /norm (InstanceNormalization)  None, 32)                        conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu (LeakyReLU)        (None, None, None,   0           ['conv_blocks_context/0/blocks/0/\n",
      "                                None, 32)                        norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout (Dropout)              (None, None, None,   0           ['leaky_re_lu[0][0]']            \n",
      "                                None, 32)                                                         \n",
      "                                                                                                  \n",
      " zero_padding3d_1 (ZeroPadding3  (None, None, None,   0          ['dropout[0][0]']                \n",
      " D)                             None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/0/blocks/1  (None, None, None,   27680      ['zero_padding3d_1[0][0]']       \n",
      " /conv (Conv3D)                 None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/0/blocks/1  (None, None, None,   64         ['conv_blocks_context/0/blocks/1/\n",
      " /norm (InstanceNormalization)  None, 32)                        conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_1 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/0/blocks/1/\n",
      "                                None, 32)                        norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_1 (Dropout)            (None, None, None,   0           ['leaky_re_lu_1[0][0]']          \n",
      "                                None, 32)                                                         \n",
      "                                                                                                  \n",
      " zero_padding3d_2 (ZeroPadding3  (None, None, None,   0          ['dropout_1[0][0]']              \n",
      " D)                             None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/1/blocks/0  (None, None, None,   55360      ['zero_padding3d_2[0][0]']       \n",
      " /conv (Conv3D)                 None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/1/blocks/0  (None, None, None,   128        ['conv_blocks_context/1/blocks/0/\n",
      " /norm (InstanceNormalization)  None, 64)                        conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_2 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/1/blocks/0/\n",
      "                                None, 64)                        norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_2 (Dropout)            (None, None, None,   0           ['leaky_re_lu_2[0][0]']          \n",
      "                                None, 64)                                                         \n",
      "                                                                                                  \n",
      " zero_padding3d_3 (ZeroPadding3  (None, None, None,   0          ['dropout_2[0][0]']              \n",
      " D)                             None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/1/blocks/1  (None, None, None,   110656     ['zero_padding3d_3[0][0]']       \n",
      " /conv (Conv3D)                 None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/1/blocks/1  (None, None, None,   128        ['conv_blocks_context/1/blocks/1/\n",
      " /norm (InstanceNormalization)  None, 64)                        conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_3 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/1/blocks/1/\n",
      "                                None, 64)                        norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_3 (Dropout)            (None, None, None,   0           ['leaky_re_lu_3[0][0]']          \n",
      "                                None, 64)                                                         \n",
      "                                                                                                  \n",
      " zero_padding3d_4 (ZeroPadding3  (None, None, None,   0          ['dropout_3[0][0]']              \n",
      " D)                             None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_context/2/blocks/0  (None, None, None,   221312     ['zero_padding3d_4[0][0]']       \n",
      " /conv (Conv3D)                 None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/2/blocks/0  (None, None, None,   256        ['conv_blocks_context/2/blocks/0/\n",
      " /norm (InstanceNormalization)  None, 128)                       conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_4 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/2/blocks/0/\n",
      "                                None, 128)                       norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_4 (Dropout)            (None, None, None,   0           ['leaky_re_lu_4[0][0]']          \n",
      "                                None, 128)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_5 (ZeroPadding3  (None, None, None,   0          ['dropout_4[0][0]']              \n",
      " D)                             None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/2/blocks/1  (None, None, None,   442496     ['zero_padding3d_5[0][0]']       \n",
      " /conv (Conv3D)                 None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/2/blocks/1  (None, None, None,   256        ['conv_blocks_context/2/blocks/1/\n",
      " /norm (InstanceNormalization)  None, 128)                       conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_5 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/2/blocks/1/\n",
      "                                None, 128)                       norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_5 (Dropout)            (None, None, None,   0           ['leaky_re_lu_5[0][0]']          \n",
      "                                None, 128)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_6 (ZeroPadding3  (None, None, None,   0          ['dropout_5[0][0]']              \n",
      " D)                             None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/3/blocks/0  (None, None, None,   884992     ['zero_padding3d_6[0][0]']       \n",
      " /conv (Conv3D)                 None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/3/blocks/0  (None, None, None,   512        ['conv_blocks_context/3/blocks/0/\n",
      " /norm (InstanceNormalization)  None, 256)                       conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_6 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/3/blocks/0/\n",
      "                                None, 256)                       norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_6 (Dropout)            (None, None, None,   0           ['leaky_re_lu_6[0][0]']          \n",
      "                                None, 256)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_7 (ZeroPadding3  (None, None, None,   0          ['dropout_6[0][0]']              \n",
      " D)                             None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/3/blocks/1  (None, None, None,   1769728    ['zero_padding3d_7[0][0]']       \n",
      " /conv (Conv3D)                 None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/3/blocks/1  (None, None, None,   512        ['conv_blocks_context/3/blocks/1/\n",
      " /norm (InstanceNormalization)  None, 256)                       conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_7 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/3/blocks/1/\n",
      "                                None, 256)                       norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_7 (Dropout)            (None, None, None,   0           ['leaky_re_lu_7[0][0]']          \n",
      "                                None, 256)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_8 (ZeroPadding3  (None, None, None,   0          ['dropout_7[0][0]']              \n",
      " D)                             None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/4/blocks/0  (None, None, None,   2212160    ['zero_padding3d_8[0][0]']       \n",
      " /conv (Conv3D)                 None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/4/blocks/0  (None, None, None,   640        ['conv_blocks_context/4/blocks/0/\n",
      " /norm (InstanceNormalization)  None, 320)                       conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_8 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/4/blocks/0/\n",
      "                                None, 320)                       norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_8 (Dropout)            (None, None, None,   0           ['leaky_re_lu_8[0][0]']          \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_9 (ZeroPadding3  (None, None, None,   0          ['dropout_8[0][0]']              \n",
      " D)                             None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/4/blocks/1  (None, None, None,   2765120    ['zero_padding3d_9[0][0]']       \n",
      " /conv (Conv3D)                 None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/4/blocks/1  (None, None, None,   640        ['conv_blocks_context/4/blocks/1/\n",
      " /norm (InstanceNormalization)  None, 320)                       conv[0][0]']                     \n",
      "                                                                                                  \n",
      " leaky_re_lu_9 (LeakyReLU)      (None, None, None,   0           ['conv_blocks_context/4/blocks/1/\n",
      "                                None, 320)                       norm[0][0]']                     \n",
      "                                                                                                  \n",
      " dropout_9 (Dropout)            (None, None, None,   0           ['leaky_re_lu_9[0][0]']          \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_10 (ZeroPadding  (None, None, None,   0          ['dropout_9[0][0]']              \n",
      " 3D)                            None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/5/0/blocks  (None, None, None,   2765120    ['zero_padding3d_10[0][0]']      \n",
      " /0/conv (Conv3D)               None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/5/0/blocks  (None, None, None,   640        ['conv_blocks_context/5/0/blocks/\n",
      " /0/norm (InstanceNormalization  None, 320)                      0/conv[0][0]']                   \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " leaky_re_lu_10 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_context/5/0/blocks/\n",
      "                                None, 320)                       0/norm[0][0]']                   \n",
      "                                                                                                  \n",
      " dropout_10 (Dropout)           (None, None, None,   0           ['leaky_re_lu_10[0][0]']         \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_11 (ZeroPadding  (None, None, None,   0          ['dropout_10[0][0]']             \n",
      " 3D)                            None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/5/1/blocks  (None, None, None,   2765120    ['zero_padding3d_11[0][0]']      \n",
      " /1/conv (Conv3D)               None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_context/5/1/blocks  (None, None, None,   640        ['conv_blocks_context/5/1/blocks/\n",
      " /1/norm (InstanceNormalization  None, 320)                      1/conv[0][0]']                   \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " leaky_re_lu_11 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_context/5/1/blocks/\n",
      "                                None, 320)                       1/norm[0][0]']                   \n",
      "                                                                                                  \n",
      " dropout_11 (Dropout)           (None, None, None,   0           ['leaky_re_lu_11[0][0]']         \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " tu/0 (Conv3DTranspose)         (None, None, None,   102400      ['dropout_11[0][0]']             \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " concatenate (Concatenate)      (None, None, None,   0           ['tu/0[0][0]',                   \n",
      "                                None, 640)                        'dropout_9[0][0]']              \n",
      "                                                                                                  \n",
      " zero_padding3d_12 (ZeroPadding  (None, None, None,   0          ['concatenate[0][0]']            \n",
      " 3D)                            None, 640)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/0/0/b  (None, None, None,   5529920    ['zero_padding3d_12[0][0]']      \n",
      " locks/0/conv (Conv3D)          None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/0/0/b  (None, None, None,   640        ['conv_blocks_localization/0/0/bl\n",
      " locks/0/norm (InstanceNormaliz  None, 320)                      ocks/0/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_12 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/0/0/bl\n",
      "                                None, 320)                       ocks/0/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_12 (Dropout)           (None, None, None,   0           ['leaky_re_lu_12[0][0]']         \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_13 (ZeroPadding  (None, None, None,   0          ['dropout_12[0][0]']             \n",
      " 3D)                            None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/0/1/b  (None, None, None,   2765120    ['zero_padding3d_13[0][0]']      \n",
      " locks/1/conv (Conv3D)          None, 320)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/0/1/b  (None, None, None,   640        ['conv_blocks_localization/0/1/bl\n",
      " locks/1/norm (InstanceNormaliz  None, 320)                      ocks/1/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_13 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/0/1/bl\n",
      "                                None, 320)                       ocks/1/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_13 (Dropout)           (None, None, None,   0           ['leaky_re_lu_13[0][0]']         \n",
      "                                None, 320)                                                        \n",
      "                                                                                                  \n",
      " tu/1 (Conv3DTranspose)         (None, None, None,   655360      ['dropout_13[0][0]']             \n",
      "                                None, 256)                                                        \n",
      "                                                                                                  \n",
      " concatenate_1 (Concatenate)    (None, None, None,   0           ['tu/1[0][0]',                   \n",
      "                                None, 512)                        'dropout_7[0][0]']              \n",
      "                                                                                                  \n",
      " zero_padding3d_14 (ZeroPadding  (None, None, None,   0          ['concatenate_1[0][0]']          \n",
      " 3D)                            None, 512)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/1/0/b  (None, None, None,   3539200    ['zero_padding3d_14[0][0]']      \n",
      " locks/0/conv (Conv3D)          None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/1/0/b  (None, None, None,   512        ['conv_blocks_localization/1/0/bl\n",
      " locks/0/norm (InstanceNormaliz  None, 256)                      ocks/0/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_14 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/1/0/bl\n",
      "                                None, 256)                       ocks/0/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_14 (Dropout)           (None, None, None,   0           ['leaky_re_lu_14[0][0]']         \n",
      "                                None, 256)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_15 (ZeroPadding  (None, None, None,   0          ['dropout_14[0][0]']             \n",
      " 3D)                            None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/1/1/b  (None, None, None,   1769728    ['zero_padding3d_15[0][0]']      \n",
      " locks/1/conv (Conv3D)          None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/1/1/b  (None, None, None,   512        ['conv_blocks_localization/1/1/bl\n",
      " locks/1/norm (InstanceNormaliz  None, 256)                      ocks/1/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_15 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/1/1/bl\n",
      "                                None, 256)                       ocks/1/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_15 (Dropout)           (None, None, None,   0           ['leaky_re_lu_15[0][0]']         \n",
      "                                None, 256)                                                        \n",
      "                                                                                                  \n",
      " tu/2 (Conv3DTranspose)         (None, None, None,   262144      ['dropout_15[0][0]']             \n",
      "                                None, 128)                                                        \n",
      "                                                                                                  \n",
      " concatenate_2 (Concatenate)    (None, None, None,   0           ['tu/2[0][0]',                   \n",
      "                                None, 256)                        'dropout_5[0][0]']              \n",
      "                                                                                                  \n",
      " zero_padding3d_16 (ZeroPadding  (None, None, None,   0          ['concatenate_2[0][0]']          \n",
      " 3D)                            None, 256)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/2/0/b  (None, None, None,   884864     ['zero_padding3d_16[0][0]']      \n",
      " locks/0/conv (Conv3D)          None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/2/0/b  (None, None, None,   256        ['conv_blocks_localization/2/0/bl\n",
      " locks/0/norm (InstanceNormaliz  None, 128)                      ocks/0/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_16 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/2/0/bl\n",
      "                                None, 128)                       ocks/0/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_16 (Dropout)           (None, None, None,   0           ['leaky_re_lu_16[0][0]']         \n",
      "                                None, 128)                                                        \n",
      "                                                                                                  \n",
      " zero_padding3d_17 (ZeroPadding  (None, None, None,   0          ['dropout_16[0][0]']             \n",
      " 3D)                            None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/2/1/b  (None, None, None,   442496     ['zero_padding3d_17[0][0]']      \n",
      " locks/1/conv (Conv3D)          None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/2/1/b  (None, None, None,   256        ['conv_blocks_localization/2/1/bl\n",
      " locks/1/norm (InstanceNormaliz  None, 128)                      ocks/1/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_17 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/2/1/bl\n",
      "                                None, 128)                       ocks/1/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_17 (Dropout)           (None, None, None,   0           ['leaky_re_lu_17[0][0]']         \n",
      "                                None, 128)                                                        \n",
      "                                                                                                  \n",
      " tu/3 (Conv3DTranspose)         (None, None, None,   65536       ['dropout_17[0][0]']             \n",
      "                                None, 64)                                                         \n",
      "                                                                                                  \n",
      " concatenate_3 (Concatenate)    (None, None, None,   0           ['tu/3[0][0]',                   \n",
      "                                None, 128)                        'dropout_3[0][0]']              \n",
      "                                                                                                  \n",
      " zero_padding3d_18 (ZeroPadding  (None, None, None,   0          ['concatenate_3[0][0]']          \n",
      " 3D)                            None, 128)                                                        \n",
      "                                                                                                  \n",
      " conv_blocks_localization/3/0/b  (None, None, None,   221248     ['zero_padding3d_18[0][0]']      \n",
      " locks/0/conv (Conv3D)          None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/3/0/b  (None, None, None,   128        ['conv_blocks_localization/3/0/bl\n",
      " locks/0/norm (InstanceNormaliz  None, 64)                       ocks/0/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_18 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/3/0/bl\n",
      "                                None, 64)                        ocks/0/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_18 (Dropout)           (None, None, None,   0           ['leaky_re_lu_18[0][0]']         \n",
      "                                None, 64)                                                         \n",
      "                                                                                                  \n",
      " zero_padding3d_19 (ZeroPadding  (None, None, None,   0          ['dropout_18[0][0]']             \n",
      " 3D)                            None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/3/1/b  (None, None, None,   110656     ['zero_padding3d_19[0][0]']      \n",
      " locks/1/conv (Conv3D)          None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/3/1/b  (None, None, None,   128        ['conv_blocks_localization/3/1/bl\n",
      " locks/1/norm (InstanceNormaliz  None, 64)                       ocks/1/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_19 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/3/1/bl\n",
      "                                None, 64)                        ocks/1/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_19 (Dropout)           (None, None, None,   0           ['leaky_re_lu_19[0][0]']         \n",
      "                                None, 64)                                                         \n",
      "                                                                                                  \n",
      " tu/4 (Conv3DTranspose)         (None, None, None,   16384       ['dropout_19[0][0]']             \n",
      "                                None, 32)                                                         \n",
      "                                                                                                  \n",
      " concatenate_4 (Concatenate)    (None, None, None,   0           ['tu/4[0][0]',                   \n",
      "                                None, 64)                         'dropout_1[0][0]']              \n",
      "                                                                                                  \n",
      " zero_padding3d_20 (ZeroPadding  (None, None, None,   0          ['concatenate_4[0][0]']          \n",
      " 3D)                            None, 64)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/4/0/b  (None, None, None,   55328      ['zero_padding3d_20[0][0]']      \n",
      " locks/0/conv (Conv3D)          None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/4/0/b  (None, None, None,   64         ['conv_blocks_localization/4/0/bl\n",
      " locks/0/norm (InstanceNormaliz  None, 32)                       ocks/0/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_20 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/4/0/bl\n",
      "                                None, 32)                        ocks/0/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_20 (Dropout)           (None, None, None,   0           ['leaky_re_lu_20[0][0]']         \n",
      "                                None, 32)                                                         \n",
      "                                                                                                  \n",
      " zero_padding3d_21 (ZeroPadding  (None, None, None,   0          ['dropout_20[0][0]']             \n",
      " 3D)                            None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/4/1/b  (None, None, None,   27680      ['zero_padding3d_21[0][0]']      \n",
      " locks/1/conv (Conv3D)          None, 32)                                                         \n",
      "                                                                                                  \n",
      " conv_blocks_localization/4/1/b  (None, None, None,   64         ['conv_blocks_localization/4/1/bl\n",
      " locks/1/norm (InstanceNormaliz  None, 32)                       ocks/1/conv[0][0]']              \n",
      " ation)                                                                                           \n",
      "                                                                                                  \n",
      " leaky_re_lu_21 (LeakyReLU)     (None, None, None,   0           ['conv_blocks_localization/4/1/bl\n",
      "                                None, 32)                        ocks/1/norm[0][0]']              \n",
      "                                                                                                  \n",
      " dropout_21 (Dropout)           (None, None, None,   0           ['leaky_re_lu_21[0][0]']         \n",
      "                                None, 32)                                                         \n",
      "                                                                                                  \n",
      " seg_outputs/4 (Conv3D)         (None, None, None,   1024        ['dropout_21[0][0]']             \n",
      "                                None, 32)                                                         \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 30,477,408\n",
      "Trainable params: 30,477,408\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n",
      "\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model instanciation + dataset loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model restoration...\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n",
      "The given value for groups will be overwritten.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-26 10:31:34.316543: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-05-26 10:31:34.701018: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1616] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14783 MB memory:  -> device: 0, name: Quadro RTX 5000, pci bus id: 0000:17:00.0, compute capability: 7.5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializing submodel : `model` !\n",
      "Successfully restored model from pretrained_models/cluster_unet_euclidian_v2/saving/model.json !\n",
      "Model cluster_unet_euclidian_v2 initialized successfully !\n",
      "Optimizer 'model_optimizer' initilized successfully !\n",
      "Submodel model compiled !\n",
      "  Loss : {'reduction': 'none', 'name': 'ge2e_seg_loss', 'mode': 'softmax', 'init_w': -1.0, 'init_b': 0.0, 'loss_averaging': 'micro', 'distance_metric': 'euclidian', 'background_mode': 'concat'}\n",
      "  Optimizer : {'name': 'Adam', 'learning_rate': {'class_name': 'DivideByStep', 'config': {'factor': <tf.Tensor: shape=(), dtype=float32, numpy=1.0>, 'minval': <tf.Tensor: shape=(), dtype=float32, numpy=1e-04>, 'maxval': <tf.Tensor: shape=(), dtype=float32, numpy=0.001>}}, 'decay': 0.0, 'beta_1': 0.9, 'beta_2': 0.999, 'epsilon': 1e-07, 'amsgrad': False}\n",
      "  Metrics : []\n",
      "\n",
      "========== cluster_unet_euclidian_v2 ==========\n",
      "Sub model model\n",
      "- Inputs \t: (None, None, None, None, 1)\n",
      "- Outputs \t: (None, None, None, None, 32)\n",
      "- Number of layers \t: 123\n",
      "- Number of parameters \t: 30.477 Millions\n",
      "- Optimizer \t: {'name': 'Adam', 'learning_rate': {'class_name': 'DivideByStep', 'config': {'factor': <tf.Tensor: shape=(), dtype=float32, numpy=1.0>, 'minval': <tf.Tensor: shape=(), dtype=float32, numpy=1e-04>, 'maxval': <tf.Tensor: shape=(), dtype=float32, numpy=0.001>}}, 'decay': 0.0, 'beta_1': 0.9, 'beta_2': 0.999, 'epsilon': 1e-07, 'amsgrad': False}\n",
      "- Loss \t : {'reduction': 'none', 'name': 'ge2e_seg_loss', 'mode': 'softmax', 'init_w': -1.0, 'init_b': 0.0, 'loss_averaging': 'micro', 'distance_metric': 'euclidian', 'background_mode': 'concat'}\n",
      "- Metrics\t : []\n",
      "\n",
      "Transfer-learning from : totalsegmentator\n",
      "Already trained on 0 epochs (0 steps)\n",
      "\n",
      "- Image size : (None, None, 1)\n",
      "- Normalization style : mean\n",
      "- Voxel dims : [1.5, 1.5, 1.5]\n",
      "- # frames : variable\n",
      "- Embedding dim   : 32\n",
      "- Distance metric : euclidian\n",
      "\n"
     ]
    }
   ],
   "source": [
    "model = get_pretrained(model_name)\n",
    "\n",
    "if 'scratch' in model_name:\n",
    "    lr = {'name' : 'DivideByStep', 'maxval' : 1e-2,'minval' : 1e-4}\n",
    "else:\n",
    "    lr = {'name' : 'DivideByStep', 'maxval' : 1e-3,'minval' : 1e-4}\n",
    "\n",
    "loss_config = {\n",
    "    'init_w' : -1 if model.distance_metric == 'euclidian' else 1.,\n",
    "    'loss_averaging'  : 'micro',\n",
    "    'background_mode' : 'concat'\n",
    "}\n",
    "model.compile(\n",
    "    optimizer = 'adam', optimizer_config = {'lr' : lr}, loss_config = loss_config\n",
    ")\n",
    "\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading dataset total_segmentator...\n",
      "Dataset length (53 data skipped, 1202 ids) :\n",
      "  Train size : 17969 (1081 ids)\n",
      "  Valid size : 2066 (121 ids)\n",
      "# ids in valid that are also in train : 0\n"
     ]
    }
   ],
   "source": [
    "dataset_name = 'total_segmentator'\n",
    "dataset = get_dataset(dataset_name, slice_step = 16, slice_size = 32)\n",
    "\n",
    "keep_mask = dataset['segmentation'].apply(lambda f: f.endswith('.npz'))\n",
    "dataset   = dataset[keep_mask]\n",
    "\n",
    "if isinstance(dataset, dict):\n",
    "    train, valid = dataset['train'], dataset['valid']\n",
    "else:\n",
    "    train, valid = train_test_split(\n",
    "        dataset, train_size = 0.9, shuffle = True, random_state = 10, split_by_unique = True, min_occurence = 0\n",
    "    )\n",
    "    train = sklearn_shuffle(train, random_state = model.epochs)\n",
    "\n",
    "print('Dataset length ({} data skipped, {} ids) :\\n  Train size : {} ({} ids)\\n  Valid size : {} ({} ids)'.format(\n",
    "    len(keep_mask) - np.sum(keep_mask.values), len(dataset['id'].unique()), \n",
    "    len(train), len(train['id'].unique()), len(valid), len(valid['id'].unique())\n",
    "))\n",
    "print('# ids in valid that are also in train : {}'.format(len([id_i for id_i in valid['id'].unique() if id_i in train['id'].values])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>subject_id</th>\n",
       "      <th>thickness</th>\n",
       "      <th>images</th>\n",
       "      <th>segmentation</th>\n",
       "      <th>label</th>\n",
       "      <th>start_frame</th>\n",
       "      <th>end_frame</th>\n",
       "      <th>id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>128</th>\n",
       "      <td>s0012</td>\n",
       "      <td>-1</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s0012/ct.nii.gz</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s0012/masks.npz</td>\n",
       "      <td>[adrenal_gland_left, adrenal_gland_right, aort...</td>\n",
       "      <td>0</td>\n",
       "      <td>32</td>\n",
       "      <td>s0012</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>129</th>\n",
       "      <td>s0012</td>\n",
       "      <td>-1</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s0012/ct.nii.gz</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s0012/masks.npz</td>\n",
       "      <td>[adrenal_gland_left, adrenal_gland_right, aort...</td>\n",
       "      <td>16</td>\n",
       "      <td>48</td>\n",
       "      <td>s0012</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>130</th>\n",
       "      <td>s0012</td>\n",
       "      <td>-1</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s0012/ct.nii.gz</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s0012/masks.npz</td>\n",
       "      <td>[adrenal_gland_left, adrenal_gland_right, aort...</td>\n",
       "      <td>32</td>\n",
       "      <td>64</td>\n",
       "      <td>s0012</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>131</th>\n",
       "      <td>s0012</td>\n",
       "      <td>-1</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s0012/ct.nii.gz</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s0012/masks.npz</td>\n",
       "      <td>[adrenal_gland_left, adrenal_gland_right, aort...</td>\n",
       "      <td>48</td>\n",
       "      <td>80</td>\n",
       "      <td>s0012</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>132</th>\n",
       "      <td>s0012</td>\n",
       "      <td>-1</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s0012/ct.nii.gz</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s0012/masks.npz</td>\n",
       "      <td>[adrenal_gland_left, adrenal_gland_right, aort...</td>\n",
       "      <td>64</td>\n",
       "      <td>96</td>\n",
       "      <td>s0012</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20083</th>\n",
       "      <td>s1403</td>\n",
       "      <td>-1</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s1403/ct.nii.gz</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s1403/masks.npz</td>\n",
       "      <td>[adrenal_gland_left, adrenal_gland_right, aort...</td>\n",
       "      <td>240</td>\n",
       "      <td>272</td>\n",
       "      <td>s1403</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20084</th>\n",
       "      <td>s1403</td>\n",
       "      <td>-1</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s1403/ct.nii.gz</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s1403/masks.npz</td>\n",
       "      <td>[adrenal_gland_left, adrenal_gland_right, aort...</td>\n",
       "      <td>256</td>\n",
       "      <td>288</td>\n",
       "      <td>s1403</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20085</th>\n",
       "      <td>s1403</td>\n",
       "      <td>-1</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s1403/ct.nii.gz</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s1403/masks.npz</td>\n",
       "      <td>[adrenal_gland_left, adrenal_gland_right, aort...</td>\n",
       "      <td>272</td>\n",
       "      <td>304</td>\n",
       "      <td>s1403</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20086</th>\n",
       "      <td>s1403</td>\n",
       "      <td>-1</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s1403/ct.nii.gz</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s1403/masks.npz</td>\n",
       "      <td>[adrenal_gland_left, adrenal_gland_right, aort...</td>\n",
       "      <td>288</td>\n",
       "      <td>320</td>\n",
       "      <td>s1403</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20087</th>\n",
       "      <td>s1403</td>\n",
       "      <td>-1</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s1403/ct.nii.gz</td>\n",
       "      <td>/storage/Totalsegmentator_dataset/s1403/masks.npz</td>\n",
       "      <td>[adrenal_gland_left, adrenal_gland_right, aort...</td>\n",
       "      <td>304</td>\n",
       "      <td>336</td>\n",
       "      <td>s1403</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>2066 rows × 8 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      subject_id  thickness                                             images  \\\n",
       "128        s0012         -1  /storage/Totalsegmentator_dataset/s0012/ct.nii.gz   \n",
       "129        s0012         -1  /storage/Totalsegmentator_dataset/s0012/ct.nii.gz   \n",
       "130        s0012         -1  /storage/Totalsegmentator_dataset/s0012/ct.nii.gz   \n",
       "131        s0012         -1  /storage/Totalsegmentator_dataset/s0012/ct.nii.gz   \n",
       "132        s0012         -1  /storage/Totalsegmentator_dataset/s0012/ct.nii.gz   \n",
       "...          ...        ...                                                ...   \n",
       "20083      s1403         -1  /storage/Totalsegmentator_dataset/s1403/ct.nii.gz   \n",
       "20084      s1403         -1  /storage/Totalsegmentator_dataset/s1403/ct.nii.gz   \n",
       "20085      s1403         -1  /storage/Totalsegmentator_dataset/s1403/ct.nii.gz   \n",
       "20086      s1403         -1  /storage/Totalsegmentator_dataset/s1403/ct.nii.gz   \n",
       "20087      s1403         -1  /storage/Totalsegmentator_dataset/s1403/ct.nii.gz   \n",
       "\n",
       "                                            segmentation                                              label  start_frame  \\\n",
       "128    /storage/Totalsegmentator_dataset/s0012/masks.npz  [adrenal_gland_left, adrenal_gland_right, aort...            0   \n",
       "129    /storage/Totalsegmentator_dataset/s0012/masks.npz  [adrenal_gland_left, adrenal_gland_right, aort...           16   \n",
       "130    /storage/Totalsegmentator_dataset/s0012/masks.npz  [adrenal_gland_left, adrenal_gland_right, aort...           32   \n",
       "131    /storage/Totalsegmentator_dataset/s0012/masks.npz  [adrenal_gland_left, adrenal_gland_right, aort...           48   \n",
       "132    /storage/Totalsegmentator_dataset/s0012/masks.npz  [adrenal_gland_left, adrenal_gland_right, aort...           64   \n",
       "...                                                  ...                                                ...          ...   \n",
       "20083  /storage/Totalsegmentator_dataset/s1403/masks.npz  [adrenal_gland_left, adrenal_gland_right, aort...          240   \n",
       "20084  /storage/Totalsegmentator_dataset/s1403/masks.npz  [adrenal_gland_left, adrenal_gland_right, aort...          256   \n",
       "20085  /storage/Totalsegmentator_dataset/s1403/masks.npz  [adrenal_gland_left, adrenal_gland_right, aort...          272   \n",
       "20086  /storage/Totalsegmentator_dataset/s1403/masks.npz  [adrenal_gland_left, adrenal_gland_right, aort...          288   \n",
       "20087  /storage/Totalsegmentator_dataset/s1403/masks.npz  [adrenal_gland_left, adrenal_gland_right, aort...          304   \n",
       "\n",
       "       end_frame     id  \n",
       "128           32  s0012  \n",
       "129           48  s0012  \n",
       "130           64  s0012  \n",
       "131           80  s0012  \n",
       "132           96  s0012  \n",
       "...          ...    ...  \n",
       "20083        272  s1403  \n",
       "20084        288  s1403  \n",
       "20085        304  s1403  \n",
       "20086        320  s1403  \n",
       "20087        336  s1403  \n",
       "\n",
       "[2066 rows x 8 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "valid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Training + history analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training config :\n",
      "HParams :\n",
      "- augment_prct\t: 0.25\n",
      "- augment_methods\t: ['noise']\n",
      "- max_size\t: (192, 192)\n",
      "- max_frames\t: 32\n",
      "- crop_mode\t: ['random_center_80', 'random_center_80', 'random']\n",
      "- skip_empty_frames\t: False\n",
      "- skip_empty_labels\t: True\n",
      "- batch_size\t: 1\n",
      "- train_batch_size\t: None\n",
      "- valid_batch_size\t: None\n",
      "- test_batch_size\t: 1\n",
      "- shuffle_size\t: 0\n",
      "- epochs\t: 1\n",
      "- verbose\t: 1\n",
      "- train_times\t: 1\n",
      "- valid_times\t: 1\n",
      "- train_size\t: None\n",
      "- valid_size\t: None\n",
      "- test_size\t: 4\n",
      "- pred_step\t: -1\n",
      "\n",
      "Running on 1 GPU\n",
      "\n",
      "Epoch 1 / 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-26 10:32:02.168117: W tensorflow/core/common_runtime/forward_type_inference.cc:332] Type inference failed. This indicates an invalid graph that escaped type checking. Error message: INVALID_ARGUMENT: expected compatible input types, but input 1:\n",
      "type_id: TFT_OPTIONAL\n",
      "args {\n",
      "  type_id: TFT_PRODUCT\n",
      "  args {\n",
      "    type_id: TFT_TENSOR\n",
      "    args {\n",
      "      type_id: TFT_INT32\n",
      "    }\n",
      "  }\n",
      "}\n",
      " is neither a subtype nor a supertype of the combined inputs preceding it:\n",
      "type_id: TFT_OPTIONAL\n",
      "args {\n",
      "  type_id: TFT_PRODUCT\n",
      "  args {\n",
      "    type_id: TFT_TENSOR\n",
      "    args {\n",
      "      type_id: TFT_FLOAT\n",
      "    }\n",
      "  }\n",
      "}\n",
      "\n",
      "\twhile inferring type of node 'ge2e_seg_loss/cond/else/_621/ge2e_seg_loss/cond/cond_2/output/_1908'\n",
      "2023-05-26 10:32:04.250383: I tensorflow/stream_executor/cuda/cuda_dnn.cc:384] Loaded cuDNN version 8204\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     21/Unknown - 42s 977ms/step - loss: nan - foreground_loss: nan - background_loss: nan         "
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "NaN loss at batch : 20\n  Logs : {'loss': nan, 'foreground_loss': nan, 'background_loss': nan}",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [5], line 29\u001b[0m\n\u001b[1;32m     26\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     27\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m l \u001b[38;5;129;01min\u001b[39;00m model\u001b[38;5;241m.\u001b[39mlayers: l\u001b[38;5;241m.\u001b[39mtrainable \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[0;32m---> 29\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     30\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtrain\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalidation_data\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mvalid\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mepochs\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mepochs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     31\u001b[0m \n\u001b[1;32m     32\u001b[0m \u001b[43m    \u001b[49m\u001b[43maugment_prct\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43maugment_prct\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshuffle_size\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mshuffle_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     33\u001b[0m \u001b[43m    \u001b[49m\u001b[43mis_rectangular\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcache\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     34\u001b[0m \n\u001b[1;32m     35\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmax_size\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mmax_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmax_frames\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mmax_frames\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcrop_mode\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mcrop_mode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mrun_eagerly\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\n\u001b[1;32m     36\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/phd_segmentation/models/interfaces/base_model.py:1118\u001b[0m, in \u001b[0;36mBaseModel.train\u001b[0;34m(self, additional_infos, verbose, eval_epoch, verbose_step, tqdm, strategy, run_eagerly, train_step, eval_step, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1113\u001b[0m metrics \u001b[38;5;241m=\u001b[39m train_function(batch)\n\u001b[1;32m   1114\u001b[0m metrics \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m   1115\u001b[0m     n : m \u001b[38;5;28;01mfor\u001b[39;00m n, m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcompiled_metrics\u001b[38;5;241m.\u001b[39mmetric_names, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcompiled_metrics\u001b[38;5;241m.\u001b[39mresult())\n\u001b[1;32m   1116\u001b[0m }\n\u001b[0;32m-> 1118\u001b[0m \u001b[43mcallbacks\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mon_train_batch_end\u001b[49m\u001b[43m(\u001b[49m\u001b[43mi\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlogs\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mmetrics\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1120\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mupdate_train_config(\n\u001b[1;32m   1121\u001b[0m     step    \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcurrent_step \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39m,\n\u001b[1;32m   1122\u001b[0m     epoch   \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcurrent_epoch\n\u001b[1;32m   1123\u001b[0m )\n\u001b[1;32m   1124\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstop_training: \u001b[38;5;28;01mbreak\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/tensorflow/python/keras/callbacks.py:456\u001b[0m, in \u001b[0;36mCallbackList.on_train_batch_end\u001b[0;34m(self, batch, logs)\u001b[0m\n\u001b[1;32m    449\u001b[0m \u001b[38;5;124;03m\"\"\"Calls the `on_train_batch_end` methods of its callbacks.\u001b[39;00m\n\u001b[1;32m    450\u001b[0m \n\u001b[1;32m    451\u001b[0m \u001b[38;5;124;03mArgs:\u001b[39;00m\n\u001b[1;32m    452\u001b[0m \u001b[38;5;124;03m    batch: Integer, index of batch within the current epoch.\u001b[39;00m\n\u001b[1;32m    453\u001b[0m \u001b[38;5;124;03m    logs: Dict. Aggregated metric results up until this batch.\u001b[39;00m\n\u001b[1;32m    454\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    455\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_should_call_train_batch_hooks:\n\u001b[0;32m--> 456\u001b[0m   \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_batch_hook\u001b[49m\u001b[43m(\u001b[49m\u001b[43mModeKeys\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mTRAIN\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mend\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlogs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlogs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/tensorflow/python/keras/callbacks.py:316\u001b[0m, in \u001b[0;36mCallbackList._call_batch_hook\u001b[0;34m(self, mode, hook, batch, logs)\u001b[0m\n\u001b[1;32m    314\u001b[0m   \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_batch_begin_hook(mode, batch, logs)\n\u001b[1;32m    315\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m hook \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mend\u001b[39m\u001b[38;5;124m'\u001b[39m:\n\u001b[0;32m--> 316\u001b[0m   \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_batch_end_hook\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlogs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    317\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    318\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mUnrecognized hook: \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(hook))\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/tensorflow/python/keras/callbacks.py:336\u001b[0m, in \u001b[0;36mCallbackList._call_batch_end_hook\u001b[0;34m(self, mode, batch, logs)\u001b[0m\n\u001b[1;32m    333\u001b[0m   batch_time \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime() \u001b[38;5;241m-\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_batch_start_time\n\u001b[1;32m    334\u001b[0m   \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_batch_times\u001b[38;5;241m.\u001b[39mappend(batch_time)\n\u001b[0;32m--> 336\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_batch_hook_helper\u001b[49m\u001b[43m(\u001b[49m\u001b[43mhook_name\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlogs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    338\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_batch_times) \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_batches_for_timing_check:\n\u001b[1;32m    339\u001b[0m   end_hook_name \u001b[38;5;241m=\u001b[39m hook_name\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/tensorflow/python/keras/callbacks.py:374\u001b[0m, in \u001b[0;36mCallbackList._call_batch_hook_helper\u001b[0;34m(self, hook_name, batch, logs)\u001b[0m\n\u001b[1;32m    372\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m callback \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcallbacks:\n\u001b[1;32m    373\u001b[0m   hook \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(callback, hook_name)\n\u001b[0;32m--> 374\u001b[0m   \u001b[43mhook\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlogs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    376\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_timing:\n\u001b[1;32m    377\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m hook_name \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_hook_times:\n",
      "File \u001b[0;32m~/phd_segmentation/custom_train_objects/callbacks/terminate_on_nan.py:24\u001b[0m, in \u001b[0;36mTerminateOnNaN.on_train_batch_end\u001b[0;34m(self, batch, logs)\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mloss\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;129;01min\u001b[39;00m logs:\n\u001b[1;32m     23\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m np\u001b[38;5;241m.\u001b[39misnan(logs[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mloss\u001b[39m\u001b[38;5;124m'\u001b[39m]) \u001b[38;5;129;01mor\u001b[39;00m np\u001b[38;5;241m.\u001b[39misinf(logs[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mloss\u001b[39m\u001b[38;5;124m'\u001b[39m]):\n\u001b[0;32m---> 24\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mNaN loss at batch : \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m  Logs : \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(batch, logs))\n",
      "\u001b[0;31mValueError\u001b[0m: NaN loss at batch : 20\n  Logs : {'loss': nan, 'foreground_loss': nan, 'background_loss': nan}"
     ]
    }
   ],
   "source": [
    "for epochs in [3]:\n",
    "    batch_size = 1\n",
    "\n",
    "    if 'highres' in model_name:\n",
    "        max_size, max_frames = 512, 32\n",
    "    elif 'lowres' in model_name:\n",
    "        max_size, max_frames = 128, 128\n",
    "    else:\n",
    "        max_size, max_frames = 256 - 64, 32\n",
    "\n",
    "    if not isinstance(max_size, tuple):\n",
    "        max_size = (max_size, max_size)\n",
    "    \n",
    "    augment_prct = 0.25\n",
    "    shuffle_size = 0 if epochs + model.epochs < 5 else batch_size * 8\n",
    "\n",
    "    crop_mode    = ['random_center_80', 'random_center_80', 'random']\n",
    "\n",
    "    if 'test' in model_name:\n",
    "        train = train.sample(10, random_state = 0)\n",
    "        valid = valid.sample(10, random_state = 0)\n",
    "\n",
    "    if 'scratch' not in model_name and model.epochs < 1 and not is_model_name(model.pretrained_name):\n",
    "        for l in model.layers[:-2]: l.trainable = False\n",
    "        epochs = 1\n",
    "    else:\n",
    "        for l in model.layers: l.trainable = True\n",
    "\n",
    "    model.train(\n",
    "        train, validation_data = valid, epochs = epochs, batch_size = batch_size,\n",
    "\n",
    "        augment_prct = augment_prct, shuffle_size = shuffle_size,\n",
    "        is_rectangular = True, cache = False,\n",
    "\n",
    "        max_size = max_size, max_frames = max_frames, crop_mode = crop_mode, run_eagerly = False\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.plot_history()\n",
    "print(model.history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'start': datetime.datetime(2023, 5, 24, 11, 26, 58, 477270), 'end': datetime.datetime(2023, 5, 24, 14, 24, 44, 408554), 'time': 10665.931284, 'interrupted': False, 'start_epoch': -1, 'final_epoch': 0}, {'start': datetime.datetime(2023, 5, 24, 14, 29, 53, 904649), 'end': datetime.datetime(2023, 5, 25, 9, 37, 51, 689042), 'time': 68877.784393, 'interrupted': False, 'start_epoch': 0, 'final_epoch': 3}]\n"
     ]
    }
   ],
   "source": [
    "pd.DataFrame(model.history.trainings_infos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "19h 7min 57sec\n"
     ]
    }
   ],
   "source": [
    "from utils import *\n",
    "print(time_to_string(68877.784393))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_memory()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-15 09:20:11.324760: I tensorflow/stream_executor/cuda/cuda_dnn.cc:384] Loaded cuDNN version 8204\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image shape : (1, 249, 188, 213, 1) - Mask shape : (1, 249, 188, 213, 104) - Output shape : (1, 249, 188, 213, 32)\n"
     ]
    }
   ],
   "source": [
    "model.max_size   = (None, None)\n",
    "model.max_frames = -1\n",
    "\n",
    "for idx, row in dataset.iloc[:1].iterrows():\n",
    "    image, target = model.encode_data(row)\n",
    "    image  = tf.expand_dims(image, axis = 0)\n",
    "    target = tf.sparse.expand_dims(target, axis = 0)\n",
    "    pred   = model.infer(image, 64)\n",
    "    \n",
    "    print('Image shape : {} - Mask shape : {} - Output shape : {}'.format(image.shape, target.shape, pred.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Test prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading checkpoint pretrained_models/cluster_unet_cosine/saving/ckpt-161\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.checkpoint.checkpoint.CheckpointLoadStatus at 0x7f914c023100>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_checkpoint(directory = model_name[:-2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if 'highres' in model_name:\n",
    "    image_size   = 512 - 64\n",
    "    model.max_size     = (image_size, image_size)\n",
    "    model.max_frames   = 32\n",
    "elif 'lowres' in model_name:\n",
    "    image_size   = 128\n",
    "    model.max_size     = (image_size, image_size)\n",
    "    model.max_frames   = 128\n",
    "else:\n",
    "    image_size   = 256\n",
    "    model.max_size     = (image_size, image_size)\n",
    "    model.max_frames   = 32\n",
    "\n",
    "model.pad_value = 0.\n",
    "print(model.max_size, model.max_frames)\n",
    "config = model.get_dataset_config(is_validation = False, batch_size = 0, shuffle_size = 0)\n",
    "\n",
    "set_level('debug', 'datasets')\n",
    "\n",
    "ds = prepare_dataset(train.iloc[:1], ** config)\n",
    "\n",
    "set_level('info', 'datasets')\n",
    "\n",
    "model.get_loss().skip_empty_frames.assign(True)\n",
    "model.get_loss().skip_empty_labels.assign(True)\n",
    "\n",
    "for inp, out in ds:\n",
    "    print('Input shape : {} - output shape : {}'.format(inp.shape, out.shape))\n",
    "    #plot_mask(inp[..., 0], out, n = 4)\n",
    "    pred = model(tf.expand_dims(inp, axis = 0))[0]\n",
    "    intersect = tf.cast(out, tf.float32) * pred\n",
    "    print(model.get_loss()(tf.sparse.expand_dims(out, 0), pred[tf.newaxis]))\n",
    "    print(tf.sparse.reduce_sum(intersect) / len(intersect.indices))\n",
    "    print(intersect)\n",
    "    plot_mask(inp[..., 0], tf.cast(pred > 0.01, tf.uint8), n = 4)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask   = (pred.numpy() > 0.1).astype(np.uint8)\n",
    "labels = np.argmax(mask[..., 1:], axis = -1) + 1\n",
    "voxels = np.any(mask[..., 1:], axis = -1)\n",
    "\n",
    "organs = model.labels\n",
    "\n",
    "show_organs = [o for o in organs if (o is not None) and ('rib' in o or 'vertebr' in o )]\n",
    "\n",
    "skip_indexes = [i for i, organ in enumerate(organs) if organ not in show_organs]\n",
    "\n",
    "#voxels[np.any(mask[..., skip_indexes], axis = -1)] = 0\n",
    "#labels[np.any(mask[..., skip_indexes], axis = -1)] = 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.any(mask, axis = -1).sum())\n",
    "print(skip_indexes)\n",
    "print(voxels.sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from utils import plot_utils\n",
    "\n",
    "importlib.reload(plot_utils)\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from loggers import set_level\n",
    "\n",
    "set_level('debug', 'utils.plot_utils')\n",
    "\n",
    "def add_color_axis(labels, cmap = None):\n",
    "    mapper = plt.cm.ScalarMappable(cmap = cmap)\n",
    "    return np.reshape(mapper.to_rgba(np.reshape(labels, [-1]).tolist()), list(labels.shape) + [4])\n",
    "\n",
    "sx, sy, sz = -3, 3, 1\n",
    "\n",
    "plot_utils.plot(\n",
    "    voxels[::sy, ::sx, ::sz].astype(bool),\n",
    "    figsize    = (10, 10),\n",
    "    facecolors = add_color_axis(labels[::sy, ::sx, ::sz], cmap = 'magma'),\n",
    "    plot_type  = 'voxels',\n",
    "    color      = None,\n",
    "    is_3d      = True,\n",
    "    with_legend = True,\n",
    "    with_colorbar = True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Test dataset performances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original dataset : <TensorSliceDataset element_spec={'subject_id': TensorSpec(shape=(), dtype=tf.string, name=None), 'thickness': TensorSpec(shape=(), dtype=tf.int32, name=None), 'images': TensorSpec(shape=(), dtype=tf.string, name=None), 'segmentation': TensorSpec(shape=(), dtype=tf.string, name=None), 'label': TensorSpec(shape=(104,), dtype=tf.string, name=None), 'start_frame': TensorSpec(shape=(), dtype=tf.int32, name=None), 'end_frame': TensorSpec(shape=(), dtype=tf.int32, name=None), 'id': TensorSpec(shape=(), dtype=tf.string, name=None)}>\n",
      "- Dataset after encoding : <ParallelMapDataset element_spec=(TensorSpec(shape=(None, None, None, 1), dtype=tf.float32, name=None), SparseTensorSpec(TensorShape([None, None, None, None]), tf.uint8))>\n",
      "- Dataset after filtering : <FilterDataset element_spec=(TensorSpec(shape=(None, None, None, 1), dtype=tf.float32, name=None), SparseTensorSpec(TensorShape([None, None, None, None]), tf.uint8))>\n",
      "- Dataset after batch : <BatchDataset element_spec=(TensorSpec(shape=(None, None, None, None, 1), dtype=tf.float32, name=None), SparseTensorSpec(TensorShape([None, None, None, None, None]), tf.uint8))>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "9it [00:20,  2.31s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "10 batchs in 20.750 sec sec (0.482 batch / sec)\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "unsupported format string passed to SparseTensor.__format__",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [5], line 14\u001b[0m\n\u001b[1;32m     10\u001b[0m ds_train \u001b[38;5;241m=\u001b[39m prepare_dataset(train\u001b[38;5;241m.\u001b[39msample(\u001b[38;5;241m10\u001b[39m), \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m config)\n\u001b[1;32m     12\u001b[0m set_level(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124minfo\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdatasets\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m---> 14\u001b[0m \u001b[43mtest_dataset_time\u001b[49m\u001b[43m(\u001b[49m\u001b[43mds_train\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msteps\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/phd_segmentation/datasets/dataset_utils.py:166\u001b[0m, in \u001b[0;36mtest_dataset_time\u001b[0;34m(dataset, steps, batch_size, **kwargs)\u001b[0m\n\u001b[1;32m    161\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m size \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m    162\u001b[0m     logger\u001b[38;5;241m.\u001b[39minfo(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTime estimated for all dataset (\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m batch) : \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[1;32m    163\u001b[0m         size, time_to_string(size \u001b[38;5;241m*\u001b[39m (i \u001b[38;5;241m/\u001b[39m temps))\n\u001b[1;32m    164\u001b[0m     ))\n\u001b[0;32m--> 166\u001b[0m logger\u001b[38;5;241m.\u001b[39minfo(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mBatch infos : \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\u001b[43m_get_infos\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m)\u001b[49m))\n\u001b[1;32m    168\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m temps\n",
      "File \u001b[0;32m~/phd_segmentation/datasets/dataset_utils.py:48\u001b[0m, in \u001b[0;36m_get_infos\u001b[0;34m(tensor, level)\u001b[0m\n\u001b[1;32m     46\u001b[0m indent \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m \u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m*\u001b[39m level\n\u001b[1;32m     47\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(tensor, (\u001b[38;5;28mlist\u001b[39m, \u001b[38;5;28mtuple\u001b[39m)):\n\u001b[0;32m---> 48\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin([\n\u001b[1;32m     49\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124mItem \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m : \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(indent, i, _get_infos(t, level\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m1\u001b[39m)) \n\u001b[1;32m     50\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m i, t \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(tensor)\n\u001b[1;32m     51\u001b[0m     ])\n\u001b[1;32m     52\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(tensor, \u001b[38;5;28mdict\u001b[39m):\n\u001b[1;32m     53\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin([\n\u001b[1;32m     54\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124mItem \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m : \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(indent, k, _get_infos(t, level\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m1\u001b[39m)) \n\u001b[1;32m     55\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m k, t \u001b[38;5;129;01min\u001b[39;00m tensor\u001b[38;5;241m.\u001b[39mitems()\n\u001b[1;32m     56\u001b[0m     ])\n",
      "File \u001b[0;32m~/phd_segmentation/datasets/dataset_utils.py:49\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     46\u001b[0m indent \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m \u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m*\u001b[39m level\n\u001b[1;32m     47\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(tensor, (\u001b[38;5;28mlist\u001b[39m, \u001b[38;5;28mtuple\u001b[39m)):\n\u001b[1;32m     48\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin([\n\u001b[0;32m---> 49\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124mItem \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m : \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(indent, i, \u001b[43m_get_infos\u001b[49m\u001b[43m(\u001b[49m\u001b[43mt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlevel\u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m) \n\u001b[1;32m     50\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m i, t \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(tensor)\n\u001b[1;32m     51\u001b[0m     ])\n\u001b[1;32m     52\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(tensor, \u001b[38;5;28mdict\u001b[39m):\n\u001b[1;32m     53\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin([\n\u001b[1;32m     54\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124mItem \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m : \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(indent, k, _get_infos(t, level\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m1\u001b[39m)) \n\u001b[1;32m     55\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m k, t \u001b[38;5;129;01min\u001b[39;00m tensor\u001b[38;5;241m.\u001b[39mitems()\n\u001b[1;32m     56\u001b[0m     ])\n",
      "File \u001b[0;32m~/phd_segmentation/datasets/dataset_utils.py:60\u001b[0m, in \u001b[0;36m_get_infos\u001b[0;34m(tensor, level)\u001b[0m\n\u001b[1;32m     58\u001b[0m infos \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mshape : \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m - type : \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(tensor\u001b[38;5;241m.\u001b[39mshape, tensor\u001b[38;5;241m.\u001b[39mdtype\u001b[38;5;241m.\u001b[39mname)\n\u001b[1;32m     59\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m tensor\u001b[38;5;241m.\u001b[39mdtype \u001b[38;5;241m!=\u001b[39m tf\u001b[38;5;241m.\u001b[39mstring:\n\u001b[0;32m---> 60\u001b[0m     infos \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43m- min : \u001b[39;49m\u001b[38;5;132;43;01m{:.3f}\u001b[39;49;00m\u001b[38;5;124;43m - max : \u001b[39;49m\u001b[38;5;132;43;01m{:.3f}\u001b[39;49;00m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mformat\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmin\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtensor\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmax\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtensor\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     61\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m infos\n",
      "\u001b[0;31mTypeError\u001b[0m: unsupported format string passed to SparseTensor.__format__"
     ]
    }
   ],
   "source": [
    "from loggers import set_level\n",
    "\n",
    "model.max_frames = 32\n",
    "model.max_size   = (256, 256)\n",
    "\n",
    "set_level('debug', 'datasets')\n",
    "\n",
    "config = model.get_dataset_config(is_validation = True, batch_size = 1, cache = False, prefetch = False, shuffle_size = 0)\n",
    "\n",
    "ds_train = prepare_dataset(train.sample(10), ** config)\n",
    "\n",
    "set_level('info', 'datasets')\n",
    "\n",
    "test_dataset_time(ds_train, steps = 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Dataset visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if 'highres' in model_name:\n",
    "    image_size   = 512 - 64\n",
    "    model.max_size     = (image_size, image_size)\n",
    "    model.max_frames   = 32\n",
    "elif 'lowres' in model_name:\n",
    "    image_size   = 128\n",
    "    model.max_size     = (image_size, image_size)\n",
    "    model.max_frames   = 128\n",
    "else:\n",
    "    image_size   = 256\n",
    "    model.max_size     = (image_size, image_size)\n",
    "    model.max_frames   = 32\n",
    "\n",
    "model.pad_value = 0.\n",
    "print(model.max_size, model.max_frames)\n",
    "config = model.get_dataset_config(is_validation = False, batch_size = 0, shuffle_size = 0)\n",
    "\n",
    "set_level('debug', 'datasets')\n",
    "\n",
    "ds = prepare_dataset(train.iloc[:5], ** config)\n",
    "\n",
    "set_level('info', 'datasets')\n",
    "\n",
    "for inp, out in ds:\n",
    "    print('Input shape : {} - output shape : {}'.format(inp.shape, out.shape))\n",
    "    plot_mask(inp[..., 0], out, n = 4)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Test processing functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "n = 5\n",
    "model.max_frames = 32\n",
    "model.max_size   = (512, 512, 32)\n",
    "\n",
    "for _, row in train.sample(n, random_state = 0).iterrows():\n",
    "    inp, _ = model.get_input(row, False)\n",
    "    print('Input shape : {} - {}'.format(inp.shape, inp.dtype))\n",
    "    print(\"Original   : {} - {}\".format(tf.reduce_min(inp), tf.reduce_max(inp)))\n",
    "    inp = model.normalize_image(inp)\n",
    "    print(\"Normalized : {} - {}\".format(tf.reduce_min(inp), tf.reduce_max(inp)))\n",
    "    inp = model.augment_input(inp)\n",
    "    print(\"Augmented  : {} - {}\".format(tf.reduce_min(inp), tf.reduce_max(inp)))\n",
    "\n",
    "\n",
    "for _, row in train.sample(n, random_state = 0).iterrows():\n",
    "    out = model.get_output_fn(row['segmentation'], row['label'])\n",
    "    print('Output shape : {} - {} - type : {}'.format(out.shape, out.dtype, out.__class__.__name__))\n",
    "\n",
    "for _, row in tqdm(train.sample(n, random_state = 0).iterrows()):\n",
    "    inp, out = model.encode_data(row)\n",
    "    print('Input shape : {} - output shape : {}'.format(inp.shape, out.shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from loggers import set_level\n",
    "\n",
    "model.max_frames = 1\n",
    "model.max_size   = (256, 256)\n",
    "\n",
    "set_level('debug', 'datasets')\n",
    "\n",
    "config = model.get_dataset_config(is_validation = True, batch_size = 0, cache = False, prefetch = False, shuffle_size = 0)\n",
    "\n",
    "ds_train = prepare_dataset(train, ** config)\n",
    "\n",
    "set_level('info', 'datasets')\n",
    "\n",
    "for i, (inp, out) in enumerate(ds_train):\n",
    "    print('Batch #{} : input shape = {} - output shape = {}'.format(i, inp.shape, out.shape))\n",
    "    tf.sparse.to_dense(out)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model.voxel_dims)\n",
    "\n",
    "img, vox_dims = load_medical_image(train.loc[0, 'images'], target_voxel_dims = model.voxel_dims)\n",
    "print(img.shape)\n",
    "\n",
    "inp = model.get_input(train.loc[0], reshape = False)\n",
    "print(inp.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Configure `learning-rate scheduler`\n",
    "\n",
    "This cell allows you to play with `learning_rate scheduler`'s parameters to get the one you want !"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = model.model_optimizer.learning_rate\n",
    "lr.factor = 1024.\n",
    "lr.warmup_steps = 1024\n",
    "lr.plot(512 * 15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}