segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Segmentation\n",
    "\n",
    "### This notebook replicates the human body segmentation experiments  in section 6.3 of \"Field Convolutions for Surface CNNs\" (Mitchel et al. 2021)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# File reading and progressbar\n",
    "import os\n",
    "import os.path as osp\n",
    "import progressbar\n",
    "\n",
    "# Numpy\n",
    "import numpy as np\n",
    "\n",
    "# PyTorch \n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch import autograd\n",
    "\n",
    "# PyTorch Geometric - used for data loading/processing\n",
    "import torch_geometric.transforms as T\n",
    "from torch_geometric.data import DataLoader\n",
    "\n",
    "# Field Convolution modules\n",
    "from nn import FCResNetBlock, LiftBlock, ECHOBlock, TangentPerceptron, LabelSmoothingLoss\n",
    "\n",
    "# Transforms\n",
    "from transforms import FCPrecomp, computeLogXPort, SupportGraph, NormalizeArea, NormalizeAxes\n",
    "\n",
    "# Load the human body segmentation dataset (Maron et al., 2017)\n",
    "from datasets import SHAPESEG\n",
    "\n",
    "# Clear your cache\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Band-limit for field convolution filters\n",
    "band_limit = 2\n",
    "\n",
    "# Number of radial samples\n",
    "n_rings = 6\n",
    "\n",
    "# Filter type (see /nn/field_conv.py)\n",
    "ftype=1\n",
    "\n",
    "# Number of channels in the network\n",
    "nf = 48\n",
    "\n",
    "# Number of ECHO descriptors to compute in the last layer of the network\n",
    "n_des = 48;\n",
    "\n",
    "# Number of descriptor bins per unit radius\n",
    "# Descriptor resolution will be approximately  PI * (n_bins + 0.5) * (n_bins + 0.5)\n",
    "n_bins = 3\n",
    "\n",
    "# Filter support radius\n",
    "epsilon = 0.2;\n",
    "\n",
    "# Number of classes for segmentation\n",
    "n_classes = 8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data pre-processing and loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path where dataset is stored\n",
    "# The path to the .zip file containing the data should be\n",
    "# /data/SHAPESEG/raw/SHAPESEG.zip\n",
    "\n",
    "path = osp.join('data', 'SHAPESEG')\n",
    "\n",
    "# Pre-processing operations\n",
    "# Normalize meshes to have unit surface area\n",
    "# Sample points on meshes and compute convolution support edges\n",
    "# Compute logarithm maps + parallel transport\n",
    "pre_transform = T.Compose((\n",
    "    NormalizeArea(),\n",
    "    SupportGraph(epsilon=epsilon, sample_n=1024),\n",
    "    computeLogXPort(),\n",
    "    NormalizeAxes()\n",
    "))\n",
    "\n",
    "# Apply a random rotation and scale every time a shape is drawn from the dataloader\n",
    "transform = T.Compose((\n",
    "    T.RandomScale((0.85, 1.15)),\n",
    "    T.RandomRotate(45, axis=0),\n",
    "    T.RandomRotate(45, axis=1),\n",
    "    T.RandomRotate(45, axis=2))\n",
    ")\n",
    "\n",
    "# Load test and train splits\n",
    "test_dataset = SHAPESEG(path, False, pre_transform=pre_transform)\n",
    "test_loader = DataLoader(test_dataset, batch_size=1)\n",
    "train_dataset = SHAPESEG(path, True, pre_transform=pre_transform, transform=transform)\n",
    "train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FCNet\n",
    "\n",
    "##### A sucession of FCResNet blocks with an ECHO block as the final layer.  A learnable gradient-like operation is used to lift scalar features to isometry-equivariant tangent vector fields at the beginning of the network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Organizes edge data at run time to expidte convolutions\n",
    "organizeEdges = FCPrecomp(band_limit=band_limit, n_rings=n_rings, epsilon=epsilon)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "\n",
    "        ## Learned 'gradient', lifting scalar features to tangent vector features\n",
    "        ## at the beginning of the network\n",
    "        self.lift = LiftBlock(3, nf, n_rings=n_rings, ftype=ftype)\n",
    "        \n",
    "        ## FCNet - four FCResNet Blocks, followed by an ECHO block \n",
    "\n",
    "        self.resnet1 = FCResNetBlock(nf, nf, band_limit=band_limit,\n",
    "                                     n_rings=n_rings, ftype=ftype)\n",
    "    \n",
    "        self.resnet2 = FCResNetBlock(nf, nf, band_limit=band_limit,\n",
    "                                     n_rings=n_rings, ftype=ftype)\n",
    "        \n",
    "        self.resnet3 = FCResNetBlock(nf, nf, band_limit=band_limit,\n",
    "                                     n_rings=n_rings, ftype=ftype)\n",
    "        \n",
    "        self.resnet4 = FCResNetBlock(nf, nf, band_limit=band_limit,\n",
    "                                     n_rings=n_rings, ftype=ftype)\n",
    "\n",
    "        self.echo = ECHOBlock(nf, n_classes, n_des=n_des, n_bins=n_bins, \n",
    "                             band_limit=band_limit, n_rings=n_rings, ftype=ftype)\n",
    "        \n",
    "        \n",
    "        #self.mlp1 = TangentPerceptron(nf, nf)\n",
    "        #self.mlp2 = TangentPerceptron(nf, nf)\n",
    "        #self.mlp3 = TangentPerceptron(nf, nf)\n",
    " \n",
    "        \n",
    "    def forward(self, data):\n",
    "        \n",
    "        \n",
    "        ##########################\n",
    "        ### Organize edge data ###\n",
    "        ##########################\n",
    "        supp_edges, supp_sten, ln, wxp = organizeEdges(data)\n",
    "        \n",
    "        attr_lift = (supp_edges, supp_sten[..., band_limit:(band_limit+2)])\n",
    "        attr_conv = (supp_edges, supp_sten)\n",
    "        attr_echo = (supp_edges, supp_sten, ln, wxp)\n",
    "        \n",
    "        \n",
    "        #############################################\n",
    "        ### Lift scalar features to vector fields ###\n",
    "        #############################################\n",
    "        \n",
    "        x = data.pos[data.sample_idx, :]\n",
    "        \n",
    "        x = self.lift(x, *attr_lift)\n",
    "        \n",
    "        #x = self.mlp3(self.mlp2(self.mlp1(x)))\n",
    "        ##########################\n",
    "        ### Field Convolutions ###\n",
    "        ##########################\n",
    "        \n",
    "        x = self.resnet1(x, *attr_conv)\n",
    "\n",
    "        x = self.resnet2(x, *attr_conv)\n",
    "        \n",
    "        x = self.resnet3(x, *attr_conv)\n",
    "        \n",
    "        x = self.resnet4(x, *attr_conv)\n",
    "\n",
    "        \n",
    "        \n",
    "        #######################################################\n",
    "        ### Compute ECHO descriptors and output predictions ###\n",
    "        #######################################################\n",
    "        \n",
    "        return self.echo(x, *attr_echo)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train on the GPU\n",
    "device = torch.device('cuda')\n",
    "model = Net().to(device)\n",
    "\n",
    "# ADAM Optimizer, lr = 0.01\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "\n",
    "## Cross entropy label smoothing loss\n",
    "loss = LabelSmoothingLoss(classes=n_classes, smoothing=0.2, dim=1)\n",
    "\n",
    "## Can also try without smoothing\n",
    "#loss = torch.nn.CrossEntropyLoss()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Training function\n",
    "## Optional batch_step parameter for gradient accumulation (not used in the paper)\n",
    "\n",
    "def train(batch_step=1):\n",
    "\n",
    "    model.train()\n",
    "    \n",
    "    # Sort out progress bar, displays average loss over last 50 samples\n",
    "    wW = 50;\n",
    "    window = torch.FloatTensor(wW).fill_(0)\n",
    "    \n",
    "    n_data = train_loader.__len__()\n",
    "    widgets = [progressbar.Percentage(), progressbar.Bar(), \n",
    "              progressbar.AdaptiveETA(), ' | ', progressbar.Variable('Loss'),]\n",
    "\n",
    "    bar = progressbar.ProgressBar(max_value=n_data, widgets=widgets)\n",
    "\n",
    "    \n",
    "    ## Zero-out\n",
    "    optimizer.zero_grad() \n",
    "\n",
    "    i = 0;\n",
    "    for data in train_loader:\n",
    "            \n",
    "        pred = model(data.to(device))\n",
    "\n",
    "        L = loss(pred, data.y.to(device)) \n",
    "        if (i < wW):\n",
    "            window[i] = L.item()\n",
    "            wAvg = torch.mean(window[:i])\n",
    "        else:\n",
    "            window = torch.cat((window[1:], torch.tensor([L.item() * batch_step])), dim=0)\n",
    "            wAvg = torch.mean(window)\n",
    "\n",
    "        # Update progress bar\n",
    "        i = i + 1\n",
    "        bar.update(i, Loss = torch.mean(window[:i]))\n",
    "        \n",
    "        ## Update loss\n",
    "        L = L / batch_step;\n",
    "        L.backward()\n",
    "\n",
    "        if ( i % batch_step == 0 or i == n_data ):\n",
    "            optimizer.step()\n",
    "            model.zero_grad()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Overally segmentation accuracy on the test dataset\n",
    "def test():\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total_num = 0\n",
    "    for i, data in enumerate(test_loader):\n",
    "        pred = F.log_softmax(model(data.to(device)),dim=1).max(1)[1]\n",
    "        correct += pred.eq(data.y).sum().item()\n",
    "        total_num += data.y.size(0)\n",
    "            \n",
    "    acc = correct / total_num\n",
    "\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train, then test\n",
    "###### We train for 15 epochs  as in the paper. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Training...')\n",
    "\n",
    "for epoch in range(15):\n",
    "    \n",
    "    print(\"Epoch {}\".format(epoch), flush=True)\n",
    "    train()\n",
    "  \n",
    "    \n",
    "acc = test()\n",
    "print(\"Test accuracy: {:06.4f}\".format(acc), flush=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}