3D_slAIcer_meshCNN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "3D_slAIcer_meshCNN.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyPnBMTTaiDIHzCp22OVmOuC",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/PascalPolygon/3D_slAIcer/blob/main/3D_slAIcer_meshCNN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x11yefwOC60c"
      },
      "source": [
        "from google.colab import drive\n",
        "import os\n",
        "import numpy as np\n",
        "from tqdm import tqdm\n",
        "import matplotlib.pyplot as plt\n",
        "import random\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "import torch\n",
        "from torch.autograd import Variable\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "import glob\n",
        "import pickle\n",
        "import pathlib\n",
        "from threading import Thread\n",
        "from heapq import heappop, heapify\n",
        "from torch.nn import ConstantPad2d\n",
        "from torch.optim import lr_scheduler\n",
        "import functools"
      ],
      "execution_count": 81,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "83rsX2jYDbur",
        "outputId": "bb2953bb-05fd-47de-fee7-61a6accee703"
      },
      "source": [
        "drive.mount('/content/drive', force_remount=True)\n",
        "print(os.getcwd())"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n",
            "/content\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nZEwb59BDppo"
      },
      "source": [
        "if torch.cuda.is_available():\n",
        "  print(f'GPU : {torch.cuda.get_device_name(0)}')\n",
        "  cuda0 = torch.device('cuda:0')\n",
        "  # net.cuda()"
      ],
      "execution_count": 83,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J4R62U17D4zI"
      },
      "source": [
        "# train_data_dir = \"/content/drive/MyDrive/MTH4320/data/train/data\"\n",
        "datasets = ['coseg_aliens', 'coseg_chairs', 'coseg_vases']\n",
        "TRAIN_DATA_DIR = \"/content/drive/MyDrive/MTH4320/data/coseg/coseg_aliens/train\"\n",
        "DATA_DIR = \"/content/drive/MyDrive/MTH4320/data/coseg\"\n",
        "TRAIN_LABELS_DIR = \"/content/drive/MyDrive/MTH4320/data/train/labels\"\n",
        "\n",
        "MODELS_DIR = \"/content/drive/MyDrive/MTH4320/models/3D_slAIcer\"\n",
        "\n",
        "EXPORT_DIR = \"/content/drive/MyDrive/MTH4320/data/export_folder\"\n",
        "# test_data_dir = \"/content/drive/MyDrive/MTH4320/data/test/data\"\n",
        "# test_labels_dir = \"/content/drive/MyDrive/MTH4320/data/test/labels\""
      ],
      "execution_count": 86,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jDkVWFZcHDAv",
        "outputId": "c51ac706-60b7-49f7-e0c5-778278df5566"
      },
      "source": [
        "\n",
        "class TrainingOptions:\n",
        "  num_aug : 1\n",
        "  export_folder : ''\n",
        "  ninput_edges : 0 \n",
        "\n",
        "class SlicerDataset(Dataset):\n",
        "    def __init__(self, opt, transform=None, target_transform=None):\n",
        "        # self.paths = os.listdir(data_dir)\n",
        "        self.opt = opt\n",
        "        self.datasets = self.opt.datasets\n",
        "        self.data_dir = self.opt.data_dir\n",
        "        self.train_dir = self.opt.train_dir\n",
        "        self.labels_dir = self.opt.labels_dir\n",
        "        self.label_paths = glob.glob(self.labels_dir+'/*.npy')\n",
        "        self.paths = []\n",
        "        for dataset in self.datasets:\n",
        "          this_data_dir = os.path.join(self.data_dir, dataset, 'train')\n",
        "          for path in glob.glob(this_data_dir+'/*obj'):\n",
        "            self.paths.append(path)\n",
        "          # if \n",
        "        \n",
        "        # self.paths = glob.glob(self.train_dir+'/*.obj')\n",
        "       \n",
        "        self.size = len(self.paths)\n",
        "        # self.export_folder = export_dir\n",
        "        \n",
        "        # self.opt.num_aug = 1\n",
        "        # self.opt.export_folder = export_dir\n",
        "        # self.opt.ninput_edges = 2280 #Typical number of edges from coseg dataset (Will this yield good results for other datasets)\n",
        "        self.get_mean_std()\n",
        "        self.max_labels_points = 0\n",
        "        # modify for network later.\n",
        "        # self.opt.input_nc = self.ninput_channels\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.paths)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        # print(self.paths[idx])\n",
        "        # print(os.path.basename(self.paths[idx]))\n",
        "        # mesh_path = os.path.join(self.train_dir, self.paths[idx])\n",
        "        mesh_path = self.paths[idx]\n",
        "        # print(mesh_path)\n",
        "\n",
        "        # path = pathlib.PurePath(mesh_path)\n",
        "        folder = mesh_path.split('/')[-3] #Third from last element of path is the dataset folder name\n",
        "        # folder = path.parent.name\n",
        "        # filename = path.name.split('.')[0]\n",
        "        filename = os.path.basename(mesh_path).split('.')[0]\n",
        "        label_path = os.path.join(self.data_dir, folder, 'slices', f'{filename}.npy')\n",
        "        # print(label_path)\n",
        "        label = np.load(label_path)[:, :2] #We don't need the z dimension. Slice is 2D all points will have the same z value\n",
        "        n_points = label.shape[0]\n",
        "        print(f'n_points = {n_points}')\n",
        "        if n_points > opt.max_labels_points:\n",
        "          print('Updating')\n",
        "          opt.max_labels_points = n_points\n",
        "\n",
        "        mesh = Mesh(file=mesh_path, opt=self.opt, hold_history=False, export_folder=self.opt.export_folder)\n",
        "        meta = {'mesh': mesh, 'label': label}\n",
        "          # get edge features\n",
        "        edge_features = mesh.extract_features()\n",
        "        edge_features = pad(edge_features, self.opt.ninput_edges)\n",
        "          #The answer: This is only supposed to run once\n",
        "        # meta['edge_features'] = edge_features #MeshCNN code does't have this line, but withouth it, this breaks. This is because edge_features doesn't exists as an attribute of meta until you add it here\n",
        "          # self.mean\n",
        "        edge_features = (edge_features - self.mean) / self.std\n",
        "        meta['edge_features'] = (edge_features - self.mean) / self.std #Normalized features\n",
        "        return meta\n",
        "\n",
        "        # return mesh_path\n",
        "    \n",
        "    def get_mean_std(self):\n",
        "        \"\"\" Computes Mean and Standard Deviation from Training Data\n",
        "        If mean/std file doesn't exist, will compute one\n",
        "        :returns\n",
        "        mean: N-dimensional mean\n",
        "        std: N-dimensional standard deviation\n",
        "        ninput_channels: N\n",
        "        (here N=5)\n",
        "        \"\"\"\n",
        "\n",
        "        mean_std_cache = os.path.join(self.train_dir, 'mean_std_cache.p')\n",
        "        if not os.path.isfile(mean_std_cache):\n",
        "            print('computing mean std from train data...')\n",
        "            # doesn't run augmentation during m/std computation\n",
        "            num_aug = self.opt.num_aug\n",
        "            self.opt.num_aug = 1\n",
        "            mean, std = np.array(0), np.array(0)\n",
        "            # print('working...')\n",
        "            for i, data in enumerate(self):\n",
        "                # print(data)\n",
        "                if i % 50 == 0: #Originall 500, but we don't have that mean files. We'll compute the average of batches of 50 instead\n",
        "                    print('{} of {}'.format(i, self.size))\n",
        "                features = data['edge_features'] #edge_features attribute was missing until I added a line in the __getitem__ functionof this loader\n",
        "                mean = mean + features.mean(axis=1)\n",
        "                std = std + features.std(axis=1)\n",
        "            mean = mean / (i + 1)\n",
        "            print(f'mean : {mean}')\n",
        "            std = std / (i + 1)\n",
        "            transform_dict = {'mean': mean[:, np.newaxis], 'std': std[:, np.newaxis],\n",
        "                              'ninput_channels': len(mean)}\n",
        "            with open(mean_std_cache, 'wb') as f:\n",
        "                pickle.dump(transform_dict, f)\n",
        "            print('saved: ', mean_std_cache)\n",
        "            self.opt.num_aug = num_aug\n",
        "        # open mean / std from file\n",
        "        with open(mean_std_cache, 'rb') as f:\n",
        "            transform_dict = pickle.load(f)\n",
        "            print('loaded mean / std from cache')\n",
        "            self.mean = transform_dict['mean']\n",
        "            self.std = transform_dict['std']\n",
        "            self.ninput_channels = transform_dict['ninput_channels']\n",
        "\n",
        "def collate_fn(batch):\n",
        "    \"\"\"Creates mini-batch tensors\n",
        "    We should build custom collate_fn rather than using default collate_fn\n",
        "    \"\"\"\n",
        "    meta = {}\n",
        "    keys = batch[0].keys()\n",
        "    for key in keys:\n",
        "        meta.update({key: np.array([d[key] for d in batch])})\n",
        "    return meta\n",
        "\n",
        "#Custom Data loader\n",
        "class DataLoader:\n",
        "    \"\"\"multi-threaded data loading\"\"\"\n",
        "\n",
        "    def __init__(self, opt):\n",
        "        self.opt = opt\n",
        "        # self.dataset = CreateDataset(opt)\n",
        "        self.dataset = SlicerDataset(opt)\n",
        "        self.max_dataset_size = len(self.dataset)\n",
        "        print(f'max dataset size : {self.max_dataset_size}')\n",
        "        self.dataloader = torch.utils.data.DataLoader(\n",
        "            self.dataset,\n",
        "            batch_size=opt.batch_size,\n",
        "            shuffle= opt.shuffle,\n",
        "            num_workers=int(opt.num_threads),\n",
        "            collate_fn=collate_fn)\n",
        "\n",
        "    def __len__(self):\n",
        "        # return min(len(self.dataset), self.opt.max_dataset_size)\n",
        "        return len(self.dataset)\n",
        "\n",
        "    def __iter__(self):\n",
        "        for i, data in enumerate(self.dataloader):\n",
        "            if i * self.opt.batch_size >= self.max_dataset_size:\n",
        "                break\n",
        "            yield data\n",
        "# custom_dataset = SlicerDataset(TRAIN_DATA_DIR, TRAIN_LABELS_DIR, EXPORT_DIR)\n",
        "# print(len(custom_dataset))\n",
        "\n",
        "#Data options\n",
        "opt = TrainingOptions()\n",
        "opt.num_aug = 1\n",
        "opt.export_folder = EXPORT_DIR\n",
        "opt.train_dir = TRAIN_DATA_DIR\n",
        "opt.labels_dir = TRAIN_LABELS_DIR\n",
        "opt.data_dir = DATA_DIR\n",
        "opt.datasets = datasets\n",
        "opt.batch_size = 4\n",
        "opt.shuffle = True\n",
        "opt.num_threads = 2\n",
        "opt.max_labels_points = 0\n",
        "\n",
        "opt.ninput_edges = 2280 #Typical number of edges from coseg dataset (Will this yield good results for other datasets)\n",
        "\n",
        "dataset = DataLoader(opt)\n",
        "dataset_size = len(dataset)\n",
        "print('#training meshes = %d' % dataset_size)\n",
        "print(f'Max points in slice = {dataset.dataset.max_labels_points}')\n",
        "# trainloader = DataLoader(custom_dataset, batch_size=4, num_workers=0)\n"
      ],
      "execution_count": 72,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loaded mean / std from cache\n",
            "max dataset size : 755\n",
            "#training meshes = 755\n",
            "Max points in slice = 0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Zn2-i7fT7lfa",
        "outputId": "21e69732-60ef-4c0d-a89e-978de4ea57e6"
      },
      "source": [
        "dataset.dataset.ninput_channels"
      ],
      "execution_count": 67,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "5"
            ]
          },
          "metadata": {},
          "execution_count": 67
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "p2UY02Xa1sZj",
        "outputId": "df11ac25-301e-4a19-b159-4dbc2c0b6d02"
      },
      "source": [
        "for i, data in enumerate(dataset):\n",
        "  continue\n",
        "  # print(data)\n",
        "print(f'Max points in slice = {opt.max_labels_points}')"
      ],
      "execution_count": 69,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "n_points = 126\n",
            "Updating\n",
            "n_points = 67\n",
            "Updating\n",
            "n_points = 78\n",
            "n_points = 68\n",
            "Updating\n",
            "n_points = 136\n",
            "Updating\n",
            "n_points = 99\n",
            "Updating\n",
            "n_points = 107\n",
            "Updating\n",
            "n_points = 49\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py:123: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
            "/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py:123: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "n_points = 91\n",
            "n_points = 252\n",
            "Updating\n",
            "n_points = 65\n",
            "n_points = 110\n",
            "n_points = 80\n",
            "n_points = 74\n",
            "n_points = 23\n",
            "n_points = 62\n",
            "n_points = 112\n",
            "Updating\n",
            "n_points = 92\n",
            "n_points = 89\n",
            "n_points = 80\n",
            "n_points = 80\n",
            "n_points = 31\n",
            "n_points = 166\n",
            "Updating\n",
            "n_points = 70\n",
            "n_points = 61\n",
            "n_points = 45\n",
            "n_points = 91\n",
            "n_points = 60\n",
            "n_points = 106\n",
            "n_points = 87\n",
            "n_points = 111\n",
            "n_points = 115\n",
            "n_points = 89\n",
            "n_points = 68\n",
            "n_points = 108\n",
            "n_points = 87\n",
            "n_points = 108\n",
            "n_points = 85\n",
            "n_points = 63\n",
            "n_points = 75\n",
            "n_points = 39\n",
            "n_points = 58\n",
            "n_points = 32\n",
            "n_points = 21\n",
            "n_points = 92\n",
            "n_points = 28\n",
            "n_points = 249\n",
            "Updating\n",
            "n_points = 125\n",
            "n_points = 79\n",
            "n_points = 69\n",
            "n_points = 79\n",
            "n_points = 231\n",
            "n_points = 97\n",
            "n_points = 99\n",
            "n_points = 283\n",
            "Updating\n",
            "n_points = 68\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Exception ignored in: <function _MultiProcessingDataLoaderIter.__del__ at 0x7f6531c389e0>\n",
            "Traceback (most recent call last):\n",
            "  File \"/usr/local/lib/python3.7/dist-packages/torch/utils/data/dataloader.py\", line 1328, in __del__\n",
            "    self._shutdown_workers()\n",
            "  File \"/usr/local/lib/python3.7/dist-packages/torch/utils/data/dataloader.py\", line 1320, in _shutdown_workers\n",
            "    if w.is_alive():\n",
            "  File \"/usr/lib/python3.7/multiprocessing/process.py\", line 151, in is_alive\n",
            "    assert self._parent_pid == os.getpid(), 'can only test a child process'\n",
            "AssertionError: can only test a child process\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "n_points = 79\n",
            "n_points = 182\n",
            "n_points = 176\n",
            "n_points = 101\n",
            "n_points = 43\n",
            "n_points = 124\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Exception ignored in: <function _MultiProcessingDataLoaderIter.__del__ at 0x7f6531c389e0>\n",
            "Traceback (most recent call last):\n",
            "  File \"/usr/local/lib/python3.7/dist-packages/torch/utils/data/dataloader.py\", line 1328, in __del__\n",
            "    self._shutdown_workers()\n",
            "  File \"/usr/local/lib/python3.7/dist-packages/torch/utils/data/dataloader.py\", line 1320, in _shutdown_workers\n",
            "    if w.is_alive():\n",
            "  File \"/usr/lib/python3.7/multiprocessing/process.py\", line 151, in is_alive\n",
            "    assert self._parent_pid == os.getpid(), 'can only test a child process'\n",
            "AssertionError: can only test a child process\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "n_points = 52\n",
            "n_points = 162\n",
            "n_points = 112\n",
            "n_points = 115\n",
            "n_points = 120\n",
            "n_points = 53\n",
            "n_points = 204\n",
            "n_points = 104\n",
            "n_points = 31\n",
            "n_points = 107\n",
            "n_points = 217\n",
            "n_points = 167\n",
            "n_points = 57\n",
            "n_points = 30\n",
            "n_points = 66\n",
            "n_points = 136\n",
            "n_points = 15\n",
            "n_points = 100\n",
            "n_points = 120\n",
            "n_points = 166\n",
            "n_points = 111\n",
            "n_points = 180\n",
            "n_points = 78\n",
            "n_points = 70\n",
            "n_points = 95\n",
            "n_points = 273\n",
            "n_points = 101\n",
            "n_points = 41\n",
            "n_points = 65\n",
            "n_points = 58\n",
            "n_points = 117\n",
            "n_points = 78\n",
            "n_points = 84\n",
            "n_points = 74\n",
            "n_points = 68\n",
            "n_points = 68\n",
            "n_points = 25\n",
            "n_points = 86\n",
            "n_points = 110\n",
            "n_points = 97\n",
            "n_points = 126\n",
            "n_points = 108\n",
            "n_points = 87\n",
            "n_points = 91\n",
            "n_points = 62\n",
            "n_points = 100\n",
            "n_points = 36\n",
            "n_points = 48\n",
            "n_points = 28\n",
            "n_points = 113\n",
            "n_points = 14\n",
            "n_points = 158\n",
            "n_points = 72\n",
            "n_points = 88\n",
            "n_points = 196\n",
            "n_points = 16\n",
            "n_points = 112\n",
            "n_points = 110\n",
            "n_points = 137\n",
            "n_points = 67\n",
            "n_points = 33\n",
            "n_points = 77\n",
            "n_points = 94\n",
            "n_points = 216\n",
            "n_points = 43\n",
            "n_points = 68\n",
            "n_points = 72\n",
            "n_points = 215\n",
            "n_points = 148\n",
            "n_points = 63\n",
            "n_points = 107\n",
            "n_points = 40\n",
            "n_points = 81\n",
            "n_points = 52\n",
            "n_points = 131\n",
            "n_points = 41\n",
            "n_points = 100\n",
            "n_points = 14\n",
            "n_points = 102\n",
            "n_points = 43\n",
            "n_points = 189\n",
            "n_points = 100\n",
            "n_points = 60\n",
            "n_points = 76\n",
            "n_points = 100\n",
            "n_points = 65\n",
            "n_points = 56\n",
            "n_points = 188\n",
            "n_points = 39\n",
            "n_points = 119\n",
            "n_points = 117\n",
            "n_points = 64\n",
            "n_points = 98\n",
            "n_points = 64\n",
            "n_points = 26\n",
            "n_points = 111\n",
            "n_points = 146\n",
            "n_points = 106\n",
            "n_points = 79\n",
            "n_points = 80\n",
            "n_points = 220\n",
            "n_points = 46\n",
            "n_points = 106\n",
            "n_points = 199\n",
            "n_points = 160\n",
            "n_points = 69\n",
            "n_points = 63\n",
            "n_points = 54\n",
            "n_points = 101\n",
            "n_points = 79\n",
            "n_points = 113\n",
            "n_points = 61\n",
            "n_points = 114\n",
            "n_points = 58\n",
            "n_points = 156\n",
            "n_points = 112\n",
            "n_points = 80\n",
            "n_points = 101\n",
            "n_points = 101\n",
            "n_points = 57\n",
            "n_points = 108\n",
            "n_points = 31\n",
            "n_points = 58\n",
            "n_points = 69\n",
            "n_points = 153\n",
            "n_points = 36\n",
            "n_points = 221\n",
            "n_points = 62\n",
            "n_points = 77\n",
            "n_points = 185\n",
            "n_points = 108\n",
            "n_points = 79\n",
            "n_points = 74\n",
            "n_points = 98\n",
            "n_points = 24\n",
            "n_points = 58\n",
            "n_points = 29\n",
            "n_points = 47\n",
            "n_points = 50\n",
            "n_points = 206\n",
            "n_points = 58\n",
            "n_points = 43\n",
            "n_points = 108\n",
            "n_points = 33\n",
            "n_points = 66\n",
            "n_points = 72\n",
            "n_points = 149\n",
            "n_points = 71\n",
            "n_points = 79\n",
            "n_points = 101\n",
            "n_points = 71\n",
            "n_points = 77\n",
            "n_points = 198\n",
            "n_points = 91\n",
            "n_points = 101\n",
            "n_points = 122\n",
            "n_points = 97\n",
            "n_points = 61\n",
            "n_points = 115\n",
            "n_points = 9\n",
            "n_points = 82\n",
            "n_points = 66\n",
            "n_points = 106\n",
            "n_points = 97\n",
            "n_points = 44\n",
            "n_points = 97\n",
            "n_points = 57\n",
            "n_points = 64\n",
            "n_points = 86\n",
            "n_points = 105\n",
            "n_points = 171\n",
            "n_points = 39\n",
            "n_points = 36\n",
            "n_points = 22\n",
            "n_points = 43\n",
            "n_points = 144\n",
            "n_points = 78\n",
            "n_points = 133\n",
            "n_points = 7\n",
            "n_points = 72\n",
            "n_points = 38\n",
            "n_points = 145\n",
            "n_points = 67\n",
            "n_points = 229\n",
            "n_points = 52\n",
            "n_points = 95\n",
            "n_points = 83\n",
            "n_points = 32\n",
            "n_points = 69\n",
            "n_points = 59\n",
            "n_points = 163\n",
            "n_points = 97\n",
            "n_points = 67\n",
            "n_points = 103\n",
            "n_points = 89\n",
            "n_points = 229\n",
            "n_points = 72\n",
            "n_points = 94\n",
            "n_points = 109\n",
            "n_points = 39\n",
            "n_points = 68\n",
            "n_points = 67\n",
            "n_points = 101\n",
            "n_points = 25\n",
            "n_points = 42\n",
            "n_points = 106\n",
            "n_points = 5\n",
            "n_points = 112\n",
            "n_points = 72\n",
            "n_points = 111\n",
            "n_points = 113\n",
            "n_points = 195\n",
            "n_points = 80\n",
            "n_points = 42\n",
            "n_points = 121\n",
            "n_points = 129\n",
            "n_points = 152\n",
            "n_points = 98\n",
            "n_points = 126\n",
            "n_points = 92\n",
            "n_points = 87\n",
            "n_points = 98\n",
            "n_points = 145\n",
            "n_points = 99\n",
            "n_points = 117\n",
            "n_points = 88\n",
            "n_points = 109\n",
            "n_points = 76\n",
            "n_points = 192\n",
            "n_points = 61\n",
            "n_points = 69\n",
            "n_points = 176\n",
            "n_points = 75\n",
            "n_points = 33\n",
            "n_points = 108\n",
            "n_points = 85\n",
            "n_points = 34\n",
            "n_points = 115\n",
            "n_points = 75\n",
            "n_points = 91\n",
            "n_points = 100\n",
            "n_points = 116\n",
            "n_points = 151\n",
            "n_points = 76\n",
            "n_points = 93\n",
            "n_points = 120\n",
            "n_points = 115\n",
            "n_points = 177\n",
            "n_points = 60\n",
            "n_points = 123\n",
            "n_points = 98\n",
            "n_points = 83\n",
            "n_points = 116\n",
            "n_points = 113\n",
            "n_points = 178\n",
            "n_points = 43\n",
            "n_points = 78\n",
            "n_points = 142\n",
            "n_points = 119\n",
            "n_points = 77\n",
            "n_points = 158\n",
            "n_points = 45\n",
            "n_points = 54\n",
            "n_points = 103\n",
            "n_points = 119\n",
            "n_points = 99\n",
            "n_points = 135\n",
            "n_points = 116\n",
            "n_points = 103\n",
            "n_points = 126\n",
            "n_points = 151\n",
            "n_points = 118\n",
            "n_points = 104\n",
            "n_points = 57\n",
            "n_points = 14\n",
            "n_points = 73\n",
            "n_points = 62\n",
            "n_points = 100\n",
            "n_points = 96\n",
            "n_points = 120\n",
            "n_points = 104\n",
            "n_points = 79\n",
            "n_points = 122\n",
            "n_points = 91\n",
            "n_points = 56\n",
            "n_points = 158\n",
            "n_points = 118\n",
            "n_points = 57\n",
            "n_points = 118\n",
            "n_points = 120\n",
            "n_points = 144\n",
            "n_points = 153\n",
            "n_points = 14\n",
            "n_points = 78\n",
            "n_points = 36\n",
            "n_points = 121\n",
            "n_points = 83\n",
            "n_points = 28\n",
            "n_points = 8\n",
            "n_points = 18\n",
            "n_points = 54\n",
            "n_points = 61\n",
            "n_points = 52\n",
            "n_points = 113\n",
            "n_points = 11\n",
            "n_points = 147\n",
            "n_points = 87\n",
            "n_points = 80\n",
            "n_points = 186\n",
            "n_points = 89\n",
            "n_points = 144\n",
            "n_points = 83\n",
            "n_points = 41\n",
            "n_points = 110\n",
            "n_points = 118\n",
            "n_points = 108\n",
            "n_points = 106\n",
            "n_points = 91\n",
            "n_points = 96\n",
            "n_points = 61\n",
            "n_points = 27\n",
            "n_points = 79\n",
            "n_points = 127\n",
            "n_points = 54\n",
            "n_points = 92\n",
            "n_points = 87\n",
            "n_points = 56\n",
            "n_points = 112\n",
            "n_points = 198\n",
            "n_points = 14\n",
            "n_points = 87\n",
            "n_points = 95\n",
            "n_points = 20\n",
            "n_points = 258\n",
            "n_points = 111\n",
            "n_points = 72\n",
            "n_points = 52\n",
            "n_points = 47\n",
            "n_points = 63\n",
            "n_points = 125\n",
            "n_points = 131\n",
            "n_points = 69\n",
            "n_points = 106\n",
            "n_points = 66\n",
            "n_points = 78\n",
            "n_points = 247\n",
            "n_points = 87\n",
            "n_points = 58\n",
            "n_points = 242\n",
            "n_points = 97\n",
            "n_points = 67\n",
            "n_points = 85\n",
            "n_points = 57\n",
            "n_points = 189\n",
            "n_points = 108\n",
            "n_points = 112\n",
            "n_points = 88\n",
            "n_points = 70\n",
            "n_points = 126\n",
            "n_points = 20\n",
            "n_points = 115\n",
            "n_points = 113\n",
            "n_points = 132\n",
            "n_points = 61\n",
            "n_points = 74\n",
            "n_points = 123\n",
            "n_points = 69\n",
            "n_points = 74\n",
            "n_points = 66\n",
            "n_points = 47\n",
            "n_points = 100\n",
            "n_points = 50\n",
            "n_points = 64\n",
            "n_points = 129\n",
            "n_points = 109\n",
            "n_points = 81\n",
            "n_points = 81\n",
            "n_points = 106\n",
            "n_points = 101\n",
            "n_points = 65\n",
            "n_points = 96\n",
            "n_points = 78\n",
            "n_points = 95\n",
            "n_points = 72\n",
            "n_points = 144\n",
            "n_points = 188\n",
            "n_points = 103\n",
            "n_points = 48\n",
            "n_points = 124\n",
            "n_points = 118\n",
            "n_points = 66\n",
            "n_points = 97\n",
            "n_points = 168\n",
            "n_points = 100\n",
            "n_points = 106\n",
            "n_points = 61\n",
            "n_points = 98\n",
            "n_points = 87\n",
            "n_points = 55\n",
            "n_points = 110\n",
            "n_points = 36\n",
            "n_points = 37\n",
            "n_points = 174\n",
            "n_points = 90\n",
            "n_points = 37\n",
            "n_points = 66\n",
            "n_points = 35\n",
            "n_points = 33\n",
            "n_points = 51\n",
            "n_points = 91\n",
            "n_points = 255\n",
            "n_points = 78\n",
            "n_points = 11\n",
            "n_points = 139\n",
            "n_points = 81\n",
            "n_points = 70\n",
            "n_points = 80\n",
            "n_points = 70\n",
            "n_points = 195\n",
            "n_points = 81\n",
            "n_points = 94\n",
            "n_points = 157\n",
            "n_points = 112\n",
            "n_points = 112\n",
            "n_points = 85\n",
            "n_points = 103\n",
            "n_points = 85\n",
            "n_points = 61\n",
            "n_points = 136\n",
            "n_points = 96\n",
            "n_points = 84\n",
            "n_points = 66\n",
            "n_points = 135\n",
            "n_points = 60\n",
            "n_points = 20\n",
            "n_points = 82\n",
            "n_points = 108\n",
            "n_points = 142\n",
            "n_points = 57\n",
            "n_points = 100\n",
            "n_points = 85\n",
            "n_points = 87\n",
            "n_points = 57\n",
            "n_points = 21\n",
            "n_points = 62\n",
            "n_points = 137\n",
            "n_points = 178\n",
            "n_points = 109\n",
            "n_points = 98\n",
            "n_points = 65\n",
            "n_points = 51\n",
            "n_points = 26\n",
            "n_points = 136\n",
            "n_points = 53\n",
            "n_points = 99\n",
            "n_points = 80\n",
            "n_points = 84\n",
            "n_points = 217\n",
            "n_points = 95\n",
            "n_points = 86\n",
            "n_points = 222\n",
            "n_points = 66\n",
            "n_points = 110\n",
            "n_points = 87\n",
            "n_points = 63\n",
            "n_points = 60\n",
            "n_points = 93\n",
            "n_points = 178\n",
            "n_points = 168\n",
            "n_points = 92\n",
            "n_points = 28\n",
            "n_points = 266\n",
            "Updating\n",
            "n_points = 79\n",
            "n_points = 78\n",
            "n_points = 111\n",
            "n_points = 144\n",
            "n_points = 23\n",
            "n_points = 36\n",
            "n_points = 83\n",
            "n_points = 82\n",
            "n_points = 78\n",
            "n_points = 79\n",
            "n_points = 74\n",
            "n_points = 34\n",
            "n_points = 62\n",
            "n_points = 139\n",
            "n_points = 82\n",
            "n_points = 97\n",
            "n_points = 228\n",
            "n_points = 133\n",
            "n_points = 63\n",
            "n_points = 51\n",
            "n_points = 56\n",
            "n_points = 100\n",
            "n_points = 43\n",
            "n_points = 105\n",
            "n_points = 91\n",
            "n_points = 82\n",
            "n_points = 230\n",
            "n_points = 67\n",
            "n_points = 33\n",
            "n_points = 61\n",
            "n_points = 120\n",
            "n_points = 63\n",
            "n_points = 66\n",
            "n_points = 15\n",
            "n_points = 63\n",
            "n_points = 43\n",
            "n_points = 24\n",
            "n_points = 83\n",
            "n_points = 23\n",
            "n_points = 187\n",
            "n_points = 99\n",
            "n_points = 88\n",
            "n_points = 68\n",
            "n_points = 100\n",
            "n_points = 69\n",
            "n_points = 66\n",
            "n_points = 62\n",
            "n_points = 56\n",
            "n_points = 156\n",
            "n_points = 183\n",
            "n_points = 141\n",
            "n_points = 126\n",
            "n_points = 38\n",
            "n_points = 40\n",
            "n_points = 56\n",
            "n_points = 68\n",
            "n_points = 45\n",
            "n_points = 68\n",
            "n_points = 124\n",
            "n_points = 61\n",
            "n_points = 43\n",
            "n_points = 26\n",
            "n_points = 142\n",
            "n_points = 41\n",
            "n_points = 102\n",
            "n_points = 73\n",
            "n_points = 93\n",
            "n_points = 158\n",
            "n_points = 63\n",
            "n_points = 113\n",
            "n_points = 200\n",
            "n_points = 80\n",
            "n_points = 42\n",
            "n_points = 108\n",
            "n_points = 55\n",
            "n_points = 99\n",
            "n_points = 70\n",
            "n_points = 34\n",
            "n_points = 100\n",
            "n_points = 27\n",
            "n_points = 56\n",
            "n_points = 56\n",
            "n_points = 100\n",
            "n_points = 78\n",
            "n_points = 191\n",
            "n_points = 32\n",
            "n_points = 118\n",
            "n_points = 151\n",
            "n_points = 33\n",
            "n_points = 75\n",
            "n_points = 112\n",
            "n_points = 40\n",
            "n_points = 59\n",
            "n_points = 103\n",
            "n_points = 101\n",
            "n_points = 45\n",
            "n_points = 145\n",
            "n_points = 70\n",
            "n_points = 71\n",
            "n_points = 89\n",
            "n_points = 89\n",
            "n_points = 80\n",
            "n_points = 81\n",
            "n_points = 140\n",
            "n_points = 216\n",
            "n_points = 50\n",
            "n_points = 111\n",
            "n_points = 95\n",
            "n_points = 75\n",
            "n_points = 37\n",
            "n_points = 96\n",
            "n_points = 66\n",
            "n_points = 78\n",
            "n_points = 33\n",
            "n_points = 63\n",
            "n_points = 126\n",
            "n_points = 56\n",
            "n_points = 103\n",
            "n_points = 73\n",
            "n_points = 101\n",
            "n_points = 55\n",
            "n_points = 69\n",
            "n_points = 105\n",
            "n_points = 42\n",
            "n_points = 68\n",
            "n_points = 91\n",
            "n_points = 62\n",
            "n_points = 121\n",
            "n_points = 113\n",
            "n_points = 63\n",
            "n_points = 115\n",
            "n_points = 115\n",
            "n_points = 193\n",
            "n_points = 111\n",
            "n_points = 66\n",
            "n_points = 110\n",
            "n_points = 80\n",
            "n_points = 83\n",
            "n_points = 107\n",
            "n_points = 97\n",
            "n_points = 72\n",
            "n_points = 88\n",
            "n_points = 107\n",
            "n_points = 23\n",
            "n_points = 68\n",
            "n_points = 67\n",
            "n_points = 81\n",
            "n_points = 90\n",
            "n_points = 99\n",
            "n_points = 97\n",
            "n_points = 78\n",
            "n_points = 73\n",
            "n_points = 68\n",
            "n_points = 26\n",
            "n_points = 25\n",
            "n_points = 204\n",
            "n_points = 113\n",
            "n_points = 157\n",
            "n_points = 48\n",
            "n_points = 154\n",
            "n_points = 106\n",
            "n_points = 36\n",
            "n_points = 97\n",
            "n_points = 94\n",
            "n_points = 80\n",
            "n_points = 81\n",
            "n_points = 45\n",
            "n_points = 83\n",
            "n_points = 99\n",
            "n_points = 14\n",
            "n_points = 119\n",
            "n_points = 91\n",
            "n_points = 102\n",
            "n_points = 141\n",
            "n_points = 54\n",
            "n_points = 171\n",
            "n_points = 35\n",
            "n_points = 71\n",
            "n_points = 41\n",
            "n_points = 44\n",
            "n_points = 132\n",
            "n_points = 72\n",
            "n_points = 222\n",
            "n_points = 63\n",
            "n_points = 114\n",
            "n_points = 43\n",
            "n_points = 131\n",
            "n_points = 80\n",
            "n_points = 32\n",
            "n_points = 85\n",
            "n_points = 93\n",
            "n_points = 88\n",
            "n_points = 104\n",
            "n_points = 99\n",
            "n_points = 67\n",
            "n_points = 77\n",
            "n_points = 41\n",
            "n_points = 54\n",
            "n_points = 99\n",
            "n_points = 87\n",
            "n_points = 32\n",
            "n_points = 83\n",
            "n_points = 49\n",
            "n_points = 62\n",
            "n_points = 57\n",
            "n_points = 61\n",
            "n_points = 10\n",
            "n_points = 87\n",
            "n_points = 65\n",
            "n_points = 123\n",
            "n_points = 36\n",
            "n_points = 33\n",
            "n_points = 85\n",
            "n_points = 101\n",
            "n_points = 70\n",
            "n_points = 65\n",
            "n_points = 67\n",
            "n_points = 87\n",
            "n_points = 92\n",
            "n_points = 77\n",
            "n_points = 41\n",
            "Max points in slice = 0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1xG-v-qCRhe3",
        "outputId": "46477163-f25e-480f-a312-adb2ea1feb9f"
      },
      "source": [
        "net = None\n",
        "opt.norm = 'batch' #Use bactch normalization\n",
        "\n",
        "# norm_layer = get_norm_layer(norm_type=opt.norm, num_groups=opt.num_groups)\n",
        "opt.norm_layer = get_norm_layer(norm_type=opt.norm)\n",
        "opt.input_nc = 5 # n input channels (1/edge_feature)\n",
        "opt.ncf = [16, 32, 32] #n conv filters for each conv layer\n",
        "opt.ninput_edges = 2280\n",
        "opt.pool_res = [1140, 780, 580] # Pooling resolution after each conv layer (n edges to keep at each pooling layer)\n",
        "opt.fc_n = 100 # number between fully connected and nclasses. n neurons in this fc layer\n",
        "opt.nresblocks = 0\n",
        "opt.lr_policy = 'lambda' #other optiosn are step, and plateau\n",
        "opt.gpu_ids = ['0']\n",
        "opt.is_train = True\n",
        "opt.checkpoints_dir = os.path.join(MODELS_DIR, 'checkpoints')\n",
        "opt.name = '3D_slAIcer'\n",
        "opt.nclasses = 300 # size of output layer\n",
        "opt.arch = 'mconvnet'\n",
        "opt.init_type = 'normal' #Network initialization [normal|xavier|kaiming|orthogonal]\n",
        "opt.init_gain = 0.2 #Scaling factor for normal, xvier and ortho\n",
        "opt.lr = 0.001\n",
        "opt.beta1 = 0.9 #First momentum term\n",
        "opt.epoch_count = 10\n",
        "opt.niter = 100\n",
        "opt.niter_decay = 500\n",
        "opt.continue_train = False\n",
        "\n",
        "SlicerModel(opt)\n",
        "# net = MeshConvNet(norm_layer, input_nc, ncf, nclasses, ninput_edges, opt.pool_res, opt.fc_n, opt.resblocks)\n",
        "#   MeshConvnET(norm_layer, nf0, conv_res, nclasses, input_res, pool_res, fc_n,\n",
        "#                  nresblocks=3):\n",
        "#  MeshConvnET(opt.norm_layer, opt.input_nc, opt.ncf, opt.nclasses, opt.ninput_edges, opt.pool_res, opt.fc_n,\n",
        "#                  opt.nresblocks=3):\n",
        "   \n",
        "# norm_layer, example batch norm\n",
        "# input_nc, n input channels (5) : 1/edge_feature\n",
        "# ncf: number of convolutional filters example [16, 32, 32] This would describe 4 conv layers\n",
        "# ninput_edges: all input meshes have their edges padded to 2280\n",
        "# pool_res: Pooling resolution after each conv layer, example [1140, 780, 580], describes n edges to keep at each pooling layer\n",
        "# fc_n: number between fully connected and nclasses. n neurons in this fc layer\n",
        "# resblocks: number of resblocks\n"
      ],
      "execution_count": 104,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "---------- Network initialized -------------\n",
            "[Network] Total number of parameters : 0.042 M\n",
            "-----------------------------------------------\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<__main__.SlicerModel at 0x7f652c6d4490>"
            ]
          },
          "metadata": {},
          "execution_count": 104
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b2GRghYsRTkZ"
      },
      "source": [
        "## MeshConvNet"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hT0CNuHWRR5c"
      },
      "source": [
        "def print_network(net):\n",
        "    \"\"\"Print the total number of parameters in the network\n",
        "    Parameters:\n",
        "        network\n",
        "    \"\"\"\n",
        "    print('---------- Network initialized -------------')\n",
        "    num_params = 0\n",
        "    for param in net.parameters():\n",
        "        num_params += param.numel()\n",
        "    print('[Network] Total number of parameters : %.3f M' % (num_params / 1e6))\n",
        "    print('-----------------------------------------------')\n",
        "    \n",
        "class MeshConvNet(nn.Module):\n",
        "    \"\"\"Network for learning a global shape descriptor (classification)\n",
        "    \"\"\"\n",
        "    def __init__(self, norm_layer, nf0, conv_res, nclasses, input_res, pool_res, fc_n,\n",
        "                 nresblocks=3):\n",
        "        super(MeshConvNet, self).__init__()\n",
        "        self.k = [nf0] + conv_res\n",
        "        self.res = [input_res] + pool_res\n",
        "        norm_args = get_norm_args(norm_layer, self.k[1:])\n",
        "\n",
        "        for i, ki in enumerate(self.k[:-1]):\n",
        "            setattr(self, 'conv{}'.format(i), MResConv(ki, self.k[i + 1], nresblocks))\n",
        "            setattr(self, 'norm{}'.format(i), norm_layer(**norm_args[i]))\n",
        "            setattr(self, 'pool{}'.format(i), MeshPool(self.res[i + 1]))\n",
        "\n",
        "\n",
        "        self.gp = torch.nn.AvgPool1d(self.res[-1])\n",
        "        # self.gp = torch.nn.MaxPool1d(self.res[-1])\n",
        "        self.fc1 = nn.Linear(self.k[-1], fc_n)\n",
        "        self.fc2 = nn.Linear(fc_n, nclasses)\n",
        "\n",
        "    def forward(self, x, mesh):\n",
        "\n",
        "        for i in range(len(self.k) - 1):\n",
        "            x = getattr(self, 'conv{}'.format(i))(x, mesh)\n",
        "            x = F.relu(getattr(self, 'norm{}'.format(i))(x))\n",
        "            x = getattr(self, 'pool{}'.format(i))(x, mesh)\n",
        "\n",
        "        x = self.gp(x)\n",
        "        x = x.view(-1, self.k[-1])\n",
        "\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = self.fc2(x)\n",
        "        return x\n",
        "\n",
        "def get_scheduler(optimizer, opt):\n",
        "    if opt.lr_policy == 'lambda':\n",
        "        def lambda_rule(epoch):\n",
        "            lr_l = 1.0 - max(0, epoch + 1 + opt.epoch_count - opt.niter) / float(opt.niter_decay + 1)\n",
        "            return lr_l\n",
        "        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)\n",
        "    elif opt.lr_policy == 'step':\n",
        "        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)\n",
        "    elif opt.lr_policy == 'plateau':\n",
        "        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)\n",
        "    else:\n",
        "        return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)\n",
        "    return scheduler\n",
        "\n",
        "class SlicerModel:\n",
        "    \"\"\" Class for training Model weights\n",
        "    :args opt: structure containing configuration params\n",
        "    e.g.,\n",
        "    --dataset_mode -> classification / segmentation)\n",
        "    --arch -> network type\n",
        "    \"\"\"\n",
        "    def __init__(self, opt):\n",
        "        self.opt = opt\n",
        "        self.gpu_ids = opt.gpu_ids\n",
        "        self.is_train = opt.is_train\n",
        "        self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')\n",
        "        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)\n",
        "        self.optimizer = None\n",
        "        self.edge_features = None\n",
        "        self.labels = None\n",
        "        self.mesh = None\n",
        "        self.soft_label = None\n",
        "        self.loss = None\n",
        "\n",
        "        #\n",
        "        self.nclasses = opt.nclasses\n",
        "\n",
        "        # load/define networks\n",
        "        self.net = MeshConvNet(opt.norm_layer, opt.input_nc, opt.ncf, opt.nclasses, opt.ninput_edges, opt.pool_res, opt.fc_n, opt.nresblocks)\n",
        "        # self.net = MeshConvNet(opt.input_nc, opt.ncf, opt.ninput_edges, opt.nclasses, opt,\n",
        "        #                                       self.gpu_ids, opt.arch, opt.init_type, opt.init_gain)\n",
        "        \n",
        "     \n",
        "        # self.net = MeshConvNet(opt.input_nc, opt.ncf, opt.ninput_edges, opt.nclasses, opt,\n",
        "        #                                       self.gpu_ids, opt.arch, opt.init_type, opt.init_gain)\n",
        "        self.net.train(self.is_train)\n",
        "        # self.criterion = networks.define_loss(opt).to(self.device)\n",
        "        self.criterion = torch.nn.MSELoss().to(self.device)\n",
        "\n",
        "        if self.is_train:\n",
        "            self.optimizer = torch.optim.Adam(self.net.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))\n",
        "            self.scheduler = get_scheduler(self.optimizer, opt)\n",
        "            print_network(self.net)\n",
        "\n",
        "        if not self.is_train or opt.continue_train:\n",
        "            self.load_network(opt.which_epoch)\n",
        "\n",
        "    def set_input(self, data):\n",
        "        input_edge_features = torch.from_numpy(data['edge_features']).float()\n",
        "        labels = torch.from_numpy(data['label']).long()\n",
        "        # set inputs\n",
        "        self.edge_features = input_edge_features.to(self.device).requires_grad_(self.is_train)\n",
        "        self.labels = labels.to(self.device)\n",
        "        self.mesh = data['mesh']\n",
        "        if self.opt.dataset_mode == 'segmentation' and not self.is_train:\n",
        "            self.soft_label = torch.from_numpy(data['soft_label'])\n",
        "\n",
        "\n",
        "    def forward(self):\n",
        "        out = self.net(self.edge_features, self.mesh)\n",
        "        return out\n",
        "\n",
        "    def backward(self, out):\n",
        "        self.loss = self.criterion(out, self.labels)\n",
        "        self.loss.backward()\n",
        "\n",
        "    def optimize_parameters(self):\n",
        "        self.optimizer.zero_grad()\n",
        "        out = self.forward()\n",
        "        self.backward(out)\n",
        "        self.optimizer.step()\n",
        "\n",
        "\n",
        "##################\n",
        "\n",
        "    def load_network(self, which_epoch):\n",
        "        \"\"\"load model from disk\"\"\"\n",
        "        save_filename = '%s_net.pth' % which_epoch\n",
        "        load_path = join(self.save_dir, save_filename)\n",
        "        net = self.net\n",
        "        if isinstance(net, torch.nn.DataParallel):\n",
        "            net = net.module\n",
        "        print('loading the model from %s' % load_path)\n",
        "        # PyTorch newer than 0.4 (e.g., built from\n",
        "        # GitHub source), you can remove str() on self.device\n",
        "        state_dict = torch.load(load_path, map_location=str(self.device))\n",
        "        if hasattr(state_dict, '_metadata'):\n",
        "            del state_dict._metadata\n",
        "        net.load_state_dict(state_dict)\n",
        "\n",
        "\n",
        "    def save_network(self, which_epoch):\n",
        "        \"\"\"save model to disk\"\"\"\n",
        "        save_filename = '%s_net.pth' % (which_epoch)\n",
        "        save_path = join(self.save_dir, save_filename)\n",
        "        if len(self.gpu_ids) > 0 and torch.cuda.is_available():\n",
        "            torch.save(self.net.module.cpu().state_dict(), save_path)\n",
        "            self.net.cuda(self.gpu_ids[0])\n",
        "        else:\n",
        "            torch.save(self.net.cpu().state_dict(), save_path)\n",
        "\n",
        "    def update_learning_rate(self):\n",
        "        \"\"\"update learning rate (called once every epoch)\"\"\"\n",
        "        self.scheduler.step()\n",
        "        lr = self.optimizer.param_groups[0]['lr']\n",
        "        print('learning rate = %.7f' % lr)\n",
        "\n",
        "    def test(self):\n",
        "        \"\"\"tests model\n",
        "        returns: number correct and total number\n",
        "        \"\"\"\n",
        "        with torch.no_grad():\n",
        "            out = self.forward()\n",
        "            # compute number of correct\n",
        "            pred_class = out.data.max(1)[1]\n",
        "            label_class = self.labels\n",
        "            self.export_segmentation(pred_class.cpu())\n",
        "            correct = self.get_accuracy(pred_class, label_class)\n",
        "        return correct, len(label_class)\n",
        "\n",
        "    def get_accuracy(self, pred, labels):\n",
        "        \"\"\"computes accuracy for classification / segmentation \"\"\"\n",
        "        if self.opt.dataset_mode == 'classification':\n",
        "            correct = pred.eq(labels).sum()\n",
        "        elif self.opt.dataset_mode == 'segmentation':\n",
        "            correct = seg_accuracy(pred, self.soft_label, self.mesh)\n",
        "        return correct\n",
        "\n",
        "    def export_segmentation(self, pred_seg):\n",
        "        if self.opt.dataset_mode == 'segmentation':\n",
        "            for meshi, mesh in enumerate(self.mesh):\n",
        "                mesh.export_segments(pred_seg[meshi, :])"
      ],
      "execution_count": 102,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vj9eYi_1PTZp"
      },
      "source": [
        "## Classes and functions to perform mesh pooling"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nD0WVhN0PSHT"
      },
      "source": [
        "class MeshUnion:\n",
        "    def __init__(self, n, device=torch.device('cpu')):\n",
        "        self.__size = n\n",
        "        self.rebuild_features = self.rebuild_features_average\n",
        "        self.groups = torch.eye(n, device=device)\n",
        "\n",
        "    def union(self, source, target):\n",
        "        self.groups[target, :] += self.groups[source, :]\n",
        "\n",
        "    def remove_group(self, index):\n",
        "        return\n",
        "\n",
        "    def get_group(self, edge_key):\n",
        "        return self.groups[edge_key, :]\n",
        "\n",
        "    def get_occurrences(self):\n",
        "        return torch.sum(self.groups, 0)\n",
        "\n",
        "    def get_groups(self, tensor_mask):\n",
        "        self.groups = torch.clamp(self.groups, 0, 1)\n",
        "        return self.groups[tensor_mask, :]\n",
        "\n",
        "    def rebuild_features_average(self, features, mask, target_edges):\n",
        "        self.prepare_groups(features, mask)\n",
        "        fe = torch.matmul(features.squeeze(-1), self.groups)\n",
        "        occurrences = torch.sum(self.groups, 0).expand(fe.shape)\n",
        "        fe = fe / occurrences\n",
        "        padding_b = target_edges - fe.shape[1]\n",
        "        if padding_b > 0:\n",
        "            padding_b = ConstantPad2d((0, padding_b, 0, 0), 0)\n",
        "            fe = padding_b(fe)\n",
        "        return fe\n",
        "\n",
        "    def prepare_groups(self, features, mask):\n",
        "        tensor_mask = torch.from_numpy(mask)\n",
        "        self.groups = torch.clamp(self.groups[tensor_mask, :], 0, 1).transpose_(1, 0)\n",
        "        padding_a = features.shape[1] - self.groups.shape[0]\n",
        "        if padding_a > 0:\n",
        "            padding_a = ConstantPad2d((0, 0, 0, padding_a), 0)\n",
        "            self.groups = padding_a(self.groups)\n",
        "            \n",
        "#Mesh pool layer\n",
        "class MeshPool(nn.Module):\n",
        "    \n",
        "    def __init__(self, target, multi_thread=False):\n",
        "        super(MeshPool, self).__init__()\n",
        "        self.__out_target = target\n",
        "        self.__multi_thread = multi_thread\n",
        "        self.__fe = None\n",
        "        self.__updated_fe = None\n",
        "        self.__meshes = None\n",
        "        self.__merge_edges = [-1, -1]\n",
        "\n",
        "    def __call__(self, fe, meshes):\n",
        "        return self.forward(fe, meshes)\n",
        "\n",
        "    def forward(self, fe, meshes):\n",
        "        self.__updated_fe = [[] for _ in range(len(meshes))]\n",
        "        pool_threads = []\n",
        "        self.__fe = fe\n",
        "        self.__meshes = meshes\n",
        "        # iterate over batch\n",
        "        for mesh_index in range(len(meshes)):\n",
        "            if self.__multi_thread:\n",
        "                pool_threads.append(Thread(target=self.__pool_main, args=(mesh_index,)))\n",
        "                pool_threads[-1].start()\n",
        "            else:\n",
        "                self.__pool_main(mesh_index)\n",
        "        if self.__multi_thread:\n",
        "            for mesh_index in range(len(meshes)):\n",
        "                pool_threads[mesh_index].join()\n",
        "        out_features = torch.cat(self.__updated_fe).view(len(meshes), -1, self.__out_target)\n",
        "        return out_features\n",
        "\n",
        "    def __pool_main(self, mesh_index):\n",
        "        mesh = self.__meshes[mesh_index]\n",
        "        queue = self.__build_queue(self.__fe[mesh_index, :, :mesh.edges_count], mesh.edges_count)\n",
        "        # recycle = []\n",
        "        # last_queue_len = len(queue)\n",
        "        last_count = mesh.edges_count + 1\n",
        "        mask = np.ones(mesh.edges_count, dtype=np.bool)\n",
        "        edge_groups = MeshUnion(mesh.edges_count, self.__fe.device)\n",
        "        while mesh.edges_count > self.__out_target:\n",
        "            value, edge_id = heappop(queue)\n",
        "            edge_id = int(edge_id)\n",
        "            if mask[edge_id]:\n",
        "                self.__pool_edge(mesh, edge_id, mask, edge_groups)\n",
        "        mesh.clean(mask, edge_groups)\n",
        "        fe = edge_groups.rebuild_features(self.__fe[mesh_index], mask, self.__out_target)\n",
        "        self.__updated_fe[mesh_index] = fe\n",
        "\n",
        "    def __pool_edge(self, mesh, edge_id, mask, edge_groups):\n",
        "        if self.has_boundaries(mesh, edge_id):\n",
        "            return False\n",
        "        elif self.__clean_side(mesh, edge_id, mask, edge_groups, 0)\\\n",
        "            and self.__clean_side(mesh, edge_id, mask, edge_groups, 2) \\\n",
        "            and self.__is_one_ring_valid(mesh, edge_id):\n",
        "            self.__merge_edges[0] = self.__pool_side(mesh, edge_id, mask, edge_groups, 0)\n",
        "            self.__merge_edges[1] = self.__pool_side(mesh, edge_id, mask, edge_groups, 2)\n",
        "            mesh.merge_vertices(edge_id)\n",
        "            mask[edge_id] = False\n",
        "            MeshPool.__remove_group(mesh, edge_groups, edge_id)\n",
        "            mesh.edges_count -= 1\n",
        "            return True\n",
        "        else:\n",
        "            return False\n",
        "\n",
        "    def __clean_side(self, mesh, edge_id, mask, edge_groups, side):\n",
        "        if mesh.edges_count <= self.__out_target:\n",
        "            return False\n",
        "        invalid_edges = MeshPool.__get_invalids(mesh, edge_id, edge_groups, side)\n",
        "        while len(invalid_edges) != 0 and mesh.edges_count > self.__out_target:\n",
        "            self.__remove_triplete(mesh, mask, edge_groups, invalid_edges)\n",
        "            if mesh.edges_count <= self.__out_target:\n",
        "                return False\n",
        "            if self.has_boundaries(mesh, edge_id):\n",
        "                return False\n",
        "            invalid_edges = self.__get_invalids(mesh, edge_id, edge_groups, side)\n",
        "        return True\n",
        "\n",
        "    @staticmethod\n",
        "    def has_boundaries(mesh, edge_id):\n",
        "        for edge in mesh.gemm_edges[edge_id]:\n",
        "            if edge == -1 or -1 in mesh.gemm_edges[edge]:\n",
        "                return True\n",
        "        return False\n",
        "\n",
        "\n",
        "    @staticmethod\n",
        "    def __is_one_ring_valid(mesh, edge_id):\n",
        "        v_a = set(mesh.edges[mesh.ve[mesh.edges[edge_id, 0]]].reshape(-1))\n",
        "        v_b = set(mesh.edges[mesh.ve[mesh.edges[edge_id, 1]]].reshape(-1))\n",
        "        shared = v_a & v_b - set(mesh.edges[edge_id])\n",
        "        return len(shared) == 2\n",
        "\n",
        "    def __pool_side(self, mesh, edge_id, mask, edge_groups, side):\n",
        "        info = MeshPool.__get_face_info(mesh, edge_id, side)\n",
        "        key_a, key_b, side_a, side_b, _, other_side_b, _, other_keys_b = info\n",
        "        self.__redirect_edges(mesh, key_a, side_a - side_a % 2, other_keys_b[0], mesh.sides[key_b, other_side_b])\n",
        "        self.__redirect_edges(mesh, key_a, side_a - side_a % 2 + 1, other_keys_b[1], mesh.sides[key_b, other_side_b + 1])\n",
        "        MeshPool.__union_groups(mesh, edge_groups, key_b, key_a)\n",
        "        MeshPool.__union_groups(mesh, edge_groups, edge_id, key_a)\n",
        "        mask[key_b] = False\n",
        "        MeshPool.__remove_group(mesh, edge_groups, key_b)\n",
        "        mesh.remove_edge(key_b)\n",
        "        mesh.edges_count -= 1\n",
        "        return key_a\n",
        "\n",
        "    @staticmethod\n",
        "    def __get_invalids(mesh, edge_id, edge_groups, side):\n",
        "        info = MeshPool.__get_face_info(mesh, edge_id, side)\n",
        "        key_a, key_b, side_a, side_b, other_side_a, other_side_b, other_keys_a, other_keys_b = info\n",
        "        shared_items = MeshPool.__get_shared_items(other_keys_a, other_keys_b)\n",
        "        if len(shared_items) == 0:\n",
        "            return []\n",
        "        else:\n",
        "            assert (len(shared_items) == 2)\n",
        "            middle_edge = other_keys_a[shared_items[0]]\n",
        "            update_key_a = other_keys_a[1 - shared_items[0]]\n",
        "            update_key_b = other_keys_b[1 - shared_items[1]]\n",
        "            update_side_a = mesh.sides[key_a, other_side_a + 1 - shared_items[0]]\n",
        "            update_side_b = mesh.sides[key_b, other_side_b + 1 - shared_items[1]]\n",
        "            MeshPool.__redirect_edges(mesh, edge_id, side, update_key_a, update_side_a)\n",
        "            MeshPool.__redirect_edges(mesh, edge_id, side + 1, update_key_b, update_side_b)\n",
        "            MeshPool.__redirect_edges(mesh, update_key_a, MeshPool.__get_other_side(update_side_a), update_key_b, MeshPool.__get_other_side(update_side_b))\n",
        "            MeshPool.__union_groups(mesh, edge_groups, key_a, edge_id)\n",
        "            MeshPool.__union_groups(mesh, edge_groups, key_b, edge_id)\n",
        "            MeshPool.__union_groups(mesh, edge_groups, key_a, update_key_a)\n",
        "            MeshPool.__union_groups(mesh, edge_groups, middle_edge, update_key_a)\n",
        "            MeshPool.__union_groups(mesh, edge_groups, key_b, update_key_b)\n",
        "            MeshPool.__union_groups(mesh, edge_groups, middle_edge, update_key_b)\n",
        "            return [key_a, key_b, middle_edge]\n",
        "\n",
        "    @staticmethod\n",
        "    def __redirect_edges(mesh, edge_a_key, side_a, edge_b_key, side_b):\n",
        "        mesh.gemm_edges[edge_a_key, side_a] = edge_b_key\n",
        "        mesh.gemm_edges[edge_b_key, side_b] = edge_a_key\n",
        "        mesh.sides[edge_a_key, side_a] = side_b\n",
        "        mesh.sides[edge_b_key, side_b] = side_a\n",
        "\n",
        "    @staticmethod\n",
        "    def __get_shared_items(list_a, list_b):\n",
        "        shared_items = []\n",
        "        for i in range(len(list_a)):\n",
        "            for j in range(len(list_b)):\n",
        "                if list_a[i] == list_b[j]:\n",
        "                    shared_items.extend([i, j])\n",
        "        return shared_items\n",
        "\n",
        "    @staticmethod\n",
        "    def __get_other_side(side):\n",
        "        return side + 1 - 2 * (side % 2)\n",
        "\n",
        "    @staticmethod\n",
        "    def __get_face_info(mesh, edge_id, side):\n",
        "        key_a = mesh.gemm_edges[edge_id, side]\n",
        "        key_b = mesh.gemm_edges[edge_id, side + 1]\n",
        "        side_a = mesh.sides[edge_id, side]\n",
        "        side_b = mesh.sides[edge_id, side + 1]\n",
        "        other_side_a = (side_a - (side_a % 2) + 2) % 4\n",
        "        other_side_b = (side_b - (side_b % 2) + 2) % 4\n",
        "        other_keys_a = [mesh.gemm_edges[key_a, other_side_a], mesh.gemm_edges[key_a, other_side_a + 1]]\n",
        "        other_keys_b = [mesh.gemm_edges[key_b, other_side_b], mesh.gemm_edges[key_b, other_side_b + 1]]\n",
        "        return key_a, key_b, side_a, side_b, other_side_a, other_side_b, other_keys_a, other_keys_b\n",
        "\n",
        "    @staticmethod\n",
        "    def __remove_triplete(mesh, mask, edge_groups, invalid_edges):\n",
        "        vertex = set(mesh.edges[invalid_edges[0]])\n",
        "        for edge_key in invalid_edges:\n",
        "            vertex &= set(mesh.edges[edge_key])\n",
        "            mask[edge_key] = False\n",
        "            MeshPool.__remove_group(mesh, edge_groups, edge_key)\n",
        "        mesh.edges_count -= 3\n",
        "        vertex = list(vertex)\n",
        "        assert(len(vertex) == 1)\n",
        "        mesh.remove_vertex(vertex[0])\n",
        "\n",
        "    def __build_queue(self, features, edges_count):\n",
        "        # delete edges with smallest norm\n",
        "        squared_magnitude = torch.sum(features * features, 0)\n",
        "        if squared_magnitude.shape[-1] != 1:\n",
        "            squared_magnitude = squared_magnitude.unsqueeze(-1)\n",
        "        edge_ids = torch.arange(edges_count, device=squared_magnitude.device, dtype=torch.float32).unsqueeze(-1)\n",
        "        heap = torch.cat((squared_magnitude, edge_ids), dim=-1).tolist()\n",
        "        heapify(heap)\n",
        "        return heap\n",
        "\n",
        "    @staticmethod\n",
        "    def __union_groups(mesh, edge_groups, source, target):\n",
        "        edge_groups.union(source, target)\n",
        "        mesh.union_groups(source, target)\n",
        "\n",
        "    @staticmethod\n",
        "    def __remove_group(mesh, edge_groups, index):\n",
        "        edge_groups.remove_group(index)\n",
        "        mesh.remove_group(index)"
      ],
      "execution_count": 74,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W-nBUZ9vNrOD"
      },
      "source": [
        "## MeshConv classes and functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "grVX9bxWNpRN"
      },
      "source": [
        "def get_norm_layer(norm_type='instance', num_groups=1):\n",
        "    if norm_type == 'batch':\n",
        "        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)\n",
        "    elif norm_type == 'instance':\n",
        "        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)\n",
        "    elif norm_type == 'group':\n",
        "        norm_layer = functools.partial(nn.GroupNorm, affine=True, num_groups=num_groups)\n",
        "    elif norm_type == 'none':\n",
        "        norm_layer = NoNorm\n",
        "    else:\n",
        "        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)\n",
        "    return norm_layer\n",
        "\n",
        "def get_norm_args(norm_layer, nfeats_list):\n",
        "    if hasattr(norm_layer, '__name__') and norm_layer.__name__ == 'NoNorm':\n",
        "        norm_args = [{'fake': True} for f in nfeats_list]\n",
        "    elif norm_layer.func.__name__ == 'GroupNorm':\n",
        "        norm_args = [{'num_channels': f} for f in nfeats_list]\n",
        "    elif norm_layer.func.__name__ == 'BatchNorm2d':\n",
        "        norm_args = [{'num_features': f} for f in nfeats_list]\n",
        "    else:\n",
        "        raise NotImplementedError('normalization layer [%s] is not found' % norm_layer.func.__name__)\n",
        "    return norm_args\n",
        "\n",
        "#MeshConv layer\n",
        "class MeshConv(nn.Module):\n",
        "    \"\"\" Computes convolution between edges and 4 incident (1-ring) edge neighbors\n",
        "    in the forward pass takes:\n",
        "    x: edge features (Batch x Features x Edges)\n",
        "    mesh: list of mesh data-structure (len(mesh) == Batch)\n",
        "    and applies convolution\n",
        "    \"\"\"\n",
        "    def __init__(self, in_channels, out_channels, k=5, bias=True):\n",
        "        super(MeshConv, self).__init__()\n",
        "        self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(1, k), bias=bias)\n",
        "        self.k = k\n",
        "\n",
        "    def __call__(self, edge_f, mesh):\n",
        "        return self.forward(edge_f, mesh)\n",
        "\n",
        "    def forward(self, x, mesh):\n",
        "        x = x.squeeze(-1)\n",
        "        G = torch.cat([self.pad_gemm(i, x.shape[2], x.device) for i in mesh], 0)\n",
        "        # build 'neighborhood image' and apply convolution\n",
        "        G = self.create_GeMM(x, G)\n",
        "        x = self.conv(G)\n",
        "        return x\n",
        "\n",
        "    def flatten_gemm_inds(self, Gi):\n",
        "        (b, ne, nn) = Gi.shape\n",
        "        ne += 1\n",
        "        batch_n = torch.floor(torch.arange(b * ne, device=Gi.device).float() / ne).view(b, ne)\n",
        "        add_fac = batch_n * ne\n",
        "        add_fac = add_fac.view(b, ne, 1)\n",
        "        add_fac = add_fac.repeat(1, 1, nn)\n",
        "        # flatten Gi\n",
        "        Gi = Gi.float() + add_fac[:, 1:, :]\n",
        "        return Gi\n",
        "\n",
        "    def create_GeMM(self, x, Gi):\n",
        "        \"\"\" gathers the edge features (x) with from the 1-ring indices (Gi)\n",
        "        applys symmetric functions to handle order invariance\n",
        "        returns a 'fake image' which can use 2d convolution on\n",
        "        output dimensions: Batch x Channels x Edges x 5\n",
        "        \"\"\"\n",
        "        Gishape = Gi.shape\n",
        "        # pad the first row of  every sample in batch with zeros\n",
        "        padding = torch.zeros((x.shape[0], x.shape[1], 1), requires_grad=True, device=x.device)\n",
        "        # padding = padding.to(x.device)\n",
        "        x = torch.cat((padding, x), dim=2)\n",
        "        Gi = Gi + 1 #shift\n",
        "\n",
        "        # first flatten indices\n",
        "        Gi_flat = self.flatten_gemm_inds(Gi)\n",
        "        Gi_flat = Gi_flat.view(-1).long()\n",
        "        #\n",
        "        odim = x.shape\n",
        "        x = x.permute(0, 2, 1).contiguous()\n",
        "        x = x.view(odim[0] * odim[2], odim[1])\n",
        "\n",
        "        f = torch.index_select(x, dim=0, index=Gi_flat)\n",
        "        f = f.view(Gishape[0], Gishape[1], Gishape[2], -1)\n",
        "        f = f.permute(0, 3, 1, 2)\n",
        "\n",
        "        # apply the symmetric functions for an equivariant conv\n",
        "        x_1 = f[:, :, :, 1] + f[:, :, :, 3]\n",
        "        x_2 = f[:, :, :, 2] + f[:, :, :, 4]\n",
        "        x_3 = torch.abs(f[:, :, :, 1] - f[:, :, :, 3])\n",
        "        x_4 = torch.abs(f[:, :, :, 2] - f[:, :, :, 4])\n",
        "        f = torch.stack([f[:, :, :, 0], x_1, x_2, x_3, x_4], dim=3)\n",
        "        return f\n",
        "\n",
        "    def pad_gemm(self, m, xsz, device):\n",
        "        \"\"\" extracts one-ring neighbors (4x) -> m.gemm_edges\n",
        "        which is of size #edges x 4\n",
        "        add the edge_id itself to make #edges x 5\n",
        "        then pad to desired size e.g., xsz x 5\n",
        "        \"\"\"\n",
        "        padded_gemm = torch.tensor(m.gemm_edges, device=device).float()\n",
        "        padded_gemm = padded_gemm.requires_grad_()\n",
        "        padded_gemm = torch.cat((torch.arange(m.edges_count, device=device).float().unsqueeze(1), padded_gemm), dim=1)\n",
        "        # pad using F\n",
        "        padded_gemm = F.pad(padded_gemm, (0, 0, 0, xsz - m.edges_count), \"constant\", 0)\n",
        "        padded_gemm = padded_gemm.unsqueeze(0)\n",
        "        return padded_gemm\n",
        "    \n",
        "class MResConv(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, skips=1):\n",
        "        super(MResConv, self).__init__()\n",
        "        self.in_channels = in_channels\n",
        "        self.out_channels = out_channels\n",
        "        self.skips = skips\n",
        "        self.conv0 = MeshConv(self.in_channels, self.out_channels, bias=False)\n",
        "        for i in range(self.skips):\n",
        "            setattr(self, 'bn{}'.format(i + 1), nn.BatchNorm2d(self.out_channels))\n",
        "            setattr(self, 'conv{}'.format(i + 1),\n",
        "                    MeshConv(self.out_channels, self.out_channels, bias=False))\n",
        "\n",
        "    def forward(self, x, mesh):\n",
        "        x = self.conv0(x, mesh)\n",
        "        x1 = x\n",
        "        for i in range(self.skips):\n",
        "            x = getattr(self, 'bn{}'.format(i + 1))(F.relu(x))\n",
        "            x = getattr(self, 'conv{}'.format(i + 1))(x, mesh)\n",
        "        x += x1\n",
        "        x = F.relu(x)\n",
        "        return x\n",
        "    "
      ],
      "execution_count": 95,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MtaLHTRkMyjd"
      },
      "source": [
        "## SlicerModel: adapted from meshCNN ClassifierModel"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4KQmPI5JMvyR"
      },
      "source": [
        ""
      ],
      "execution_count": 76,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9HVhpuLLdKSz"
      },
      "source": [
        "## Mesh utils from MeshCNN (with slight modifications)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XenKgzPPciFH"
      },
      "source": [
        "def pad(input_arr, target_length, val=0, dim=1):\n",
        "    shp = input_arr.shape\n",
        "    npad = [(0, 0) for _ in range(len(shp))]\n",
        "    npad[dim] = (0, target_length - shp[dim])\n",
        "    return np.pad(input_arr, pad_width=npad, mode='constant', constant_values=val)\n",
        "\n",
        "def fill_mesh(mesh2fill, file: str, opt):\n",
        "    load_path = get_mesh_path(file, opt.num_aug)\n",
        "    if os.path.exists(load_path):\n",
        "        mesh_data = np.load(load_path, encoding='latin1', allow_pickle=True)\n",
        "    else:\n",
        "        mesh_data = from_scratch(file, opt)\n",
        "        np.savez_compressed(load_path, gemm_edges=mesh_data.gemm_edges, vs=mesh_data.vs, edges=mesh_data.edges,\n",
        "                            edges_count=mesh_data.edges_count, ve=mesh_data.ve, v_mask=mesh_data.v_mask,\n",
        "                            filename=mesh_data.filename, sides=mesh_data.sides,\n",
        "                            edge_lengths=mesh_data.edge_lengths, edge_areas=mesh_data.edge_areas,\n",
        "                            features=mesh_data.features)\n",
        "    mesh2fill.vs = mesh_data['vs']\n",
        "    mesh2fill.edges = mesh_data['edges']\n",
        "    mesh2fill.gemm_edges = mesh_data['gemm_edges']\n",
        "    mesh2fill.edges_count = int(mesh_data['edges_count'])\n",
        "    mesh2fill.ve = mesh_data['ve']\n",
        "    mesh2fill.v_mask = mesh_data['v_mask']\n",
        "    mesh2fill.filename = str(mesh_data['filename'])\n",
        "    mesh2fill.edge_lengths = mesh_data['edge_lengths']\n",
        "    mesh2fill.edge_areas = mesh_data['edge_areas']\n",
        "    mesh2fill.features = mesh_data['features']\n",
        "    mesh2fill.sides = mesh_data['sides']\n",
        "\n",
        "def get_mesh_path(file: str, num_aug: int):\n",
        "    filename, _ = os.path.splitext(file)\n",
        "    dir_name = os.path.dirname(filename)\n",
        "    prefix = os.path.basename(filename)\n",
        "    load_dir = os.path.join(dir_name, 'cache')\n",
        "    load_file = os.path.join(load_dir, '%s_%03d.npz' % (prefix, np.random.randint(0, num_aug)))\n",
        "    if not os.path.isdir(load_dir):\n",
        "        os.makedirs(load_dir, exist_ok=True)\n",
        "    return load_file\n",
        "\n",
        "def from_scratch(file, opt):\n",
        "\n",
        "    class MeshPrep:\n",
        "        def __getitem__(self, item):\n",
        "            return eval('self.' + item)\n",
        "\n",
        "    mesh_data = MeshPrep()\n",
        "    mesh_data.vs = mesh_data.edges = None\n",
        "    mesh_data.gemm_edges = mesh_data.sides = None\n",
        "    mesh_data.edges_count = None\n",
        "    mesh_data.ve = None\n",
        "    mesh_data.v_mask = None\n",
        "    mesh_data.filename = 'unknown'\n",
        "    mesh_data.edge_lengths = None\n",
        "    mesh_data.edge_areas = []\n",
        "    mesh_data.vs, faces = fill_from_file(mesh_data, file)\n",
        "    mesh_data.v_mask = np.ones(len(mesh_data.vs), dtype=bool)\n",
        "    faces, face_areas = remove_non_manifolds(mesh_data, faces)\n",
        "    if opt.num_aug > 1:\n",
        "        faces = augmentation(mesh_data, opt, faces)\n",
        "    build_gemm(mesh_data, faces, face_areas)\n",
        "    if opt.num_aug > 1:\n",
        "        post_augmentation(mesh_data, opt)\n",
        "    mesh_data.features = extract_features(mesh_data)\n",
        "    return mesh_data\n",
        "\n",
        "def fill_from_file(mesh, file):\n",
        "    # mesh.filename = ntpath.split(file)[1]\n",
        "    mesh.filename = os.path.basename(file)\n",
        "    mesh.fullfilename = file\n",
        "    vs, faces = [], []\n",
        "    f = open(file)\n",
        "    for line in f:\n",
        "        line = line.strip()\n",
        "        splitted_line = line.split()\n",
        "        if not splitted_line:\n",
        "            continue\n",
        "        elif splitted_line[0] == 'v':\n",
        "            vs.append([float(v) for v in splitted_line[1:4]])\n",
        "        elif splitted_line[0] == 'f':\n",
        "            face_vertex_ids = [int(c.split('/')[0]) for c in splitted_line[1:]]\n",
        "            assert len(face_vertex_ids) == 3\n",
        "            face_vertex_ids = [(ind - 1) if (ind >= 0) else (len(vs) + ind)\n",
        "                               for ind in face_vertex_ids]\n",
        "            faces.append(face_vertex_ids)\n",
        "    f.close()\n",
        "    vs = np.asarray(vs)\n",
        "    faces = np.asarray(faces, dtype=int)\n",
        "    assert np.logical_and(faces >= 0, faces < len(vs)).all()\n",
        "    return vs, faces\n",
        "\n",
        "\n",
        "def remove_non_manifolds(mesh, faces):\n",
        "    mesh.ve = [[] for _ in mesh.vs]\n",
        "    edges_set = set()\n",
        "    mask = np.ones(len(faces), dtype=bool)\n",
        "    _, face_areas = compute_face_normals_and_areas(mesh, faces)\n",
        "    for face_id, face in enumerate(faces):\n",
        "        if face_areas[face_id] == 0:\n",
        "            mask[face_id] = False\n",
        "            continue\n",
        "        faces_edges = []\n",
        "        is_manifold = False\n",
        "        for i in range(3):\n",
        "            cur_edge = (face[i], face[(i + 1) % 3])\n",
        "            if cur_edge in edges_set:\n",
        "                is_manifold = True\n",
        "                break\n",
        "            else:\n",
        "                faces_edges.append(cur_edge)\n",
        "        if is_manifold:\n",
        "            mask[face_id] = False\n",
        "        else:\n",
        "            for idx, edge in enumerate(faces_edges):\n",
        "                edges_set.add(edge)\n",
        "    return faces[mask], face_areas[mask]\n",
        "\n",
        "\n",
        "def build_gemm(mesh, faces, face_areas):\n",
        "    \"\"\"\n",
        "    gemm_edges: array (#E x 4) of the 4 one-ring neighbors for each edge\n",
        "    sides: array (#E x 4) indices (values of: 0,1,2,3) indicating where an edge is in the gemm_edge entry of the 4 neighboring edges\n",
        "    for example edge i -> gemm_edges[gemm_edges[i], sides[i]] == [i, i, i, i]\n",
        "    \"\"\"\n",
        "    mesh.ve = [[] for _ in mesh.vs]\n",
        "    edge_nb = []\n",
        "    sides = []\n",
        "    edge2key = dict()\n",
        "    edges = []\n",
        "    edges_count = 0\n",
        "    nb_count = []\n",
        "    for face_id, face in enumerate(faces):\n",
        "        faces_edges = []\n",
        "        for i in range(3):\n",
        "            cur_edge = (face[i], face[(i + 1) % 3])\n",
        "            faces_edges.append(cur_edge)\n",
        "        for idx, edge in enumerate(faces_edges):\n",
        "            edge = tuple(sorted(list(edge)))\n",
        "            faces_edges[idx] = edge\n",
        "            if edge not in edge2key:\n",
        "                edge2key[edge] = edges_count\n",
        "                edges.append(list(edge))\n",
        "                edge_nb.append([-1, -1, -1, -1])\n",
        "                sides.append([-1, -1, -1, -1])\n",
        "                mesh.ve[edge[0]].append(edges_count)\n",
        "                mesh.ve[edge[1]].append(edges_count)\n",
        "                mesh.edge_areas.append(0)\n",
        "                nb_count.append(0)\n",
        "                edges_count += 1\n",
        "            mesh.edge_areas[edge2key[edge]] += face_areas[face_id] / 3\n",
        "        for idx, edge in enumerate(faces_edges):\n",
        "            edge_key = edge2key[edge]\n",
        "            edge_nb[edge_key][nb_count[edge_key]] = edge2key[faces_edges[(idx + 1) % 3]]\n",
        "            edge_nb[edge_key][nb_count[edge_key] + 1] = edge2key[faces_edges[(idx + 2) % 3]]\n",
        "            nb_count[edge_key] += 2\n",
        "        for idx, edge in enumerate(faces_edges):\n",
        "            edge_key = edge2key[edge]\n",
        "            sides[edge_key][nb_count[edge_key] - 2] = nb_count[edge2key[faces_edges[(idx + 1) % 3]]] - 1\n",
        "            sides[edge_key][nb_count[edge_key] - 1] = nb_count[edge2key[faces_edges[(idx + 2) % 3]]] - 2\n",
        "    mesh.edges = np.array(edges, dtype=np.int32)\n",
        "    mesh.gemm_edges = np.array(edge_nb, dtype=np.int64)\n",
        "    mesh.sides = np.array(sides, dtype=np.int64)\n",
        "    mesh.edges_count = edges_count\n",
        "    mesh.edge_areas = np.array(mesh.edge_areas, dtype=np.float32) / np.sum(face_areas) #todo whats the difference between edge_areas and edge_lenghts?\n",
        "\n",
        "\n",
        "def compute_face_normals_and_areas(mesh, faces):\n",
        "    face_normals = np.cross(mesh.vs[faces[:, 1]] - mesh.vs[faces[:, 0]],\n",
        "                            mesh.vs[faces[:, 2]] - mesh.vs[faces[:, 1]])\n",
        "    face_areas = np.sqrt((face_normals ** 2).sum(axis=1))\n",
        "    face_normals /= face_areas[:, np.newaxis]\n",
        "    assert (not np.any(face_areas[:, np.newaxis] == 0)), 'has zero area face: %s' % mesh.filename\n",
        "    face_areas *= 0.5\n",
        "    return face_normals, face_areas\n",
        "\n",
        "\n",
        "# Data augmentation methods\n",
        "def augmentation(mesh, opt, faces=None):\n",
        "    if hasattr(opt, 'scale_verts') and opt.scale_verts:\n",
        "        scale_verts(mesh)\n",
        "    if hasattr(opt, 'flip_edges') and opt.flip_edges:\n",
        "        faces = flip_edges(mesh, opt.flip_edges, faces)\n",
        "    return faces\n",
        "\n",
        "\n",
        "def post_augmentation(mesh, opt):\n",
        "    if hasattr(opt, 'slide_verts') and opt.slide_verts:\n",
        "        slide_verts(mesh, opt.slide_verts)\n",
        "\n",
        "\n",
        "def slide_verts(mesh, prct):\n",
        "    edge_points = get_edge_points(mesh)\n",
        "    dihedral = dihedral_angle(mesh, edge_points).squeeze() #todo make fixed_division epsilon=0\n",
        "    thr = np.mean(dihedral) + np.std(dihedral)\n",
        "    vids = np.random.permutation(len(mesh.ve))\n",
        "    target = int(prct * len(vids))\n",
        "    shifted = 0\n",
        "    for vi in vids:\n",
        "        if shifted < target:\n",
        "            edges = mesh.ve[vi]\n",
        "            if min(dihedral[edges]) > 2.65:\n",
        "                edge = mesh.edges[np.random.choice(edges)]\n",
        "                vi_t = edge[1] if vi == edge[0] else edge[0]\n",
        "                nv = mesh.vs[vi] + np.random.uniform(0.2, 0.5) * (mesh.vs[vi_t] - mesh.vs[vi])\n",
        "                mesh.vs[vi] = nv\n",
        "                shifted += 1\n",
        "        else:\n",
        "            break\n",
        "    mesh.shifted = shifted / len(mesh.ve)\n",
        "\n",
        "\n",
        "def scale_verts(mesh, mean=1, var=0.1):\n",
        "    for i in range(mesh.vs.shape[1]):\n",
        "        mesh.vs[:, i] = mesh.vs[:, i] * np.random.normal(mean, var)\n",
        "\n",
        "\n",
        "def angles_from_faces(mesh, edge_faces, faces):\n",
        "    normals = [None, None]\n",
        "    for i in range(2):\n",
        "        edge_a = mesh.vs[faces[edge_faces[:, i], 2]] - mesh.vs[faces[edge_faces[:, i], 1]]\n",
        "        edge_b = mesh.vs[faces[edge_faces[:, i], 1]] - mesh.vs[faces[edge_faces[:, i], 0]]\n",
        "        normals[i] = np.cross(edge_a, edge_b)\n",
        "        div = fixed_division(np.linalg.norm(normals[i], ord=2, axis=1), epsilon=0)\n",
        "        normals[i] /= div[:, np.newaxis]\n",
        "    dot = np.sum(normals[0] * normals[1], axis=1).clip(-1, 1)\n",
        "    angles = np.pi - np.arccos(dot)\n",
        "    return angles\n",
        "\n",
        "\n",
        "def flip_edges(mesh, prct, faces):\n",
        "    edge_count, edge_faces, edges_dict = get_edge_faces(faces)\n",
        "    dihedral = angles_from_faces(mesh, edge_faces[:, 2:], faces)\n",
        "    edges2flip = np.random.permutation(edge_count)\n",
        "    # print(dihedral.min())\n",
        "    # print(dihedral.max())\n",
        "    target = int(prct * edge_count)\n",
        "    flipped = 0\n",
        "    for edge_key in edges2flip:\n",
        "        if flipped == target:\n",
        "            break\n",
        "        if dihedral[edge_key] > 2.7:\n",
        "            edge_info = edge_faces[edge_key]\n",
        "            if edge_info[3] == -1:\n",
        "                continue\n",
        "            new_edge = tuple(sorted(list(set(faces[edge_info[2]]) ^ set(faces[edge_info[3]]))))\n",
        "            if new_edge in edges_dict:\n",
        "                continue\n",
        "            new_faces = np.array(\n",
        "                [[edge_info[1], new_edge[0], new_edge[1]], [edge_info[0], new_edge[0], new_edge[1]]])\n",
        "            if check_area(mesh, new_faces):\n",
        "                del edges_dict[(edge_info[0], edge_info[1])]\n",
        "                edge_info[:2] = [new_edge[0], new_edge[1]]\n",
        "                edges_dict[new_edge] = edge_key\n",
        "                rebuild_face(faces[edge_info[2]], new_faces[0])\n",
        "                rebuild_face(faces[edge_info[3]], new_faces[1])\n",
        "                for i, face_id in enumerate([edge_info[2], edge_info[3]]):\n",
        "                    cur_face = faces[face_id]\n",
        "                    for j in range(3):\n",
        "                        cur_edge = tuple(sorted((cur_face[j], cur_face[(j + 1) % 3])))\n",
        "                        if cur_edge != new_edge:\n",
        "                            cur_edge_key = edges_dict[cur_edge]\n",
        "                            for idx, face_nb in enumerate(\n",
        "                                    [edge_faces[cur_edge_key, 2], edge_faces[cur_edge_key, 3]]):\n",
        "                                if face_nb == edge_info[2 + (i + 1) % 2]:\n",
        "                                    edge_faces[cur_edge_key, 2 + idx] = face_id\n",
        "                flipped += 1\n",
        "    # print(flipped)\n",
        "    return faces\n",
        "\n",
        "\n",
        "def rebuild_face(face, new_face):\n",
        "    new_point = list(set(new_face) - set(face))[0]\n",
        "    for i in range(3):\n",
        "        if face[i] not in new_face:\n",
        "            face[i] = new_point\n",
        "            break\n",
        "    return face\n",
        "\n",
        "def check_area(mesh, faces):\n",
        "    face_normals = np.cross(mesh.vs[faces[:, 1]] - mesh.vs[faces[:, 0]],\n",
        "                            mesh.vs[faces[:, 2]] - mesh.vs[faces[:, 1]])\n",
        "    face_areas = np.sqrt((face_normals ** 2).sum(axis=1))\n",
        "    face_areas *= 0.5\n",
        "    return face_areas[0] > 0 and face_areas[1] > 0\n",
        "\n",
        "\n",
        "def get_edge_faces(faces):\n",
        "    edge_count = 0\n",
        "    edge_faces = []\n",
        "    edge2keys = dict()\n",
        "    for face_id, face in enumerate(faces):\n",
        "        for i in range(3):\n",
        "            cur_edge = tuple(sorted((face[i], face[(i + 1) % 3])))\n",
        "            if cur_edge not in edge2keys:\n",
        "                edge2keys[cur_edge] = edge_count\n",
        "                edge_count += 1\n",
        "                edge_faces.append(np.array([cur_edge[0], cur_edge[1], -1, -1]))\n",
        "            edge_key = edge2keys[cur_edge]\n",
        "            if edge_faces[edge_key][2] == -1:\n",
        "                edge_faces[edge_key][2] = face_id\n",
        "            else:\n",
        "                edge_faces[edge_key][3] = face_id\n",
        "    return edge_count, np.array(edge_faces), edge2keys\n",
        "\n",
        "\n",
        "def set_edge_lengths(mesh, edge_points=None):\n",
        "    if edge_points is not None:\n",
        "        edge_points = get_edge_points(mesh)\n",
        "    edge_lengths = np.linalg.norm(mesh.vs[edge_points[:, 0]] - mesh.vs[edge_points[:, 1]], ord=2, axis=1)\n",
        "    mesh.edge_lengths = edge_lengths\n",
        "\n",
        "\n",
        "def extract_features(mesh):\n",
        "    features = []\n",
        "    edge_points = get_edge_points(mesh)\n",
        "    set_edge_lengths(mesh, edge_points)\n",
        "    with np.errstate(divide='raise'):\n",
        "        try:\n",
        "            for extractor in [dihedral_angle, symmetric_opposite_angles, symmetric_ratios]:\n",
        "                feature = extractor(mesh, edge_points)\n",
        "                features.append(feature)\n",
        "            return np.concatenate(features, axis=0)\n",
        "        except Exception as e:\n",
        "            print(e)\n",
        "            raise ValueError(mesh.filename, 'bad features')\n",
        "\n",
        "\n",
        "def dihedral_angle(mesh, edge_points):\n",
        "    normals_a = get_normals(mesh, edge_points, 0)\n",
        "    normals_b = get_normals(mesh, edge_points, 3)\n",
        "    dot = np.sum(normals_a * normals_b, axis=1).clip(-1, 1)\n",
        "    angles = np.expand_dims(np.pi - np.arccos(dot), axis=0)\n",
        "    return angles\n",
        "\n",
        "\n",
        "def symmetric_opposite_angles(mesh, edge_points):\n",
        "    \"\"\" computes two angles: one for each face shared between the edge\n",
        "        the angle is in each face opposite the edge\n",
        "        sort handles order ambiguity\n",
        "    \"\"\"\n",
        "    angles_a = get_opposite_angles(mesh, edge_points, 0)\n",
        "    angles_b = get_opposite_angles(mesh, edge_points, 3)\n",
        "    angles = np.concatenate((np.expand_dims(angles_a, 0), np.expand_dims(angles_b, 0)), axis=0)\n",
        "    angles = np.sort(angles, axis=0)\n",
        "    return angles\n",
        "\n",
        "\n",
        "def symmetric_ratios(mesh, edge_points):\n",
        "    \"\"\" computes two ratios: one for each face shared between the edge\n",
        "        the ratio is between the height / base (edge) of each triangle\n",
        "        sort handles order ambiguity\n",
        "    \"\"\"\n",
        "    ratios_a = get_ratios(mesh, edge_points, 0)\n",
        "    ratios_b = get_ratios(mesh, edge_points, 3)\n",
        "    ratios = np.concatenate((np.expand_dims(ratios_a, 0), np.expand_dims(ratios_b, 0)), axis=0)\n",
        "    return np.sort(ratios, axis=0)\n",
        "\n",
        "\n",
        "def get_edge_points(mesh):\n",
        "    \"\"\" returns: edge_points (#E x 4) tensor, with four vertex ids per edge\n",
        "        for example: edge_points[edge_id, 0] and edge_points[edge_id, 1] are the two vertices which define edge_id \n",
        "        each adjacent face to edge_id has another vertex, which is edge_points[edge_id, 2] or edge_points[edge_id, 3]\n",
        "    \"\"\"\n",
        "    edge_points = np.zeros([mesh.edges_count, 4], dtype=np.int32)\n",
        "    for edge_id, edge in enumerate(mesh.edges):\n",
        "        edge_points[edge_id] = get_side_points(mesh, edge_id)\n",
        "        # edge_points[edge_id, 3:] = mesh.get_side_points(edge_id, 2)\n",
        "    return edge_points\n",
        "\n",
        "\n",
        "def get_side_points(mesh, edge_id):\n",
        "    # if mesh.gemm_edges[edge_id, side] == -1:\n",
        "    #     return mesh.get_side_points(edge_id, ((side + 2) % 4))\n",
        "    # else:\n",
        "    edge_a = mesh.edges[edge_id]\n",
        "\n",
        "    if mesh.gemm_edges[edge_id, 0] == -1:\n",
        "        edge_b = mesh.edges[mesh.gemm_edges[edge_id, 2]]\n",
        "        edge_c = mesh.edges[mesh.gemm_edges[edge_id, 3]]\n",
        "    else:\n",
        "        edge_b = mesh.edges[mesh.gemm_edges[edge_id, 0]]\n",
        "        edge_c = mesh.edges[mesh.gemm_edges[edge_id, 1]]\n",
        "    if mesh.gemm_edges[edge_id, 2] == -1:\n",
        "        edge_d = mesh.edges[mesh.gemm_edges[edge_id, 0]]\n",
        "        edge_e = mesh.edges[mesh.gemm_edges[edge_id, 1]]\n",
        "    else:\n",
        "        edge_d = mesh.edges[mesh.gemm_edges[edge_id, 2]]\n",
        "        edge_e = mesh.edges[mesh.gemm_edges[edge_id, 3]]\n",
        "    first_vertex = 0\n",
        "    second_vertex = 0\n",
        "    third_vertex = 0\n",
        "    if edge_a[1] in edge_b:\n",
        "        first_vertex = 1\n",
        "    if edge_b[1] in edge_c:\n",
        "        second_vertex = 1\n",
        "    if edge_d[1] in edge_e:\n",
        "        third_vertex = 1\n",
        "    return [edge_a[first_vertex], edge_a[1 - first_vertex], edge_b[second_vertex], edge_d[third_vertex]]\n",
        "\n",
        "\n",
        "def get_normals(mesh, edge_points, side):\n",
        "    edge_a = mesh.vs[edge_points[:, side // 2 + 2]] - mesh.vs[edge_points[:, side // 2]]\n",
        "    edge_b = mesh.vs[edge_points[:, 1 - side // 2]] - mesh.vs[edge_points[:, side // 2]]\n",
        "    normals = np.cross(edge_a, edge_b)\n",
        "    div = fixed_division(np.linalg.norm(normals, ord=2, axis=1), epsilon=0.1)\n",
        "    normals /= div[:, np.newaxis]\n",
        "    return normals\n",
        "\n",
        "def get_opposite_angles(mesh, edge_points, side):\n",
        "    edges_a = mesh.vs[edge_points[:, side // 2]] - mesh.vs[edge_points[:, side // 2 + 2]]\n",
        "    edges_b = mesh.vs[edge_points[:, 1 - side // 2]] - mesh.vs[edge_points[:, side // 2 + 2]]\n",
        "\n",
        "    edges_a /= fixed_division(np.linalg.norm(edges_a, ord=2, axis=1), epsilon=0.1)[:, np.newaxis]\n",
        "    edges_b /= fixed_division(np.linalg.norm(edges_b, ord=2, axis=1), epsilon=0.1)[:, np.newaxis]\n",
        "    dot = np.sum(edges_a * edges_b, axis=1).clip(-1, 1)\n",
        "    return np.arccos(dot)\n",
        "\n",
        "\n",
        "def get_ratios(mesh, edge_points, side):\n",
        "    edges_lengths = np.linalg.norm(mesh.vs[edge_points[:, side // 2]] - mesh.vs[edge_points[:, 1 - side // 2]],\n",
        "                                   ord=2, axis=1)\n",
        "    point_o = mesh.vs[edge_points[:, side // 2 + 2]]\n",
        "    point_a = mesh.vs[edge_points[:, side // 2]]\n",
        "    point_b = mesh.vs[edge_points[:, 1 - side // 2]]\n",
        "    line_ab = point_b - point_a\n",
        "    projection_length = np.sum(line_ab * (point_o - point_a), axis=1) / fixed_division(\n",
        "        np.linalg.norm(line_ab, ord=2, axis=1), epsilon=0.1)\n",
        "    closest_point = point_a + (projection_length / edges_lengths)[:, np.newaxis] * line_ab\n",
        "    d = np.linalg.norm(point_o - closest_point, ord=2, axis=1)\n",
        "    return d / edges_lengths\n",
        "\n",
        "def fixed_division(to_div, epsilon):\n",
        "    if epsilon == 0:\n",
        "        to_div[to_div == 0] = 0.1\n",
        "    else:\n",
        "        to_div += epsilon\n",
        "    return to_div\n",
        "\n",
        "\n",
        "class MeshUnion:\n",
        "    def __init__(self, n, device=torch.device('cpu')):\n",
        "        self.__size = n\n",
        "        self.rebuild_features = self.rebuild_features_average\n",
        "        self.groups = torch.eye(n, device=device)\n",
        "\n",
        "    def union(self, source, target):\n",
        "        self.groups[target, :] += self.groups[source, :]\n",
        "\n",
        "    def remove_group(self, index):\n",
        "        return\n",
        "\n",
        "    def get_group(self, edge_key):\n",
        "        return self.groups[edge_key, :]\n",
        "\n",
        "    def get_occurrences(self):\n",
        "        return torch.sum(self.groups, 0)\n",
        "\n",
        "    def get_groups(self, tensor_mask):\n",
        "        self.groups = torch.clamp(self.groups, 0, 1)\n",
        "        return self.groups[tensor_mask, :]\n",
        "\n",
        "    def rebuild_features_average(self, features, mask, target_edges):\n",
        "        self.prepare_groups(features, mask)\n",
        "        fe = torch.matmul(features.squeeze(-1), self.groups)\n",
        "        occurrences = torch.sum(self.groups, 0).expand(fe.shape)\n",
        "        fe = fe / occurrences\n",
        "        padding_b = target_edges - fe.shape[1]\n",
        "        if padding_b > 0:\n",
        "            padding_b = nn.ConstantPad2d((0, padding_b, 0, 0), 0)\n",
        "            fe = padding_b(fe)\n",
        "        return fe\n",
        "\n",
        "    def prepare_groups(self, features, mask):\n",
        "        tensor_mask = torch.from_numpy(mask)\n",
        "        self.groups = torch.clamp(self.groups[tensor_mask, :], 0, 1).transpose_(1, 0)\n",
        "        padding_a = features.shape[1] - self.groups.shape[0]\n",
        "        if padding_a > 0:\n",
        "            padding_a = nn.ConstantPad2d((0, 0, 0, padding_a), 0)\n",
        "            self.groups = padding_a(self.groups)\n",
        "\n",
        "class Mesh:\n",
        "    def __init__(self, file=None, opt=None, hold_history=False, export_folder=''):\n",
        "        self.vs = self.v_mask = self.filename = self.features = self.edge_areas = None\n",
        "        self.edges = self.gemm_edges = self.sides = None\n",
        "        self.pool_count = 0\n",
        "        fill_mesh(self, file, opt)\n",
        "        self.export_folder = export_folder\n",
        "        self.history_data = None\n",
        "        if hold_history:\n",
        "            self.init_history()\n",
        "        self.export()\n",
        "\n",
        "    def extract_features(self):\n",
        "        return self.features\n",
        "\n",
        "    def merge_vertices(self, edge_id):\n",
        "        self.remove_edge(edge_id)\n",
        "        edge = self.edges[edge_id]\n",
        "        v_a = self.vs[edge[0]]\n",
        "        v_b = self.vs[edge[1]]\n",
        "        # update pA\n",
        "        v_a.__iadd__(v_b)\n",
        "        v_a.__itruediv__(2)\n",
        "        self.v_mask[edge[1]] = False\n",
        "        mask = self.edges == edge[1]\n",
        "        self.ve[edge[0]].extend(self.ve[edge[1]])\n",
        "        self.edges[mask] = edge[0]\n",
        "\n",
        "    def remove_vertex(self, v):\n",
        "        self.v_mask[v] = False\n",
        "\n",
        "    def remove_edge(self, edge_id):\n",
        "        vs = self.edges[edge_id]\n",
        "        for v in vs:\n",
        "            if edge_id not in self.ve[v]:\n",
        "                print(self.ve[v])\n",
        "                print(self.filename)\n",
        "            self.ve[v].remove(edge_id)\n",
        "\n",
        "    def clean(self, edges_mask, groups):\n",
        "        edges_mask = edges_mask.astype(bool)\n",
        "        torch_mask = torch.from_numpy(edges_mask.copy())\n",
        "        self.gemm_edges = self.gemm_edges[edges_mask]\n",
        "        self.edges = self.edges[edges_mask]\n",
        "        self.sides = self.sides[edges_mask]\n",
        "        new_ve = []\n",
        "        edges_mask = np.concatenate([edges_mask, [False]])\n",
        "        new_indices = np.zeros(edges_mask.shape[0], dtype=np.int32)\n",
        "        new_indices[-1] = -1\n",
        "        new_indices[edges_mask] = np.arange(0, np.ma.where(edges_mask)[0].shape[0])\n",
        "        self.gemm_edges[:, :] = new_indices[self.gemm_edges[:, :]]\n",
        "        for v_index, ve in enumerate(self.ve):\n",
        "            update_ve = []\n",
        "            # if self.v_mask[v_index]:\n",
        "            for e in ve:\n",
        "                update_ve.append(new_indices[e])\n",
        "            new_ve.append(update_ve)\n",
        "        self.ve = new_ve\n",
        "        self.__clean_history(groups, torch_mask)\n",
        "        self.pool_count += 1\n",
        "        self.export()\n",
        "\n",
        "\n",
        "    def export(self, file=None, vcolor=None):\n",
        "        if file is None:\n",
        "            if self.export_folder:\n",
        "                filename, file_extension = os.path.splitext(self.filename)\n",
        "                file = '%s/%s_%d%s' % (self.export_folder, filename, self.pool_count, file_extension)\n",
        "            else:\n",
        "                return\n",
        "        faces = []\n",
        "        vs = self.vs[self.v_mask]\n",
        "        gemm = np.array(self.gemm_edges)\n",
        "        new_indices = np.zeros(self.v_mask.shape[0], dtype=np.int32)\n",
        "        new_indices[self.v_mask] = np.arange(0, np.ma.where(self.v_mask)[0].shape[0])\n",
        "        for edge_index in range(len(gemm)):\n",
        "            cycles = self.__get_cycle(gemm, edge_index)\n",
        "            for cycle in cycles:\n",
        "                faces.append(self.__cycle_to_face(cycle, new_indices))\n",
        "        with open(file, 'w+') as f:\n",
        "            for vi, v in enumerate(vs):\n",
        "                vcol = ' %f %f %f' % (vcolor[vi, 0], vcolor[vi, 1], vcolor[vi, 2]) if vcolor is not None else ''\n",
        "                f.write(\"v %f %f %f%s\\n\" % (v[0], v[1], v[2], vcol))\n",
        "            for face_id in range(len(faces) - 1):\n",
        "                f.write(\"f %d %d %d\\n\" % (faces[face_id][0] + 1, faces[face_id][1] + 1, faces[face_id][2] + 1))\n",
        "            f.write(\"f %d %d %d\" % (faces[-1][0] + 1, faces[-1][1] + 1, faces[-1][2] + 1))\n",
        "            for edge in self.edges:\n",
        "                f.write(\"\\ne %d %d\" % (new_indices[edge[0]] + 1, new_indices[edge[1]] + 1))\n",
        "\n",
        "    def export_segments(self, segments):\n",
        "        if not self.export_folder:\n",
        "            return\n",
        "        cur_segments = segments\n",
        "        for i in range(self.pool_count + 1):\n",
        "            filename, file_extension = os.path.splitext(self.filename)\n",
        "            file = '%s/%s_%d%s' % (self.export_folder, filename, i, file_extension)\n",
        "            fh, abs_path = mkstemp()\n",
        "            edge_key = 0\n",
        "            with os.fdopen(fh, 'w') as new_file:\n",
        "                with open(file) as old_file:\n",
        "                    for line in old_file:\n",
        "                        if line[0] == 'e':\n",
        "                            new_file.write('%s %d' % (line.strip(), cur_segments[edge_key]))\n",
        "                            if edge_key < len(cur_segments):\n",
        "                                edge_key += 1\n",
        "                                new_file.write('\\n')\n",
        "                        else:\n",
        "                            new_file.write(line)\n",
        "            os.remove(file)\n",
        "            move(abs_path, file)\n",
        "            if i < len(self.history_data['edges_mask']):\n",
        "                cur_segments = segments[:len(self.history_data['edges_mask'][i])]\n",
        "                cur_segments = cur_segments[self.history_data['edges_mask'][i]]\n",
        "\n",
        "    def __get_cycle(self, gemm, edge_id):\n",
        "        cycles = []\n",
        "        for j in range(2):\n",
        "            next_side = start_point = j * 2\n",
        "            next_key = edge_id\n",
        "            if gemm[edge_id, start_point] == -1:\n",
        "                continue\n",
        "            cycles.append([])\n",
        "            for i in range(3):\n",
        "                tmp_next_key = gemm[next_key, next_side]\n",
        "                tmp_next_side = self.sides[next_key, next_side]\n",
        "                tmp_next_side = tmp_next_side + 1 - 2 * (tmp_next_side % 2)\n",
        "                gemm[next_key, next_side] = -1\n",
        "                gemm[next_key, next_side + 1 - 2 * (next_side % 2)] = -1\n",
        "                next_key = tmp_next_key\n",
        "                next_side = tmp_next_side\n",
        "                cycles[-1].append(next_key)\n",
        "        return cycles\n",
        "\n",
        "    def __cycle_to_face(self, cycle, v_indices):\n",
        "        face = []\n",
        "        for i in range(3):\n",
        "            v = list(set(self.edges[cycle[i]]) & set(self.edges[cycle[(i + 1) % 3]]))[0]\n",
        "            face.append(v_indices[v])\n",
        "        return face\n",
        "\n",
        "    def init_history(self):\n",
        "        self.history_data = {\n",
        "                               'groups': [],\n",
        "                               'gemm_edges': [self.gemm_edges.copy()],\n",
        "                               'occurrences': [],\n",
        "                               'old2current': np.arange(self.edges_count, dtype=np.int32),\n",
        "                               'current2old': np.arange(self.edges_count, dtype=np.int32),\n",
        "                               'edges_mask': [torch.ones(self.edges_count,dtype=torch.bool)],\n",
        "                               'edges_count': [self.edges_count],\n",
        "                              }\n",
        "        if self.export_folder:\n",
        "            self.history_data['collapses'] = MeshUnion(self.edges_count)\n",
        "\n",
        "    def union_groups(self, source, target):\n",
        "        if self.export_folder and self.history_data:\n",
        "            self.history_data['collapses'].union(self.history_data['current2old'][source], self.history_data['current2old'][target])\n",
        "        return\n",
        "\n",
        "    def remove_group(self, index):\n",
        "        if self.history_data is not None:\n",
        "            self.history_data['edges_mask'][-1][self.history_data['current2old'][index]] = 0\n",
        "            self.history_data['old2current'][self.history_data['current2old'][index]] = -1\n",
        "            if self.export_folder:\n",
        "                self.history_data['collapses'].remove_group(self.history_data['current2old'][index])\n",
        "\n",
        "    def get_groups(self):\n",
        "        return self.history_data['groups'].pop()\n",
        "\n",
        "    def get_occurrences(self):\n",
        "        return self.history_data['occurrences'].pop()\n",
        "    \n",
        "    def __clean_history(self, groups, pool_mask):\n",
        "        if self.history_data is not None:\n",
        "            mask = self.history_data['old2current'] != -1\n",
        "            self.history_data['old2current'][mask] = np.arange(self.edges_count, dtype=np.int32)\n",
        "            self.history_data['current2old'][0: self.edges_count] = np.ma.where(mask)[0]\n",
        "            if self.export_folder != '':\n",
        "                self.history_data['edges_mask'].append(self.history_data['edges_mask'][-1].clone())\n",
        "            self.history_data['occurrences'].append(groups.get_occurrences())\n",
        "            self.history_data['groups'].append(groups.get_groups(pool_mask))\n",
        "            self.history_data['gemm_edges'].append(self.gemm_edges.copy())\n",
        "            self.history_data['edges_count'].append(self.edges_count)\n",
        "    \n",
        "    def unroll_gemm(self):\n",
        "        self.history_data['gemm_edges'].pop()\n",
        "        self.gemm_edges = self.history_data['gemm_edges'][-1]\n",
        "        self.history_data['edges_count'].pop()\n",
        "        self.edges_count = self.history_data['edges_count'][-1]\n",
        "\n",
        "    def get_edge_areas(self):\n",
        "        return self.edge_areas"
      ],
      "execution_count": 22,
      "outputs": []
    }
  ]
}