GstreamerVulkanWriter.h

#pragma once
#include <vulkan/vulkan.h>
#include <gst/gst.h>
#include <gst/app/gstappsrc.h>
#include <string>
#include <stdexcept>
#include <cstdint>
#include <vector>
#include <cstring>
 
#ifdef _MSC_VER
#pragma comment(lib, "gstreamer-1.0.lib")
#pragma comment(lib, "gstapp-1.0.lib")
#pragma comment(lib, "gstbase-1.0.lib")
#pragma comment(lib, "gstvideo-1.0.lib")
#pragma comment(lib, "gobject-2.0.lib")
#pragma comment(lib, "glib-2.0.lib")
#endif
 
namespace NDEVR
{
    class GStreamerVulkanWriter
    {
    public:
        GStreamerVulkanWriter(VkPhysicalDevice phys, VkDevice dev, VkQueue q, uint32_t queueFamilyIndex, VkCommandPool commandPool, uint32_t w, uint32_t h, VkFormat fmt, const std::string& outFile, uint32_t fpsNumer = 30, uint32_t fpsDenom = 1)
            : m_phys(phys), m_dev(dev), m_queue(q), m_qfam(queueFamilyIndex), m_cmdPool(commandPool), m_w(w), m_h(h), m_fmt(fmt), m_fpsN(fpsNumer), m_fpsD(fpsDenom)
        {
            if (m_fmt != VK_FORMAT_B8G8R8A8_UNORM && m_fmt != VK_FORMAT_R8G8B8A8_UNORM) 
                throw std::runtime_error("unsupported format");
            m_stride = 4;
            m_frameBytes = m_w * m_h * m_stride;
            createStaging();
            initGStreamer(outFile);
        }
 
        ~GStreamerVulkanWriter()
        {
            if (m_pipeline)
            {
                if (m_appsrc) gst_app_src_end_of_stream(GST_APP_SRC(m_appsrc));
                GstBus* bus = gst_element_get_bus(m_pipeline);
                if (bus)
                {
                    for (;;)
                    {
                        GstMessage* msg = gst_bus_timed_pop_filtered(bus, 3 * GST_SECOND, (GstMessageType)(GST_MESSAGE_EOS | GST_MESSAGE_ERROR | GST_MESSAGE_STATE_CHANGED));
                        if (!msg) break;
                        gst_message_unref(msg);
                        if (msg && (GST_MESSAGE_TYPE(msg) == GST_MESSAGE_EOS || GST_MESSAGE_TYPE(msg) == GST_MESSAGE_ERROR)) break;
                    }
                    gst_object_unref(bus);
                }
                gst_element_set_state(m_pipeline, GST_STATE_NULL);
                gst_object_unref(m_pipeline);
            }
            if (m_map) vkUnmapMemory(m_dev, m_mem);
            if (m_buf) vkDestroyBuffer(m_dev, m_buf, nullptr);
            if (m_mem) vkFreeMemory(m_dev, m_mem, nullptr);
        }

        void pushFrame(VkImage srcImage, VkImageLayout currentLayout, uint64_t frameIndex)
        {
            VkCommandBuffer cmd = beginCmd();
            VkImageMemoryBarrier pre{};
            pre.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
            pre.oldLayout = currentLayout;
            pre.newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
            pre.srcQueueFamilyIndex = m_qfam;
            pre.dstQueueFamilyIndex = m_qfam;
            pre.image = srcImage;
            pre.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            pre.subresourceRange.baseMipLevel = 0;
            pre.subresourceRange.levelCount = 1;
            pre.subresourceRange.baseArrayLayer = 0;
            pre.subresourceRange.layerCount = 1;
            pre.srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT | VK_ACCESS_TRANSFER_READ_BIT;
            pre.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
            vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0, nullptr, 1, &pre);
 
            VkBufferImageCopy region{};
            region.bufferOffset = 0;
            region.bufferRowLength = 0;
            region.bufferImageHeight = 0;
            region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            region.imageSubresource.mipLevel = 0;
            region.imageSubresource.baseArrayLayer = 0;
            region.imageSubresource.layerCount = 1;
            region.imageOffset = { 0, 0, 0 };
            region.imageExtent = { m_w, m_h, 1 };
            vkCmdCopyImageToBuffer(cmd, srcImage, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, m_buf, 1, &region);
 
            VkImageMemoryBarrier post{};
            post.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
            post.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
            post.newLayout = currentLayout;
            post.srcQueueFamilyIndex = m_qfam;
            post.dstQueueFamilyIndex = m_qfam;
            post.image = srcImage;
            post.subresourceRange = pre.subresourceRange;
            post.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
            post.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT | VK_ACCESS_TRANSFER_READ_BIT | VK_ACCESS_SHADER_READ_BIT;
            vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT | VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT, 0, 0, nullptr, 0, nullptr, 1, &post);
 
            endCmdAndWait(cmd);
 
            GstBuffer* buffer = gst_buffer_new_and_alloc(m_frameBytes);
            GstMapInfo info{};
            gst_buffer_map(buffer, &info, GST_MAP_WRITE);
            std::memcpy(info.data, m_map, m_frameBytes);
            gst_buffer_unmap(buffer, &info);
            GST_BUFFER_PTS(buffer) = gst_util_uint64_scale(frameIndex * (uint64_t)m_fpsD, GST_SECOND, m_fpsN);
            GST_BUFFER_DURATION(buffer) = gst_util_uint64_scale((uint64_t)m_fpsD, GST_SECOND, m_fpsN);
            gst_app_src_push_buffer(GST_APP_SRC(m_appsrc), buffer);
        }

        void finish()
        {
            if (m_appsrc) gst_app_src_end_of_stream(GST_APP_SRC(m_appsrc));
        }
 
    private:
        VkPhysicalDevice m_phys{};       
        VkDevice m_dev{};                
        VkQueue m_queue{};               
        uint32_t m_qfam{};               
        VkCommandPool m_cmdPool{};       
        uint32_t m_w{}, m_h{}, m_stride{};  
        VkFormat m_fmt{};                
        VkBuffer m_buf{};                
        VkDeviceMemory m_mem{};          
        void* m_map{};                   
        VkDeviceSize m_frameBytes{};     
        GstElement* m_pipeline{};        
        GstElement* m_appsrc{};          
        uint32_t m_fpsN{}, m_fpsD{};     
 
        uint32_t findMemoryType(uint32_t bits, VkMemoryPropertyFlags flags)
        {
            VkPhysicalDeviceMemoryProperties mp{};
            vkGetPhysicalDeviceMemoryProperties(m_phys, &mp);
            for (uint32_t i = 0; i < mp.memoryTypeCount; ++i)
            {
                if ((bits & (1u << i)) && (mp.memoryTypes[i].propertyFlags & flags) == flags) return i;
            }
            throw std::runtime_error("no memory type");
        }
 
        void createStaging()
        {
            VkBufferCreateInfo bi{ VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO };
            bi.size = m_frameBytes;
            bi.usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT;
            bi.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
            if (vkCreateBuffer(m_dev, &bi, nullptr, &m_buf) != VK_SUCCESS) throw std::runtime_error("vkCreateBuffer");
            VkMemoryRequirements req{};
            vkGetBufferMemoryRequirements(m_dev, m_buf, &req);
            VkMemoryAllocateInfo ai{ VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO };
            ai.allocationSize = req.size;
            ai.memoryTypeIndex = findMemoryType(req.memoryTypeBits, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
            if (vkAllocateMemory(m_dev, &ai, nullptr, &m_mem) != VK_SUCCESS) throw std::runtime_error("vkAllocateMemory");
            if (vkBindBufferMemory(m_dev, m_buf, m_mem, 0) != VK_SUCCESS) throw std::runtime_error("vkBindBufferMemory");
            if (vkMapMemory(m_dev, m_mem, 0, m_frameBytes, 0, &m_map) != VK_SUCCESS) throw std::runtime_error("vkMapMemory");
        }
 
        VkCommandBuffer beginCmd()
        {
            VkCommandBufferAllocateInfo ai{ VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO };
            ai.commandPool = m_cmdPool;
            ai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
            ai.commandBufferCount = 1;
            VkCommandBuffer cmd{};
            if (vkAllocateCommandBuffers(m_dev, &ai, &cmd) != VK_SUCCESS) throw std::runtime_error("vkAllocateCommandBuffers");
            VkCommandBufferBeginInfo bi{ VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO };
            bi.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
            if (vkBeginCommandBuffer(cmd, &bi) != VK_SUCCESS) throw std::runtime_error("vkBeginCommandBuffer");
            return cmd;
        }
 
        void endCmdAndWait(VkCommandBuffer cmd)
        {
            if (vkEndCommandBuffer(cmd) != VK_SUCCESS) throw std::runtime_error("vkEndCommandBuffer");
            VkFenceCreateInfo fi{ VK_STRUCTURE_TYPE_FENCE_CREATE_INFO };
            VkFence fence{};
            if (vkCreateFence(m_dev, &fi, nullptr, &fence) != VK_SUCCESS) throw std::runtime_error("vkCreateFence");
            VkSubmitInfo si{ VK_STRUCTURE_TYPE_SUBMIT_INFO };
            si.commandBufferCount = 1;
            si.pCommandBuffers = &cmd;
            if (vkQueueSubmit(m_queue, 1, &si, fence) != VK_SUCCESS) throw std::runtime_error("vkQueueSubmit");
            vkWaitForFences(m_dev, 1, &fence, VK_TRUE, UINT64_C(0xffffffffffffffff));
            vkDestroyFence(m_dev, fence, nullptr);
            vkFreeCommandBuffers(m_dev, m_cmdPool, 1, &cmd);
        }
 
        void initGStreamer(const std::string& outFile)
        {
            static bool inited = false;
            if (!inited) { int argc = 0; char** argv = nullptr; gst_init(&argc, &argv); inited = true; }
            const char* vfmt = (m_fmt == VK_FORMAT_B8G8R8A8_UNORM) ? "BGRA" : "RGBA";
            std::string pipe =
                "appsrc name=src is-live=false format=time do-timestamp=false "
                "! videoconvert n-threads=0 "
                "! x264enc tune=zerolatency speed-preset=ultrafast bitrate=5000 key-int-max=60 "
                "! h264parse "
                "! mp4mux faststart=true "
                "! filesink location=\"" + outFile + "\"";
            GError* err = nullptr;
            m_pipeline = gst_parse_launch(pipe.c_str(), &err);
            if (!m_pipeline || err) throw std::runtime_error("gst_parse_launch");
            m_appsrc = gst_bin_get_by_name(GST_BIN(m_pipeline), "src");
            if (!m_appsrc) throw std::runtime_error("appsrc not found");
            GstCaps* caps = gst_caps_new_simple(
                "video/x-raw",
                "format", G_TYPE_STRING, vfmt,
                "width", G_TYPE_INT, (int)m_w,
                "height", G_TYPE_INT, (int)m_h,
                "framerate", GST_TYPE_FRACTION, (int)m_fpsN, (int)m_fpsD,
                NULL);
            g_object_set(G_OBJECT(m_appsrc),
                "caps", caps,
                "stream-type", 0,
                "format", GST_FORMAT_TIME,
                "block", TRUE,
                NULL);
            gst_caps_unref(caps);
            gst_element_set_state(m_pipeline, GST_STATE_PLAYING);
        }
    };
}