Rollback add point clouds summary writer to tensorboard interface and metric writer.

PiperOrigin-RevId: 654045032

Author: CLU Authors

clu/metric_writers/async_writer.py

@@ -23,15 +23,14 @@
 
 from collections.abc import Mapping, Sequence
 import contextlib
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from clu import asynclib
 
 from clu.metric_writers import interface
 from clu.metric_writers import multi_writer
 import wrapt
 
-
 Array = interface.Array
 Scalar = interface.Scalar
 
@@ -97,44 +96,21 @@ def write_videos(self, step: int, videos: Mapping[str, Array]):
 
   @_wrap_exceptions
   def write_audios(
-      self, step: int, audios: Mapping[str, Array], *, sample_rate: int
-  ):
+      self, step: int, audios: Mapping[str, Array], *, sample_rate: int):
     self._pool(self._writer.write_audios)(
-        step=step, audios=audios, sample_rate=sample_rate
-    )
+        step=step, audios=audios, sample_rate=sample_rate)
 
   @_wrap_exceptions
   def write_texts(self, step: int, texts: Mapping[str, str]):
     self._pool(self._writer.write_texts)(step=step, texts=texts)
 
   @_wrap_exceptions
-  def write_histograms(
-      self,
-      step: int,
-      arrays: Mapping[str, Array],
-      num_buckets: Optional[Mapping[str, int]] = None,
-  ):
+  def write_histograms(self,
+                       step: int,
+                       arrays: Mapping[str, Array],
+                       num_buckets: Optional[Mapping[str, int]] = None):
     self._pool(self._writer.write_histograms)(
-        step=step, arrays=arrays, num_buckets=num_buckets
-    )
-
-  @_wrap_exceptions
-  def write_pointcloud(
-      self,
-      step: int,
-      point_clouds: Mapping[str, Array],
-      *,
-      point_colors: Optional[Array] = None,
-      configs: Optional[
-          Mapping[str, Union[str, int, float, bool, None]]
-      ] = None,
-  ):
-    self._pool(self._writer.write_pointcloud)(
-        step=step,
-        point_clouds=point_clouds,
-        point_colors=point_colors,
-        configs=configs,
-    )
+        step=step, arrays=arrays, num_buckets=num_buckets)
 
   @_wrap_exceptions
   def write_hparams(self, hparams: Mapping[str, Any]):

clu/metric_writers/async_writer_test.py

@@ -72,27 +72,6 @@ def test_write_videos(self):
     self.sync_writer.write_videos.assert_called_with(4,
                                                      {"input_videos": mock.ANY})
 
-  def test_write_pointcloud(self):
-    point_clouds = np.random.normal(0, 1, (1, 1024, 3)).astype(np.float32)
-    point_colors = np.random.uniform(0, 1, (1, 1024, 3)).astype(np.float32)
-    config = {
-        "material": "PointCloudMaterial",
-        "size": 0.09,
-    }
-    self.writer.write_pointcloud(
-        step=0,
-        point_clouds={"pcd": point_clouds},
-        point_colors={"pcd": point_colors},
-        configs={"config": config},
-    )
-    self.writer.flush()
-    self.sync_writer.write_pointcloud.assert_called_with(
-        step=0,
-        point_clouds={"pcd": mock.ANY},
-        point_colors={"pcd": mock.ANY},
-        configs={"config": mock.ANY},
-    )
-
   def test_write_texts(self):
     self.writer.write_texts(4, {"samples": "bla"})
     self.writer.flush()

clu/metric_writers/interface.py

@@ -153,28 +153,6 @@ def write_histograms(self,
         of the MetricWriter.
     """
 
-  @abc.abstractmethod
-  def write_pointcloud(
-      self,
-      step: int,
-      point_clouds: Mapping[str, Array],
-      *,
-      point_colors: Optional[Mapping[str, Array]] = None,
-      configs: Optional[
-          Mapping[str, Union[str, int, float, bool, None]]
-      ] = None,
-  ):
-    """Writes point cloud summaries.
-
-    Args:
-      step: Step at which the point cloud was generated.
-      point_clouds: Mapping from point clouds key to point cloud of shape [N, 3]
-        array of point coordinates.
-      point_colors: Mapping from point colors key to [N, 3] array of point
-        colors.
-      configs: A dictionary of configuration options for the point cloud.
-    """
-
   @abc.abstractmethod
   def write_hparams(self, hparams: Mapping[str, Any]):
     """Write hyper parameters.

clu/metric_writers/logging_writer.py

@@ -15,7 +15,7 @@
 """MetricWriter that writes all values to INFO log."""
 
 from collections.abc import Mapping
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from absl import logging
 from clu.metric_writers import interface
@@ -77,29 +77,6 @@ def write_histograms(self,
                      self._collection_str, key,
                      _get_histogram_as_string(histo, bins))
 
-  def write_pointcloud(
-      self,
-      step: int,
-      point_clouds: Mapping[str, Array],
-      *,
-      point_colors: Optional[Mapping[str, Any]] = None,
-      configs: Optional[
-          Mapping[str, Union[str, int, float, bool, None]]
-      ] = None,
-  ):
-    logging.info(
-        "[%d]%s Got point clouds: %s, point_colors: %s, configs: %s.",
-        step,
-        self._collection_str,
-        {k: v.shape for k, v in point_clouds.items()},
-        (
-            {k: v.shape for k, v in point_colors.items()}
-            if point_colors is not None
-            else None
-        ),
-        configs,
-    )
-
   def write_hparams(self, hparams: Mapping[str, Any]):
     logging.info("[Hyperparameters]%s %s", self._collection_str, hparams)
 

clu/metric_writers/logging_writer_test.py

@@ -80,42 +80,14 @@ def test_write_histogram(self):
         "INFO:absl:[4] Histogram for 'c' = {[-0.4, 0.6]: 5}",
     ])
 
-  def test_write_pointcloud(self):
-    point_clouds = np.random.normal(0, 1, (1, 1024, 3)).astype(np.float32)
-    point_colors = np.random.uniform(0, 1, (1, 1024, 3)).astype(np.float32)
-    config = {
-        "material": "PointCloudMaterial",
-        "size": 0.09,
-    }
-    with self.assertLogs(level="INFO") as logs:
-      self.writer.write_pointcloud(
-          step=4,
-          point_clouds={"pcd": point_clouds},
-          point_colors={"pcd": point_colors},
-          configs={"configs": config},
-      )
-    self.assertEqual(
-        logs.output,
-        [
-            "INFO:absl:[4] Got point clouds: {'pcd': (1, 1024, 3)},"
-            " point_colors: {'pcd': (1, 1024, 3)}, configs: {'configs':"
-            " {'material': 'PointCloudMaterial', 'size': 0.09}}."
-        ],
-    )
-
   def test_write_hparams(self):
     with self.assertLogs(level="INFO") as logs:
       self.writer.write_hparams({"learning_rate": 0.1, "batch_size": 128})
-    self.assertEqual(
-        logs.output,
-        [
-            "INFO:absl:[Hyperparameters] {'learning_rate': 0.1, 'batch_size':"
-            " 128}"
-        ],
-    )
+    self.assertEqual(logs.output, [
+        "INFO:absl:[Hyperparameters] {'learning_rate': 0.1, 'batch_size': 128}"
+    ])
 
   def test_collection(self):
-    writer = logging_writer.LoggingWriter(collection="train")
     writer = logging_writer.LoggingWriter(collection="train")
     with self.assertLogs(level="INFO") as logs:
       writer.write_scalars(0, {"a": 3, "b": 0.15})

clu/metric_writers/multi_writer.py

@@ -15,7 +15,7 @@
 """MetricWriter that writes to multiple MetricWriters."""
 
 from collections.abc import Mapping, Sequence
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from clu.metric_writers import interface
 
@@ -57,27 +57,13 @@ def write_texts(self, step: int, texts: Mapping[str, str]):
     for w in self._writers:
       w.write_texts(step, texts)
 
-  def write_histograms(
-      self,
-      step: int,
-      arrays: Mapping[str, Array],
-      num_buckets: Optional[Mapping[str, int]] = None):
+  def write_histograms(self,
+                       step: int,
+                       arrays: Mapping[str, Array],
+                       num_buckets: Optional[Mapping[str, int]] = None):
     for w in self._writers:
       w.write_histograms(step, arrays, num_buckets)
 
-  def write_pointcloud(
-      self,
-      step: int,
-      point_clouds: Mapping[str, Array],
-      *,
-      point_colors: Optional[Mapping[str, Array]] = None,
-      configs: Optional[Mapping[str, Union[str, int, float, bool, None]]] = None
-  ):
-    for w in self._writers:
-      w.write_pointcloud(
-          step, point_clouds, point_colors=point_colors, configs=configs
-      )
-
   def write_hparams(self, hparams: Mapping[str, Any]):
     for w in self._writers:
       w.write_hparams(hparams)

clu/metric_writers/multi_writer_test.py

@@ -18,7 +18,6 @@
 
 from clu.metric_writers import interface
 from clu.metric_writers import multi_writer
-import numpy as np
 import tensorflow as tf
 
 
@@ -49,29 +48,6 @@ def test_write_scalars(self):
       ])
       w.flush.assert_called()
 
-  def test_write_pointcloud(self):
-    point_clouds = np.random.normal(0, 1, (1, 1024, 3)).astype(np.float32)
-    point_colors = np.random.uniform(0, 1, (1, 1024, 3)).astype(np.float32)
-    config = {
-        "material": "PointCloudMaterial",
-        "size": 0.09,
-    }
-    self.writer.write_pointcloud(
-        step=0,
-        point_clouds={"pcd": point_clouds},
-        point_colors={"pcd": point_colors},
-        configs={"config": config},
-    )
-    self.writer.flush()
-    for w in self.writers:
-      w.write_pointcloud.assert_called_with(
-          step=0,
-          point_clouds={"pcd": point_clouds},
-          point_colors={"pcd": point_colors},
-          configs={"config": config},
-      )
-      w.flush.assert_called()
-
 
 if __name__ == "__main__":
   tf.test.main()

clu/metric_writers/tf/summary_writer.py

@@ -19,7 +19,7 @@
 """
 
 from collections.abc import Mapping
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from absl import logging
 
@@ -31,7 +31,6 @@
 with epy.lazy_imports():
   # pylint: disable=g-import-not-at-top
   from tensorboard.plugins.hparams import api as hparams_api
-  from tensorboard.plugins.mesh import summary as mesh_summary  # pylint: disable=line-too-long
   # pylint: enable=g-import-not-at-top
 
 
@@ -98,28 +97,6 @@ def write_histograms(
         buckets = None if num_buckets is None else num_buckets.get(key)
         tf.summary.histogram(key, value, step=step, buckets=buckets)
 
-  def write_pointcloud(
-      self,
-      step: int,
-      point_clouds: Mapping[str, Array],
-      *,
-      point_colors: Optional[Mapping[str, Array]] = None,
-      configs: Optional[
-          Mapping[str, Union[str, int, float, bool, None]]
-      ] = None,
-  ):
-    with self._summary_writer.as_default():
-      for key, vertices in point_clouds.items():
-        colors = None if point_colors is None else point_colors.get(key)
-        config = None if configs is None else configs.get(key)
-        mesh_summary.mesh(
-            key,
-            vertices=vertices,
-            colors=colors,
-            step=step,
-            config_dict=config,
-        )
-
   def write_hparams(self, hparams: Mapping[str, Any]):
     with self._summary_writer.as_default():
       hparams_api.hparams(dict(utils.flatten_dict(hparams)))

clu/metric_writers/tf/summary_writer_test.py

@@ -83,21 +83,6 @@ def _load_hparams(logdir: str):
   return hparams
 
 
-def _load_pointcloud_data(logdir: str):
-  """Loads pointcloud summaries from events in a logdir."""
-  paths = tf.io.gfile.glob(os.path.join(logdir, "events.out.tfevents.*"))
-  data = collections.defaultdict(dict)
-  for path in paths:
-    for event in tf.compat.v1.train.summary_iterator(path):
-      for value in event.summary.value:
-        if value.metadata.plugin_data.plugin_name == "mesh":
-          if "config" not in value.tag:
-            data[event.step][value.tag] = tf.make_ndarray(value.tensor)
-          else:
-            data[event.step][value.tag] = value.metadata.plugin_data.content
-  return data
-
-
 class SummaryWriterTest(tf.test.TestCase):
 
   def setUp(self):
@@ -157,24 +142,6 @@ def test_write_histograms(self):
     ]
     self.assertAllClose(data["b"], ([0, 2], expected_histograms_b))
 
-  def test_write_pointcloud(self):
-    point_clouds = np.random.normal(0, 1, (1, 1024, 3)).astype(np.float32)
-    point_colors = np.random.uniform(0, 1, (1, 1024, 3)).astype(np.float32)
-    config = {
-        "material": "PointCloudMaterial",
-        "size": 0.09,
-    }
-    self.writer.write_pointcloud(
-        step=0,
-        point_clouds={"pcd": point_clouds},
-        point_colors={"pcd": point_colors},
-        configs={"config": config},
-    )
-    self.writer.flush()
-    data = _load_pointcloud_data(self.logdir)
-    self.assertAllClose(data[0]["pcd_VERTEX"], point_clouds)
-    self.assertAllClose(data[0]["pcd_COLOR"], point_colors)
-
   def test_hparams(self):
     self.writer.write_hparams(dict(batch_size=512, num_epochs=90))
     hparams = _load_hparams(self.logdir)