<jupyter_start><jupyter_text>Copyright 2019 Google LLC.<jupyter_code>#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.<jupyter_output><empty_output><jupyter_text>Mesh Segmentation using Feature Steered Graph ConvolutionsSegmenting a mesh to its semantic parts is an important problem for 3D shapeunderstanding. This colab demonstrates how to build a semantic mesh segmentationmodel for deformable shapes using graph convolution layers defined in[Tensorflow Graphics](https://github.com/tensorflow/graphics).![](https://storage.googleapis.com/tensorflow-graphics/notebooks/mesh_segmentation/mesh_segmentation_demo.png)This notebook covers the following key topics:* How to use graph-convolutional layers to define a CNN for mesh segmentation.* How to setup a data pipeline to represent mesh connectivity with SparseTensors.Note: The easiest way to use this tutorial is as a Colab notebook, which allowsyou to dive in with no setup. Image Convolutions vs Graph ConvolutionsImages are represented by uniform grids of pixels. Running convolutions onuniform grids is a well understood process and is at the core of a significantamount of products and academic publications.![](https://storage.googleapis.com/tensorflow-graphics/notebooks/mesh_segmentation/cat_image_convolutions.png)However, things become a bit more complicated when dealing with threedimensional objects such as meshes or point clouds since these are not definedon regular grids. A convolution operation for meshes or point clouds mustoperate on irregular data structures. This makes convolutional neuralnetworks based on them harder to implement.![](https://storage.googleapis.com/tensorflow-graphics/notebooks/mesh_segmentation/cat_mesh_convolutions.png)Any general mesh can be denoted as a graph that is not constrained to a regular grid. Many graph-convolutional operators have been published inthe recent years. In this demo we use the method described in[Feature Steered Graph Convolutions](https://arxiv.org/abs/1706.05206). Similarto it's image counterpart, this basic building block can be used do solve aplethora of problems. This Colab focusses on segmenting deformable meshes ofhuman bodies into parts (e.g. head, right foot, etc.).  Setup & ImportsTo run this Colab optimally, please update the runtime type to use a GPUhardware accelerator. - click on the 'Runtime' menu, then 'Change runtime type':![](https://storage.googleapis.com/tensorflow-graphics/notebooks/non_rigid_deformation/change_runtime.jpg)-   finally, set the 'Hardware accelerator' to 'GPU'.![](https://storage.googleapis.com/tensorflow-graphics/notebooks/mesh_segmentation/gpu_runtime.png)If Tensorflow Graphics is not installed on your system, the following cell willinstall the Tensorflow Graphics package for you.<jupyter_code>!pip install tensorflow_graphics<jupyter_output><empty_output><jupyter_text>Now that Tensorflow Graphics and dependencies are installed, let's import everything needed to run the demos contained in this notebook.<jupyter_code>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import glob
import os
import tensorflow as tf

from tensorflow_graphics.nn.layer import graph_convolution as graph_conv
from tensorflow_graphics.notebooks import mesh_segmentation_dataio as dataio
from tensorflow_graphics.notebooks import mesh_viewer<jupyter_output><empty_output><jupyter_text>Note this notebook works best in Graph mode.  Fetch model files and dataFor 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).<jupyter_code>path_to_model_zip = tf.keras.utils.get_file(
    'model.zip',
    origin='https://storage.googleapis.com/tensorflow-graphics/notebooks/mesh_segmentation/model.zip',
    extract=True)

path_to_data_zip = tf.keras.utils.get_file(
    'data.zip',
    origin='https://storage.googleapis.com/tensorflow-graphics/notebooks/mesh_segmentation/data.zip',
    extract=True)

local_model_dir = os.path.join(os.path.dirname(path_to_model_zip), 'model')
test_data_files = [
    os.path.join(
        os.path.dirname(path_to_data_zip),
        'data/Dancer_test_sequence.tfrecords')
]<jupyter_output><empty_output><jupyter_text>Load and visualize test dataFor graph convolutions, we need a *weighted adjacency matrix* denoting the meshconnectivity. Feature-steered graph convolutions expect self-edges in the meshconnectivity for each vertex, i.e. the diagonal of the weighted adjacency matrixshould be non-zero. This matrix is defined as:```A[i, j] = w[i,j] if vertex i and vertex j share an edge,A[i, i] = w[i,i] for each vertex i,A[i, j] = 0 otherwise.where, w[i, j] = 1/(degree(vertex i)), and sum(j)(w[i,j]) = 1```Here degree(vertex i) is the number of edges incident on a vertex (including theself-edge). This weighted adjacency matrix is stored as a SparseTensor.We will load the test meshes from the test [tf.data.TFRecordDataset](https://www.tensorflow.org/api_docs/python/tf/data/TFRecordDataset)downloaded above. Each mesh is stored as a[tf.Example](https://www.tensorflow.org/api_docs/python/tf/train/Example), withthe following fields:*   'num_vertices': Number of vertices in each mesh*   'num_triangles': Number of triangles in each mesh.*   'vertices': A [V, 3] float tensor of vertex positions.*   'triangles': A [T, 3] integer tensor of vertex indices for each triangle.*   'labels': A [V] integer tensor with segmentation class label for each    vertex.where 'V' is number of vertices and 'T' is number of triangles in the mesh. Aseach mesh may have a varying number of vertices and faces (and the correspondingconnectivity matrix), we pad the data tensors with '0's in each batch.For details on the dataset pipeline implementation, take a look atmesh_segmentation_dataio.py.Lets try to load a batch from the test TFRecordDataset, and visualize the firstmesh with each vertex colored by the part label.<jupyter_code>test_io_params = {
    'is_training': False,
    'sloppy': False,
    'shuffle': True,
}
test_tfrecords = test_data_files


input_graph = tf.Graph()
with input_graph.as_default():
  mesh_load_op = dataio.create_input_from_dataset(
      dataio.create_dataset_from_tfrecords, test_tfrecords, test_io_params)
  with tf.Session() as sess:
    test_mesh_data, test_labels = sess.run(mesh_load_op)

input_mesh_data = {
    'vertices': test_mesh_data['vertices'][0, ...],
    'faces': test_mesh_data['triangles'][0, ...],
    'vertex_colors': mesh_viewer.SEGMENTATION_COLORMAP[test_labels[0, ...]],
}
input_viewer = mesh_viewer.Viewer(input_mesh_data)<jupyter_output><empty_output><jupyter_text>Model DefinitionGiven a mesh with V vertices and D-dimensional per-vertex input features (e.g.vertex position, normal), we would like to create a network capable ofclassifying each vertex to a part label. Lets first create a mesh encoder thatencodes each vertex in the mesh into C-dimensional logits, where C is the numberof parts. First we use 1x1 convolutions to change input feature dimensions,followed by a sequence of feature steered graph convolutions and ReLUnon-linearities, and finally 1x1 convolutions to logits, which are used forcomputing softmax cross entropy as described below.Note that this model does not use any form of pooling, which is outside the scope of this notebook.![](https://storage.googleapis.com/tensorflow-graphics/notebooks/mesh_segmentation/mesh_segmentation_model_def.png)<jupyter_code>MODEL_PARAMS = {
    'num_filters': 8,
    'num_classes': 16,
    'encoder_filter_dims': [32, 64, 128],
}


def mesh_encoder(batch_mesh_data, num_filters, output_dim, conv_layer_dims):
  """A mesh encoder using feature steered graph convolutions.

    The shorthands used below are
      `B`: Batch size.
      `V`: The maximum number of vertices over all meshes in the batch.
      `D`: The number of dimensions of input vertex features, D=3 if vertex
        positions are used as features.

  Args:
    batch_mesh_data: A mesh_data dict with following keys
      'vertices': A [B, V, D] `float32` tensor of vertex features, possibly
        0-padded.
      'neighbors': A [B, V, V] `float32` sparse tensor of edge weights.
      'num_vertices': A [B] `int32` tensor of number of vertices per mesh.
    num_filters: The number of weight matrices to be used in feature steered
      graph conv.
    output_dim: A dimension of output per vertex features.
    conv_layer_dims: A list of dimensions used in graph convolution layers.

  Returns:
    vertex_features: A [B, V, output_dim] `float32` tensor of per vertex
      features.
  """
  batch_vertices = batch_mesh_data['vertices']

  # Linear: N x D --> N x 16.
  vertex_features = tf.keras.layers.Conv1D(16, 1, name='lin16')(batch_vertices)

  # graph convolution layers
  for dim in conv_layer_dims:
    with tf.variable_scope('conv_%d' % dim):
      vertex_features = graph_conv.feature_steered_convolution_layer(
          vertex_features,
          batch_mesh_data['neighbors'],
          batch_mesh_data['num_vertices'],
          num_weight_matrices=num_filters,
          num_output_channels=dim)
    vertex_features = tf.nn.relu(vertex_features)

  # Linear: N x 128 --> N x 256.
  vertex_features = tf.keras.layers.Conv1D(
      256, 1, name='lin256')(
          vertex_features)
  vertex_features = tf.nn.relu(vertex_features)

  # Linear: N x 256 --> N x output_dim.
  vertex_features = tf.keras.layers.Conv1D(
      output_dim, 1, name='lin_output')(
          vertex_features)

  return vertex_features<jupyter_output><empty_output><jupyter_text>Given a mesh encoder, let's define a model_fn for a custom[tf.Estimator](https://www.tensorflow.org/api_docs/python/tf/estimator/Estimator)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 thevertex tensor may be zero-padded (see Dataset Pipeline above), hence we must mask out the contribution from the padded values.<jupyter_code>def get_learning_rate(params):
  """Returns a decaying learning rate."""
  global_step = tf.train.get_or_create_global_step()
  learning_rate = tf.train.exponential_decay(
      params['init_learning_rate'],
      global_step,
      params['lr_decay_steps'],
      params['lr_decay_rate'])
  return learning_rate

def model_fn(features, labels, mode, params):
  """Returns a mesh segmentation model_fn for use with tf.Estimator."""
  logits = mesh_encoder(features, params['num_filters'], params['num_classes'],
                        params['encoder_filter_dims'])
  predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
  outputs = {
      'vertices': features['vertices'],
      'triangles': features['triangles'],
      'num_vertices': features['num_vertices'],
      'num_triangles': features['num_triangles'],
      'predictions': predictions,
  }
  # For predictions, return the outputs.
  if mode == tf.estimator.ModeKeys.PREDICT:
    outputs['labels'] = features['labels']
    return tf.estimator.EstimatorSpec(mode=mode, predictions=outputs)
  # Loss
  # Weight the losses by masking out padded vertices/labels.
  vertex_ragged_sizes = features['num_vertices']
  mask = tf.sequence_mask(vertex_ragged_sizes, tf.shape(labels)[-1])
  loss_weights = tf.cast(mask, dtype=tf.float32)
  loss = tf.losses.sparse_softmax_cross_entropy(
      logits=logits, labels=labels, weights=loss_weights)
  # For training, build the optimizer.
  if mode == tf.estimator.ModeKeys.TRAIN:
    optimizer = tf.train.AdamOptimizer(
        learning_rate=get_learning_rate(params),
        beta1=params['beta'],
        epsilon=params['adam_epsilon'])
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
      train_op = optimizer.minimize(
          loss=loss, global_step=tf.train.get_global_step())
    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

  # For eval, return eval metrics.
  eval_ops = {
      'mean_loss':
          tf.metrics.mean(loss),
      'accuracy':
          tf.metrics.accuracy(
              labels=labels, predictions=predictions, weights=loss_weights)
  }
  return tf.estimator.EstimatorSpec(
      mode=mode, loss=loss, eval_metric_ops=eval_ops)<jupyter_output><empty_output><jupyter_text>Test model & visualize resultsNow that we have defined the model, let's load the weights from the trained model downloaded above and use tf.Estimator.predict to predict the part labels for meshes in the test dataset.<jupyter_code>test_io_params = {
    'is_training': False,
    'sloppy': False,
    'shuffle': True,
    'repeat': False
}
test_tfrecords = test_data_files

def predict_fn():
  return dataio.create_input_from_dataset(dataio.create_dataset_from_tfrecords,
                                          test_tfrecords,
                                          test_io_params)


estimator = tf.estimator.Estimator(model_fn=model_fn,
                                   model_dir=local_model_dir,
                                   params=MODEL_PARAMS)
test_predictions = estimator.predict(input_fn=predict_fn)<jupyter_output><empty_output><jupyter_text>Run the following cell repeatedly to cycle through the meshes in the test sequence. The left view shows the input mesh, and the right view shows the predicted part labels.<jupyter_code>prediction = next(test_predictions)
input_mesh_data = {
    'vertices': prediction['vertices'],
    'faces': prediction['triangles'],
}
predicted_mesh_data = {
    'vertices': prediction['vertices'],
    'faces': prediction['triangles'],
    'vertex_colors': mesh_viewer.SEGMENTATION_COLORMAP[prediction['predictions']],
}

input_viewer = mesh_viewer.Viewer(input_mesh_data)
prediction_viewer = mesh_viewer.Viewer(predicted_mesh_data)<jupyter_output><empty_output><jupyter_text>Train the model from scratch.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.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.<jupyter_code>path_to_train_data_zip = tf.keras.utils.get_file(
    'train_data.zip',
    origin='https://storage.googleapis.com/tensorflow-graphics/notebooks/mesh_segmentation/train_data.zip',
    extract=True)

train_data_files = glob.glob(
    os.path.join(os.path.dirname(path_to_train_data_zip), '*train*.tfrecords'))

retrain_model_dir = os.path.join(local_model_dir, 'retrain')
train_io_params = {
    'batch_size': 8,
    'parallel_threads': 8,
    'is_training': True,
    'shuffle': True,
    'sloppy': True,
}

eval_io_params = {
    'batch_size': 8,
    'parallel_threads': 8,
    'is_training': False,
    'shuffle': False
}


def train_fn():
  return dataio.create_input_from_dataset(dataio.create_dataset_from_tfrecords,
                                          train_data_files, train_io_params)


def eval_fn():
  return dataio.create_input_from_dataset(dataio.create_dataset_from_tfrecords,
                                          test_data_files, eval_io_params)


train_params = {
    'beta': 0.9,
    'adam_epsilon': 1e-8,
    'init_learning_rate': 0.001,
    'lr_decay_steps': 10000,
    'lr_decay_rate': 0.95,
}

train_params.update(MODEL_PARAMS)

checkpoint_delay = 120  # Checkpoint every 2 minutes.
max_steps = 100000  # Number of training steps.

config = tf.estimator.RunConfig(
    log_step_count_steps=1,
    save_checkpoints_secs=checkpoint_delay,
    keep_checkpoint_max=3)

classifier = tf.estimator.Estimator(
    model_fn=model_fn,
    model_dir=retrain_model_dir,
    config=config,
    params=train_params)
train_spec = tf.estimator.TrainSpec(input_fn=train_fn, max_steps=max_steps)
eval_spec = tf.estimator.EvalSpec(
    input_fn=eval_fn,
    steps=None,
    start_delay_secs=2 * checkpoint_delay,
    throttle_secs=checkpoint_delay)

print('Start training & eval.')
tf.estimator.train_and_evaluate(classifier, train_spec, eval_spec)
print('Train and eval done.')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08771201518206619, 'fraction_numerical': 0.007413118254062389, 'mean_word_length': 4.574545454545454, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 44, 'https://': 16, 'lorem ipsum': 0, 'www.': 4, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '133301', 'license': 'Apache-2.0', 'n_tokens_mistral': 5148, 'n_tokens_neox': 4927, 'n_words': 1497, 'path': 'tensorflow_graphics/notebooks/mesh_segmentation_demo.ipynb', 'repo_name': 'kristenrq/graphics', 'size': '16862'}