add_mesh (#448)

tensorboardX/proto/plugin_mesh.proto

@@ -0,0 +1,29 @@
+syntax = "proto3";
+
+package tensorboard.mesh;
+
+// A MeshPluginData encapsulates information on which plugins are able to make
+// use of a certain summary value.
+message MeshPluginData {
+  enum ContentType {
+    UNDEFINED = 0;
+    VERTEX = 1;
+    FACE = 2;  // Triangle face.
+    COLOR = 3;
+  }
+
+  // Version `0` is the only supported version.
+  int32 version = 1;
+
+  // The name of the mesh summary this particular summary belongs to.
+  string name = 2;
+
+  // Type of data in the summary.
+  ContentType content_type = 3;
+
+  // JSON-serialized dictionary of ThreeJS classes configuration.
+  string json_config = 5;
+
+  // Shape of underlying data. Cache it here for performance reasons.
+  repeated int32 shape = 6;
+}

tensorboardX/summary.py

@@ -17,6 +17,7 @@
 from .proto.tensor_shape_pb2 import TensorShapeProto
 from .proto.plugin_pr_curve_pb2 import PrCurvePluginData
 from .proto.plugin_text_pb2 import TextPluginData
+from .proto.plugin_mesh_pb2 import MeshPluginData
 from .proto import layout_pb2
 from .x2num import make_np
 from .utils import _prepare_video, convert_to_HWC
@@ -457,3 +458,51 @@ def compute_curve(labels, predictions, num_thresholds=None, weights=None):
     precision = tp / np.maximum(_MINIMUM_COUNT, tp + fp)
     recall = tp / np.maximum(_MINIMUM_COUNT, tp + fn)
     return np.stack((tp, fp, tn, fn, precision, recall))
+
+
+def _get_tensor_summary(tag, tensor, content_type, json_config):
+    mesh_plugin_data = MeshPluginData(
+        version=0,
+        name=tag,
+        content_type=content_type,
+        json_config=json_config,
+        shape=tensor.shape,
+    )
+    content = mesh_plugin_data.SerializeToString()
+    smd = SummaryMetadata(
+        plugin_data=[SummaryMetadata.PluginData(
+            plugin_name='mesh',
+            content=content)])
+
+    tensor = TensorProto(dtype='DT_FLOAT',
+                         float_val=tensor.reshape(-1).tolist(),
+                         tensor_shape=TensorShapeProto(dim=[
+                             TensorShapeProto.Dim(size=tensor.shape[0]),
+                             TensorShapeProto.Dim(size=tensor.shape[1]),
+                             TensorShapeProto.Dim(size=tensor.shape[2]),
+                         ]))
+    tensor_summary = Summary.Value(
+        tag='{}_{}'.format(tag, content_type),
+        tensor=tensor,
+        metadata=smd,
+    )
+    return tensor_summary
+
+
+def mesh(tag, vertices, colors, faces, config_dict=None):
+
+    import json
+    summaries = []
+    tensors = [
+        (make_np(vertices), 1),
+        (make_np(faces), 2),
+        (make_np(colors), 3)
+    ]
+
+    for tensor, content_type in tensors:
+        if tensor is None:
+            continue
+        summaries.append(
+            _get_tensor_summary(tag, tensor, content_type, json.dumps(config_dict, sort_keys=True)))
+
+    return Summary(value=summaries)

tensorboardX/writer.py

@@ -21,7 +21,7 @@
 from .utils import figure_to_image
 from .summary import (
     scalar, histogram, histogram_raw, image, audio, text,
-    pr_curve, pr_curve_raw, video, custom_scalars, image_boxes
+    pr_curve, pr_curve_raw, video, custom_scalars, image_boxes, mesh
 )
 
 
@@ -648,7 +648,8 @@ def add_video(self, tag, vid_tensor, global_step=None, fps=4, walltime=None):
             fps (float or int): Frames per second
             walltime (float): Optional override default walltime (time.time()) of event
         Shape:
-            vid_tensor: :math:`(N, T, C, H, W)`. The values should lie in [0, 255] for type `uint8` or [0, 1] for type `float`.
+            vid_tensor: :math:`(N, T, C, H, W)`. The values should lie in [0, 255] for type
+              `uint8` or [0, 1] for type `float`.
         """
         self._get_file_writer().add_summary(
             video(tag, vid_tensor, fps), global_step, walltime)
@@ -936,6 +937,59 @@ def add_custom_scalars(self, layout):
         """
         self._get_file_writer().add_summary(custom_scalars(layout))
 
+    def add_mesh(self, tag, vertices, colors=None, faces=None, config_dict=None, global_step=None, walltime=None):
+        """Add meshes or 3D point clouds to TensorBoard. The visualization is based on Three.js,
+        so it allows users to interact with the rendered object. Besides the basic definitions
+        such as vertices, faces, users can further provide camera parameter, lighting condition, etc.
+        Please check https://threejs.org/docs/index.html#manual/en/introduction/Creating-a-scene for
+        advanced usage. Note that currently this depends on tb-nightly to show.
+
+        Args:
+            tag (string): Data identifier
+            vertices (torch.Tensor): List of the 3D coordinates of vertices.
+            colors (torch.Tensor): Colors for each vertex
+            faces (torch.Tensor): Indices of vertices within each triangle. (Optional)
+            config_dict: Dictionary with ThreeJS classes names and configuration.
+            global_step (int): Global step value to record
+            walltime (float): Optional override default walltime (time.time())
+              seconds after epoch of event
+
+        Shape:
+            vertices: :math:`(B, N, 3)`. (batch, number_of_vertices, channels)
+
+            colors: :math:`(B, N, 3)`. The values should lie in [0, 255] for type `uint8` or [0, 1] for type `float`.
+
+            faces: :math:`(B, N, 3)`. The values should lie in [0, number_of_vertices] for type `uint8`.
+
+        Examples::
+
+            from tensorboardX import SummaryWriter
+            vertices_tensor = np.array([[
+                [1, 1, 1],
+                [-1, -1, 1],
+                [1, -1, -1],
+                [-1, 1, -1],
+            ]], dtype=float)
+            colors_tensor = np.array([[
+                [255, 0, 0],
+                [0, 255, 0],
+                [0, 0, 255],
+                [255, 0, 255],
+            ]], dtype=int)
+            faces_tensor = np.array([[
+                [0, 2, 3],
+                [0, 3, 1],
+                [0, 1, 2],
+                [1, 3, 2],
+            ]], dtype=int)
+
+            writer = SummaryWriter()
+            writer.add_mesh('my_mesh', vertices=vertices_tensor, colors=colors_tensor, faces=faces_tensor)
+
+            writer.close()
+        """
+        self._get_file_writer().add_summary(mesh(tag, vertices, colors, faces, config_dict), global_step, walltime)
+
     def close(self):
         if self.all_writers is None:
             return  # ignore double close