Skip to content

Commit c1b69fd

Browse files
anirulanirul
committed
Fix Vulkan screenshot capture
1 parent 179963a commit c1b69fd

4 files changed

Lines changed: 336 additions & 26 deletions

File tree

‎frame/file/image.cpp‎

Lines changed: 66 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
#include <algorithm>
99
#include <cctype>
1010
#include <chrono>
11+
#include <cstdint>
1112
#include <cstdlib>
1213
#include <cstring>
1314
#include <fstream>
@@ -103,6 +104,50 @@ std::filesystem::path StripAssetPrefix(std::filesystem::path input)
103104
return input;
104105
}
105106

107+
const std::uint8_t* PreparePngBytes(
108+
const void* image,
109+
glm::ivec2 size,
110+
const proto::PixelStructure& pixel_structure,
111+
std::vector<std::uint8_t>& converted)
112+
{
113+
const auto* source = static_cast<const std::uint8_t*>(image);
114+
if (!source)
115+
{
116+
return nullptr;
117+
}
118+
119+
const auto pixel_count = static_cast<std::size_t>(size.x) *
120+
static_cast<std::size_t>(size.y);
121+
switch (pixel_structure.value())
122+
{
123+
case proto::PixelStructure::BGR: {
124+
converted.resize(pixel_count * 3);
125+
for (std::size_t i = 0; i < pixel_count; ++i)
126+
{
127+
const auto src = i * 3;
128+
converted[src + 0] = source[src + 2];
129+
converted[src + 1] = source[src + 1];
130+
converted[src + 2] = source[src + 0];
131+
}
132+
return converted.data();
133+
}
134+
case proto::PixelStructure::BGR_ALPHA: {
135+
converted.resize(pixel_count * 4);
136+
for (std::size_t i = 0; i < pixel_count; ++i)
137+
{
138+
const auto src = i * 4;
139+
converted[src + 0] = source[src + 2];
140+
converted[src + 1] = source[src + 1];
141+
converted[src + 2] = source[src + 0];
142+
converted[src + 3] = source[src + 3];
143+
}
144+
return converted.data();
145+
}
146+
default:
147+
return source;
148+
}
149+
}
150+
106151
} // namespace
107152

108153
Image::Image(
@@ -485,19 +530,35 @@ Image::Image(
485530

486531
void Image::SaveImageToFile(const std::string& file) const
487532
{
488-
// For OpenGL it seams...
533+
if (pixel_element_size_.value() != proto::PixelElementSize::BYTE)
534+
{
535+
throw std::runtime_error(
536+
"SaveImageToFile only supports BYTE images.");
537+
}
538+
539+
// Frame stores image rows bottom-up, so flip back for PNG output.
489540
stbi_flip_vertically_on_write(true);
490541
const auto& logger = frame::Logger::GetInstance();
491542
logger->info("Saving [{}]...", file);
492543
if (!image_)
493544
throw std::runtime_error("no pointer to be saved?");
494-
stbi_write_png(
545+
546+
const int channels = DesiredChannels(pixel_structure_);
547+
std::vector<std::uint8_t> converted = {};
548+
const auto* output_data =
549+
PreparePngBytes(image_, size_, pixel_structure_, converted);
550+
const int stride_in_bytes = size_.x * channels;
551+
const int written = stbi_write_png(
495552
file.c_str(),
496553
size_.x,
497554
size_.y,
498-
pixel_structure_.value(),
499-
image_,
500-
size_.x * pixel_structure_.value());
555+
channels,
556+
output_data,
557+
stride_in_bytes);
558+
if (written == 0)
559+
{
560+
throw std::runtime_error("failed to write png image.");
561+
}
501562
}
502563

503564
void Image::SetData(void* data)

‎frame/vulkan/device.cpp‎

Lines changed: 230 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -17,19 +17,21 @@
1717
#include <filesystem>
1818

1919
#include <stdexcept>
20-
#include "absl/flags/flag.h"
21-
22-
#include "frame/bvh.h"
23-
#include "frame/camera.h"
24-
#include "frame/json/program_key.h"
25-
#include "frame/level.h"
26-
#include "frame/common/application.h"
27-
#include "frame/node_mesh.h"
28-
#include "frame/vulkan/buffer.h"
29-
#include "frame/vulkan/buffer_resources.h"
30-
#include "frame/vulkan/build_level.h"
31-
#include "frame/vulkan/command_resources.h"
32-
#include "frame/vulkan/command_queue.h"
20+
#include "absl/flags/flag.h"
21+
22+
#include "frame/bvh.h"
23+
#include "frame/camera.h"
24+
#include "frame/common/application.h"
25+
#include "frame/file/image.h"
26+
#include "frame/json/program_key.h"
27+
#include "frame/level.h"
28+
#include "frame/node_mesh.h"
29+
#include "frame/proto/uniform.pb.h"
30+
#include "frame/vulkan/buffer.h"
31+
#include "frame/vulkan/buffer_resources.h"
32+
#include "frame/vulkan/build_level.h"
33+
#include "frame/vulkan/command_resources.h"
34+
#include "frame/vulkan/command_queue.h"
3335
#include "frame/vulkan/gpu_memory_manager.h"
3436
#include "frame/vulkan/mesh_resources.h"
3537
#include "frame/vulkan/mesh_utils.h"
@@ -40,11 +42,10 @@
4042
#include "frame/vulkan/scoped_timer.h"
4143
#include "frame/vulkan/shader_compiler.h"
4244
#include "frame/vulkan/swapchain_resources.h"
43-
#include "frame/vulkan/sync_resources.h"
44-
#include "frame/vulkan/texture.h"
45-
#include "frame/vulkan/texture_resources.h"
45+
#include "frame/vulkan/sync_resources.h"
46+
#include "frame/vulkan/texture.h"
47+
#include "frame/vulkan/texture_resources.h"
4648
#include "frame/vulkan/skinned_mesh.h"
47-
#include "frame/proto/uniform.pb.h"
4849
#include <glm/gtc/matrix_transform.hpp>
4950
#include <glm/gtc/matrix_inverse.hpp>
5051
#include <glm/gtc/type_ptr.hpp>
@@ -77,6 +78,170 @@ double ElapsedMilliseconds(const SteadyClock::time_point& start)
7778
.count();
7879
}
7980

81+
proto::PixelStructure PixelStructureForScreenshotFormat(vk::Format format)
82+
{
83+
proto::PixelStructure pixel_structure = {};
84+
switch (format)
85+
{
86+
case vk::Format::eR8G8B8Unorm:
87+
case vk::Format::eR8G8B8Srgb:
88+
pixel_structure.set_value(proto::PixelStructure::RGB);
89+
return pixel_structure;
90+
case vk::Format::eR8G8B8A8Unorm:
91+
case vk::Format::eR8G8B8A8Srgb:
92+
pixel_structure.set_value(proto::PixelStructure::RGB_ALPHA);
93+
return pixel_structure;
94+
case vk::Format::eB8G8R8Unorm:
95+
case vk::Format::eB8G8R8Srgb:
96+
pixel_structure.set_value(proto::PixelStructure::BGR);
97+
return pixel_structure;
98+
case vk::Format::eB8G8R8A8Unorm:
99+
case vk::Format::eB8G8R8A8Srgb:
100+
pixel_structure.set_value(proto::PixelStructure::BGR_ALPHA);
101+
return pixel_structure;
102+
default:
103+
throw std::runtime_error(
104+
"Unsupported Vulkan screenshot format: " + vk::to_string(format) +
105+
".");
106+
}
107+
}
108+
109+
std::uint32_t BytesPerPixelForScreenshotFormat(vk::Format format)
110+
{
111+
switch (format)
112+
{
113+
case vk::Format::eR8G8B8Unorm:
114+
case vk::Format::eR8G8B8Srgb:
115+
case vk::Format::eB8G8R8Unorm:
116+
case vk::Format::eB8G8R8Srgb:
117+
return 3;
118+
case vk::Format::eR8G8B8A8Unorm:
119+
case vk::Format::eR8G8B8A8Srgb:
120+
case vk::Format::eB8G8R8A8Unorm:
121+
case vk::Format::eB8G8R8A8Srgb:
122+
return 4;
123+
default:
124+
throw std::runtime_error(
125+
"Unsupported Vulkan screenshot format: " + vk::to_string(format) +
126+
".");
127+
}
128+
}
129+
130+
void FlipRowsInPlace(
131+
std::vector<std::uint8_t>& pixels, std::size_t row_stride)
132+
{
133+
if (row_stride == 0)
134+
{
135+
return;
136+
}
137+
138+
const auto row_count = pixels.size() / row_stride;
139+
if (row_count < 2)
140+
{
141+
return;
142+
}
143+
144+
std::vector<std::uint8_t> swap_buffer(row_stride);
145+
for (std::size_t top = 0, bottom = row_count - 1; top < bottom;
146+
++top, --bottom)
147+
{
148+
auto* top_row = pixels.data() + top * row_stride;
149+
auto* bottom_row = pixels.data() + bottom * row_stride;
150+
std::memcpy(swap_buffer.data(), top_row, row_stride);
151+
std::memcpy(top_row, bottom_row, row_stride);
152+
std::memcpy(bottom_row, swap_buffer.data(), row_stride);
153+
}
154+
}
155+
156+
std::vector<std::uint8_t> ReadScreenshotPixels(
157+
vk::Device device,
158+
frame::vulkan::GpuMemoryManager& gpu_memory_manager,
159+
frame::vulkan::CommandQueue& command_queue,
160+
vk::Image image,
161+
vk::Format format,
162+
glm::uvec2 size)
163+
{
164+
const auto bytes_per_pixel = BytesPerPixelForScreenshotFormat(format);
165+
const auto row_stride = static_cast<std::size_t>(size.x) * bytes_per_pixel;
166+
const auto image_size = static_cast<vk::DeviceSize>(row_stride) * size.y;
167+
168+
vk::UniqueDeviceMemory staging_memory;
169+
auto staging_buffer = gpu_memory_manager.CreateBuffer(
170+
image_size,
171+
vk::BufferUsageFlagBits::eTransferDst,
172+
vk::MemoryPropertyFlagBits::eHostVisible |
173+
vk::MemoryPropertyFlagBits::eHostCoherent,
174+
staging_memory);
175+
176+
command_queue.SubmitOneTime([&](vk::CommandBuffer command_buffer) {
177+
const vk::ImageSubresourceRange subresource_range(
178+
vk::ImageAspectFlagBits::eColor,
179+
0,
180+
1,
181+
0,
182+
1);
183+
const vk::ImageMemoryBarrier to_transfer(
184+
vk::AccessFlagBits::eShaderRead,
185+
vk::AccessFlagBits::eTransferRead,
186+
vk::ImageLayout::eShaderReadOnlyOptimal,
187+
vk::ImageLayout::eTransferSrcOptimal,
188+
VK_QUEUE_FAMILY_IGNORED,
189+
VK_QUEUE_FAMILY_IGNORED,
190+
image,
191+
subresource_range);
192+
command_buffer.pipelineBarrier(
193+
vk::PipelineStageFlagBits::eFragmentShader |
194+
vk::PipelineStageFlagBits::eComputeShader,
195+
vk::PipelineStageFlagBits::eTransfer,
196+
{},
197+
nullptr,
198+
nullptr,
199+
to_transfer);
200+
201+
const vk::BufferImageCopy copy_region(
202+
0,
203+
0,
204+
0,
205+
vk::ImageSubresourceLayers{
206+
vk::ImageAspectFlagBits::eColor, 0, 0, 1},
207+
vk::Offset3D{0, 0, 0},
208+
vk::Extent3D{size.x, size.y, 1});
209+
command_buffer.copyImageToBuffer(
210+
image,
211+
vk::ImageLayout::eTransferSrcOptimal,
212+
*staging_buffer,
213+
copy_region);
214+
215+
const vk::ImageMemoryBarrier to_shader_read(
216+
vk::AccessFlagBits::eTransferRead,
217+
vk::AccessFlagBits::eShaderRead,
218+
vk::ImageLayout::eTransferSrcOptimal,
219+
vk::ImageLayout::eShaderReadOnlyOptimal,
220+
VK_QUEUE_FAMILY_IGNORED,
221+
VK_QUEUE_FAMILY_IGNORED,
222+
image,
223+
subresource_range);
224+
command_buffer.pipelineBarrier(
225+
vk::PipelineStageFlagBits::eTransfer,
226+
vk::PipelineStageFlagBits::eFragmentShader |
227+
vk::PipelineStageFlagBits::eComputeShader,
228+
{},
229+
nullptr,
230+
nullptr,
231+
to_shader_read);
232+
});
233+
234+
auto* mapped = static_cast<std::uint8_t*>(
235+
device.mapMemory(*staging_memory, 0, image_size));
236+
std::vector<std::uint8_t> pixels(static_cast<std::size_t>(image_size), 0);
237+
std::memcpy(pixels.data(), mapped, pixels.size());
238+
device.unmapMemory(*staging_memory);
239+
240+
// Vulkan readback is top-down; Image::SaveImageToFile expects bottom-up.
241+
FlipRowsInPlace(pixels, row_stride);
242+
return pixels;
243+
}
244+
80245
struct AnimatedRaytraceTimingStats
81246
{
82247
std::size_t frame_count = 0;
@@ -3321,10 +3486,54 @@ void Device::Shutdown()
33213486
vk_unique_device_.reset();
33223487
}
33233488

3324-
void Device::ScreenShot(const std::string& file) const
3325-
{
3326-
logger_->warn("Vulkan screenshot not implemented (requested: {})", file);
3327-
}
3489+
void Device::ScreenShot(const std::string& file) const
3490+
{
3491+
if (device_lost_)
3492+
{
3493+
throw std::runtime_error(
3494+
"Cannot take a Vulkan screenshot after device loss.");
3495+
}
3496+
if (!vk_unique_device_ || !swapchain_resources_ || !output_image_resources_ ||
3497+
!gpu_memory_manager_ || !command_queue_)
3498+
{
3499+
throw std::runtime_error(
3500+
"Vulkan screenshot resources are not initialized.");
3501+
}
3502+
if (!output_image_resources_->HasSwapchainPreviewImage() ||
3503+
!output_image_resources_->IsSwapchainPreviewInShaderRead())
3504+
{
3505+
throw std::runtime_error(
3506+
"No rendered Vulkan frame is available for screenshot.");
3507+
}
3508+
3509+
const auto extent = swapchain_resources_->GetExtent();
3510+
if (extent.width == 0 || extent.height == 0)
3511+
{
3512+
throw std::runtime_error("Cannot take a screenshot of an empty frame.");
3513+
}
3514+
3515+
vk_unique_device_->waitIdle();
3516+
3517+
const auto format = swapchain_resources_->GetImageFormat();
3518+
const auto pixel_structure = PixelStructureForScreenshotFormat(format);
3519+
auto pixels = ReadScreenshotPixels(
3520+
*vk_unique_device_,
3521+
*gpu_memory_manager_,
3522+
*command_queue_,
3523+
output_image_resources_->GetSwapchainPreviewImage(),
3524+
format,
3525+
glm::uvec2(extent.width, extent.height));
3526+
3527+
proto::PixelElementSize pixel_element_size = {};
3528+
pixel_element_size.set_value(proto::PixelElementSize::BYTE);
3529+
frame::file::Image output_image(
3530+
glm::uvec2(extent.width, extent.height),
3531+
pixel_element_size,
3532+
pixel_structure);
3533+
output_image.SetData(pixels.data());
3534+
output_image.SaveImageToFile(file);
3535+
logger_->info("Saved Vulkan screenshot to {}.", file);
3536+
}
33283537

33293538
std::unique_ptr<frame::BufferInterface> Device::CreatePointBuffer(
33303539
std::vector<float>&& /*vector*/)

‎frame/vulkan/output_image_resources.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,7 @@ void OutputImageResources::CreateSwapchainPreviewImage()
187187
vk::SampleCountFlagBits::e1,
188188
vk::ImageTiling::eOptimal,
189189
vk::ImageUsageFlagBits::eTransferDst |
190+
vk::ImageUsageFlagBits::eTransferSrc |
190191
vk::ImageUsageFlagBits::eSampled);
191192
swapchain_preview_image_ =
192193
device_.vk_unique_device_->createImageUnique(image_info);

0 commit comments

Comments
 (0)