tfgraphics/my_mesh_segmentation_demo_indrive.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "my_mesh_segmentation_demo_indrive.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/LittleQili/CS473/blob/main/my_mesh_segmentation_demo_indrive.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X2Fj4S3r0p1A"
      },
      "source": [
        "##### Copyright 2019 Google LLC.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WhCKWToEah3m"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "%cd drive/MyDrive/mesh_seg"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x09idYC2vsZQ"
      },
      "source": [
        "# Mesh Segmentation using Feature Steered Graph Convolutions\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/graphics/blob/master/tensorflow_graphics/notebooks/mesh_segmentation_demo.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/graphics/blob/master/tensorflow_graphics/notebooks/mesh_segmentation_demo.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "26AvKq8MJRGl"
      },
      "source": [
        "!pip install tensorflow_graphics\n",
        "# !pip install tensorboardcolab"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Gcq20GLutgqy"
      },
      "source": [
        "# from tensorboardcolab import TensorBoardColab, TensorBoardColabCallback\n",
        "# tbc=TensorBoardColab()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KlBviBxue7n0"
      },
      "source": [
        "import glob\n",
        "import os\n",
        "import tensorflow.compat.v1 as tf\n",
        "tf.disable_v2_behavior()\n",
        "\n",
        "from tensorflow_graphics.nn.layer import graph_convolution as graph_conv\n",
        "from tensorflow_graphics.notebooks import mesh_segmentation_dataio as dataio\n",
        "from tensorflow_graphics.notebooks import mesh_viewer"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AGaDtH49dlJb"
      },
      "source": [
        "Note this notebook works best in Graph mode."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Gh-ZSwXnB-5"
      },
      "source": [
        "### Fetch model files and data\n",
        "\n",
        "For convenience, we provide a pre-trained model. Let's now download a pre-trained model checkpoint and the test data. The meshes are generated using Unity Multipurpose Avatar system [UMA](https://assetstore.unity.com/packages/3d/characters/uma-2-unity-multipurpose-avatar-35611)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_ZkB3iIcvvzJ"
      },
      "source": [
        "# mydrivedatafolder = '/content/drive/MyDrive/mesh_seg/data/'\n",
        "mydrivedatafolder = ''\n",
        "path_to_model_zip = tf.keras.utils.get_file(\n",
        "    mydrivedatafolder+'model.zip',\n",
        "    origin='https://storage.googleapis.com/tensorflow-graphics/notebooks/mesh_segmentation/model.zip',\n",
        "    extract=True)\n",
        "\n",
        "path_to_data_zip = tf.keras.utils.get_file(\n",
        "    mydrivedatafolder+'data.zip',\n",
        "    origin='https://storage.googleapis.com/tensorflow-graphics/notebooks/mesh_segmentation/data.zip',\n",
        "    extract=True)\n",
        "\n",
        "local_model_dir = os.path.join(os.path.dirname(path_to_model_zip), 'model')\n",
        "# print result: ['/root/.keras/datasets/data/Dancer_test_sequence.tfrecords'] 只是一个文件路径而已\n",
        "test_data_files = [\n",
        "    os.path.join(\n",
        "        os.path.dirname(path_to_data_zip),\n",
        "        'data/Dancer_test_sequence.tfrecords')\n",
        "]\n",
        "print(test_data_files)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmh4b6VKcATt"
      },
      "source": [
        "## Load and visualize test data\n",
        "\n",
        "For graph convolutions, we need a *weighted adjacency matrix* denoting the mesh\n",
        "connectivity. Feature-steered graph convolutions expect self-edges in the mesh\n",
        "connectivity for each vertex, i.e. the diagonal of the weighted adjacency matrix\n",
        "should be non-zero. This matrix is defined as:\n",
        "```\n",
        "A[i, j] = w[i,j] if vertex i and vertex j share an edge,\n",
        "A[i, i] = w[i,i] for each vertex i,\n",
        "A[i, j] = 0 otherwise.\n",
        "where, w[i, j] = 1/(degree(vertex i)), and sum(j)(w[i,j]) = 1\n",
        "```\n",
        "Here degree(vertex i) is the number of edges incident on a vertex (including the\n",
        "self-edge). This weighted adjacency matrix is stored as a SparseTensor.\n",
        "\n",
        "We will load the test meshes from the test [tf.data.TFRecordDataset](https://www.tensorflow.org/api_docs/python/tf/data/TFRecordDataset)\n",
        "downloaded above. Each mesh is stored as a\n",
        "[tf.Example](https://www.tensorflow.org/api_docs/python/tf/train/Example), with\n",
        "the following fields:\n",
        "\n",
        "*   'num_vertices', V: Number of vertices in each mesh.\n",
        "*   'num_triangles', T: Number of triangles in each mesh.\n",
        "*   'vertices': A [V, 3] float tensor of vertex positions.\n",
        "*   'triangles': A [T, 3] integer tensor of vertex indices for each triangle.\n",
        "*   'labels': A [V] integer tensor with segmentation class label for each\n",
        "    vertex.\n",
        "\n",
        "As\n",
        "each mesh may have a varying number of vertices and faces (and the corresponding\n",
        "connectivity matrix), we pad the data tensors with '0's in each batch.\n",
        "\n",
        "For details on the dataset pipeline implementation, take a look at\n",
        "mesh_segmentation_dataio.py.\n",
        "\n",
        "Let's try to load a batch from the test TFRecordDataset, and visualize the first\n",
        "mesh with each vertex colored by the part label."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aqV6vkCkWB7J"
      },
      "source": [
        "## Model Definition\n",
        "\n",
        "Given a mesh with V vertices and D-dimensional per-vertex input features (e.g.\n",
        "vertex position, normal), we would like to create a network capable of\n",
        "classifying each vertex to a part label. Let's first create a mesh encoder that\n",
        "encodes each vertex in the mesh into C-dimensional logits, where C is the number\n",
        "of parts. First we use 1x1 convolutions to change input feature dimensions,\n",
        "followed by a sequence of feature steered graph convolutions and ReLU\n",
        "non-linearities, and finally 1x1 convolutions to logits, which are used for\n",
        "computing softmax cross entropy as described below.\n",
        "\n",
        "Note that this model does not use any form of pooling, which is outside the scope of this notebook.\n",
        "\n",
        "![](https://storage.googleapis.com/tensorflow-graphics/notebooks/mesh_segmentation/mesh_segmentation_model_def.png)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fQVeuGazM0LK"
      },
      "source": [
        "MODEL_PARAMS = {\n",
        "    'num_filters': 8,\n",
        "    'num_classes': 16,\n",
        "    'encoder_filter_dims': [32, 64, 128],\n",
        "}\n",
        "\n",
        "def my_inceptionv0(vertex_features):\n",
        "  vf_256 = tf.keras.layers.Conv1D(\n",
        "      256, 1, name='inceptionv0lin256')(\n",
        "          vertex_features)\n",
        "  vf_128 = tf.keras.layers.Conv1D(\n",
        "      128, 1, name='inceptionv0lin128')(\n",
        "          vertex_features)\n",
        "  vf_64 = tf.keras.layers.Conv1D(\n",
        "      64, 1, name='inceptionv0lin64')(\n",
        "          vertex_features)\n",
        "  vf_32 = tf.keras.layers.Conv1D(\n",
        "      32, 1, name='inceptionv0lin32')(\n",
        "          vertex_features)\n",
        "\n",
        "  res_inc = tf.keras.layers.Concatenate(axis=-1)([vf_256,vf_128,vf_64,vf_32])\n",
        "  res_inc = tf.nn.relu(res_inc)\n",
        "  return res_inc\n",
        "\n",
        "def my_inceptionv1(vertex_features,num_filters,batch_mesh_data):\n",
        "  dim0 = [32, 64, 128]\n",
        "  dim1 = [64, 64, 128]\n",
        "  dim2 = [32, 32, 128]\n",
        "\n",
        "  vf_0 = vertex_features\n",
        "  vf_1 = vertex_features\n",
        "  vf_2 = vertex_features\n",
        "\n",
        "  for dim in dim0:\n",
        "    # with tf.variable_scope('conv_%d' % dim):\n",
        "    vf_0 = graph_conv.feature_steered_convolution_layer(\n",
        "          vf_0,\n",
        "          batch_mesh_data['neighbors'],\n",
        "          batch_mesh_data['num_vertices'],\n",
        "          num_weight_matrices=num_filters,\n",
        "          num_output_channels=dim,\n",
        "          name='graph_convolution_feature_steered_convolution0'+str(dim))\n",
        "    vf_0 = tf.nn.relu(vf_0)\n",
        "      \n",
        "  for dim in dim1:\n",
        "    # with tf.variable_scope('conv_%d' % dim):\n",
        "    vf_1 = graph_conv.feature_steered_convolution_layer(\n",
        "          vf_1,\n",
        "          batch_mesh_data['neighbors'],\n",
        "          batch_mesh_data['num_vertices'],\n",
        "          num_weight_matrices=num_filters,\n",
        "          num_output_channels=dim,\n",
        "          name='graph_convolution_feature_steered_convolution1'+str(dim),)\n",
        "          # var_name = 'graph_convolution_feature_steered_convolutionweight1'+'dim')\n",
        "    vf_1 = tf.nn.relu(vf_1)\n",
        "  for dim in dim2:\n",
        "    # with tf.variable_scope('conv_%d' % dim):\n",
        "    vf_2 = graph_conv.feature_steered_convolution_layer(\n",
        "          vf_2,\n",
        "          batch_mesh_data['neighbors'],\n",
        "          batch_mesh_data['num_vertices'],\n",
        "          num_weight_matrices=num_filters,\n",
        "          num_output_channels=dim,\n",
        "          name='graph_convolution_feature_steered_convolution2'+str(dim))\n",
        "    vf_2 = tf.nn.relu(vf_2)\n",
        "  \n",
        "  res_inc = tf.keras.layers.Concatenate(axis=-1)([vf_0,vf_1,vf_2])\n",
        "  res_inc = tf.keras.layers.Conv1D(\n",
        "      128, 1, name='inceptionv1lin128')(\n",
        "          res_inc)\n",
        "  return res_inc\n",
        "\n",
        "def my_inceptionv1_5(vertex_features,num_filters,batch_mesh_data):\n",
        "  dim0 = [32, 64]\n",
        "  dim1 = [64, 64]\n",
        "  dim2 = [32, 32]\n",
        "\n",
        "  vf_0 = vertex_features\n",
        "  vf_1 = vertex_features\n",
        "  vf_2 = vertex_features\n",
        "\n",
        "  for dim in dim0:\n",
        "    # with tf.variable_scope('conv_%d' % dim):\n",
        "    vf_0 = graph_conv.feature_steered_convolution_layer(\n",
        "          vf_0,\n",
        "          batch_mesh_data['neighbors'],\n",
        "          batch_mesh_data['num_vertices'],\n",
        "          num_weight_matrices=num_filters,\n",
        "          num_output_channels=dim,\n",
        "          name='graph_convolution_feature_steered_convolution0'+str(dim))\n",
        "    vf_0 = tf.nn.relu(vf_0)\n",
        "      \n",
        "  for dim in dim1:\n",
        "    # with tf.variable_scope('conv_%d' % dim):\n",
        "    vf_1 = graph_conv.feature_steered_convolution_layer(\n",
        "          vf_1,\n",
        "          batch_mesh_data['neighbors'],\n",
        "          batch_mesh_data['num_vertices'],\n",
        "          num_weight_matrices=num_filters,\n",
        "          num_output_channels=dim,\n",
        "          name='graph_convolution_feature_steered_convolution1'+str(dim))\n",
        "          # var_name = 'graph_convolution_feature_steered_convolutionweight1'+'dim')\n",
        "    vf_1 = tf.nn.relu(vf_1)\n",
        "  for dim in dim2:\n",
        "    # with tf.variable_scope('conv_%d' % dim):\n",
        "    vf_2 = graph_conv.feature_steered_convolution_layer(\n",
        "          vf_2,\n",
        "          batch_mesh_data['neighbors'],\n",
        "          batch_mesh_data['num_vertices'],\n",
        "          num_weight_matrices=num_filters,\n",
        "          num_output_channels=dim,\n",
        "          name='graph_convolution_feature_steered_convolution2'+str(dim))\n",
        "    vf_2 = tf.nn.relu(vf_2)\n",
        "  \n",
        "  res_inc = tf.keras.layers.Concatenate(axis=-1)([vf_0,vf_1,vf_2])\n",
        "  res_inc = graph_conv.feature_steered_convolution_layer(\n",
        "        res_inc,\n",
        "        batch_mesh_data['neighbors'],\n",
        "        batch_mesh_data['num_vertices'],\n",
        "        num_weight_matrices=num_filters,\n",
        "        num_output_channels=128,\n",
        "        name='graph_convolution_feature_steered_convolution'+str(128))\n",
        "  res_inc = tf.nn.relu(res_inc)\n",
        "  return res_inc\n",
        "\n",
        "def my_res2_inception1_5_conv(vertex_features,num_filters,batch_mesh_data):\n",
        "  dim0 = [32, 64]\n",
        "  dim1 = [64, 64]\n",
        "  dim2 = [32, 32]\n",
        "\n",
        "  vf_0 = vertex_features\n",
        "  vf_1 = vertex_features\n",
        "  vf_2 = vertex_features\n",
        "\n",
        "  for dim in dim0:\n",
        "    # with tf.variable_scope('conv_%d' % dim):\n",
        "    res_0 = tf.keras.layers.Conv1D(dim, 1, name='linres0'+str(dim))(vf_0)\n",
        "    vf_0 = graph_conv.feature_steered_convolution_layer(\n",
        "          vf_0,\n",
        "          batch_mesh_data['neighbors'],\n",
        "          batch_mesh_data['num_vertices'],\n",
        "          num_weight_matrices=num_filters,\n",
        "          num_output_channels=dim,\n",
        "          name='graph_convolution_feature_steered_convolution0'+str(dim))\n",
        "    vf_0 = tf.nn.relu(vf_0)\n",
        "    vf_0 = vf_0 + res_0\n",
        "      \n",
        "  for dim in dim1:\n",
        "    # with tf.variable_scope('conv_%d' % dim):\n",
        "    res_1 = tf.keras.layers.Conv1D(dim, 1, name='linres1'+str(dim))(vf_1)\n",
        "    vf_1 = graph_conv.feature_steered_convolution_layer(\n",
        "          vf_1,\n",
        "          batch_mesh_data['neighbors'],\n",
        "          batch_mesh_data['num_vertices'],\n",
        "          num_weight_matrices=num_filters,\n",
        "          num_output_channels=dim,\n",
        "          name='graph_convolution_feature_steered_convolution1'+str(dim))\n",
        "          # var_name = 'graph_convolution_feature_steered_convolutionweight1'+'dim')\n",
        "    vf_1 = tf.nn.relu(vf_1)\n",
        "    vf_1 = vf_1 + res_1\n",
        "  for dim in dim2:\n",
        "    # with tf.variable_scope('conv_%d' % dim):\n",
        "    res_2 = tf.keras.layers.Conv1D(dim, 1, name='linres2'+str(dim))(vf_2)\n",
        "    vf_2 = graph_conv.feature_steered_convolution_layer(\n",
        "          vf_2,\n",
        "          batch_mesh_data['neighbors'],\n",
        "          batch_mesh_data['num_vertices'],\n",
        "          num_weight_matrices=num_filters,\n",
        "          num_output_channels=dim,\n",
        "          name='graph_convolution_feature_steered_convolution2'+str(dim))\n",
        "    vf_2 = tf.nn.relu(vf_2)\n",
        "    vf_2 = vf_2 + res_2\n",
        "  \n",
        "  res_inc = tf.keras.layers.Concatenate(axis=-1)([vf_0,vf_1,vf_2])\n",
        "  res_out = tf.keras.layers.Conv1D(128, 1, name='linresout'+str(dim))(res_inc)\n",
        "  res_inc = graph_conv.feature_steered_convolution_layer(\n",
        "        res_inc,\n",
        "        batch_mesh_data['neighbors'],\n",
        "        batch_mesh_data['num_vertices'],\n",
        "        num_weight_matrices=num_filters,\n",
        "        num_output_channels=128,\n",
        "        name='graph_convolution_feature_steered_convolution'+str(128))\n",
        "  res_inc = tf.nn.relu(res_inc)\n",
        "  res_inc = res_inc + res_out\n",
        "  return res_inc\n",
        "\n",
        "def my_resgraph_conv(vertex_features,num_filters,batch_mesh_data):\n",
        "  arrdim = [32, 64, 128]\n",
        "  vf_res = vertex_features\n",
        "  \n",
        "  for dim in arrdim:\n",
        "    res_vertex_features = tf.keras.layers.Conv1D(dim, 1, name='linres'+str(dim))(vf_res)\n",
        "    # with tf.variable_scope('conv_%d' % dim):\n",
        "    vf_res = graph_conv.feature_steered_convolution_layer(\n",
        "          vf_res,\n",
        "          batch_mesh_data['neighbors'],\n",
        "          batch_mesh_data['num_vertices'],\n",
        "          num_weight_matrices=num_filters,\n",
        "          num_output_channels=dim,\n",
        "          name='graph_convolution_feature_steered_convolutionres'+str(dim))\n",
        "    vf_res = tf.nn.relu(vf_res)\n",
        "    vf_res = vf_res + res_vertex_features\n",
        "  return vf_res\n",
        "\n",
        "\n",
        "def mesh_encoder(batch_mesh_data, num_filters, output_dim, conv_layer_dims):\n",
        "  \"\"\"A mesh encoder using feature steered graph convolutions.\n",
        "\n",
        "    The shorthands used below are\n",
        "      `B`: Batch size.\n",
        "      `V`: The maximum number of vertices over all meshes in the batch.\n",
        "      `D`: The number of dimensions of input vertex features, D=3 if vertex\n",
        "        positions are used as features.\n",
        "\n",
        "  Args:\n",
        "    batch_mesh_data: A mesh_data dict with following keys\n",
        "      'vertices': A [B, V, D] `float32` tensor of vertex features, possibly\n",
        "        0-padded.\n",
        "      'neighbors': A [B, V, V] `float32` sparse tensor of edge weights.\n",
        "      'num_vertices': A [B] `int32` tensor of number of vertices per mesh.\n",
        "    num_filters: The number of weight matrices to be used in feature steered\n",
        "      graph conv.\n",
        "    output_dim: A dimension of output per vertex features.\n",
        "    conv_layer_dims: A list of dimensions used in graph convolution layers.\n",
        "\n",
        "  Returns:\n",
        "    vertex_features: A [B, V, output_dim] `float32` tensor of per vertex\n",
        "      features.\n",
        "  \"\"\"\n",
        "  batch_vertices = batch_mesh_data['vertices']\n",
        "\n",
        "  # Linear: N x D --> N x 16.\n",
        "  vertex_features = tf.keras.layers.Conv1D(16, 1, name='lin16')(batch_vertices)\n",
        "  # res_vertex_features = vertex_features\n",
        "\n",
        "  # graph convolution layers\n",
        "  # for dim in conv_layer_dims:\n",
        "  #   with tf.variable_scope('conv_%d' % dim):\n",
        "  #     vertex_features = graph_conv.feature_steered_convolution_layer(\n",
        "  #         vertex_features,\n",
        "  #         batch_mesh_data['neighbors'],\n",
        "  #         batch_mesh_data['num_vertices'],\n",
        "  #         num_weight_matrices=num_filters,\n",
        "  #         num_output_channels=dim)\n",
        "  #   vertex_features = tf.nn.relu(vertex_features)\n",
        "\n",
        "  # inception v1 v1.5\n",
        "  # vertex_features = my_inceptionv1(vertex_features,num_filters,batch_mesh_data)\n",
        "  # vertex_features = my_inceptionv1_5(vertex_features,num_filters,batch_mesh_data)\n",
        "\n",
        "  # resudal v1\n",
        "  # res_vertex_features = tf.keras.layers.Conv1D(128, 1, name='reslin16')(res_vertex_features)\n",
        "  # vertex_features = res_vertex_features + vertex_features\n",
        "\n",
        "  # resudal v2\n",
        "  # vertex_features = my_resgraph_conv(vertex_features,num_filters,batch_mesh_data)\n",
        "\n",
        "  # resudal v2 inception 1.5\n",
        "  vertex_features = my_res2_inception1_5_conv(vertex_features,num_filters,batch_mesh_data)\n",
        "  \n",
        "  # resudal v1.5 latter part\n",
        "  res_vertex_features = tf.keras.layers.Conv1D(output_dim, 1, name='linresout')(vertex_features)\n",
        "\n",
        "  # Linear: N x 128 --> N x 256.\n",
        "  # vertex_features = tf.keras.layers.Conv1D(\n",
        "  #     256, 1, name='lin256')(\n",
        "  #         vertex_features)\n",
        "  # vertex_features = tf.nn.relu(vertex_features)\n",
        "\n",
        "  # inception v0\n",
        "  vertex_features = my_inceptionv0(vertex_features)\n",
        "\n",
        "  # Linear: N x 256 --> N x output_dim.\n",
        "  vertex_features = tf.keras.layers.Conv1D(\n",
        "      output_dim, 1, name='lin_output')(\n",
        "          vertex_features)\n",
        "  vertex_features = vertex_features + res_vertex_features\n",
        "\n",
        "  return vertex_features"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6c2pz4r_79F_"
      },
      "source": [
        "Given a mesh encoder, let's define a model_fn for a custom\n",
        "[tf.Estimator](https://www.tensorflow.org/api_docs/python/tf/estimator/Estimator)\n",
        "for vertex classification using softmax cross entropy loss. A tf.Estimator model_fn returns the ops necessary to perform training, evaluation, or predictions given inputs and a number of other parameters. Recall that the\n",
        "vertex tensor may be zero-padded (see Dataset Pipeline above), hence we must mask out the contribution from the padded values."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WE-cuv0i78ak"
      },
      "source": [
        "def get_learning_rate(params):\n",
        "  \"\"\"Returns a decaying learning rate.\"\"\"\n",
        "  global_step = tf.train.get_or_create_global_step()\n",
        "  learning_rate = tf.train.exponential_decay(\n",
        "      params['init_learning_rate'],\n",
        "      global_step,\n",
        "      params['lr_decay_steps'],\n",
        "      params['lr_decay_rate'])\n",
        "  return learning_rate\n",
        "\n",
        "def model_fn(features, labels, mode, params):\n",
        "  \"\"\"Returns a mesh segmentation model_fn for use with tf.Estimator.\"\"\"\n",
        "  logits = mesh_encoder(features, params['num_filters'], params['num_classes'],\n",
        "                        params['encoder_filter_dims'])\n",
        "  predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)\n",
        "  outputs = {\n",
        "      'vertices': features['vertices'],\n",
        "      'triangles': features['triangles'],\n",
        "      'num_vertices': features['num_vertices'],\n",
        "      'num_triangles': features['num_triangles'],\n",
        "      'predictions': predictions,\n",
        "  }\n",
        "  # For predictions, return the outputs.\n",
        "  if mode == tf.estimator.ModeKeys.PREDICT:\n",
        "    outputs['labels'] = features['labels']\n",
        "    return tf.estimator.EstimatorSpec(mode=mode, predictions=outputs)\n",
        "  # Loss\n",
        "  # Weight the losses by masking out padded vertices/labels.\n",
        "  vertex_ragged_sizes = features['num_vertices']\n",
        "  mask = tf.sequence_mask(vertex_ragged_sizes, tf.shape(labels)[-1])\n",
        "  loss_weights = tf.cast(mask, dtype=tf.float32)\n",
        "  loss = tf.losses.sparse_softmax_cross_entropy(\n",
        "      logits=logits, labels=labels, weights=loss_weights)\n",
        "  # For training, build the optimizer.\n",
        "  if mode == tf.estimator.ModeKeys.TRAIN:\n",
        "    optimizer = tf.train.AdamOptimizer(\n",
        "        learning_rate=get_learning_rate(params),\n",
        "        beta1=params['beta'],\n",
        "        epsilon=params['adam_epsilon'])\n",
        "    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)\n",
        "    with tf.control_dependencies(update_ops):\n",
        "      train_op = optimizer.minimize(\n",
        "          loss=loss, global_step=tf.train.get_global_step())\n",
        "    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)\n",
        "\n",
        "  # For eval, return eval metrics.\n",
        "  eval_ops = {\n",
        "      'mean_loss':\n",
        "          tf.metrics.mean(loss),\n",
        "      'accuracy':\n",
        "          tf.metrics.accuracy(\n",
        "              labels=labels, predictions=predictions, weights=loss_weights)\n",
        "  }\n",
        "  return tf.estimator.EstimatorSpec(\n",
        "      mode=mode, loss=loss, eval_metric_ops=eval_ops)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QE8Meyi5GKWA"
      },
      "source": [
        "## Train the model from scratch\n",
        "\n",
        "Now let's train the mesh segmentation model from scratch. First we will download the train dataset files, and use tf.Estimator.train_and_evaluate to train a model.\n",
        "\n",
        "Note: Training code is provided inside colab for demonstration, and may be slow. For optimal performance, consider running the training process as a command line process, and a tensorboard process to track."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S6vxJ8t5HRcS"
      },
      "source": [
        "path_to_train_data_zip = tf.keras.utils.get_file(\n",
        "    mydrivedatafolder+'train_data.zip',\n",
        "    origin='https://storage.googleapis.com/tensorflow-graphics/notebooks/mesh_segmentation/train_data.zip',\n",
        "    extract=True)\n",
        "\n",
        "train_data_files = glob.glob(\n",
        "    os.path.join(os.path.dirname(path_to_train_data_zip), '*train*.tfrecords'))\n",
        "\n",
        "retrain_model_dir = os.path.join(local_model_dir, 'retrain')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "inNSvIN-HcBg"
      },
      "source": [
        "train_io_params = {\n",
        "    'batch_size': 8,\n",
        "    'parallel_threads': 8,\n",
        "    'is_training': True,\n",
        "    'shuffle': True,\n",
        "    'sloppy': True,\n",
        "}\n",
        "\n",
        "eval_io_params = {\n",
        "    'batch_size': 8,\n",
        "    'parallel_threads': 8,\n",
        "    'is_training': False,\n",
        "    'shuffle': False\n",
        "}\n",
        "\n",
        "\n",
        "def train_fn():\n",
        "  return dataio.create_input_from_dataset(dataio.create_dataset_from_tfrecords,\n",
        "                                          train_data_files, train_io_params)\n",
        "\n",
        "\n",
        "def eval_fn():\n",
        "  return dataio.create_input_from_dataset(dataio.create_dataset_from_tfrecords,\n",
        "                                          test_data_files, eval_io_params)\n",
        "\n",
        "\n",
        "train_params = {\n",
        "    'beta': 0.9,\n",
        "    'adam_epsilon': 1e-8,\n",
        "    'init_learning_rate': 0.001,\n",
        "    'lr_decay_steps': 1000,\n",
        "    'lr_decay_rate': 0.95,\n",
        "}\n",
        "\n",
        "train_params.update(MODEL_PARAMS)\n",
        "\n",
        "# about 15 steps one save and evaluation.\n",
        "# checkpoint_delay = 20  # Checkpoint every 2 minutes.\n",
        "eval_spec_delay_secs = 20\n",
        "eval_steps = 15\n",
        "save_checkpoints_steps = 100\n",
        "# save_summary_steps = 20\n",
        "max_steps = 2000  # Number of training steps.\n",
        "log_step_count_steps = 2\n",
        "\n",
        "config = tf.estimator.RunConfig(\n",
        "    log_step_count_steps=log_step_count_steps,\n",
        "    save_checkpoints_secs=eval_spec_delay_secs,\n",
        "    # save_checkpoints_steps=save_checkpoints_steps,\n",
        "    # save_summary_steps = save_summary_steps, # I cannot find it where.\n",
        "    keep_checkpoint_max=5)\n",
        "\n",
        "classifier = tf.estimator.Estimator(\n",
        "    model_fn=model_fn,\n",
        "    model_dir=retrain_model_dir,\n",
        "    config=config,\n",
        "    params=train_params)\n",
        "train_spec = tf.estimator.TrainSpec(input_fn=train_fn, max_steps=max_steps)\n",
        "eval_spec = tf.estimator.EvalSpec(\n",
        "    input_fn=eval_fn,\n",
        "    steps=eval_steps,\n",
        "    start_delay_secs=eval_spec_delay_secs,# evaluation time delays\n",
        "    throttle_secs=eval_spec_delay_secs/2 # do not evaluate after within this parameter seconds after the latest evaluation.\n",
        "    )\n",
        "\n",
        "print('Start training & eval.')\n",
        "tf.estimator.train_and_evaluate(classifier, train_spec, eval_spec)\n",
        "print('Train and eval done.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H6_aJMQAvu0g"
      },
      "source": [
        "%cd /root/.keras/datasets/model/retrain\n",
        "%ls -a\n",
        "%cp /root/.keras/datasets/model/retrain/model.ckpt-2000.meta /content/drive/MyDrive/trained_model\n",
        "%cp /root/.keras/datasets/model/retrain/graph.pbtxt /content/drive/MyDrive/trained_model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3pizREREwPLu"
      },
      "source": [
        "%rm -rf /root/.keras/datasets/model\n",
        "%cd /root/.keras/datasets\n",
        "%ls -a"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}