[gui] [vulkan] Surpport for getting depth information for python users.

python/taichi/ui/window.py

@@ -126,6 +126,14 @@ def write_image(self, filename):
         """
         return self.window.write_image(filename)
 
+    def get_depth_buffer(self):
+        """Get the depth information of current scene to numpy array.
+
+        Returns:
+            2d numpy array: [width, height] with (0.0~1.0) float-format.
+        """
+        return self.window.get_depth_buffer()
+
     def get_image_buffer(self):
         """Get the window content to numpy array.
 

taichi/python/export_ggui.cpp

@@ -291,6 +291,16 @@ struct PyWindow {
     window->write_image(filename);
   }
 
+  py::array_t<float> get_depth_buffer() {
+    uint32_t w, h;
+    auto &depth_buffer = window->get_depth_buffer(w, h);
+
+    return py::array_t<float>(
+        py::detail::any_container<ssize_t>({w, h}),
+        py::detail::any_container<ssize_t>({sizeof(float) * h, sizeof(float)}),
+        depth_buffer.data(), nullptr);
+  }
+
   py::array_t<float> get_image_buffer() {
     uint32_t w, h;
     auto &img_buffer = window->get_image_buffer(w, h);
@@ -388,6 +398,7 @@ void export_ggui(py::module &m) {
       .def("get_canvas", &PyWindow::get_canvas)
       .def("show", &PyWindow::show)
       .def("write_image", &PyWindow::write_image)
+      .def("get_depth_buffer", &PyWindow::get_depth_buffer)
       .def("get_image_buffer", &PyWindow::get_image_buffer)
       .def("is_pressed", &PyWindow::is_pressed)
       .def("get_cursor_pos", &PyWindow::py_get_cursor_pos)

taichi/rhi/device.h

@@ -286,6 +286,7 @@ struct BufferImageCopyParams {
   } image_extent;
   uint32_t image_base_layer{0};
   uint32_t image_layer_count{1};
+  uint32_t image_aspect_flag{1};
 };
 
 struct ImageCopyParams {
@@ -531,6 +532,7 @@ class Surface {
   virtual int get_image_count() = 0;
   virtual BufferFormat image_format() = 0;
   virtual void resize(uint32_t width, uint32_t height) = 0;
+  virtual DeviceAllocation get_depth_data(DeviceAllocation &depth_alloc) = 0;
   virtual DeviceAllocation get_image_data() {
     TI_NOT_IMPLEMENTED
   }

taichi/rhi/vulkan/vulkan_device.cpp

@@ -1147,8 +1147,11 @@ inline void buffer_image_copy_ti_to_vk(VkBufferImageCopy &copy_info,
   copy_info.imageOffset.y = params.image_offset.y;
   copy_info.imageOffset.z = params.image_offset.z;
   copy_info.imageSubresource.aspectMask =
-      VK_IMAGE_ASPECT_COLOR_BIT;  // FIXME: add option in BufferImageCopyParams
-                                  // to support copying depth images
+      params.image_aspect_flag;  // FIXME: add option in BufferImageCopyParams
+                                 // to support copying depth images
+                                 // FIXED: added an option in
+                                 // BufferImageCopyParams as image_aspect_flag
+                                 // by yuhaoLong(mocki)
   copy_info.imageSubresource.baseArrayLayer = params.image_base_layer;
   copy_info.imageSubresource.layerCount = params.image_layer_count;
   copy_info.imageSubresource.mipLevel = params.image_mip_level;
@@ -2377,6 +2380,9 @@ VulkanSurface::~VulkanSurface() {
     }
     swapchain_images_.clear();
   }
+  if (depth_buffer_ != kDeviceNullAllocation) {
+    device_->dealloc_memory(depth_buffer_);
+  }
   if (screenshot_buffer_ != kDeviceNullAllocation) {
     device_->dealloc_memory(screenshot_buffer_);
   }
@@ -2444,6 +2450,38 @@ void VulkanSurface::present_image(
   device_->wait_idle();
 }
 
+DeviceAllocation VulkanSurface::get_depth_data(DeviceAllocation &depth_alloc) {
+  auto *stream = device_->get_graphics_stream();
+
+  auto [w, h] = get_size();
+  size_t size_bytes = w * h * 4;
+
+  if (depth_buffer_ == kDeviceNullAllocation) {
+    Device::AllocParams params{size_bytes, /*host_wrtie*/ false,
+                               /*host_read*/ true, /*export_sharing*/ false,
+                               AllocUsage::Uniform};
+    depth_buffer_ = device_->allocate_memory(params);
+  }
+
+  device_->image_transition(depth_alloc, ImageLayout::present_src,
+                            ImageLayout::transfer_src);
+
+  std::unique_ptr<CommandList> cmd_list{nullptr};
+
+  BufferImageCopyParams copy_params;
+  copy_params.image_extent.x = w;
+  copy_params.image_extent.y = h;
+  copy_params.image_aspect_flag = VK_IMAGE_ASPECT_DEPTH_BIT;
+  cmd_list = stream->new_command_list();
+  cmd_list->image_to_buffer(depth_buffer_.get_ptr(), depth_alloc,
+                            ImageLayout::transfer_src, copy_params);
+  cmd_list->image_transition(depth_alloc, ImageLayout::transfer_src,
+                             ImageLayout::present_src);
+  stream->submit_synced(cmd_list.get());
+
+  return depth_buffer_;
+}
+
 DeviceAllocation VulkanSurface::get_image_data() {
   auto *stream = device_->get_graphics_stream();
   DeviceAllocation img_alloc = swapchain_images_[image_index_];
@@ -2491,6 +2529,7 @@ DeviceAllocation VulkanSurface::get_image_data() {
   BufferImageCopyParams copy_params;
   copy_params.image_extent.x = w;
   copy_params.image_extent.y = h;
+  copy_params.image_aspect_flag = VK_IMAGE_ASPECT_COLOR_BIT;
   cmd_list = stream->new_command_list();
   // TODO: directly map the image to cpu memory
   cmd_list->image_to_buffer(screenshot_buffer_.get_ptr(), img_alloc,

taichi/rhi/vulkan/vulkan_device.h

@@ -449,6 +449,7 @@ class VulkanSurface : public Surface {
   BufferFormat image_format() override;
   void resize(uint32_t width, uint32_t height) override;
 
+  DeviceAllocation get_depth_data(DeviceAllocation &depth_alloc) override;
   DeviceAllocation get_image_data() override;
 
  private:
@@ -473,6 +474,7 @@ class VulkanSurface : public Surface {
   std::vector<DeviceAllocation> swapchain_images_;
 
   // DeviceAllocation screenshot_image_{kDeviceNullAllocation};
+  DeviceAllocation depth_buffer_{kDeviceNullAllocation};
   DeviceAllocation screenshot_buffer_{kDeviceNullAllocation};
 };
 

taichi/ui/backends/vulkan/swap_chain.cpp

@@ -63,6 +63,23 @@ uint32_t SwapChain::height() {
 taichi::lang::Surface &SwapChain::surface() {
   return *(surface_.get());
 }
+std::vector<float> &SwapChain::dump_depth_buffer() {
+  auto [w, h] = surface_->get_size();
+  curr_width_ = w;
+  curr_height_ = h;
+  depth_buffer_data_.clear();
+  depth_buffer_data_.resize(w * h);
+  DeviceAllocation depth_buffer = surface_->get_depth_data(depth_allocation_);
+  float *ptr = (float *)app_context_->device().map(depth_buffer);
+
+  for (int i = 0; i < w; i++) {
+    for (int j = 0; j < h; j++) {
+      depth_buffer_data_[i * h + (h - j - 1)] = ptr[j * w + i];
+    }
+  }
+  app_context_->device().unmap(depth_buffer);
+  return depth_buffer_data_;
+}
 std::vector<uint32_t> &SwapChain::dump_image_buffer() {
   auto [w, h] = surface_->get_size();
   curr_width_ = w;

taichi/ui/backends/vulkan/swap_chain.h

@@ -16,6 +16,8 @@ class TI_DLL_EXPORT SwapChain {
 
   void resize(uint32_t width, uint32_t height);
 
+  std::vector<float> &dump_depth_buffer();
+
   std::vector<uint32_t> &dump_image_buffer();
 
   void write_image(const std::string &filename);
@@ -29,6 +31,8 @@ class TI_DLL_EXPORT SwapChain {
 
   std::unique_ptr<taichi::lang::Surface> surface_;
 
+  std::vector<float> depth_buffer_data_;
+
   std::vector<uint32_t> image_buffer_data_;
 
   class AppContext *app_context_;

taichi/ui/backends/vulkan/window.cpp

@@ -102,6 +102,19 @@ void Window::write_image(const std::string &filename) {
   }
 }
 
+std::vector<float> &Window::get_depth_buffer(uint32_t &w, uint32_t &h) {
+  if (!drawn_frame_) {
+    draw_frame();
+  }
+  w = renderer_->swap_chain().width();
+  h = renderer_->swap_chain().height();
+  auto &depth_buffer = renderer_->swap_chain().dump_depth_buffer();
+  if (!config_.show_window) {
+    prepare_for_next_frame();
+  }
+  return depth_buffer;
+}
+
 std::vector<uint32_t> &Window::get_image_buffer(uint32_t &w, uint32_t &h) {
   if (!drawn_frame_) {
     draw_frame();

taichi/ui/backends/vulkan/window.h

@@ -42,6 +42,8 @@ class Window final : public WindowBase {
 
   void write_image(const std::string &filename) override;
 
+  std::vector<float> &get_depth_buffer(uint32_t &w, uint32_t &h) override;
+
   std::vector<uint32_t> &get_image_buffer(uint32_t &w, uint32_t &h) override;
 
   ~Window();

taichi/ui/common/window_base.h

@@ -41,6 +41,8 @@ class WindowBase {
 
   virtual void write_image(const std::string &filename) = 0;
 
+  virtual std::vector<float> &get_depth_buffer(uint32_t &w, uint32_t &h) = 0;
+
   virtual std::vector<uint32_t> &get_image_buffer(uint32_t &w, uint32_t &h) = 0;
 
   virtual GuiBase *GUI();