diff --git a/python/taichi/ui/window.py b/python/taichi/ui/window.py index 5e200247bdbf4..cd8e70146f674 100644 --- a/python/taichi/ui/window.py +++ b/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. diff --git a/taichi/python/export_ggui.cpp b/taichi/python/export_ggui.cpp index 656ef014e93cd..a02957b0f553e 100644 --- a/taichi/python/export_ggui.cpp +++ b/taichi/python/export_ggui.cpp @@ -291,6 +291,16 @@ struct PyWindow { window->write_image(filename); } + py::array_t get_depth_buffer() { + uint32_t w, h; + auto &depth_buffer = window->get_depth_buffer(w, h); + + return py::array_t( + py::detail::any_container({w, h}), + py::detail::any_container({sizeof(float) * h, sizeof(float)}), + depth_buffer.data(), nullptr); + } + py::array_t 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) diff --git a/taichi/rhi/device.h b/taichi/rhi/device.h index e23a77f438b26..ca4dc5abef5eb 100644 --- a/taichi/rhi/device.h +++ b/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 } diff --git a/taichi/rhi/vulkan/vulkan_device.cpp b/taichi/rhi/vulkan/vulkan_device.cpp index 80a23fc5c3e39..a7de6a4981f66 100644 --- a/taichi/rhi/vulkan/vulkan_device.cpp +++ b/taichi/rhi/vulkan/vulkan_device.cpp @@ -1147,8 +1147,11 @@ inline void buffer_image_copy_ti_to_vk(VkBufferImageCopy ©_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 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, diff --git a/taichi/rhi/vulkan/vulkan_device.h b/taichi/rhi/vulkan/vulkan_device.h index 32adc61650108..bcfcfe61ee488 100644 --- a/taichi/rhi/vulkan/vulkan_device.h +++ b/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 swapchain_images_; // DeviceAllocation screenshot_image_{kDeviceNullAllocation}; + DeviceAllocation depth_buffer_{kDeviceNullAllocation}; DeviceAllocation screenshot_buffer_{kDeviceNullAllocation}; }; diff --git a/taichi/ui/backends/vulkan/swap_chain.cpp b/taichi/ui/backends/vulkan/swap_chain.cpp index 78c1390ad7e1c..387caec0d0526 100644 --- a/taichi/ui/backends/vulkan/swap_chain.cpp +++ b/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 &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 &SwapChain::dump_image_buffer() { auto [w, h] = surface_->get_size(); curr_width_ = w; diff --git a/taichi/ui/backends/vulkan/swap_chain.h b/taichi/ui/backends/vulkan/swap_chain.h index d9a91257ca367..435e12bbcd306 100644 --- a/taichi/ui/backends/vulkan/swap_chain.h +++ b/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 &dump_depth_buffer(); + std::vector &dump_image_buffer(); void write_image(const std::string &filename); @@ -29,6 +31,8 @@ class TI_DLL_EXPORT SwapChain { std::unique_ptr surface_; + std::vector depth_buffer_data_; + std::vector image_buffer_data_; class AppContext *app_context_; diff --git a/taichi/ui/backends/vulkan/window.cpp b/taichi/ui/backends/vulkan/window.cpp index 45f63480e1b25..f96d1706cafef 100644 --- a/taichi/ui/backends/vulkan/window.cpp +++ b/taichi/ui/backends/vulkan/window.cpp @@ -102,6 +102,19 @@ void Window::write_image(const std::string &filename) { } } +std::vector &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 &Window::get_image_buffer(uint32_t &w, uint32_t &h) { if (!drawn_frame_) { draw_frame(); diff --git a/taichi/ui/backends/vulkan/window.h b/taichi/ui/backends/vulkan/window.h index 19f8c374d66e4..a03d4707a91b7 100644 --- a/taichi/ui/backends/vulkan/window.h +++ b/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 &get_depth_buffer(uint32_t &w, uint32_t &h) override; + std::vector &get_image_buffer(uint32_t &w, uint32_t &h) override; ~Window(); diff --git a/taichi/ui/common/window_base.h b/taichi/ui/common/window_base.h index 2eca2145426e4..9a8d0242da342 100644 --- a/taichi/ui/common/window_base.h +++ b/taichi/ui/common/window_base.h @@ -41,6 +41,8 @@ class WindowBase { virtual void write_image(const std::string &filename) = 0; + virtual std::vector &get_depth_buffer(uint32_t &w, uint32_t &h) = 0; + virtual std::vector &get_image_buffer(uint32_t &w, uint32_t &h) = 0; virtual GuiBase *GUI();