Allow use of mesh_demo_v2 with PLY files lacking vertex colors

tensorboard/plugins/mesh/demo_utils.py

@@ -26,12 +26,13 @@ def _parse_vertex(vertex_row):
   """Parses a line in a PLY file which encodes a vertex coordinates.
 
   Args:
-    vertex_row: string with vertex coordinates and color.
+    vertex_row: string with vertex coordinates and optional color.
 
   Returns:
     2-tuple containing a length-3 array of vertex coordinates (as
-    floats) and a length-3 array of RGB color values (as ints between 0
-    and 255, inclusive).
+    floats), and a length-3 array of RGB color values (as ints between 0
+    and 255, inclusive) if the vertex line contained color information,
+    otherwise None.
   """
   vertex = vertex_row.strip().split()
   # The row must contain coordinates with RGB/RGBA color in addition to that.
@@ -40,7 +41,9 @@ def _parse_vertex(vertex_row):
     # TODO(b/129298103): add support of RGBA in .ply files.
     return ([float(coord) for coord in vertex[:3]],
             [int(channel) for channel in vertex[3:6]])
-  raise ValueError('PLY file must contain vertices with colors.')
+  elif len(vertex) == 3:
+    return ([float(coord) for coord in vertex], None)
+  raise ValueError('PLY file must have at least 3 vertex properties')
 
 
 def _parse_face(face_row):
@@ -64,7 +67,8 @@ def read_ascii_ply(filename):
     filename: path to a PLY file to read.
 
   Returns:
-    numpy `[dim_1, 3]` array of vertices, `[dim_1, 3]` array of colors and
+    numpy `[dim_1, 3]` array of vertices, `[dim_1, 3]` array of colors
+    (or None if the file did not contain color information), and a
     `[dim_1, 3]` array of faces of the mesh.
   """
   with tf.io.gfile.GFile(filename) as ply_file:
@@ -77,10 +81,14 @@ def read_ascii_ply(filename):
         face_count = int(line.split()[-1])
     # Read vertices and their colors.
     vertex_data = [_parse_vertex(next(ply_file)) for _ in range(vert_count)]
-    vertices = [datum[0] for datum in vertex_data]
-    colors = [datum[1] for datum in vertex_data]
+    vertices = np.array([datum[0] for datum in vertex_data]).astype(np.float32)
+    colors_raw = [datum[1] for datum in vertex_data if datum[1] is not None]
+    if len(colors_raw) == vert_count:
+      colors = np.array(colors_raw).astype(np.uint8)
+    elif len(colors_raw) == 0:
+      colors = None
+    else:
+      raise ValueError('Missing colors for %d vertices' % (vert_count - len(colors)))
     # Read faces.
     faces = [_parse_face(next(ply_file)) for _ in range(face_count)]
-    return (np.array(vertices).astype(np.float32),
-            np.array(colors).astype(np.uint8),
-            np.array(faces).astype(np.int32))
+    return (vertices, colors, np.array(faces).astype(np.int32))

tensorboard/plugins/mesh/demo_utils_test.py

@@ -35,11 +35,13 @@ def test_parse_vertex(self):
     self.assertListEqual(coords, vertex_data[:3])
     self.assertListEqual(colors, vertex_data[3:6])
 
-  def test_prase_vertex_expects_colors(self):
-    """Tests that method will throw error if color is not poresent."""
-    with self.assertRaisesRegexp(ValueError,
-        'PLY file must contain vertices with colors'):
-      demo_utils._parse_vertex('1 2 3')
+  def test_parse_vertex_no_colors(self):
+    """Tests vertex parsing without colors."""
+    # Vertex 3D coordinates with RGBA color.
+    vertex_data = [-0.249245, 1.119303, 0.3095566]
+    coords, colors = demo_utils._parse_vertex(' '.join(map(str, vertex_data)))
+    self.assertListEqual(coords, vertex_data)
+    self.assertIsNone(colors)
 
   def test_parse_face(self):
     """Tests face line parsing."""

tensorboard/plugins/mesh/mesh_demo.py

@@ -55,6 +55,8 @@ def run():
 
   # Read sample PLY file.
   vertices, colors, faces = demo_utils.read_ascii_ply(FLAGS.mesh_path)
+  if colors is None:
+    raise ValueError('Demo requires PLY file containing vertex colors')
 
   # Add batch dimension.
   vertices = np.expand_dims(vertices, 0)

tensorboard/plugins/mesh/mesh_demo_v2.py

@@ -30,7 +30,11 @@
 
 flags.DEFINE_string('logdir', '/tmp/mesh_demo',
                     'Directory to write event logs to.')
-flags.DEFINE_string('mesh_path', None, 'Path to PLY file to visualize.')
+flags.DEFINE_string('mesh_path',
+                    None,
+                    'Path to PLY file to visualize. For an example, download '
+                    'https://people.sc.fsu.edu/~jburkardt/data/ply/teapot.ply')
+flags.DEFINE_string('tag_name', 'mesh_color_tensor', 'Summary tag to use.')
 
 FLAGS = flags.FLAGS
 
@@ -42,11 +46,14 @@
 
 def train_step(vertices, faces, colors, config_dict, step):
   """Executes summary as a train step."""
-  # Change colors over time.
-  t = float(step) / _MAX_STEPS
-  transformed_colors = t * (255 - colors) + (1 - t) * colors
+  if colors is not None:
+    # Change colors over time.
+    t = float(step) / _MAX_STEPS
+    transformed_colors = t * (255 - colors) + (1 - t) * colors
+  else:
+    transformed_colors = None
   mesh_summary.mesh(
-      'mesh_color_tensor', vertices=vertices, faces=faces,
+      FLAGS.tag_name, vertices=vertices, faces=faces,
       colors=transformed_colors, config_dict=config_dict, step=step)
 
 
@@ -70,7 +77,8 @@ def run():
   # Add batch dimension.
   vertices = np.expand_dims(vertices, 0)
   faces = np.expand_dims(faces, 0)
-  colors = np.expand_dims(colors, 0)
+  if colors is not None:
+    colors = np.expand_dims(colors, 0)
 
   # Create summary writer.
   writer = tf.summary.create_file_writer(FLAGS.logdir)