"""Framebuffer capture — reads swapchain image pixels back to CPU."""
from __future__ import annotations
from typing import Any
import numpy as np
import vulkan as vk
from .gpu.memory import begin_single_time_commands, create_buffer, end_single_time_commands
__all__ = ["capture_swapchain_frame"]
[docs]
def capture_swapchain_frame(
device: Any,
physical_device: Any,
graphics_queue: Any,
command_pool: Any,
swapchain_images: list,
image_index: int,
extent: tuple[int, int],
image_format: int,
) -> np.ndarray:
"""Capture a swapchain image as an RGBA numpy array.
Returns (height, width, 4) uint8 array. Must be called after the frame has been rendered.
"""
vk.vkDeviceWaitIdle(device)
w, h = extent
buf_size = w * h * 4
staging_buf, staging_mem = create_buffer(
device,
physical_device,
buf_size,
vk.VK_BUFFER_USAGE_TRANSFER_DST_BIT,
vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
)
image = swapchain_images[image_index]
cmd = begin_single_time_commands(device, command_pool)
# PRESENT_SRC_KHR -> TRANSFER_SRC_OPTIMAL
barrier = vk.VkImageMemoryBarrier(
srcAccessMask=vk.VK_ACCESS_MEMORY_READ_BIT,
dstAccessMask=vk.VK_ACCESS_TRANSFER_READ_BIT,
oldLayout=vk.VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
newLayout=vk.VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
srcQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
dstQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
image=image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=1,
),
)
vk.vkCmdPipelineBarrier(
cmd,
vk.VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
0, 0, None, 0, None, 1, [barrier],
)
region = vk.VkBufferImageCopy(
bufferOffset=0,
bufferRowLength=0,
bufferImageHeight=0,
imageSubresource=vk.VkImageSubresourceLayers(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
mipLevel=0,
baseArrayLayer=0,
layerCount=1,
),
imageOffset=vk.VkOffset3D(x=0, y=0, z=0),
imageExtent=vk.VkExtent3D(width=w, height=h, depth=1),
)
vk.vkCmdCopyImageToBuffer(
cmd, image, vk.VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, staging_buf, 1, [region],
)
# TRANSFER_SRC_OPTIMAL -> PRESENT_SRC_KHR
barrier2 = vk.VkImageMemoryBarrier(
srcAccessMask=vk.VK_ACCESS_TRANSFER_READ_BIT,
dstAccessMask=vk.VK_ACCESS_MEMORY_READ_BIT,
oldLayout=vk.VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
newLayout=vk.VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
srcQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
dstQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
image=image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=1,
),
)
vk.vkCmdPipelineBarrier(
cmd,
vk.VK_PIPELINE_STAGE_TRANSFER_BIT,
vk.VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,
0, 0, None, 0, None, 1, [barrier2],
)
end_single_time_commands(device, graphics_queue, command_pool, cmd)
# Read pixels — vkMapMemory returns a CFFI buffer, cast to readable bytes
mapped = vk.vkMapMemory(device, staging_mem, 0, buf_size, 0)
if isinstance(mapped, vk.ffi.buffer):
pixels = np.frombuffer(bytes(mapped), dtype=np.uint8).reshape(h, w, 4).copy()
else:
pixels = np.frombuffer(vk.ffi.buffer(mapped, buf_size), dtype=np.uint8).reshape(h, w, 4).copy()
vk.vkUnmapMemory(device, staging_mem)
vk.vkDestroyBuffer(device, staging_buf, None)
vk.vkFreeMemory(device, staging_mem, None)
# Swizzle BGRA -> RGBA if needed
if image_format in (vk.VK_FORMAT_B8G8R8A8_SRGB, vk.VK_FORMAT_B8G8R8A8_UNORM):
pixels[:, :, [0, 2]] = pixels[:, :, [2, 0]]
return pixels