Tooth_acg_neural_volume_compression.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dWIMc0lIB4K1"
      },
      "source": [
        "# ACG Seminar Assignment -- Neural Volume Compression  \n",
        "\n",
        "Student: Anže Kristan  \n",
        "class: Advanced Computer Graphics  \n",
        "semester: Spring 2023/24  \n",
        "school: Faculty of Computer and Information Science, University of Ljubljana  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2M8dmjIDYC3u"
      },
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5CEuX--IZ-ke"
      },
      "source": [
        "#### imports and setup:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "VunWvEZAlI8R"
      },
      "outputs": [],
      "source": [
        "# from google.colab import drive\n",
        "# drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "wK91ZzKZlWBs"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torchsummary import summary\n",
        "import numpy as np\n",
        "import seaborn as sns\n",
        "import pandas as pd\n",
        "from scipy.ndimage import gaussian_filter\n",
        "from os import path as ospath"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "dLsUTnT3lXOc"
      },
      "outputs": [],
      "source": [
        "folder_path = None #change path\n",
        "files = {\n",
        "    \"tooth\": {\n",
        "        \"file_name\": \"tooth_103x94x161_1x1x1_uint8.raw\",\n",
        "        \"shape\": [103, 94, 161],\n",
        "        \"dtype\": np.uint8\n",
        "    }\n",
        "}\n",
        "\n",
        "# from https://www.geeksforgeeks.org/reading-binary-files-in-python/\n",
        "\n",
        "# Open the file in binary mode\n",
        "def readRawFile(file_path, datatype, shape, doReshape=True, divideBy=2**8, gaussianFilter=False, cutoffLow=(0,0), cutoffHigh=(255,255)):\n",
        "  array = np.fromfile(file_path, dtype=datatype).astype(np.float32) # read from file based on datatype\n",
        "  array[array<cutoffLow[0]] = cutoffLow[1] # set values below cutoff[0] to the value of cutoff[1]\n",
        "  array[array>cutoffHigh[0]] = cutoffHigh[1]\n",
        "  array = np.divide(array, divideBy) # divide input by specified (to get input into range 0-1)\n",
        "  if (doReshape):\n",
        "    array = np.reshape(array, list(reversed(shape))).transpose() # need to make sure reshape takes the same order as the file all x->all y->all z\n",
        "    if (gaussianFilter):\n",
        "      array = gaussian_filter(array, sigma=0.1, radius=3)\n",
        "  return array"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vss-WbLAYObS"
      },
      "source": [
        "#### Dataset class:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Pgh-GIJaC8f"
      },
      "source": [
        "Dataset class to help with machine learning; sample it like an array with an index (int) and returns a sample corresponding to that index"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "QIXSbByvl5e2"
      },
      "outputs": [],
      "source": [
        "class Dataset:\n",
        "  def __init__(self,samples,stride=[4,4,4],kernel_size=[8,8,8], returnSurroundingIndices=True):\n",
        "    self.stride = stride # [x,y,z]\n",
        "    self.kernel_size = kernel_size # single value treated as length of the side of a cube\n",
        "    self.returnSurroundingIndices = returnSurroundingIndices\n",
        "    self.orig_shape = samples.shape\n",
        "    self.samples = self.pad_to_divisible_and_kernel(samples)\n",
        "\n",
        "    self.unfold_dims = np.array([ # num samples in each dimension\n",
        "        int(np.ceil((self.orig_shape[0] - self.kernel_size[0])/self.stride[0])+1),\n",
        "        int(np.ceil((self.orig_shape[1] - self.kernel_size[1])/self.stride[1])+1),\n",
        "        int(np.ceil((self.orig_shape[2] - self.kernel_size[2])/self.stride[2])+1)\n",
        "    ])\n",
        "    self.length =  np.prod(self.unfold_dims) # total num samples\n",
        "\n",
        "  def pad_to_divisible_and_kernel(self, arr):\n",
        "    # pad the array to be divisible by kernel size\n",
        "    to_pad = np.remainder(arr.shape, self.kernel_size)\n",
        "    to_pad = [self.kernel_size[x] - to_pad[x] if to_pad[x] > 0 else 0 for x in range(3)]\n",
        "    arr_pd = np.pad(arr, ( (0, to_pad[0]), (0, to_pad[1]), (0, to_pad[2]) ), mode='constant', constant_values='0')\n",
        "\n",
        "    if self.returnSurroundingIndices:\n",
        "      # pad an extra kernel_size around array\n",
        "      extra_pad = [[self.kernel_size[x],self.kernel_size[x]] for x in range(3)]\n",
        "      arr_pd2 = torch.tensor(np.pad(arr_pd, extra_pad, mode='constant', constant_values=0))\n",
        "    else:\n",
        "      arr_pd2 = arr_pd\n",
        "    return torch.Tensor(arr_pd2)\n",
        "\n",
        "\n",
        "  def __getitem__(self,index):\n",
        "    arr_coord = np.multiply(np.unravel_index(index, self.unfold_dims), self.stride)\n",
        "    if self.returnSurroundingIndices:\n",
        "      arr_coord = np.add(arr_coord, self.kernel_size) # skip initial kernel_size of padding\n",
        "    ix = 0\n",
        "    if self.returnSurroundingIndices:\n",
        "      surrounding_indices = [-1, 0, 1]\n",
        "      ret_arr = torch.empty((27, *self.kernel_size))\n",
        "    else:\n",
        "      surrounding_indices = [0]\n",
        "      ret_arr = torch.empty((1, *self.kernel_size))\n",
        "    for z in surrounding_indices:\n",
        "      for y in surrounding_indices:\n",
        "        for x in surrounding_indices:\n",
        "          tmp_coords = np.add(arr_coord, np.multiply([x, y, z], self.kernel_size))\n",
        "          # print(ix, tmp_coords, z, y, x)\n",
        "          ret_arr[ix] = self.get_sub_matrix_from_coords(tmp_coords)\n",
        "          ix += 1\n",
        "    return ret_arr\n",
        "\n",
        "  def get_sub_matrix_from_coords(self, coords):\n",
        "    batch_start = coords\n",
        "    # print(batch_start)\n",
        "    batch_end = np.add(coords, self.kernel_size)\n",
        "    # print(batch_end)\n",
        "    ret_arr = self.samples[batch_start[0]:batch_end[0], batch_start[1]:batch_end[1], batch_start[2]:batch_end[2]]\n",
        "    # print(ret_arr.shape)\n",
        "    return ret_arr\n",
        "\n",
        "  def __len__(self):\n",
        "    return self.length\n",
        "\n",
        "  def shuffle(self):\n",
        "    pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nWjVMP_BZ7vW"
      },
      "source": [
        "#### Test read write:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KlbeDJ179jTr"
      },
      "source": [
        "Pass model, to see if reading/writing file gives same output as input:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "TK9YzcHsClC0"
      },
      "outputs": [],
      "source": [
        "tooth_kernel_size = [8]*3\n",
        "tooth_stride = [8]*3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DCNMmiNq9evm",
        "outputId": "775b31aa-cf41-49e0-fb38-aa56604086b5"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(103, 94, 161)"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "tooth_path=ospath.join(folder_path, files[\"tooth\"][\"file_name\"])\n",
        "tooth_array=readRawFile(tooth_path, files[\"tooth\"][\"dtype\"], files[\"tooth\"][\"shape\"]) #, cutoffLow=(105,0), gaussianFilter=True) # cutoff = (105,0) and gaussianFilter are good denoising parameters for the tooth model\n",
        "tooth_array.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bnoy4JroIoep",
        "outputId": "8bd79932-f7ae-4ca7-f99a-a00be67e8d90"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "count    1.558802e+06\n",
              "mean     1.031652e+02\n",
              "std      4.591717e+01\n",
              "min      0.000000e+00\n",
              "25%      7.600000e+01\n",
              "50%      7.900000e+01\n",
              "75%      1.220000e+02\n",
              "max      2.550000e+02\n",
              "dtype: float64"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "pd.Series(readRawFile(tooth_path, files[\"tooth\"][\"dtype\"], files[\"tooth\"][\"shape\"], doReshape=False, divideBy=1)).describe()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_tPQNKv9EGdN",
        "outputId": "8e9bb4e8-fba9-482f-e584-05044d2178f2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "3276 torch.Size([104, 96, 168]) [13 12 21] (103, 94, 161) [8, 8, 8] [7 6 1]\n"
          ]
        }
      ],
      "source": [
        "toothPassDset = Dataset(tooth_array, stride=tooth_stride, kernel_size=tooth_kernel_size, returnSurroundingIndices=False)\n",
        "print(toothPassDset.length, toothPassDset.samples.shape, toothPassDset.unfold_dims, toothPassDset.orig_shape, tooth_kernel_size, np.remainder((toothPassDset.orig_shape), (tooth_kernel_size)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "qVtV1QnkCasx"
      },
      "outputs": [],
      "source": [
        "pass_array_out = torch.empty(toothPassDset.samples.shape)\n",
        "for ix in range(toothPassDset.length):\n",
        "  outIst = np.multiply(np.unravel_index(ix, toothPassDset.unfold_dims), toothPassDset.stride)\n",
        "  outIed = np.add(outIst, toothPassDset.kernel_size)\n",
        "  pass_array_out[outIst[0]:outIed[0], outIst[1]:outIed[1], outIst[2]:outIed[2]] = toothPassDset[ix]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "t0PMIZ27n3Kt",
        "outputId": "4171f748-c982-48c2-a6e1-d3de2d1cb210"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "3275 [ 96  88 160] [104  96 168]\n"
          ]
        }
      ],
      "source": [
        "print(ix, outIst, outIed)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "l0NJdTIdMsni",
        "outputId": "f1fbb291-2631-46ff-8832-59ebd31c9000"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "1.0"
            ]
          },
          "execution_count": 11,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "np.sum((pass_array_out == toothPassDset.samples).flatten().numpy())/np.prod(pass_array_out.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HiVkT0onfkkk",
        "outputId": "0f19f9f2-2b8c-4f65-f90e-2c869ad9362c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "MSE reconstructed - input: 0.0\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "array([73, 75, 78, 77, 81], dtype=uint8)"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "write_array = pass_array_out[0:toothPassDset.orig_shape[0], 0:toothPassDset.orig_shape[1], 0:toothPassDset.orig_shape[2]].numpy()\n",
        "mse = np.mean((write_array - tooth_array)**2) # from https://stackoverflow.com/a/18047482\n",
        "print(f\"MSE reconstructed - input: {mse}\")\n",
        "write_array = (write_array.transpose().reshape(-1)*255).astype(np.uint8)\n",
        "write_array[:5]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SNAa08wMevHv",
        "outputId": "c8236df7-f882-426f-cca5-2aea34a8c73c"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "9.622774412657926e-06"
            ]
          },
          "execution_count": 13,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "np.sum(readRawFile(tooth_path, files[\"tooth\"][\"dtype\"], files[\"tooth\"][\"shape\"], False) == write_array)/np.prod(files[\"tooth\"][\"shape\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "quqU3h27gQyI"
      },
      "outputs": [],
      "source": [
        "write_array.tofile(ospath.join(folder_path, \"pass_out_test.raw\"), format=\"%<hhf\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PLIxWwNjgpcc",
        "outputId": "3d71cf74-f896-44c0-838a-657dca458cd0"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(1558802,)"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "read_array = np.fromfile(ospath.join(folder_path, \"pass_out_test.raw\"), dtype=np.uint8)\n",
        "read_array.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cvTzipvvhWUg",
        "outputId": "a940c2ca-4e41-46c8-b6ef-573041da8ede"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "1558802"
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "np.sum(write_array == read_array)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K0wR4xxaigKo"
      },
      "source": [
        "As it seems like the dataset, reading, writing and such works, we can move on to (neural network) models for compression"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gs7Bdey6Yco8"
      },
      "source": [
        "#### Fit function:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z9AIamtnuAXw"
      },
      "source": [
        "Define the fit function:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "YGKp15QM1jBj"
      },
      "outputs": [],
      "source": [
        "from torch.optim import SGD, Adam\n",
        "from torch.nn import MSELoss, CrossEntropyLoss, L1Loss\n",
        "import datetime\n",
        "\n",
        "def fit(model, dset_train, num_epochs=1, train_batches=1, optimFunc=Adam, lossFunc=MSELoss, folder_path=folder_path, verbose=2):\n",
        "  # with help from https://pytorch.org/tutorials/beginner/introyt/trainingyt.html\n",
        "\n",
        "  curr_best_loss = np.Infinity\n",
        "  training_losses = []\n",
        "  save_path = ospath.join(folder_path, model.name)\n",
        "\n",
        "  optimizer = optimFunc(model.parameters(), lr=0.1)\n",
        "  loss = lossFunc()\n",
        "  model.train(True)\n",
        "  start_time = datetime.datetime.now()\n",
        "  if (verbose > 0):\n",
        "    print(f\"starting training for model '{model.name}' at: {datetime.datetime.now()}\")\n",
        "  for i in range(num_epochs):\n",
        "    t_losses = []\n",
        "\n",
        "    for j in range(train_batches):\n",
        "      x_train = dset_train[j]\n",
        "      p = model(x_train)\n",
        "      l = loss(p, x_train)\n",
        "\n",
        "      optimizer.zero_grad()\n",
        "      l.backward()\n",
        "      optimizer.step()\n",
        "\n",
        "      t_losses.append(l.item())\n",
        "\n",
        "    curr_loss = np.mean(t_losses)\n",
        "    training_losses.append(curr_loss)\n",
        "\n",
        "    if curr_loss < curr_best_loss:\n",
        "      curr_best_loss = curr_loss\n",
        "      torch.save(model.state_dict(), save_path)\n",
        "    if (verbose > 1):\n",
        "      print(f\"epoch {i+1}/{num_epochs} at time {datetime.datetime.now().time()}; current loss: {curr_loss}\")\n",
        "    elif (verbose == 1):\n",
        "      print(f\"epoch {i+1}/{num_epochs: <{10}}\", end=\"\\r\")\n",
        "\n",
        "  end_time = datetime.datetime.now()\n",
        "  if (verbose > 0):\n",
        "    print(f\"end of training {num_epochs} epochs: {datetime.datetime.now()}; elapsed time: {end_time - start_time}\")\n",
        "  best_model = model\n",
        "  best_model.load_state_dict(torch.load(save_path))\n",
        "\n",
        "  return best_model, training_losses"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DsIbRSAHhkOc"
      },
      "source": [
        "#### Linear models:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QO4fmrHPUX5r"
      },
      "source": [
        "##### AELinPass  \n",
        "Try a pass linear model to make sure it works"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "8xP4Wsd4S3UQ"
      },
      "outputs": [],
      "source": [
        "class autoEncLinPass(nn.Module):\n",
        "\n",
        "  def __init__(self, name=\"linae1\", kernel_size=[8,8,8]):\n",
        "    super(autoEncLinPass, self).__init__()\n",
        "    sample_size = np.prod(kernel_size)\n",
        "    self.encoder = nn.Sequential(\n",
        "      nn.Flatten(),\n",
        "      )\n",
        "    self.decoder = nn.Sequential(\n",
        "      nn.Unflatten(1, [1, *kernel_size]),\n",
        "      )\n",
        "    self.name = name\n",
        "\n",
        "  def get_encoding(self, x):\n",
        "    return self.encoder(x)\n",
        "\n",
        "  def get_decoding(self, c):\n",
        "    return self.decoder(c)\n",
        "\n",
        "  def forward(self, x):\n",
        "    c = self.encoder(x)\n",
        "    y = self.decoder(c)\n",
        "    return y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "DHZ-U-IlURhc"
      },
      "outputs": [],
      "source": [
        "def trainAndOutputAutoEncPassModel(folder_path=folder_path, model=autoEncLinPass, kernel_size=[4,4,4], stride_size=[4,4,4], epochs=4, name=\"model1lin\", vol_name=\"tooth\", verbose=True):\n",
        "  # to train model\n",
        "  tooth_path=ospath.join(folder_path, files[vol_name][\"file_name\"])\n",
        "  tooth_array=readRawFile(tooth_path, files[vol_name][\"dtype\"], files[vol_name][\"shape\"])#, cutoffLow=(105,0))#, gaussianFilter=True)\n",
        "  model1toothDset = Dataset(tooth_array, stride=stride_size, kernel_size=kernel_size, returnSurroundingIndices=False)\n",
        "  model1 = model(name=name, kernel_size=kernel_size)\n",
        "\n",
        "  # to output encoded representation\n",
        "  if (verbose):\n",
        "    print(f\"starting encoding at: {datetime.datetime.now()}\")\n",
        "  toothEncDset = Dataset(tooth_array, stride=kernel_size, kernel_size=kernel_size, returnSurroundingIndices=False)\n",
        "  model1_encoded_reps = np.array([model1.encoder(toothEncDset[x]).detach().numpy().flatten() for x in range(toothEncDset.length)])\n",
        "  model1_encoded_reps.reshape(-1).tofile(ospath.join(folder_path, f\"{model1.name}.enc\"), format=\"%<f\")\n",
        "\n",
        "  # output reconstructed model\n",
        "  if (verbose):\n",
        "      print(f\"starting decoding at: {datetime.datetime.now()}\")\n",
        "  model1_reconstructed = np.empty(toothEncDset.samples.shape)\n",
        "  for ix, enc in enumerate(model1_encoded_reps):\n",
        "    tmp_dec = model1.decoder(torch.tensor(enc).reshape((1,-1))).detach().numpy()\n",
        "    ixst = np.multiply(np.unravel_index(ix, toothEncDset.unfold_dims), toothEncDset.stride)\n",
        "    ixed = np.add(ixst,toothEncDset.kernel_size)\n",
        "    # print(ix, ixst, ixed)\n",
        "    model1_reconstructed[ixst[0]:ixed[0], ixst[1]:ixed[1], ixst[2]:ixed[2]] = tmp_dec\n",
        "\n",
        "  model1_reconstructed = model1_reconstructed[0:toothEncDset.orig_shape[0], 0:toothEncDset.orig_shape[1], 0:toothEncDset.orig_shape[2]]\n",
        "  # model1_reconstructed = gaussian_filter(model1_reconstructed, sigma=1, radius=np.divide(kernel_size,2).astype(np.int32)) # filter output to get rid of some of the blockiness\n",
        "  mse = np.mean(np.power((model1_reconstructed - tooth_array),2)) # from https://stackoverflow.com/a/18047482\n",
        "  print(f\"MSE reconstructed - input: {mse}\")\n",
        "  model1_reconstructed = (model1_reconstructed*255).astype(np.uint8).transpose().reshape(-1)\n",
        "  model1_reconstructed.tofile(ospath.join(folder_path, f\"{model1.name}.raw\"), format=\"%<hhf\")\n",
        "\n",
        "  return"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "VCklSMQ2TR-x"
      },
      "outputs": [],
      "source": [
        "model1_kernel = [8]*3\n",
        "model1_stride = [8]*3\n",
        "model1_name = \"linae1pass_8_8\"\n",
        "model1 = autoEncLinPass\n",
        "\n",
        "model1_epochs = 20 # 0 epochs just runs the encoder/decoder without training\n",
        "\n",
        "# trainAndOutputAutoEncPassModel(folder_path=folder_path, model=model1, kernel_size=model1_kernel, stride_size=model1_stride, epochs=model1_epochs, name=model1_name, vol_name=\"tooth\", verbose=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v_pxDvGoCYgl"
      },
      "source": [
        "##### AELin1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "qaVQCC8GYfHV"
      },
      "outputs": [],
      "source": [
        "class autoEncLin1(nn.Module):\n",
        "\n",
        "  def __init__(self, name=\"linae1\", kernel_size=[8,8,8], enc_size=9):\n",
        "    super(autoEncLin1, self).__init__()\n",
        "    sample_size = np.prod(kernel_size)\n",
        "    self.encoder = nn.Sequential(\n",
        "      nn.Flatten(),\n",
        "      nn.Linear(in_features=sample_size, out_features=int(sample_size/2)),\n",
        "      nn.PReLU(), #use parametric relu for activation functions\n",
        "      nn.Linear(int(sample_size/2), int(sample_size/4)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/4), int(sample_size/8)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/8), enc_size),\n",
        "      nn.PReLU(),\n",
        "      )\n",
        "    self.decoder = nn.Sequential(\n",
        "      nn.Linear(enc_size, int(sample_size/8)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/8), int(sample_size/4)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/4), int(sample_size/2)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/2), int(sample_size)),\n",
        "      nn.Unflatten(1, kernel_size),\n",
        "      nn.Sigmoid(),\n",
        "      )\n",
        "    self.name = name\n",
        "\n",
        "  def get_encoding(self, x):\n",
        "    return self.encoder(x)\n",
        "\n",
        "  def get_decoding(self, c):\n",
        "    return self.decoder(c)\n",
        "\n",
        "  def forward(self, x):\n",
        "    c = self.encoder(x)\n",
        "    y = self.decoder(c)\n",
        "    return y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "id": "CwVExSsnAVa2"
      },
      "outputs": [],
      "source": [
        "def trainAndOutputAutoEnc1Model(folder_path=folder_path, model=autoEncLin1, kernel_size=[4,4,4], stride_size=[4,4,4], enc_size=4, epochs=4, name=\"model1lin\", optim=SGD, loss=MSELoss, vol_name=\"tooth\", verbose=0, load_model_weights=False):\n",
        "  # to train model\n",
        "  tooth_path=ospath.join(folder_path, files[vol_name][\"file_name\"])\n",
        "  tooth_array=readRawFile(tooth_path, files[vol_name][\"dtype\"], files[vol_name][\"shape\"], cutoffLow=(100,0))#, gaussianFilter=True)\n",
        "  model1toothDset = Dataset(tooth_array, stride=stride_size, kernel_size=kernel_size, returnSurroundingIndices=False)\n",
        "  model1 = model(name=name, kernel_size=kernel_size, enc_size=enc_size)\n",
        "  if (load_model_weights):\n",
        "    save_path = ospath.join(folder_path, model1.name)\n",
        "    model1.load_state_dict(torch.load(save_path))\n",
        "  model1_best, model1_losses = fit(model1, model1toothDset, num_epochs=epochs, train_batches=model1toothDset.length-1, optimFunc=optim, lossFunc=loss, folder_path=folder_path, verbose=verbose)\n",
        "\n",
        "  # to output encoded representation\n",
        "  if (verbose > 0):\n",
        "    print(f\"starting encoding at: {datetime.datetime.now()}\")\n",
        "  toothEncDset = Dataset(tooth_array, stride=kernel_size, kernel_size=kernel_size, returnSurroundingIndices=False)\n",
        "  model1_encoded_reps = np.array([model1_best.encoder(toothEncDset[x]).detach().numpy().flatten().astype(np.float16) for x in range(toothEncDset.length)])\n",
        "  model1_encoded_reps.reshape(-1).tofile(ospath.join(folder_path, f\"{model1.name}.enc\"), format=\"%<hf\")\n",
        "\n",
        "  # output reconstructed model\n",
        "  if (verbose>0):\n",
        "      print(f\"starting decoding at: {datetime.datetime.now()}\")\n",
        "  model1_reconstructed = np.empty(toothEncDset.samples.shape)\n",
        "  for ix, enc in enumerate(model1_encoded_reps):\n",
        "    tmp_dec = model1_best.decoder(torch.tensor(enc.astype(np.float32)).reshape((1,-1))).detach().numpy()\n",
        "    ixst = np.multiply(np.unravel_index(ix, toothEncDset.unfold_dims), toothEncDset.stride)\n",
        "    ixed = np.add(ixst,toothEncDset.kernel_size)\n",
        "    # print(ix, ixst, ixed)\n",
        "    model1_reconstructed[ixst[0]:ixed[0], ixst[1]:ixed[1], ixst[2]:ixed[2]] = tmp_dec\n",
        "  if (verbose>0):\n",
        "      print(f\"finish decoding at: {datetime.datetime.now()}\")\n",
        "\n",
        "  model1_reconstructed = model1_reconstructed[0:toothEncDset.orig_shape[0], 0:toothEncDset.orig_shape[1], 0:toothEncDset.orig_shape[2]]\n",
        "  # model1_reconstructed = gaussian_filter(model1_reconstructed, sigma=1, radius=np.divide(kernel_size,2).astype(np.float32)) # filter output to get rid of some of the blockiness\n",
        "  mse = np.mean((model1_reconstructed - tooth_array)**2) # from https://stackoverflow.com/a/18047482\n",
        "  mae = np.mean(np.abs((model1_reconstructed - tooth_array)))\n",
        "  print(f\"MSE reconstructed - input: {mse}; MAE: {mae}\")\n",
        "  model1_reconstructed = (model1_reconstructed*255).astype(np.uint8).transpose().reshape(-1)\n",
        "  model1_reconstructed.tofile(ospath.join(folder_path, f\"{model1.name}.raw\"), format=\"%<hhf\")\n",
        "\n",
        "  model1_decoder = model1_best\n",
        "  model1_decoder.encoder = None\n",
        "  torch.save(model1_decoder.state_dict(), ospath.join(folder_path, f\"{model1.name}.dec\"))\n",
        "\n",
        "  return model1_losses"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2t_minHgkciU",
        "outputId": "1ea7e2b1-d905-45b9-bebd-d0f9532b8864"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "starting training for model 'linae1_8_8_8' at: 2024-05-26 15:38:03.882953\n",
            "epoch 1/1 at time 15:38:14.897525; current loss: 0.0012940410842699234\n",
            "end of training 1 epochs: 2024-05-26 15:38:14.897757; elapsed time: 0:00:11.014809\n",
            "starting encoding at: 2024-05-26 15:38:14.909588\n",
            "starting decoding at: 2024-05-26 15:38:16.133603\n",
            "finish decoding at: 2024-05-26 15:38:17.270399\n",
            "MSE reconstructed - input: 0.0017210463625708953; MAE: 0.015702462532426957\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "count    1.000000\n",
              "mean     0.001294\n",
              "std           NaN\n",
              "min      0.001294\n",
              "25%      0.001294\n",
              "50%      0.001294\n",
              "75%      0.001294\n",
              "max      0.001294\n",
              "dtype: float64"
            ]
          },
          "execution_count": 41,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model1_kernel = [8]*3\n",
        "model1_stride = [8]*3\n",
        "model1_enc_size = 8\n",
        "model1_name = \"linae1_8_8_8\"\n",
        "model1 = autoEncLin1\n",
        "\n",
        "model1_epochs = 1 # 0 epochs just runs the encoder/decoder without training\n",
        "model1_optim = SGD  #Adam#SGD\n",
        "model1_loss = MSELoss   #CrossEntropyLoss#MSELoss\n",
        "model1_load_weights = True # if changing anything above this line (especially kernel,stride,enc_size,name), set load_model_weights to False\n",
        "\n",
        "model1_losses = trainAndOutputAutoEnc1Model(folder_path, model1, model1_kernel, model1_stride, model1_enc_size, model1_epochs, model1_name, model1_optim, model1_loss, vol_name=\"tooth\", verbose=2, load_model_weights=model1_load_weights)\n",
        "pd.Series(model1_losses).describe()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rdPi-wUqc4nk",
        "outputId": "5e838633-7e3e-4969-9b57-cff42fa477eb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "starting training for model 'linae1_4_4_8' at: 2024-05-26 15:38:17.387903\n",
            "epoch 1/1 at time 15:39:12.868775; current loss: 0.0006420350716156618\n",
            "end of training 1 epochs: 2024-05-26 15:39:12.869028; elapsed time: 0:00:55.481128\n",
            "starting encoding at: 2024-05-26 15:39:12.877005\n",
            "starting decoding at: 2024-05-26 15:39:20.718999\n",
            "finish decoding at: 2024-05-26 15:39:29.386775\n",
            "MSE reconstructed - input: 0.000742868928268601; MAE: 0.009005612449657726\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "count    1.000000\n",
              "mean     0.000642\n",
              "std           NaN\n",
              "min      0.000642\n",
              "25%      0.000642\n",
              "50%      0.000642\n",
              "75%      0.000642\n",
              "max      0.000642\n",
              "dtype: float64"
            ]
          },
          "execution_count": 42,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model1_kernel = [4]*3\n",
        "model1_stride = [4]*3\n",
        "model1_enc_size = 8\n",
        "model1_name = \"linae1_4_4_8\"\n",
        "model1 = autoEncLin1\n",
        "\n",
        "model1_epochs = 1 # 0 epochs just runs the encoder/decoder without training\n",
        "model1_optim = SGD  #Adam#SGD\n",
        "model1_loss = MSELoss   #CrossEntropyLoss#MSELoss\n",
        "model1_load_weights = True # if changing anything above this line (especially kernel,stride,enc_size,name), set load_model_weights to False\n",
        "\n",
        "model1_losses = trainAndOutputAutoEnc1Model(folder_path, model1, model1_kernel, model1_stride, model1_enc_size, model1_epochs, model1_name, model1_optim, model1_loss, vol_name=\"tooth\", verbose=2, load_model_weights=model1_load_weights)\n",
        "pd.Series(model1_losses).describe()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GPM284W_loUw"
      },
      "source": [
        "##### AELin2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "id": "_ZQiIbCNe6Ks"
      },
      "outputs": [],
      "source": [
        "class autoEncLin2(nn.Module):\n",
        "\n",
        "  def __init__(self, name=\"linae2\", kernel_size=[8,8,8], enc_size=9):\n",
        "    super(autoEncLin2, self).__init__()\n",
        "    sample_size = np.prod(kernel_size)\n",
        "    self.encoder = nn.Sequential(\n",
        "      nn.Flatten(),\n",
        "      nn.Linear(in_features=sample_size, out_features=int(sample_size/8)),\n",
        "      nn.PReLU(), #use parametric relu for activation functions\n",
        "      nn.Linear(int(sample_size/8), int(sample_size/8)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/8), int(sample_size/12)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/12), enc_size),\n",
        "      nn.PReLU(),\n",
        "      )\n",
        "    self.decoder = nn.Sequential(\n",
        "      nn.Linear(enc_size, int(sample_size/12)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/12), int(sample_size/8)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/8), int(sample_size/8)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/8), int(sample_size)),\n",
        "      nn.Unflatten(1, kernel_size),\n",
        "      nn.Sigmoid(),\n",
        "      )\n",
        "    self.name = name\n",
        "\n",
        "  def get_encoding(self, x):\n",
        "    return self.encoder(x)\n",
        "\n",
        "  def get_decoding(self, c):\n",
        "    return self.decoder(c)\n",
        "\n",
        "  def forward(self, x):\n",
        "    c = self.encoder(x)\n",
        "    y = self.decoder(c)\n",
        "    return y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DcY0OG-TeTeE",
        "outputId": "df0f13f4-127b-4a90-c3a3-03f101894669"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "starting training for model 'linae2_8_8_8' at: 2024-05-26 15:43:32.260149\n",
            "epoch 1/1 at time 15:43:39.872894; current loss: 0.0022808388583393066\n",
            "end of training 1 epochs: 2024-05-26 15:43:39.874846; elapsed time: 0:00:07.614695\n",
            "starting encoding at: 2024-05-26 15:43:39.882852\n",
            "starting decoding at: 2024-05-26 15:43:41.070351\n",
            "finish decoding at: 2024-05-26 15:43:42.214508\n",
            "MSE reconstructed - input: 0.0029273420156157907; MAE: 0.02294337057319233\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "count    1.000000\n",
              "mean     0.002281\n",
              "std           NaN\n",
              "min      0.002281\n",
              "25%      0.002281\n",
              "50%      0.002281\n",
              "75%      0.002281\n",
              "max      0.002281\n",
              "dtype: float64"
            ]
          },
          "execution_count": 44,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model1_kernel = [8]*3\n",
        "model1_stride = [8]*3\n",
        "model1_enc_size = 8\n",
        "model1_name = f\"linae2_8_8_{model1_enc_size}\"\n",
        "model1 = autoEncLin2\n",
        "\n",
        "model1_epochs = 1 # 0 epochs just runs the encoder/decoder without training\n",
        "model1_optim = SGD  #Adam#SGD\n",
        "model1_loss = MSELoss   #CrossEntropyLoss#MSELoss\n",
        "model1_load_weights = True # if changing anything above this line (especially kernel,stride,enc_size,name), set load_model_weights to False\n",
        "\n",
        "model1_losses = trainAndOutputAutoEnc1Model(folder_path, model1, model1_kernel, model1_stride, model1_enc_size, model1_epochs, model1_name, model1_optim, model1_loss, vol_name=\"tooth\", verbose=2, load_model_weights=model1_load_weights)\n",
        "pd.Series(model1_losses).describe()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "v9oGsw7-coId",
        "outputId": "ee675a6a-2faf-478c-a4df-ed51c04d74f1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "starting training for model 'linae2_4_4_8' at: 2024-05-26 15:43:42.378243\n",
            "epoch 1/1 at time 15:44:35.730839; current loss: 0.0017531924037302\n",
            "end of training 1 epochs: 2024-05-26 15:44:35.731926; elapsed time: 0:00:53.353684\n",
            "starting encoding at: 2024-05-26 15:44:35.737713\n",
            "starting decoding at: 2024-05-26 15:44:45.032692\n",
            "finish decoding at: 2024-05-26 15:44:52.306596\n",
            "MSE reconstructed - input: 0.0021151287194519485; MAE: 0.016955845855715797\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "count    1.000000\n",
              "mean     0.001753\n",
              "std           NaN\n",
              "min      0.001753\n",
              "25%      0.001753\n",
              "50%      0.001753\n",
              "75%      0.001753\n",
              "max      0.001753\n",
              "dtype: float64"
            ]
          },
          "execution_count": 45,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model1_kernel = [4]*3\n",
        "model1_stride = [4]*3\n",
        "model1_enc_size = 8\n",
        "model1_name = f\"linae2_4_4_{model1_enc_size}\"\n",
        "model1 = autoEncLin2\n",
        "\n",
        "model1_epochs = 1 # 0 epochs just runs the encoder/decoder without training\n",
        "model1_optim = SGD  #Adam#SGD\n",
        "model1_loss = MSELoss   #CrossEntropyLoss#MSELoss\n",
        "model1_load_weights = True # if changing anything above this line (especially kernel,stride,enc_size,name), set load_model_weights to False\n",
        "\n",
        "model1_losses = trainAndOutputAutoEnc1Model(folder_path, model1, model1_kernel, model1_stride, model1_enc_size, model1_epochs, model1_name, model1_optim, model1_loss, vol_name=\"tooth\", verbose=2, load_model_weights=model1_load_weights)\n",
        "pd.Series(model1_losses).describe()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MdCaWq-IhrFW"
      },
      "source": [
        "#### Non-neural encoders (or mixed models)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k3gsSGFME9q9"
      },
      "source": [
        "Try a model that uses coordinates, variances and singular values"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 50,
      "metadata": {
        "id": "TbKVr6kJjwcu"
      },
      "outputs": [],
      "source": [
        "from torch.optim import SGD, Adam\n",
        "from torch.nn import MSELoss, CrossEntropyLoss\n",
        "import datetime\n",
        "\n",
        "def fitWithIndex(model, dset_train, num_epochs=1, train_batches=1, optimFunc=Adam, lossFunc=MSELoss, folder_path=folder_path, verbose=2):\n",
        "  # with help from https://pytorch.org/tutorials/beginner/introyt/trainingyt.html\n",
        "\n",
        "  curr_best_loss = np.Infinity\n",
        "  training_losses = []\n",
        "  save_path = ospath.join(folder_path, model.name)\n",
        "\n",
        "  optimizer = optimFunc(model.parameters(), lr=0.1)\n",
        "  loss = lossFunc()\n",
        "  model.train(True)\n",
        "  start_time = datetime.datetime.now()\n",
        "  if (verbose > 0):\n",
        "    print(f\"starting training for model '{model.name}' at: {datetime.datetime.now()}\")\n",
        "  for i in range(num_epochs):\n",
        "    t_losses = []\n",
        "\n",
        "    for j in range(train_batches):\n",
        "      x_train = dset_train[j]\n",
        "      p = model(x_train, j)\n",
        "      l = loss(p, x_train)\n",
        "\n",
        "      optimizer.zero_grad()\n",
        "      l.backward()\n",
        "      optimizer.step()\n",
        "\n",
        "      t_losses.append(l.item())\n",
        "\n",
        "    curr_loss = np.mean(t_losses)\n",
        "    training_losses.append(curr_loss)\n",
        "\n",
        "    if curr_loss < curr_best_loss:\n",
        "      curr_best_loss = curr_loss\n",
        "      torch.save(model.state_dict(), save_path)\n",
        "    if (verbose > 1):\n",
        "      print(f\"epoch {i+1}/{num_epochs} at time {datetime.datetime.now().time()}; current loss: {curr_loss}\")\n",
        "    elif (verbose == 1):\n",
        "      print(f\"epoch {i+1}/{num_epochs: <{10}}\", end=\"\\r\")\n",
        "\n",
        "  end_time = datetime.datetime.now()\n",
        "  if (verbose > 0):\n",
        "    print(f\"end of training {num_epochs} epochs: {datetime.datetime.now()}; elapsed time: {end_time - start_time}\")\n",
        "  best_model = model\n",
        "  best_model.load_state_dict(torch.load(save_path))\n",
        "\n",
        "  return best_model, training_losses"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ikY5_a1_cboi"
      },
      "source": [
        "##### Mixmodel1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 51,
      "metadata": {
        "id": "EBqgjHsziQJ4"
      },
      "outputs": [],
      "source": [
        "class MixedModel1(nn.Module):\n",
        "\n",
        "  def get_encoding(self, x, index):\n",
        "    arr_coord = np.multiply(np.unravel_index(index, self.dset.unfold_dims), self.dset.stride)+self.kernel_size # add one kernel_size to skip padded 0s\n",
        "    ret_coords = np.divide(arr_coord, self.dset.samples.shape)\n",
        "    retCodeInfo = torch.empty(9, dtype=torch.float32)\n",
        "    # get normalized coords\n",
        "    retCodeInfo[0:3] = torch.tensor(ret_coords[0:3], dtype=torch.float32)\n",
        "    # get mean variances\n",
        "    retCodeInfo[3] = torch.mean(torch.var(torch.reshape(x, self.kernel_size), axis=(0)))\n",
        "    retCodeInfo[4] = torch.mean(torch.var(torch.reshape(x, self.kernel_size), axis=(1)))\n",
        "    retCodeInfo[5] = torch.mean(torch.var(torch.reshape(x, self.kernel_size), axis=(2)))\n",
        "    # get svds\n",
        "    retCodeInfo[6:9] = torch.linalg.svdvals(torch.reshape(x, self.kernel_size))[0:3,0]\n",
        "    #combine into ret_array\n",
        "\n",
        "    return retCodeInfo.reshape((1,-1))\n",
        "\n",
        "  def __init__(self, dset, name=\"mxae1\"):\n",
        "    super(MixedModel1, self).__init__()\n",
        "    self.dset = dset\n",
        "    kernel_size = dset.kernel_size\n",
        "    self.kernel_size = kernel_size\n",
        "    sample_size = np.prod(kernel_size)\n",
        "    self.encoder = self.get_encoding\n",
        "    self.decoder = nn.Sequential(\n",
        "      nn.Linear(9, int(sample_size/8)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/8), int(sample_size/4)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/4), int(sample_size/2)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/2), int(sample_size)),\n",
        "      nn.Unflatten(1, kernel_size),\n",
        "      nn.Sigmoid(),\n",
        "      )\n",
        "    self.name = name\n",
        "\n",
        "\n",
        "  def get_decoding(self, c):\n",
        "    return self.decoder(c)\n",
        "\n",
        "  def forward(self, x, i):\n",
        "    c = self.encoder(x, i)\n",
        "    y = self.decoder(c)\n",
        "    return y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 52,
      "metadata": {
        "id": "iEZiV4pvDlfO"
      },
      "outputs": [],
      "source": [
        "def trainAndOutputMixedModel(folder_path=folder_path, model=MixedModel1, kernel_size=[4,4,4], stride_size=[4,4,4], epochs=4, name=\"model1mix\", optim=SGD, loss=MSELoss, vol_name=\"tooth\", verbose=True, load_model_weights=False):\n",
        "  # to train model\n",
        "  tooth_path=ospath.join(folder_path, files[vol_name][\"file_name\"])\n",
        "  tooth_array=readRawFile(tooth_path, files[vol_name][\"dtype\"], files[vol_name][\"shape\"], cutoffLow=(105, 0))\n",
        "  model1toothDset = Dataset(tooth_array, stride=stride_size, kernel_size=kernel_size, returnSurroundingIndices=False)\n",
        "\n",
        "  model1 = model(dset=model1toothDset, name=name)\n",
        "  if (load_model_weights):\n",
        "    save_path = ospath.join(folder_path, model1.name)\n",
        "    model1.load_state_dict(torch.load(save_path))\n",
        "  model1_best, model1_losses = fitWithIndex(model1, model1toothDset, num_epochs=epochs, train_batches=model1toothDset.length-1, optimFunc=optim, lossFunc=loss, folder_path=folder_path, verbose=verbose)\n",
        "\n",
        "  # to output encoded representation\n",
        "  if (verbose):\n",
        "    print(f\"starting encoding at: {datetime.datetime.now()}\")\n",
        "  toothEncDset = Dataset(tooth_array, stride=kernel_size, kernel_size=kernel_size, returnSurroundingIndices=False)\n",
        "  model1_encoded_reps = np.array([model1_best.encoder(toothEncDset[x], x).detach().numpy().flatten().astype(np.float16) for x in range(toothEncDset.length)])\n",
        "  model1_encoded_reps.reshape(-1).tofile(ospath.join(folder_path, f\"{model1.name}.enc\"), format=\"%<hf\")\n",
        "\n",
        "  # output reconstructed model\n",
        "  if (verbose):\n",
        "      print(f\"starting decoding at: {datetime.datetime.now()}\")\n",
        "  model1_reconstructed = np.empty(toothEncDset.samples.shape)\n",
        "  for ix, enc in enumerate(model1_encoded_reps):\n",
        "    tmp_dec = model1_best.decoder(torch.tensor(enc, dtype=torch.float32).reshape((1,-1))).detach().numpy()\n",
        "    ixst = np.multiply(np.unravel_index(ix, toothEncDset.unfold_dims), toothEncDset.stride)\n",
        "    ixed = np.add(ixst,toothEncDset.kernel_size)\n",
        "    # print(ix, ixst, ixed)\n",
        "    model1_reconstructed[ixst[0]:ixed[0], ixst[1]:ixed[1], ixst[2]:ixed[2]] = tmp_dec\n",
        "  if (verbose>0):\n",
        "      print(f\"finish decoding at: {datetime.datetime.now()}\")\n",
        "\n",
        "  model1_reconstructed = model1_reconstructed[0:toothEncDset.orig_shape[0], 0:toothEncDset.orig_shape[1], 0:toothEncDset.orig_shape[2]]\n",
        "  # model1_reconstructed = gaussian_filter(model1_reconstructed, sigma=1, radius=np.divide(kernel_size,2).astype(np.float32)) # filter output to get rid of some of the blockiness\n",
        "  mse = np.mean((model1_reconstructed - tooth_array)**2) # from https://stackoverflow.com/a/18047482\n",
        "  mae = np.mean(np.abs((model1_reconstructed - tooth_array)))\n",
        "  print(f\"MSE reconstructed - input: {mse}; MAE: {mae}\")\n",
        "  model1_reconstructed = (model1_reconstructed*255).astype(np.uint8).transpose().reshape(-1)\n",
        "  model1_reconstructed.tofile(ospath.join(folder_path, f\"{model1.name}.raw\"), format=\"%<hhf\")\n",
        "\n",
        "  return model1_losses"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 53,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vLaWLrC_bSuu",
        "outputId": "ff409cc2-b15b-4ddf-f64c-9019f813dee2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "starting training for model 'coordvarsvd1_4_4' at: 2024-05-26 15:49:10.338370\n",
            "epoch 1/1 at time 15:49:56.141748; current loss: 0.001932225920083854\n",
            "end of training 1 epochs: 2024-05-26 15:49:56.141933; elapsed time: 0:00:45.803566\n",
            "starting encoding at: 2024-05-26 15:49:56.147734\n",
            "starting decoding at: 2024-05-26 15:50:05.656440\n",
            "finish decoding at: 2024-05-26 15:50:13.344371\n",
            "MSE reconstructed - input: 0.004003983363984769; MAE: 0.022566684042712035\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "count    1.000000\n",
              "mean     0.001932\n",
              "std           NaN\n",
              "min      0.001932\n",
              "25%      0.001932\n",
              "50%      0.001932\n",
              "75%      0.001932\n",
              "max      0.001932\n",
              "dtype: float64"
            ]
          },
          "execution_count": 53,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model1_kernel = [4]*3\n",
        "model1_stride = model1_kernel\n",
        "model1_name = \"coordvarsvd1_4_4\"\n",
        "model1 = MixedModel1\n",
        "\n",
        "model1_epochs = 1 # 0 epochs just runs the encoder/decoder without training\n",
        "model1_optim = SGD  #Adam#SGD\n",
        "model1_loss = MSELoss   #CrossEntropyLoss#MSELoss\n",
        "model1_load_weights = True # if changing anything above this line (especially kernel,stride,enc_size,name), set load_model_weights to False\n",
        "\n",
        "model1_losses = trainAndOutputMixedModel(folder_path, model1, model1_kernel, model1_stride, model1_epochs, model1_name, model1_optim, model1_loss, vol_name=\"tooth\", verbose=2, load_model_weights=model1_load_weights)\n",
        "pd.Series(model1_losses).describe()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5sMmJuZfcWUA"
      },
      "source": [
        "##### Mixmodel2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 54,
      "metadata": {
        "id": "RSVpbCMVepUM"
      },
      "outputs": [],
      "source": [
        "class MixedModel2(nn.Module):\n",
        "\n",
        "  def get_encoding(self, x, index):\n",
        "    arr_coord = np.multiply(np.unravel_index(index, self.dset.unfold_dims), self.dset.stride)+self.kernel_size # add one kernel_size to skip padded 0s\n",
        "    ret_coords = np.divide(arr_coord, self.dset.samples.shape)\n",
        "    retCodeInfo = torch.empty(3, dtype=torch.float32)\n",
        "    retCodeInfo = torch.linalg.svdvals(torch.reshape(x, self.kernel_size))[0:3,0]\n",
        "\n",
        "    return retCodeInfo.reshape((1,-1))\n",
        "\n",
        "  def __init__(self, dset, name=\"mxae2\"):\n",
        "    super(MixedModel2, self).__init__()\n",
        "    self.dset = dset\n",
        "    kernel_size = dset.kernel_size\n",
        "    self.kernel_size = kernel_size\n",
        "    sample_size = np.prod(kernel_size)\n",
        "    self.encoder = self.get_encoding\n",
        "    self.decoder = nn.Sequential(\n",
        "      nn.Linear(3, int(sample_size/8)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/8), int(sample_size/4)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/4), int(sample_size/2)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/2), int(sample_size)),\n",
        "      nn.Unflatten(1, kernel_size),\n",
        "      nn.Sigmoid(),\n",
        "      )\n",
        "    self.name = name\n",
        "\n",
        "\n",
        "  def get_decoding(self, c):\n",
        "    return self.decoder(c)\n",
        "\n",
        "  def forward(self, x, i):\n",
        "    c = self.encoder(x, i)\n",
        "    y = self.decoder(c)\n",
        "    return y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 55,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Sl2eXMnOglxT",
        "outputId": "eadb5bd8-af05-4f54-a38c-efb2c37f8170"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "starting training for model 'svd_4_4' at: 2024-05-26 15:50:13.463935\n",
            "epoch 1/1 at time 15:50:54.582840; current loss: 0.003970364428952854\n",
            "end of training 1 epochs: 2024-05-26 15:50:54.583052; elapsed time: 0:00:41.119118\n",
            "starting encoding at: 2024-05-26 15:50:54.588614\n",
            "starting decoding at: 2024-05-26 15:50:59.283024\n",
            "finish decoding at: 2024-05-26 15:51:08.095513\n",
            "MSE reconstructed - input: 0.0049014430961756205; MAE: 0.0259280144515157\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "count    1.00000\n",
              "mean     0.00397\n",
              "std          NaN\n",
              "min      0.00397\n",
              "25%      0.00397\n",
              "50%      0.00397\n",
              "75%      0.00397\n",
              "max      0.00397\n",
              "dtype: float64"
            ]
          },
          "execution_count": 55,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model1_kernel = [4]*3\n",
        "model1_stride = model1_kernel\n",
        "model1_name = \"svd_4_4\"\n",
        "model1 = MixedModel2\n",
        "\n",
        "model1_epochs = 1 # 0 epochs just runs the encoder/decoder without training\n",
        "model1_optim = SGD  #Adam#SGD\n",
        "model1_loss = MSELoss   #CrossEntropyLoss#MSELoss\n",
        "model1_load_weights = True # if changing anything above this line (especially kernel,stride,enc_size,name), set load_model_weights to False\n",
        "\n",
        "model1_losses = trainAndOutputMixedModel(folder_path, model1, model1_kernel, model1_stride, model1_epochs, model1_name, model1_optim, model1_loss, vol_name=\"tooth\", verbose=2, load_model_weights=model1_load_weights)\n",
        "pd.Series(model1_losses).describe()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CsjeXV-9P0Cz"
      },
      "source": [
        "##### Mixmodel3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 56,
      "metadata": {
        "id": "pmdZg5C2P5Ri"
      },
      "outputs": [],
      "source": [
        "class MixedModel3(nn.Module):\n",
        "\n",
        "  def get_encoding(self, x, index):\n",
        "    arr_coord = np.multiply(np.unravel_index(index, self.dset.unfold_dims), self.dset.stride)+self.kernel_size # add one kernel_size to skip padded 0s\n",
        "    ret_coords = np.divide(arr_coord, self.dset.samples.shape)\n",
        "    retCodeInfo = torch.empty(3, dtype=torch.float32)\n",
        "    retCodeInfo = torch.linalg.svdvals(torch.reshape(x, self.kernel_size))[0:3,0]\n",
        "\n",
        "    return retCodeInfo.reshape((1,-1))\n",
        "\n",
        "  def __init__(self, dset, name=\"mxae3\"):\n",
        "    super(MixedModel3, self).__init__()\n",
        "    self.dset = dset\n",
        "    kernel_size = dset.kernel_size\n",
        "    self.kernel_size = kernel_size\n",
        "    sample_size = np.prod(kernel_size)\n",
        "    self.encoder = self.get_encoding\n",
        "    self.decoder = nn.Sequential(\n",
        "      nn.Linear(3, int(sample_size/12)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/12), int(sample_size/8)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/8), int(sample_size/8)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/8), int(sample_size)),\n",
        "      nn.Unflatten(1, kernel_size),\n",
        "      nn.Sigmoid(),\n",
        "      )\n",
        "    self.name = name\n",
        "\n",
        "\n",
        "  def get_decoding(self, c):\n",
        "    return self.decoder(c)\n",
        "\n",
        "  def forward(self, x, i):\n",
        "    c = self.encoder(x, i)\n",
        "    y = self.decoder(c)\n",
        "    return y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 60,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dbJP5shIQGKK",
        "outputId": "328d039b-550e-4fe7-acbc-10e27a8814b2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "starting training for model 'svd2_4_4' at: 2024-05-26 15:59:16.224574\n",
            "epoch 1/1 at time 15:59:57.004942; current loss: 0.003994967876574053\n",
            "end of training 1 epochs: 2024-05-26 15:59:57.005140; elapsed time: 0:00:40.780569\n",
            "starting encoding at: 2024-05-26 15:59:57.011107\n",
            "starting decoding at: 2024-05-26 16:00:01.841116\n",
            "finish decoding at: 2024-05-26 16:00:10.194839\n",
            "MSE reconstructed - input: 0.004972182177168173; MAE: 0.02673630686496637\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "count    1.000000\n",
              "mean     0.003995\n",
              "std           NaN\n",
              "min      0.003995\n",
              "25%      0.003995\n",
              "50%      0.003995\n",
              "75%      0.003995\n",
              "max      0.003995\n",
              "dtype: float64"
            ]
          },
          "execution_count": 60,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model1_kernel = [4]*3\n",
        "model1_stride = model1_kernel\n",
        "model1_name = \"svd2_4_4\"\n",
        "model1 = MixedModel3\n",
        "\n",
        "model1_epochs = 1 # 0 epochs just runs the encoder/decoder without training\n",
        "model1_optim = SGD  #Adam#SGD\n",
        "model1_loss = MSELoss   #CrossEntropyLoss#MSELoss\n",
        "model1_load_weights = True # if changing anything above this line (especially kernel,stride,enc_size,name), set load_model_weights to False\n",
        "\n",
        "model1_losses = trainAndOutputMixedModel(folder_path, model1, model1_kernel, model1_stride, model1_epochs, model1_name, model1_optim, model1_loss, vol_name=\"tooth\", verbose=2, load_model_weights=model1_load_weights)\n",
        "pd.Series(model1_losses).describe()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wInKtCgCQe2Q"
      },
      "source": [
        "##### Mixmodel4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 58,
      "metadata": {
        "id": "D0X39tbJQe2S"
      },
      "outputs": [],
      "source": [
        "class MixedModel4(nn.Module):\n",
        "\n",
        "  def get_encoding(self, x, index):\n",
        "    arr_coord = np.multiply(np.unravel_index(index, self.dset.unfold_dims), self.dset.stride)+self.kernel_size # add one kernel_size to skip padded 0s\n",
        "    ret_coords = np.divide(arr_coord, self.dset.samples.shape)\n",
        "    retCodeInfo = torch.tensor(ret_coords, dtype=torch.float32)\n",
        "\n",
        "    return retCodeInfo.reshape((1,-1))\n",
        "\n",
        "  def __init__(self, dset, name=\"mxae4\"):\n",
        "    super(MixedModel4, self).__init__()\n",
        "    self.dset = dset\n",
        "    kernel_size = dset.kernel_size\n",
        "    self.kernel_size = kernel_size\n",
        "    sample_size = np.prod(kernel_size)\n",
        "    self.encoder = self.get_encoding\n",
        "    self.decoder = nn.Sequential(\n",
        "      nn.Linear(3, int(sample_size/12)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/12), int(sample_size/8)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/8), int(sample_size/8)),\n",
        "      nn.PReLU(),\n",
        "      nn.Linear(int(sample_size/8), int(sample_size)),\n",
        "      nn.Unflatten(1, kernel_size),\n",
        "      nn.Sigmoid(),\n",
        "      )\n",
        "    self.name = name\n",
        "\n",
        "\n",
        "  def get_decoding(self, c):\n",
        "    return self.decoder(c)\n",
        "\n",
        "  def forward(self, x, i):\n",
        "    c = self.encoder(x, i)\n",
        "    y = self.decoder(c)\n",
        "    return y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 59,
      "metadata": {
        "id": "6-Pad2t9Qe2S"
      },
      "outputs": [],
      "source": [
        "model1_kernel = [4]*3\n",
        "model1_stride = model1_kernel\n",
        "model1_name = \"coords2_4_4\"\n",
        "model1 = MixedModel4\n",
        "\n",
        "model1_epochs = 0 # 0 epochs just runs the encoder/decoder without training\n",
        "model1_optim = SGD  #Adam#SGD\n",
        "model1_loss = MSELoss   #CrossEntropyLoss#MSELoss\n",
        "model1_load_weights = True # if changing anything above this line (especially kernel,stride,enc_size,name), set load_model_weights to False\n",
        "\n",
        "# model1_losses = trainAndOutputMixedModel(folder_path, model1, model1_kernel, model1_stride, model1_epochs, model1_name, model1_optim, model1_loss, vol_name=\"tooth\", verbose=2, load_model_weights=model1_load_weights)\n",
        "# pd.Series(model1_losses).describe()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 59,
      "metadata": {
        "id": "yGE9XgbDQ70I"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "5CEuX--IZ-ke",
        "vss-WbLAYObS",
        "nWjVMP_BZ7vW",
        "gs7Bdey6Yco8"
      ],
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}