{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# UNet segmentation\n\nCredit: A Grigis\n\nA simple example on how to use the SphericalUNet architecture on the\nclassification dataset.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport matplotlib.pyplot as plt\nimport torch\nfrom torch import nn\nfrom torch.utils.data import DataLoader\nfrom surfify import utils\nfrom surfify import plotting\nfrom surfify import models\nfrom surfify import datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Inspect dataset\n\nFirst we load the classification dataset (with 3 classes) and inspect the\ngenrated labels.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "standard_ico = True\nico_order = 3\nn_classes = 3\nn_epochs = 20\nico_vertices, ico_triangles = utils.icosahedron(\n    order=ico_order, standard_ico=standard_ico)\nn_vertices = len(ico_vertices)\nX, y = datasets.make_classification(\n    ico_vertices, n_samples=40, n_classes=n_classes, scale=1, seed=42)\nprint(\"Surface:\", ico_vertices.shape, ico_triangles.shape)\nprint(\"Data:\", X.shape, y.shape)\nplotting.plot_trisurf(ico_vertices, ico_triangles, y, is_label=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Train the model\n\nWe now train the SphericalUNet model using a CrossEntropy loss and a SGD\noptimizer. As it is obvious to segment the input classification dataset\nan accuracy of 100% is expected.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "dataset = datasets.ClassificationDataset(\n    ico_vertices, n_samples=40, n_classes=n_classes, scale=1, seed=42)\nloader = DataLoader(dataset, batch_size=5, shuffle=True)\nmodel = models.SphericalUNet(\n    in_order=ico_order, in_channels=n_classes, out_channels=n_classes,\n    depth=2, start_filts=8, conv_mode=\"DiNe\", dine_size=1, up_mode=\"transpose\",\n    standard_ico=standard_ico)\nloss_fn = nn.CrossEntropyLoss()\noptimizer = torch.optim.SGD(\n    model.parameters(), lr=0.1, momentum=0.99, weight_decay=1e-4)\nsize = len(loader.dataset)\nn_batches = len(loader)\nfor epoch in range(n_epochs):\n    for batch, (X, y) in enumerate(loader):\n        pred = model(X)\n        loss = loss_fn(pred, y)\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n        loss, current = loss.item(), batch * len(X)\n        if epoch % 5 == 0:\n            print(\"loss {0}: {1:>7f}  [{2:>5d}/{3:>5d}]\".format(\n                epoch, loss, current, size))\nmodel.eval()\ntest_loss, correct = 0, 0\ny_preds = []\nwith torch.no_grad():\n    for X, y in loader:\n        pred = model(X)\n        test_loss += loss_fn(pred, y).item()\n        logit = torch.nn.functional.softmax(pred, dim=1)\n        y_pred = pred.argmax(dim=1)\n        correct += (y_pred == y).type(torch.float).sum().item()\n        y_preds.append(y_pred.numpy())\ntest_loss /= n_batches\ncorrect /= (size * n_vertices)\ny_preds = np.concatenate(y_preds, axis=0)\nprint(\"Test Error: \\n Accuracy: {0:>0.1f}%, Avg loss: {1:>8f}\".format(\n    100 * correct, test_loss))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Inspect the predicted labels\n\nFinally the predicted labels of the first sample are displayed. As expected\nthey corresspond exactly to the ground truth.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plotting.plot_trisurf(ico_vertices, ico_triangles, y_pred[0], is_label=True)\nplt.show()"
      ]
    }
  ],
  "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.6.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}