tensorbored 2.21.0rc1769983804__py3-none-any.whl

tensorbored/plugins/mesh/mesh_plugin.py

@@ -0,0 +1,292 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""TensorBoard 3D mesh visualizer plugin."""
+
+import numpy as np
+from werkzeug import wrappers
+
+from tensorbored.backend import http_util
+from tensorbored.data import provider
+from tensorbored.plugins import base_plugin
+from tensorbored.plugins.mesh import metadata
+from tensorbored.plugins.mesh import plugin_data_pb2
+from tensorbored import plugin_util
+
+_DEFAULT_DOWNSAMPLING = 100  # meshes per time series
+
+
+class MeshPlugin(base_plugin.TBPlugin):
+    """A plugin that serves 3D visualization of meshes."""
+
+    plugin_name = metadata.PLUGIN_NAME
+
+    def __init__(self, context):
+        """Instantiates a MeshPlugin via TensorBoard core.
+
+        Args:
+          context: A base_plugin.TBContext instance. A magic container that
+            TensorBoard uses to make objects available to the plugin.
+        """
+        self._data_provider = context.data_provider
+        self._downsample_to = (context.sampling_hints or {}).get(
+            self.plugin_name, _DEFAULT_DOWNSAMPLING
+        )
+        self._version_checker = plugin_util._MetadataVersionChecker(
+            data_kind="mesh",
+            latest_known_version=0,
+        )
+
+    def _instance_tag_metadata(self, ctx, experiment, run, instance_tag):
+        """Gets the `MeshPluginData` proto for an instance tag."""
+        results = self._data_provider.list_tensors(
+            ctx,
+            experiment_id=experiment,
+            plugin_name=metadata.PLUGIN_NAME,
+            run_tag_filter=provider.RunTagFilter(
+                runs=[run], tags=[instance_tag]
+            ),
+        )
+        content = results[run][instance_tag].plugin_content
+        return metadata.parse_plugin_metadata(content)
+
+    def _tag(self, ctx, experiment, run, instance_tag):
+        """Gets the user-facing tag name for an instance tag."""
+        return self._instance_tag_metadata(
+            ctx, experiment, run, instance_tag
+        ).name
+
+    def _instance_tags(self, ctx, experiment, run, tag):
+        """Gets the instance tag names for a user-facing tag."""
+        index = self._data_provider.list_tensors(
+            ctx,
+            experiment_id=experiment,
+            plugin_name=metadata.PLUGIN_NAME,
+            run_tag_filter=provider.RunTagFilter(runs=[run]),
+        )
+        return [
+            instance_tag
+            for (instance_tag, ts) in index.get(run, {}).items()
+            if tag == metadata.parse_plugin_metadata(ts.plugin_content).name
+        ]
+
+    @wrappers.Request.application
+    def _serve_tags(self, request):
+        """A route (HTTP handler) that returns a response with tags.
+
+        Args:
+          request: The werkzeug.Request object.
+
+        Returns:
+          A response that contains a JSON object. The keys of the object
+          are all the runs. Each run is mapped to a (potentially empty)
+          list of all tags that are relevant to this plugin.
+        """
+        ctx = plugin_util.context(request.environ)
+        experiment = plugin_util.experiment_id(request.environ)
+        all_runs = self._data_provider.list_tensors(
+            ctx,
+            experiment_id=experiment,
+            plugin_name=metadata.PLUGIN_NAME,
+        )
+
+        # tagToContent is itself a dictionary mapping tag name to string
+        # SummaryMetadata.plugin_data.content. Retrieve the keys of that dictionary
+        # to obtain a list of tags associated with each run. For each tag estimate
+        # number of samples.
+        response = dict()
+        for run, tags in all_runs.items():
+            response[run] = dict()
+            for instance_tag, metadatum in tags.items():
+                md = metadata.parse_plugin_metadata(metadatum.plugin_content)
+                if not self._version_checker.ok(md.version, run, instance_tag):
+                    continue
+                # Make sure we only operate on user-defined tags here.
+                tag = self._tag(ctx, experiment, run, instance_tag)
+                meta = self._instance_tag_metadata(
+                    ctx, experiment, run, instance_tag
+                )
+                # Batch size must be defined, otherwise we don't know how many
+                # samples were there.
+                response[run][tag] = {"samples": meta.shape[0]}
+        return http_util.Respond(request, response, "application/json")
+
+    def get_plugin_apps(self):
+        """Gets all routes offered by the plugin.
+
+        This method is called by TensorBoard when retrieving all the
+        routes offered by the plugin.
+
+        Returns:
+          A dictionary mapping URL path to route that handles it.
+        """
+        # Note that the methods handling routes are decorated with
+        # @wrappers.Request.application.
+        return {
+            "/tags": self._serve_tags,
+            "/meshes": self._serve_mesh_metadata,
+            "/data": self._serve_mesh_data,
+        }
+
+    def is_active(self):
+        return False  # `list_plugins` as called by TB core suffices
+
+    def frontend_metadata(self):
+        return base_plugin.FrontendMetadata(element_name="mesh-dashboard")
+
+    def _get_sample(self, tensor_datum, sample):
+        """Returns a single sample from a batch of samples."""
+        return tensor_datum.numpy[sample].tolist()
+
+    def _get_tensor_metadata(
+        self, event, content_type, components, data_shape, config
+    ):
+        """Converts a TensorDatum into a JSON-compatible response.
+
+        Args:
+          event: TensorDatum object containing data in proto format.
+          content_type: enum plugin_data_pb2.MeshPluginData.ContentType value,
+            representing content type in TensorDatum.
+          components: Bitmask representing all parts (vertices, colors, etc.) that
+            belong to the summary.
+          data_shape: list of dimensions sizes of the tensor.
+          config: rendering scene configuration as dictionary.
+
+        Returns:
+          Dictionary of transformed metadata.
+        """
+        return {
+            "wall_time": event.wall_time,
+            "step": event.step,
+            "content_type": content_type,
+            "components": components,
+            "config": config,
+            "data_shape": list(data_shape),
+        }
+
+    def _get_tensor_data(self, event, sample):
+        """Convert a TensorDatum into a JSON-compatible response."""
+        data = self._get_sample(event, sample)
+        return data
+
+    def _collect_tensor_events(self, request, step=None):
+        """Collects list of tensor events based on request."""
+        ctx = plugin_util.context(request.environ)
+        experiment = plugin_util.experiment_id(request.environ)
+        run = request.args.get("run")
+        tag = request.args.get("tag")
+
+        tensor_events = []  # List of tuples (meta, tensor) that contain tag.
+        for instance_tag in self._instance_tags(ctx, experiment, run, tag):
+            tensors = self._data_provider.read_tensors(
+                ctx,
+                experiment_id=experiment,
+                plugin_name=metadata.PLUGIN_NAME,
+                run_tag_filter=provider.RunTagFilter(
+                    runs=[run], tags=[instance_tag]
+                ),
+                downsample=self._downsample_to,
+            )[run][instance_tag]
+            meta = self._instance_tag_metadata(
+                ctx, experiment, run, instance_tag
+            )
+            tensor_events += [(meta, tensor) for tensor in tensors]
+
+        if step is not None:
+            tensor_events = [
+                event for event in tensor_events if event[1].step == step
+            ]
+        else:
+            # Make sure tensors sorted by step in ascending order.
+            tensor_events = sorted(
+                tensor_events, key=lambda tensor_data: tensor_data[1].step
+            )
+
+        return tensor_events
+
+    @wrappers.Request.application
+    def _serve_mesh_data(self, request):
+        """A route that returns data for particular summary of specified type.
+
+        Data can represent vertices coordinates, vertices indices in faces,
+        vertices colors and so on. Each mesh may have different combination of
+        abovementioned data and each type/part of mesh summary must be served as
+        separate roundtrip to the server.
+
+        Args:
+          request: werkzeug.Request containing content_type as a name of enum
+            plugin_data_pb2.MeshPluginData.ContentType.
+
+        Returns:
+          werkzeug.Response either float32 or int32 data in binary format.
+        """
+        step = float(request.args.get("step", 0.0))
+        tensor_events = self._collect_tensor_events(request, step)
+        content_type = request.args.get("content_type")
+        try:
+            content_type = plugin_data_pb2.MeshPluginData.ContentType.Value(
+                content_type
+            )
+        except ValueError:
+            return http_util.Respond(
+                request, "Bad content_type", "text/plain", 400
+            )
+        sample = int(request.args.get("sample", 0))
+
+        response = [
+            self._get_tensor_data(tensor, sample)
+            for meta, tensor in tensor_events
+            if meta.content_type == content_type
+        ]
+
+        np_type = {
+            plugin_data_pb2.MeshPluginData.VERTEX: np.float32,
+            plugin_data_pb2.MeshPluginData.FACE: np.int32,
+            plugin_data_pb2.MeshPluginData.COLOR: np.uint8,
+        }[content_type]
+
+        response = np.array(response, dtype=np_type)
+        # Looks like reshape can take around 160ms, so why not store it reshaped.
+        response = response.reshape(-1).tobytes()
+
+        return http_util.Respond(request, response, "arraybuffer")
+
+    @wrappers.Request.application
+    def _serve_mesh_metadata(self, request):
+        """A route that returns the mesh metadata associated with a tag.
+
+        Metadata consists of wall time, type of elements in tensor, scene
+        configuration and so on.
+
+        Args:
+          request: The werkzeug.Request object.
+
+        Returns:
+          A JSON list of mesh data associated with the run and tag
+          combination.
+        """
+        tensor_events = self._collect_tensor_events(request)
+
+        # We convert the tensor data to text.
+        response = [
+            self._get_tensor_metadata(
+                tensor,
+                meta.content_type,
+                meta.components,
+                meta.shape,
+                meta.json_config,
+            )
+            for meta, tensor in tensor_events
+        ]
+        return http_util.Respond(request, response, "application/json")

tensorbored/plugins/mesh/metadata.py

@@ -0,0 +1,152 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Internal information about the mesh plugin."""
+
+import dataclasses
+
+from typing import Any
+
+from tensorbored.compat.proto import summary_pb2
+from tensorbored.plugins.mesh import plugin_data_pb2
+
+PLUGIN_NAME = "mesh"
+
+# The most recent value for the `version` field of the
+# `MeshPluginData` proto.
+_PROTO_VERSION = 0
+
+
+@dataclasses.dataclass(frozen=True)
+class MeshTensor:
+    """A mesh tensor.
+
+    Attributes:
+      data: Tensor of shape `[dim_1, ..., dim_n, 3]` representing the mesh data
+        of one of the following:
+          - 3D coordinates of vertices
+          - Indices of vertices within each triangle
+          - Colors for each vertex
+      content_type: Type of the mesh plugin data content.
+      data_type: Data type of the elements in the tensor.
+    """
+
+    data: Any  # Expects type `tf.Tensor`, not specified here to avoid heavy TF dep.
+    content_type: plugin_data_pb2.MeshPluginData.ContentType
+    data_type: Any  # Expects type `tf.DType`, not specified here to avoid heavy TF dep.
+
+
+def get_components_bitmask(content_types):
+    """Creates bitmask for all existing components of the summary.
+
+    Args:
+      content_type: list of plugin_data_pb2.MeshPluginData.ContentType,
+        representing all components related to the summary.
+    Returns: bitmask based on passed tensors.
+    """
+    components = 0
+    for content_type in content_types:
+        if content_type == plugin_data_pb2.MeshPluginData.UNDEFINED:
+            raise ValueError("Cannot include UNDEFINED content type in mask.")
+        components = components | (1 << content_type)
+    return components
+
+
+def get_current_version():
+    """Returns current verions of the proto."""
+    return _PROTO_VERSION
+
+
+def get_instance_name(name, content_type):
+    """Returns a unique instance name for a given summary related to the
+    mesh."""
+    return "%s_%s" % (
+        name,
+        plugin_data_pb2.MeshPluginData.ContentType.Name(content_type),
+    )
+
+
+def create_summary_metadata(
+    name,
+    display_name,
+    content_type,
+    components,
+    shape,
+    description=None,
+    json_config=None,
+):
+    """Creates summary metadata which defined at MeshPluginData proto.
+
+    Arguments:
+      name: Original merged (summaries of different types) summary name.
+      display_name: The display name used in TensorBoard.
+      content_type: Value from MeshPluginData.ContentType enum describing data.
+      components: Bitmask representing present parts (vertices, colors, etc.) that
+        belong to the summary.
+      shape: list of dimensions sizes of the tensor.
+      description: The description to show in TensorBoard.
+      json_config: A string, JSON-serialized dictionary of ThreeJS classes
+        configuration.
+
+    Returns:
+      A `summary_pb2.SummaryMetadata` protobuf object.
+    """
+    # Shape should be at least BxNx3 where B represents the batch dimensions
+    # and N - the number of points, each with x,y,z coordinates.
+    if len(shape) != 3:
+        raise ValueError(
+            "Tensor shape should be of shape BxNx3, but got %s." % str(shape)
+        )
+    mesh_plugin_data = plugin_data_pb2.MeshPluginData(
+        version=get_current_version(),
+        name=name,
+        content_type=content_type,
+        components=components,
+        shape=shape,
+        json_config=json_config,
+    )
+    content = mesh_plugin_data.SerializeToString()
+    return summary_pb2.SummaryMetadata(
+        display_name=display_name,  # Will not be used in TensorBoard UI.
+        summary_description=description,
+        plugin_data=summary_pb2.SummaryMetadata.PluginData(
+            plugin_name=PLUGIN_NAME, content=content
+        ),
+    )
+
+
+def parse_plugin_metadata(content):
+    """Parse summary metadata to a Python object.
+
+    Arguments:
+      content: The `content` field of a `SummaryMetadata` proto
+        corresponding to the mesh plugin.
+
+    Returns:
+      A `MeshPluginData` protobuf object.
+    Raises: Error if the version of the plugin is not supported.
+    """
+    if not isinstance(content, bytes):
+        raise TypeError("Content type must be bytes.")
+    result = plugin_data_pb2.MeshPluginData.FromString(content)
+    # Add components field to older version of the proto.
+    if result.components == 0:
+        result.components = get_components_bitmask(
+            [
+                plugin_data_pb2.MeshPluginData.VERTEX,
+                plugin_data_pb2.MeshPluginData.FACE,
+                plugin_data_pb2.MeshPluginData.COLOR,
+            ]
+        )
+    return result

tensorbored/plugins/mesh/plugin_data_pb2.py

@@ -0,0 +1,37 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensorbored/plugins/mesh/plugin_data.proto
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n*tensorbored/plugins/mesh/plugin_data.proto\x12\x10tensorbored.mesh\"\xea\x01\n\x0eMeshPluginData\x12\x0f\n\x07version\x18\x01 \x01(\x05\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x42\n\x0c\x63ontent_type\x18\x03 \x01(\x0e\x32,.tensorbored.mesh.MeshPluginData.ContentType\x12\x13\n\x0bjson_config\x18\x05 \x01(\t\x12\r\n\x05shape\x18\x06 \x03(\x05\x12\x12\n\ncomponents\x18\x07 \x01(\r\"=\n\x0b\x43ontentType\x12\r\n\tUNDEFINED\x10\x00\x12\n\n\x06VERTEX\x10\x01\x12\x08\n\x04\x46\x41\x43\x45\x10\x02\x12\t\n\x05\x43OLOR\x10\x03\x62\x06proto3')
+
+
+
+_MESHPLUGINDATA = DESCRIPTOR.message_types_by_name['MeshPluginData']
+_MESHPLUGINDATA_CONTENTTYPE = _MESHPLUGINDATA.enum_types_by_name['ContentType']
+MeshPluginData = _reflection.GeneratedProtocolMessageType('MeshPluginData', (_message.Message,), {
+  'DESCRIPTOR' : _MESHPLUGINDATA,
+  '__module__' : 'tensorbored.plugins.mesh.plugin_data_pb2'
+  # @@protoc_insertion_point(class_scope:tensorbored.mesh.MeshPluginData)
+  })
+_sym_db.RegisterMessage(MeshPluginData)
+
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  _MESHPLUGINDATA._serialized_start=65
+  _MESHPLUGINDATA._serialized_end=299
+  _MESHPLUGINDATA_CONTENTTYPE._serialized_start=238
+  _MESHPLUGINDATA_CONTENTTYPE._serialized_end=299
+# @@protoc_insertion_point(module_scope)