tests.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6bc3c6d6-a0d0-45f6-a13c-33b031981033",
   "metadata": {},
   "source": [
    "# Tests"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "082658ac-5ebd-421f-8d3a-d576fd9bf468",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test labels rearranging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2d8fc4cd-7bf9-41ac-9d7f-4f6780fe57c2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-06-09 09:33:52.114958: 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-06-09 09:33:52.215281: 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-06-09 09:33:52.239627: 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",
      "<frozen importlib._bootstrap>:228: RuntimeWarning: scipy._lib.messagestream.MessageStream size changed, may indicate binary incompatibility. Expected 56 from C header, got 64 from PyObject\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3, 3) (3, 3, 3) tf.Tensor([3 3], shape=(2,), dtype=int64) tf.Tensor([3 3 3], shape=(3,), dtype=int64)\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-06-09 09:33:55.058187: 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-06-09 09:33:55.816986: 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",
      "2023-06-09 09:33:55.817858: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1616] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 14485 MB memory:  -> device: 1, name: Quadro RTX 5000, pci bus id: 0000:65:00.0, compute capability: 7.5\n"
     ]
    }
   ],
   "source": [
    "import importlib\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from utils import is_equal\n",
    "from utils.med_utils import label_utils\n",
    "\n",
    "importlib.reload(label_utils)\n",
    "\n",
    "labels  = [None, 'lung_left', 'lung_right', 'heart']\n",
    "mapping = [None, 'heart', 'lung_left', 'lung_right']\n",
    "one_hot_labels = labels[1:]\n",
    "\n",
    "array = tf.cast([\n",
    "    [1, 2, 1],\n",
    "    [0, 0, 2],\n",
    "    [3, 1, 2]\n",
    "], dtype = tf.int32)\n",
    "one_hot = tf.cast(\n",
    "    np.expand_dims(array, -1) == np.arange(np.max(array) + 1).reshape([1] * len(array.shape) + [-1]),\n",
    "    tf.uint8\n",
    ")[..., 1:]\n",
    "sparse = tf.sparse.from_dense(array)\n",
    "sparse_one_hot = tf.sparse.from_dense(one_hot)\n",
    "\n",
    "print(array.shape, one_hot.shape, sparse.dense_shape, sparse_one_hot.dense_shape)\n",
    "\n",
    "target = tf.cast([\n",
    "    [2, 3, 2],\n",
    "    [0, 0, 3],\n",
    "    [1, 2, 3]\n",
    "], dtype = tf.int32)\n",
    "\n",
    "print(is_equal(target, label_utils.rearrange_labels(array, labels, mapping, is_one_hot = False))[1])\n",
    "print(is_equal(target, label_utils.rearrange_labels(one_hot, one_hot_labels, mapping, is_one_hot = True))[1])\n",
    "print(is_equal(\n",
    "    target, tf.sparse.to_dense(label_utils.rearrange_labels(sparse, labels, mapping, is_one_hot = False))\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    target, np.argmax(tf.sparse.to_dense(\n",
    "        label_utils.rearrange_labels(sparse_one_hot, one_hot_labels, mapping, is_one_hot = True)\n",
    "    ), axis = -1).astype(np.int32)\n",
    ")[1])\n",
    "\n",
    "\n",
    "_ = \"\"\"\n",
    "array   = np.random.randint(0, len(labels), size = (256, 256, 256))\n",
    "one_hot = np.equal(np.expand_dims(array, -1), np.arange(len(labels)).reshape([1] * len(array.shape) + [-1]))\n",
    "\n",
    "%timeit rearrange_labels_one_hot(one_hot, labels, mapping)\n",
    "%timeit rearrange_labels_one_hot(np.equal(np.expand_dims(array, -1), np.arange(len(labels)).reshape([1] * len(array.shape) + [-1])), labels, mapping)\n",
    "%timeit rearrange_labels_dense(array, labels, mapping)\n",
    "#%timeit rearrange_labels_vectorize(array, labels, mapping)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d5d23729-2d7a-4de3-b229-9bf91032d67d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-06-09 09:08:59.960365: 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-06-09 09:09:00.057079: 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-06-09 09:09:00.081826: 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",
      "<frozen importlib._bootstrap>:228: RuntimeWarning: scipy._lib.messagestream.MessageStream size changed, may indicate binary incompatibility. Expected 56 from C header, got 64 from PyObject\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(True, 'Value are equals !')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-06-09 09:09:09.163988: 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-06-09 09:09:09.921810: 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",
      "2023-06-09 09:09:09.922672: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1616] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 14485 MB memory:  -> device: 1, name: Quadro RTX 5000, pci bus id: 0000:65:00.0, compute capability: 7.5\n"
     ]
    }
   ],
   "source": [
    "import importlib\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from utils import is_equal\n",
    "from utils.med_utils import label_utils\n",
    "\n",
    "importlib.reload(label_utils)\n",
    "\n",
    "labels  = [None, 'lung_left', 'lung_right', 'heart']\n",
    "mapping = [None, 'heart', 'lung_left', 'lung_right']\n",
    "one_hot_labels = labels[1:]\n",
    "\n",
    "array = tf.cast([\n",
    "    [1, 2, 1],\n",
    "    [0, 0, 2],\n",
    "    [3, 1, 2]\n",
    "], dtype = tf.int32)\n",
    "\n",
    "target = tf.cast([\n",
    "    [2, 3, 2],\n",
    "    [0, 0, 3],\n",
    "    [1, 2, 3]\n",
    "], dtype = tf.int32)\n",
    "\n",
    "array  = tf.random.uniform((128, 128, 128), minval = 0, maxval = len(labels), dtype = tf.int32)\n",
    "target = label_utils.rearrange_labels(array, labels, mapping)\n",
    "\n",
    "keys, values = [], []\n",
    "for i, label in enumerate(labels):\n",
    "    if label in mapping:\n",
    "        keys.append(i)\n",
    "        values.append(mapping.index(label))\n",
    "keys, values = tf.cast(keys, tf.int32), tf.cast(values, tf.int32)\n",
    "table = tf.lookup.KeyValueTensorInitializer(keys, values)\n",
    "table = tf.lookup.StaticHashTable(table, default_value = 0)\n",
    "\n",
    "print(is_equal(target, table.lookup(array)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b158a80b-e659-495d-ac97-2d9d886ac9ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.34 ms ± 4.58 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
      "52.7 ms ± 26.7 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n",
      "2.87 ms ± 20.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit table.lookup(array)\n",
    "%timeit label_utils.rearrange_labels(array, labels, mapping)\n",
    "%timeit label_utils.add_zero_label(tf.sparse.from_dense(array), is_one_hot = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73a547d0-974b-4503-9b86-c8560abd71e9",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test loss averaging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a268c018-bb8c-4915-a884-3af69dc8c07b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from utils import compute_centroids\n",
    "\n",
    "def average_loss(losses, batch_size, mode = 'macro', batch_idx = None, ids = None):\n",
    "    if batch_size == 1:\n",
    "        if mode == 'macro':\n",
    "            return tf.reduce_mean(losses, axis = -1, keepdims = True)\n",
    "        elif mode == 'micro':\n",
    "            return tf.reduce_mean(tf.squeeze(compute_centroids(\n",
    "                tf.expand_dims(losses, axis = 1), tf.cast(ids, tf.int32)\n",
    "            )[1], axis = -1), axis = -1, keepdims = True)\n",
    "    else:\n",
    "        if mode == 'macro':\n",
    "            return tf.squeeze(compute_centroids(\n",
    "                tf.expand_dims(losses, axis = 1), tf.cast(batch_idx, tf.int32)\n",
    "            )[1], axis = -1)\n",
    "        else:\n",
    "            n_labels    = tf.reduce_max(ids)\n",
    "            indexes     = tf.cast(batch_idx * n_labels + ids, tf.int32)\n",
    "            loss_label_batch_idx, loss_per_label_per_batch = compute_centroids(\n",
    "                tf.expand_dims(losses, axis = 1), indexes\n",
    "            )\n",
    "\n",
    "            loss_batch_idx = loss_label_batch_idx // n_labels\n",
    "            return tf.squeeze(compute_centroids(\n",
    "                loss_per_label_per_batch, loss_batch_idx\n",
    "            )[1], axis = 1)\n",
    "\n",
    "losses = tf.cast([1, 2, 1.5, 2.5, 3, 4.5, 4, 3.5], tf.float32)\n",
    "ids    = tf.cast([1, 1, 2, 2, 3, 2, 3, 2], tf.int32)\n",
    "batch  = tf.cast([1, 1, 1, 1, 2, 2, 2, 2], tf.int32)\n",
    "\n",
    "print(np.all(average_loss(losses, 1, 'macro', batch, ids) == np.mean(losses)))\n",
    "print(np.all(average_loss(losses, 1, 'micro', batch, ids) == np.mean([1.5, 3, 3.5])))\n",
    "print(np.all(average_loss(losses, 2, 'macro', batch, ids) == np.mean([[1, 2, 1.5, 2.5], [3, 3.5, 4, 4.5]], axis = -1)))\n",
    "print(np.all(average_loss(losses, 2, 'macro', batch, ids) == np.mean([[1.5, 2], [3.5, 4]], axis = -1)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9239465-4107-471a-8f08-4dd45959d361",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test sparse argmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "36d59968-d9e9-4fcb-a4c2-3644b2c1ba9e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(True, shape=(), dtype=bool)\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "mask = tf.sparse.SparseTensor(\n",
    "    indices = [\n",
    "        [1, 1, 1],\n",
    "        [2, 1, 1],\n",
    "        [2, 2, 1],\n",
    "        [1, 2, 2],\n",
    "        [2, 0, 2],\n",
    "        [3, 1, 2]\n",
    "    ],\n",
    "    values = [1] * 6,\n",
    "    dense_shape = (4, 4, 3)\n",
    ")\n",
    "mask = tf.sparse.reorder(mask)\n",
    "mask = tf.sparse.expand_dims(mask, axis = 0)\n",
    "\n",
    "res1 = tf.argmax(tf.sparse.to_dense(mask), axis = -1, output_type = tf.int32)\n",
    "res2 = tf.tensor_scatter_nd_update(\n",
    "    tf.zeros(mask.dense_shape[:-1], dtype = tf.int32),\n",
    "    mask.indices[:, :-1],\n",
    "    tf.cast(mask.indices[:, -1], tf.int32)\n",
    ")\n",
    "print(tf.reduce_all(res1 == res2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "259a8196-050c-4a5e-9750-0744b8a28231",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test `pad_or_crop`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "da272845-4628-46ea-8efc-af6ee2a5443b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[ 0  1  2  3]\n",
      " [ 4  5  6  7]\n",
      " [ 8  9 10 11]\n",
      " [12 13 14 15]], shape=(4, 4), dtype=int32)\n",
      "\n",
      "Cropping tests\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "\n",
      "Padding tests\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "\n",
      "Cropping + padding tests\n",
      "Value are equals !\n"
     ]
    }
   ],
   "source": [
    "import importlib\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "from utils import is_equal\n",
    "from utils.med_utils import pre_processing\n",
    "\n",
    "try:\n",
    "    tf.config.set_visible_devices([], 'GPU')\n",
    "except:\n",
    "    pass\n",
    "\n",
    "def _to_dense(res):\n",
    "    return res[0], tf.sparse.to_dense(res[1])\n",
    "\n",
    "importlib.reload(pre_processing)\n",
    "\n",
    "image = tf.reshape(tf.range(16), [4, 4])\n",
    "mask  = image % 3 == 0\n",
    "\n",
    "print(image)\n",
    "\n",
    "print('\\nCropping tests')\n",
    "crop_shape = (2, 2)\n",
    "print(is_equal(\n",
    "    image[:2, :2], pre_processing.pad_or_crop(image, crop_shape, crop_mode = 'start')\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    image[1:3, 1:3], pre_processing.pad_or_crop(image, crop_shape, crop_mode = 'center')\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    image[-2:, -2:], pre_processing.pad_or_crop(image, crop_shape, crop_mode = 'end')\n",
    ")[1])\n",
    "\n",
    "print(is_equal(\n",
    "    (image[:2, :2], mask[:2, :2]), pre_processing.pad_or_crop(image, crop_shape, mask = mask, crop_mode = 'start')\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    (image[1:3, 1:3], mask[1:3, 1:3]), pre_processing.pad_or_crop(image, crop_shape, mask = mask, crop_mode = 'center')\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    (image[-2:, -2:], mask[-2:, -2:]), pre_processing.pad_or_crop(image, crop_shape, mask = mask, crop_mode = 'end')\n",
    ")[1])\n",
    "\n",
    "print(is_equal(\n",
    "    (image[:2, :2], mask[:2, :2]), _to_dense(pre_processing.pad_or_crop(\n",
    "        image, crop_shape, mask = tf.sparse.from_dense(mask), crop_mode = 'start'\n",
    "    ))\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    (image[1:3, 1:3], mask[1:3, 1:3]), _to_dense(pre_processing.pad_or_crop(\n",
    "        image, crop_shape, mask = tf.sparse.from_dense(mask), crop_mode = 'center'\n",
    "    ))\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    (image[-2:, -2:], mask[-2:, -2:]), _to_dense(pre_processing.pad_or_crop(\n",
    "        image, crop_shape, mask = tf.sparse.from_dense(mask), crop_mode = 'end'\n",
    "    ))\n",
    ")[1])\n",
    "\n",
    "\n",
    "print('\\nPadding tests')\n",
    "pad_shape = (6, 6)\n",
    "print(is_equal(\n",
    "    tf.pad(image, [(2, 0)] * 2), pre_processing.pad_or_crop(image, pad_shape, pad_mode = 'before', pad_value = 0)\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    tf.pad(image, [(1, 1)] * 2), pre_processing.pad_or_crop(image, pad_shape, pad_mode = 'even', pad_value = 0)\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    tf.pad(image, [(0, 2)] * 2), pre_processing.pad_or_crop(image, pad_shape, pad_mode = 'after', pad_value = 0)\n",
    ")[1])\n",
    "\n",
    "print(is_equal(\n",
    "    (tf.pad(image, [(2, 0)] * 2), tf.pad(mask, [(2, 0)] * 2)), pre_processing.pad_or_crop(\n",
    "        image, pad_shape, mask = mask, pad_mode = 'before', pad_value = 0\n",
    "    )\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    (tf.pad(image, [(1, 1)] * 2), tf.pad(mask, [(1, 1)] * 2)), pre_processing.pad_or_crop(\n",
    "        image, pad_shape, mask = mask, pad_mode = 'even', pad_value = 0\n",
    "    )\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    (tf.pad(image, [(0, 2)] * 2), tf.pad(mask, [(0, 2)] * 2)), pre_processing.pad_or_crop(\n",
    "        image, pad_shape, mask = mask, pad_mode = 'after', pad_value = 0\n",
    "    )\n",
    ")[1])\n",
    "\n",
    "print(is_equal(\n",
    "    (tf.pad(image, [(2, 0)] * 2), tf.pad(mask, [(2, 0)] * 2)), _to_dense(pre_processing.pad_or_crop(\n",
    "        image, pad_shape, mask = tf.sparse.from_dense(mask), pad_mode = 'before', pad_value = 0\n",
    "    ))\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    (tf.pad(image, [(1, 1)] * 2), tf.pad(mask, [(1, 1)] * 2)), _to_dense(pre_processing.pad_or_crop(\n",
    "        image, pad_shape, mask = tf.sparse.from_dense(mask), pad_mode = 'even', pad_value = 0\n",
    "    ))\n",
    ")[1])\n",
    "print(is_equal(\n",
    "    (tf.pad(image, [(0, 2)] * 2), tf.pad(mask, [(0, 2)] * 2)), _to_dense(pre_processing.pad_or_crop(\n",
    "        image, pad_shape, mask = tf.sparse.from_dense(mask), pad_mode = 'after', pad_value = 0\n",
    "    ))\n",
    ")[1])\n",
    "\n",
    "\n",
    "print('\\nCropping + padding tests')\n",
    "crop_pad_shape = (6, 2)\n",
    "print(is_equal(tf.pad(image[:, 1:3], [(1, 1), (0, 0)]), pre_processing.pad_or_crop(\n",
    "    image, crop_pad_shape, crop_mode = 'center', pad_mode = 'even', pad_value = 0\n",
    "))[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc48376a-8be1-406f-a76e-8b499afa7888",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test `crop_then_reshape`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "630dea2e-c950-4822-8a5e-77cd53f24dc2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading dataset total_segmentator...\n",
      "Dataset length : 1203\n"
     ]
    }
   ],
   "source": [
    "import importlib\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from utils import is_equal\n",
    "from datasets import get_dataset\n",
    "from utils.med_utils import pre_processing, resample_volume, load_medical_data\n",
    "\n",
    "try:\n",
    "    tf.config.set_visible_devices([], 'GPU')\n",
    "except:\n",
    "    pass\n",
    "\n",
    "dataset = get_dataset('total_segmentator')\n",
    "\n",
    "print('Dataset length : {}'.format(len(dataset)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f75508e6-9da1-4008-aa05-b610ae67395a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-06-12 12:49:59.771974: 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": [
      "Image shape : (249, 188, 213) - voxel dims : [1.4999999 1.5       1.4999999]\n",
      "\n",
      "Test cropping\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "\n",
      "Test padding\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "\n",
      "Test cropping then resizing\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "\n",
      "Test padding then resizing\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "\n",
      "Test simple frame cropping\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "\n",
      "Test simple resizing\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "\n",
      "Test simple resize to multiple\n",
      "Value are equals !\n",
      "Value are equals !\n",
      "Value are equals !\n"
     ]
    }
   ],
   "source": [
    "importlib.reload(pre_processing)\n",
    "\n",
    "# All the below tests should alos work if the image has a 4th channel axis (uncomment the below line to test it)\n",
    "row = dataset.iloc[0]\n",
    "img, voxel_dims = load_medical_data(row['images'])\n",
    "#img = tf.expand_dims(img, axis = -1)\n",
    "print('Image shape : {} - voxel dims : {}'.format(img.shape, voxel_dims))\n",
    "\n",
    "print('\\nTest cropping')\n",
    "target_shape = (128, 128, 32)\n",
    "target_voxel = voxel_dims\n",
    "\n",
    "img2 = pre_processing.crop_then_reshape(img, voxel_dims, target_shape = target_shape, target_voxel_dims = target_voxel)\n",
    "tar2 = pre_processing.pad_or_crop(img, target_shape = target_shape)\n",
    "print(is_equal(target_shape, img2.shape[:len(target_shape)])[1])\n",
    "print(is_equal(tar2, img2)[1])\n",
    "\n",
    "print('\\nTest padding')\n",
    "target_shape = (256, 256, 32)\n",
    "target_voxel = voxel_dims\n",
    "img3 = pre_processing.crop_then_reshape(img, voxel_dims, target_shape = target_shape, target_voxel_dims = target_voxel)\n",
    "tar3 = pre_processing.pad_or_crop(img, target_shape = target_shape)\n",
    "print(is_equal(target_shape, img3.shape[:len(target_shape)])[1])\n",
    "print(is_equal(tar3, img3)[1])\n",
    "\n",
    "print('\\nTest cropping then resizing')\n",
    "target_shape = (128, 128, 32)\n",
    "target_voxel = (0.5, 0.5, 0.5)\n",
    "\n",
    "img4 = pre_processing.crop_then_reshape(img, voxel_dims, target_shape = target_shape, target_voxel_dims = target_voxel)\n",
    "tar4, new_voxel = resample_volume(\n",
    "    pre_processing.pad_or_crop(img, target_shape = np.round(np.array(target_shape) / 3).astype(np.int32)),\n",
    "    voxel_dims   = voxel_dims,\n",
    "    target_shape = target_shape\n",
    ")\n",
    "print(is_equal(target_shape, img4.shape[:len(target_shape)])[1])\n",
    "print(is_equal(target_voxel, new_voxel, max_err = 0.1)[1])\n",
    "print(is_equal(tar4, img4)[1])\n",
    "\n",
    "print('\\nTest padding then resizing')\n",
    "target_shape = (128, 128, 32)\n",
    "target_voxel = (3, 3, 3)\n",
    "\n",
    "img5 = pre_processing.crop_then_reshape(img, voxel_dims, target_shape = target_shape, target_voxel_dims = target_voxel)\n",
    "tar5, new_voxel = resample_volume(\n",
    "    pre_processing.pad_or_crop(img, target_shape = np.array(target_shape) * 2),\n",
    "    voxel_dims   = voxel_dims,\n",
    "    target_shape = target_shape\n",
    ")\n",
    "print(is_equal(target_shape, img5.shape[:len(target_shape)])[1])\n",
    "print(is_equal(target_voxel, new_voxel, max_err = 1e-3)[1])\n",
    "print(is_equal(tar5, img5)[1])\n",
    "\n",
    "print('\\nTest simple frame cropping')\n",
    "target_shape = (-1, -1, 32)\n",
    "target_voxel = voxel_dims\n",
    "\n",
    "img6 = pre_processing.crop_then_reshape(\n",
    "    img, voxel_dims, target_shape = target_shape, target_voxel_dims = target_voxel, crop_mode = 'start'\n",
    ")\n",
    "tar6, new_voxel = img[:, :, :target_shape[2]], target_voxel\n",
    "print(is_equal(tar6.shape, img6.shape)[1])\n",
    "print(is_equal(target_voxel, new_voxel, max_err = 0.1)[1])\n",
    "print(is_equal(tar6, img6)[1])\n",
    "\n",
    "print('\\nTest simple resizing')\n",
    "# in this scenario, the target should be a resized version of the 64 first frames\n",
    "# the resized should have a shape half of the original shape, as the new voxel dim is 2 times higher\n",
    "# if each voxel is 2 times bigger, we should have a volume 2 times smaller to cover the same space in the 3D-world space\n",
    "# It is the reason why we take the 64 first frames then resize them to 32\n",
    "target_shape = (-1, -1, 32)\n",
    "target_voxel = (3, 3, 3)\n",
    "\n",
    "img7 = pre_processing.crop_then_reshape(\n",
    "    img, voxel_dims, target_shape = target_shape, target_voxel_dims = target_voxel, crop_mode = 'start'\n",
    ")\n",
    "tar7, new_voxel = resample_volume(\n",
    "    img[:, :, :target_shape[2] * 2],\n",
    "    voxel_dims   = voxel_dims,\n",
    "    target_shape = [s if s > 0 else img.shape[i] // 2 for i, s in enumerate(target_shape)]\n",
    ")\n",
    "print(is_equal(tar7.shape, img7.shape)[1])\n",
    "print(is_equal(target_voxel, new_voxel, max_err = 0.1)[1])\n",
    "print(is_equal(tar7, img7)[1])\n",
    "\n",
    "print('\\nTest simple resize to multiple')\n",
    "target_shape = (-1, -1, 32)\n",
    "target_voxel = voxel_dims\n",
    "multiples    = np.array([32, 32, 32])\n",
    "\n",
    "img8 = pre_processing.crop_then_reshape(\n",
    "    img, voxel_dims, target_shape = target_shape, target_voxel_dims = target_voxel, crop_mode = 'start',\n",
    "    multiple_shape = multiples\n",
    ")\n",
    "tar8, new_voxel = resample_volume(\n",
    "    img[:, :, :target_shape[2]], voxel_dims, target_voxel_dims = voxel_dims,\n",
    "    target_shape = [\n",
    "        img.shape[0] // multiples[0] * multiples[0],\n",
    "        img.shape[1] // multiples[1] * multiples[1],\n",
    "        target_shape[2] // multiples[2] * multiples[2]\n",
    "    ]\n",
    ")\n",
    "print(is_equal(tar8.shape, img8.shape)[1])\n",
    "print(is_equal(target_voxel, new_voxel, max_err = 0.1)[1])\n",
    "print(is_equal(tar8, img8)[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80215eba-6a49-4185-af26-9b3a1a3b74e3",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test Clustering Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "be085920-beb3-4d4e-a74a-287e213cca98",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import importlib\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from custom_train_objects.losses import ge2e_seg_loss\n",
    "\n",
    "try:\n",
    "    tf.config.set_visible_devices([tf.config.list_physical_devices('GPU')[0]], 'GPU')\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "15863be6-32ce-4339-90b0-361e9ae8ea04",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[1. 1.]\n",
      " [1. 2.]\n",
      " [2. 0.]\n",
      " [2. 1.]\n",
      " [2. 2.]\n",
      " [3. 1.]], shape=(6, 2), dtype=float32)\n",
      "tf.Tensor([0 1 1 0 0 1], shape=(6,), dtype=int32)\n",
      "tf.Tensor(\n",
      "[[1.6666666 1.3333334]\n",
      " [2.        1.       ]], shape=(2, 2), dtype=float32)\n",
      "tf.Tensor(\n",
      "[[0. 0.]\n",
      " [0. 0.]\n",
      " [0. 0.]\n",
      " [0. 0.]\n",
      " [0. 0.]\n",
      " [0. 0.]\n",
      " [0. 0.]\n",
      " [0. 0.]\n",
      " [0. 0.]\n",
      " [0. 0.]], shape=(10, 2), dtype=float32)\n",
      "tf.Tensor([[0. 0.]], shape=(1, 2), dtype=float32)\n",
      "tf.Tensor(0.9176704, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "importlib.reload(ge2e_seg_loss)\n",
    "\n",
    "loss = ge2e_seg_loss.GE2ESegLoss()\n",
    "\n",
    "mask = tf.sparse.SparseTensor(\n",
    "    indices = [\n",
    "        [1, 1, 1],\n",
    "        [2, 1, 1],\n",
    "        [2, 2, 1],\n",
    "        [1, 2, 2],\n",
    "        [2, 0, 2],\n",
    "        [3, 1, 2]\n",
    "    ],\n",
    "    values = [1] * 6,\n",
    "    dense_shape = (4, 4, 3)\n",
    ")\n",
    "mask = tf.sparse.reorder(mask)\n",
    "mask = tf.sparse.expand_dims(mask, axis = 0)\n",
    "embeddings = np.zeros((1, 4, 4, 2))\n",
    "for i, j in mask.indices[:, 1:-1].numpy(): embeddings[:, i, j] = [i, j]\n",
    "embeddings = tf.cast(embeddings, tf.float32)\n",
    "\n",
    "fore_emb, fore_centr, fore_centr_ids = loss.compute_foreground_centroids(mask, embeddings)\n",
    "print(fore_emb)\n",
    "print(fore_centr_ids)\n",
    "print(fore_centr)\n",
    "\n",
    "back_emb, back_centr, _ = loss.compute_background_centroid(mask, embeddings)\n",
    "print(back_emb)\n",
    "print(back_centr)\n",
    "print(loss(mask, embeddings))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0827969b-38be-4e31-bd9e-741526454052",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test Dice Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f163e4c-479f-4ca2-97f4-5acfd1bfadeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import importlib\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from custom_train_objects.losses import dice_loss\n",
    "from utils.med_utils import load_medical_image, load_medical_seg\n",
    "\n",
    "try:\n",
    "    tf.config.set_visible_devices([tf.config.list_physical_devices('GPU')[0]], 'GPU')\n",
    "except:\n",
    "    pass\n",
    "\n",
    "path  = '/storage/Totalsegmentator_dataset/s0001'\n",
    "\n",
    "#image = tf.expand_dims(load_medical_image(os.path.join(path, 'ct.nii.gz'))[0], axis = 0)\n",
    "mask  = tf.sparse.expand_dims(load_medical_seg(os.path.join(path, 'masks.npz'))[0], axis = 0)\n",
    "\n",
    "print('Image shape : {} - mask shape : {}'.format((), mask.shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd7a8dff-d5ff-4a93-ba7f-9712e7bdd337",
   "metadata": {},
   "outputs": [],
   "source": [
    "importlib.reload(dice_loss)\n",
    "\n",
    "loss = dice_loss.DiceLoss(skip_empty_frames = False, skip_empty_labels = True, smoothing = 0.01)\n",
    "\n",
    "dense_mask = tf.sparse.to_dense(tf.cast(mask, tf.float32))\n",
    "\n",
    "empty_labels = tf.sparse.reduce_sum(tf.sparse.reshape(mask, [1, -1, 104]), axis = 1) == 0\n",
    "empty_labels = tf.cast(tf.reshape(empty_labels, [1, 1, 1, 1, 104]), tf.float32)\n",
    "\n",
    "print(empty_labels)\n",
    "\n",
    "print(loss(mask, dense_mask).numpy())\n",
    "print(loss(mask, dense_mask * 0.5).numpy())\n",
    "print(loss(mask, dense_mask + 0.25 * empty_labels).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c7fb9fc-e661-40d0-b66f-c633dc13f943",
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit loss(mask, dense_mask + 0.25 * empty_labels)\n",
    "%timeit loss(dense_mask, dense_mask + 0.25 * empty_labels)\n",
    "%timeit loss(tf.sparse.to_dense(mask), dense_mask + 0.25 * empty_labels)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01b48a72-7578-40f6-87c1-873bbe0da4d3",
   "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": 5
}