text: # Copyright 2020 Deepmind Technologies Limited.
  #
  # 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
  #
  #     http://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.
  
  """Mesh data utilities."""
  import matplotlib.pyplot as plt
  import modules
  from mpl_toolkits import mplot3d  # pylint: disable=unused-import
  from mpl_toolkits.mplot3d.art3d import Poly3DCollection
  import networkx as nx
  import numpy as np
  import six
  from six.moves import range
  import tensorflow.compat.v1 as tf
  import tensorflow_probability as tfp
  
  tfd = tfp.distributions
  
  class_dict = {'02691156': 0, '02747177': 1, '02773838': 2, '02801938': 3, '02808440': 4, '02818832': 5, '02828884': 6,
                '02843684': 7, '02871439': 8, '02876657': 9, '02880940': 10, '02924116': 11, '02933112': 12,
                '02942699': 13, '02946921': 14, '02954340': 15, '02958343': 16, '02992529': 17, '03001627': 18,
                '03046257': 19, '03085013': 20, '03207941': 21, '03211117': 22, '03261776': 23, '03325088': 24,
                '03337140': 25, '03467517': 26, '03513137': 27, '03593526': 28, '03624134': 29, '03636649': 30,
                '03642806': 31, '03691459': 32, '03710193': 33, '03759954': 34, '03761084': 35, '03790512': 36,
                '03797390': 37, '03928116': 38, '03938244': 39, '03948459': 40, '03991062': 41,
                '04004475': 42, '04074963': 43, '04090263': 44, '04099429': 45, '04225987': 46,
                '04256520': 47, '04330267': 48, '04379243': 49, '04401088': 50, '04460130': 51, '04468005': 52,
                '04530566': 53, '04554684': 54}
  
  
  def random_shift(vertices, shift_factor=0.25):
      """Apply random shift to vertices."""
      max_shift_pos = tf.cast(255 - tf.reduce_max(vertices, axis=0), tf.float32)
      max_shift_pos = tf.maximum(max_shift_pos, 1e-9)
  
      max_shift_neg = tf.cast(tf.reduce_min(vertices, axis=0), tf.float32)
      max_shift_neg = tf.maximum(max_shift_neg, 1e-9)
  
      shift = tfd.TruncatedNormal(
          tf.zeros([1, 3]), shift_factor * 255, -max_shift_neg,
          max_shift_pos).sample()
      shift = tf.cast(shift, tf.int32)
      vertices += shift
      return vertices
  
  
  def make_vertex_model_dataset(ds, apply_random_shift=False):
      """Prepare dataset for vertex model training."""
  
      def _vertex_model_map_fn(example):
          vertices = example['vertices']
  
          # Randomly shift vertices
          if apply_random_shift:
              vertices = random_shift(vertices)
  
          # Re-order vertex coordinates as (z, y, x).
          vertices_permuted = tf.stack(
              [vertices[:, 2], vertices[:, 1], vertices[:, 0]], axis=-1)
  
          # Flatten quantized vertices, reindex starting from 1, and pad with a
          # zero stopping token.
          vertices_flat = tf.reshape(vertices_permuted, [-1])
          example['vertices_flat'] = tf.pad(vertices_flat + 1, [[0, 1]])
  
          # Create mask to indicate valid tokens after padding and batching.
          example['vertices_flat_mask'] = tf.ones_like(
              example['vertices_flat'], dtype=tf.float32)
          return example
  
      return ds.map(_vertex_model_map_fn)
  
  
  def make_face_model_dataset(
          ds, apply_random_shift=False, shuffle_vertices=True, quantization_bits=8):
      """Prepare dataset for face model training."""
  
      def _face_model_map_fn(example):
          vertices = example['vertices']
  
          # Randomly shift vertices
          if apply_random_shift:
              vertices = random_shift(vertices)
          example['num_vertices'] = tf.shape(vertices)[0]
  
          # Optionally shuffle vertices and re-order faces to match
          if shuffle_vertices:
              permutation = tf.random_shuffle(tf.range(example['num_vertices']))
              vertices = tf.gather(vertices, permutation)
              face_permutation = tf.concat(
                  [tf.constant([0, 1], dtype=tf.int32), tf.argsort(permutation) + 2],
                  axis=0)
              example['faces'] = tf.cast(
                  tf.gather(face_permutation, example['faces']), tf.int64)
  
          # Vertices are quantized. So convert to floats for input to face model
          example['vertices'] = modules.dequantize_verts(vertices, quantization_bits)
          example['vertices_mask'] = tf.ones_like(
              example['vertices'][..., 0], dtype=tf.float32)
          example['faces_mask'] = tf.ones_like(example['faces'], dtype=tf.float32)
          return example
  
      return ds.map(_face_model_map_fn)
  
  
  def read_obj(obj_path):
      """Read vertices and faces from .obj file."""
      vertex_list = []
      flat_vertices_list = []
      flat_vertices_indices = {}
      flat_triangles = []
  
      with open(obj_path) as obj_file:
          for line in obj_file:
              tokens = line.split()
              if not tokens:
                  continue
              line_type = tokens[0]
              # We skip lines not starting with v or f.
              if line_type == 'v':
                  vertex_list.append([float(x) for x in tokens[1:]])
              elif line_type == 'f':
                  triangle = []
                  for i in range(len(tokens) - 1):
                      vertex_name = tokens[i + 1].split('/')[0]
                      if vertex_name in flat_vertices_indices:
                          triangle.append(flat_vertices_indices[vertex_name])
                          continue
                      flat_vertex = []
                      for index in six.ensure_str(vertex_name).split('/'):
                          if not index:
                              continue
                          # obj triangle indices are 1 indexed, so subtract 1 here.
                          flat_vertex += vertex_list[int(index) - 1]
                      flat_vertex_index = len(flat_vertices_list)
                      flat_vertices_list.append(flat_vertex)
                      flat_vertices_indices[vertex_name] = flat_vertex_index
                      triangle.append(flat_vertex_index)
                  flat_triangles.append(triangle)
  
      return np.array(flat_vertices_list, dtype=np.float32), flat_triangles
  
  
  def write_obj(vertices, faces, file_path, transpose=True, scale=1.):
      """Write vertices and faces to obj."""
      if transpose:
          vertices = vertices[:, [1, 2, 0]]
      vertices *= scale
      if faces is not None:
          if min(min(faces)) == 0:
              f_add = 1
          else:
              f_add = 0
      with open(file_path, 'w') as f:
          for v in vertices:
              f.write('v {} {} {}\n'.format(v[0], v[1], v[2]))
          for face in faces:
              line = 'f'
              for i in face:
                  line += ' {}'.format(i + f_add)
              line += '\n'
              f.write(line)
  
  
  def quantize_verts(verts, n_bits=8):
      """Convert vertices in [-1., 1.] to discrete values in [0, n_bits**2 - 1]."""
      min_range = -0.5
      max_range = 0.5
      range_quantize = 2 ** n_bits - 1
      verts_quantize = (verts - min_range) * range_quantize / (
              max_range - min_range)
      return verts_quantize.astype('int32')
  
  
  def dequantize_verts(verts, n_bits=8, add_noise=False):
      """Convert quantized vertices to floats."""
      min_range = -0.5
      max_range = 0.5
      range_quantize = 2 ** n_bits - 1
      verts = verts.astype('float32')
      verts = verts * (max_range - min_range) / range_quantize + min_range
      if add_noise:
          verts += np.random.uniform(size=verts.shape) * (1 / range_quantize)
      return verts
  
  
  def face_to_cycles(face):
      """Find cycles in face."""
      g = nx.Graph()
      for v in range(len(face) - 1):
          g.add_edge(face[v], face[v + 1])
      g.add_edge(face[-1], face[0])
      return list(nx.cycle_basis(g))
  
  
  def flatten_faces(faces):
      """Converts from list of faces to flat face array with stopping indices."""
      if not faces:
          return np.array([0])
      else:
          l = [f + [-1] for f in faces[:-1]]
          l += [faces[-1] + [-2]]
          return np.array([item for sublist in l for item in sublist]) + 2  # pylint: disable=g-complex-comprehension
  
  
  def unflatten_faces(flat_faces):
      """Converts from flat face sequence to a list of separate faces."""
  
      def group(seq):
          g = []
          for el in seq:
              if el == 0 or el == -1:
                  yield g
                  g = []
              else:
                  g.append(el - 1)
          yield g
  
      outputs = list(group(flat_faces - 1))[:-1]
      # Remove empty faces
      return [o for o in outputs if len(o) > 2]
  
  
  def center_vertices(vertices):
      """Translate the vertices so that bounding box is centered at zero."""
      vert_min = vertices.min(axis=0)
      vert_max = vertices.max(axis=0)
      vert_center = 0.5 * (vert_min + vert_max)
      return vertices - vert_center
  
  
  def normalize_vertices_scale(vertices):
      """Scale the vertices so that the long diagonal of the bounding box is one."""
      vert_min = vertices.min(axis=0)
      vert_max = vertices.max(axis=0)
      extents = vert_max - vert_min
      scale = np.sqrt(np.sum(extents ** 2))
      return vertices / scale
  
  
  def quantize_process_mesh(vertices, faces, tris=None, quantization_bits=8):
      """Quantize vertices, remove resulting duplicates and reindex faces."""
      vertices = quantize_verts(vertices, quantization_bits)
      vertices, inv = np.unique(vertices, axis=0, return_inverse=True)
  
      # Sort vertices by z then y then x.
      sort_inds = np.lexsort(vertices.T)
      vertices = vertices[sort_inds]
  
      # Re-index faces and tris to re-ordered vertices.
      faces = [np.argsort(sort_inds)[inv[f]] for f in faces]
      if tris is not None:
          tris = np.array([np.argsort(sort_inds)[inv[t]] for t in tris])
  
      # Merging duplicate vertices and re-indexing the faces causes some faces to
      # contain loops (e.g [2, 3, 5, 2, 4]). Split these faces into distinct
      # sub-faces.
      sub_faces = []
      for f in faces:
          cliques = face_to_cycles(f)
          for c in cliques:
              c_length = len(c)
              # Only append faces with more than two verts.
              if c_length > 2:
                  d = np.argmin(c)
                  # Cyclically permute faces just that first index is the smallest.
                  sub_faces.append([c[(d + i) % c_length] for i in range(c_length)])
      faces = sub_faces
      if tris is not None:
          tris = np.array([v for v in tris if len(set(v)) == len(v)])
  
      # Sort faces by lowest vertex indices. If two faces have the same lowest
      # index then sort by next lowest and so on.
      faces.sort(key=lambda f: tuple(sorted(f)))
      if tris is not None:
          tris = tris.tolist()
          tris.sort(key=lambda f: tuple(sorted(f)))
          tris = np.array(tris)
  
      # After removing degenerate faces some vertices are now unreferenced.
      # Remove these.
      num_verts = vertices.shape[0]
      vert_connected = np.equal(
          np.arange(num_verts)[:, None], np.hstack(faces)[None]).any(axis=-1)
      vertices = vertices[vert_connected]
  
      # Re-index faces and tris to re-ordered vertices.
      vert_indices = (
              np.arange(num_verts) - np.cumsum(1 - vert_connected.astype('int')))
      faces = [vert_indices[f].tolist() for f in faces]
      if tris is not None:
          tris = np.array([vert_indices[t].tolist() for t in tris])
  
      return vertices, faces, tris
  
  
  def load_process_mesh(mesh_obj_path, quantization_bits=8):
      """Load obj file and process."""
      # Load mesh
      vertices, faces = read_obj(mesh_obj_path)
      # Transpose so that z-axis is vertical.
      vertices = vertices[:, [2, 0, 1]]
  
      # Translate the vertices so that bounding box is centered at zero.
      vertices = center_vertices(vertices)
  
      # Scale the vertices so that the long diagonal of the bounding box is equal
      # to one.
      vertices = normalize_vertices_scale(vertices)
  
      # Quantize and sort vertices, remove resulting duplicates, sort and reindex
      # faces.
      vertices, faces, _ = quantize_process_mesh(
          vertices, faces, quantization_bits=quantization_bits)
  
      # Flatten faces and add 'new face' = 1 and 'stop' = 0 tokens.
      faces = flatten_faces(faces)
      # Discard degenerate meshes without faces.
      if len(vertices) >= 800 or len(faces) >= 2800:
          return {
                     'vertices': vertices,
                     'faces': faces,
                 }, False
      else:
          return {
                     'vertices': vertices,
                     'faces': faces,
                 }, True
  
  
  def plot_meshes(mesh_list,
                  ax_lims=0.3,
                  fig_size=4,
                  el=30,
                  rot_start=120,
                  vert_size=10,
                  vert_alpha=0.75,
                  n_cols=4):
      """Plots mesh data using matplotlib."""
  
      n_plot = len(mesh_list)
      n_cols = np.minimum(n_plot, n_cols)
      n_rows = np.ceil(n_plot / n_cols).astype('int')
      fig = plt.figure(figsize=(fig_size * n_cols, fig_size * n_rows))
      for p_inc, mesh in enumerate(mesh_list):
  
          for key in [
              'vertices', 'faces', 'vertices_conditional', 'pointcloud', 'class_name'
          ]:
              if key not in list(mesh.keys()):
                  mesh[key] = None
  
          ax = fig.add_subplot(n_rows, n_cols, p_inc + 1, projection='3d')
  
          if mesh['faces'] is not None:
              if mesh['vertices_conditional'] is not None:
                  face_verts = np.concatenate(
                      [mesh['vertices_conditional'], mesh['vertices']], axis=0)
              else:
                  face_verts = mesh['vertices']
              collection = []
              for f in mesh['faces']:
                  collection.append(face_verts[f])
              plt_mesh = Poly3DCollection(collection)
              plt_mesh.set_edgecolor((0., 0., 0., 0.3))
              plt_mesh.set_facecolor((1, 0, 0, 0.2))
              ax.add_collection3d(plt_mesh)
  
          if mesh['vertices'] is not None:
              ax.scatter3D(
                  mesh['vertices'][:, 0],
                  mesh['vertices'][:, 1],
                  mesh['vertices'][:, 2],
                  lw=0.,
                  s=vert_size,
                  c='g',
                  alpha=vert_alpha)
  
          if mesh['vertices_conditional'] is not None:
              ax.scatter3D(
                  mesh['vertices_conditional'][:, 0],
                  mesh['vertices_conditional'][:, 1],
                  mesh['vertices_conditional'][:, 2],
                  lw=0.,
                  s=vert_size,
                  c='b',
                  alpha=vert_alpha)
  
          if mesh['pointcloud'] is not None:
              ax.scatter3D(
                  mesh['pointcloud'][:, 0],
                  mesh['pointcloud'][:, 1],
                  mesh['pointcloud'][:, 2],
                  lw=0.,
                  s=2.5 * vert_size,
                  c='b',
                  alpha=1.)
  
          ax.set_xlim(-ax_lims, ax_lims)
          ax.set_ylim(-ax_lims, ax_lims)
          ax.set_zlim(-ax_lims, ax_lims)
  
          ax.view_init(el, rot_start)
  
          display_string = ''
          if mesh['faces'] is not None:
              display_string += 'Num. faces: {}\n'.format(len(collection))
          if mesh['vertices'] is not None:
              num_verts = mesh['vertices'].shape[0]
              if mesh['vertices_conditional'] is not None:
                  num_verts += mesh['vertices_conditional'].shape[0]
              display_string += 'Num. verts: {}\n'.format(num_verts)
          if mesh['class_name'] is not None:
              display_string += 'Synset: {}'.format(mesh['class_name'])
          if mesh['pointcloud'] is not None:
              display_string += 'Num. pointcloud: {}\n'.format(
                  mesh['pointcloud'].shape[0])
          ax.text2D(0.05, 0.8, display_string, transform=ax.transAxes)
      plt.subplots_adjust(
          left=0., right=1., bottom=0., top=1., wspace=0.025, hspace=0.025)
      plt.show()
  