Add point clouds summary writer to tensorboard interface and metric writer.

CLU Authors · copybara-github · commit d0909659ba8b · 2024-07-17T11:01:40.000-07:00
PiperOrigin-RevId: 653292723
diff --git a/clu/metric_writers/async_writer.py b/clu/metric_writers/async_writer.py
@@ -21,15 +21,17 @@
   thread on the next write_*() call.
 """
 
+from collections.abc import Mapping, Sequence
 import contextlib
-from typing import Any, Mapping, Optional, Sequence
+from typing import Any, Optional, Union
 
 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
 
@@ -95,21 +97,44 @@ 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)
+        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,
+    )
 
   @_wrap_exceptions
   def write_hparams(self, hparams: Mapping[str, Any]):
diff --git a/clu/metric_writers/async_writer_test.py b/clu/metric_writers/async_writer_test.py
@@ -72,6 +72,27 @@ 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()
diff --git a/clu/metric_writers/interface.py b/clu/metric_writers/interface.py
@@ -20,7 +20,8 @@
 """
 
 import abc
-from typing import Any, Mapping, Optional, Union
+from collections.abc import Mapping
+from typing import Any, Optional, Union
 
 import jax.numpy as jnp
 import numpy as np
@@ -152,6 +153,28 @@ 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.
diff --git a/clu/metric_writers/logging_writer.py b/clu/metric_writers/logging_writer.py
@@ -14,7 +14,8 @@
 
 """MetricWriter that writes all values to INFO log."""
 
-from typing import Any, Mapping, Optional, Tuple
+from collections.abc import Mapping
+from typing import Any, Optional, Union
 
 from absl import logging
 from clu.metric_writers import interface
@@ -76,6 +77,29 @@ 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)
 
@@ -89,7 +113,7 @@ def close(self):
 def _compute_histogram_as_tf(
     array: np.ndarray,
     num_buckets: Optional[int] = None
-) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
+) -> tuple[Optional[np.ndarray], Optional[np.ndarray]]:
   """Compute the histogram of the input array as TF would do.
 
   Args:
diff --git a/clu/metric_writers/logging_writer_test.py b/clu/metric_writers/logging_writer_test.py
@@ -80,14 +80,42 @@ 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})
diff --git a/clu/metric_writers/multi_writer.py b/clu/metric_writers/multi_writer.py
@@ -14,7 +14,8 @@
 
 """MetricWriter that writes to multiple MetricWriters."""
 
-from typing import Any, Mapping, Optional, Sequence
+from collections.abc import Mapping, Sequence
+from typing import Any, Optional, Union
 
 from clu.metric_writers import interface
 
@@ -56,13 +57,27 @@ 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)
diff --git a/clu/metric_writers/multi_writer_test.py b/clu/metric_writers/multi_writer_test.py
@@ -18,6 +18,7 @@
 
 from clu.metric_writers import interface
 from clu.metric_writers import multi_writer
+import numpy as np
 import tensorflow as tf
 
 
@@ -48,6 +49,29 @@ 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()
diff --git a/clu/metric_writers/tf/summary_writer.py b/clu/metric_writers/tf/summary_writer.py
@@ -19,7 +19,7 @@
 """
 
 from collections.abc import Mapping
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from absl import logging
 
@@ -31,6 +31,7 @@
 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
 
 
@@ -97,6 +98,28 @@ 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)))
diff --git a/clu/metric_writers/tf/summary_writer_test.py b/clu/metric_writers/tf/summary_writer_test.py
diff --git a/clu/metric_writers/torch_tensorboard_writer.py b/clu/metric_writers/torch_tensorboard_writer.py
