From 222b337c2e84cb51f04623c6e5753eb66720648c Mon Sep 17 00:00:00 2001
From: Yixing Lao <yixing.lao@gmail.com>
Date: Wed, 17 Apr 2024 17:30:49 +0800
Subject: [PATCH] new: ct.render.render_geometries as depth image (#57)

---
 camtools/render.py | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/camtools/render.py b/camtools/render.py
index d01f8cc..e711222 100644
--- a/camtools/render.py
+++ b/camtools/render.py
@@ -14,9 +14,10 @@ def render_geometries(
     view_status_str: str = None,
     height: int = 720,
     width: int = 1280,
-    visible: bool = False,
     point_size: float = 1.0,
     line_radius: float = None,
+    to_depth=False,
+    visible: bool = False,
 ) -> None:
     """
     Render Open3D geometries to an image. This function may require a display.
@@ -33,15 +34,17 @@ def render_geometries(
             size and the point size.
         height: int image height.
         width: int image width.
-        visible: bool whether to show the window.
         point_size: float point size for point cloud objects.
         line_radius: float line radius for line set objects, when set, the line
             sets will be converted to cylinder meshes with the given radius.
             The radius is in world metric space, not relative pixel space like
             the point size.
+        to_depth: bool whether to render a depth image instead of RGB image.
+        visible: bool whether to show the window.
 
     Returns:
-        image: (H, W, 3) float32 np.ndarray image.
+        (H, W, 3) float32 np.ndarray RGB image by default; (H, W) float32
+        depth image if to_depth is True.
     """
 
     if not isinstance(geometries, list):
@@ -100,9 +103,12 @@ def render_geometries(
 
     vis.poll_events()
     vis.update_renderer()
-    buffer = vis.capture_screen_float_buffer()
+    if to_depth:
+        buffer = vis.capture_depth_float_buffer()
+    else:
+        buffer = vis.capture_screen_float_buffer()
     vis.destroy_window()
-    im_buffer = np.asarray(buffer)
+    im_buffer = np.asarray(buffer).astype(np.float32)
 
     return im_buffer