DisplayFlow/demo/windows_sender/VideoEncoder.cpp

#include "VideoEncoder.h"
#include <iostream>

VideoEncoder::VideoEncoder() = default;

VideoEncoder::~VideoEncoder() {
    if (swsContext_) sws_freeContext(swsContext_);
    if (codecContext_) avcodec_free_context(&codecContext_);
    if (frame_) av_frame_free(&frame_);
    if (packet_) av_packet_free(&packet_);
    if (stagingTexture_) stagingTexture_.Reset();
}

bool VideoEncoder::Initialize(ID3D11Device* device, int width, int height, int fps, int bitrate) {
    device_ = device;
    device_->GetImmediateContext(&context_);
    width_ = width;
    height_ = height;

    // 1. Create Staging Texture for CPU access
    D3D11_TEXTURE2D_DESC desc = {};
    desc.Width = width;
    desc.Height = height;
    desc.MipLevels = 1;
    desc.ArraySize = 1;
    desc.Format = DXGI_FORMAT_B8G8R8A8_UNORM;
    desc.SampleDesc.Count = 1;
    desc.Usage = D3D11_USAGE_STAGING;
    desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
    
    HRESULT hr = device_->CreateTexture2D(&desc, nullptr, &stagingTexture_);
    if (FAILED(hr)) {
        std::cerr << "Failed to create staging texture" << std::endl;
        return false;
    }

    // 2. Initialize FFmpeg
    const AVCodec* codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (!codec) {
        std::cerr << "Codec H.264 not found" << std::endl;
        return false;
    }

    codecContext_ = avcodec_alloc_context3(codec);
    if (!codecContext_) {
        std::cerr << "Could not allocate video codec context" << std::endl;
        return false;
    }

    codecContext_->bit_rate = bitrate;
    codecContext_->width = width;
    codecContext_->height = height;
    codecContext_->time_base = {1, fps};
    codecContext_->framerate = {fps, 1};
    codecContext_->gop_size = 10;
    codecContext_->max_b_frames = 1;
    codecContext_->pix_fmt = AV_PIX_FMT_YUV420P;

    // Set H.264 specific options (latency vs quality)
    av_opt_set(codecContext_->priv_data, "preset", "ultrafast", 0);
    av_opt_set(codecContext_->priv_data, "tune", "zerolatency", 0);

    if (avcodec_open2(codecContext_, codec, nullptr) < 0) {
        std::cerr << "Could not open codec" << std::endl;
        return false;
    }

    frame_ = av_frame_alloc();
    if (!frame_) {
        std::cerr << "Could not allocate video frame" << std::endl;
        return false;
    }
    frame_->format = codecContext_->pix_fmt;
    frame_->width = codecContext_->width;
    frame_->height = codecContext_->height;

    if (av_frame_get_buffer(frame_, 32) < 0) {
        std::cerr << "Could not allocate the video frame data" << std::endl;
        return false;
    }

    packet_ = av_packet_alloc();
    if (!packet_) {
        std::cerr << "Could not allocate packet" << std::endl;
        return false;
    }

    return true;
}

bool VideoEncoder::EncodeFrame(ID3D11Texture2D* texture, std::vector<uint8_t>& outputData, bool& isKeyFrame) {
    if (!texture || !stagingTexture_ || !context_) return false;

    // 1. Copy GPU texture to Staging texture
    context_->CopyResource(stagingTexture_.Get(), texture);

    // 2. Map Staging texture to read data
    D3D11_MAPPED_SUBRESOURCE mapped = {};
    HRESULT hr = context_->Map(stagingTexture_.Get(), 0, D3D11_MAP_READ, 0, &mapped);
    if (FAILED(hr)) return false;

    // 3. Convert BGRA to YUV420P
    if (!swsContext_) {
        swsContext_ = sws_getContext(
            width_, height_, AV_PIX_FMT_BGRA,
            width_, height_, AV_PIX_FMT_YUV420P,
            SWS_BILINEAR, nullptr, nullptr, nullptr
        );
    }

    uint8_t* srcSlice[] = { (uint8_t*)mapped.pData };
    int srcStride[] = { (int)mapped.RowPitch };

    // We need to handle potential padding in mapped.RowPitch vs width*4
    // FFmpeg handles strides correctly.

    sws_scale(swsContext_, srcSlice, srcStride, 0, height_, frame_->data, frame_->linesize);

    context_->Unmap(stagingTexture_.Get(), 0);

    // 4. Encode
    frame_->pts = pts_++;

    int ret = avcodec_send_frame(codecContext_, frame_);
    if (ret < 0) {
        std::cerr << "Error sending a frame for encoding" << std::endl;
        return false;
    }

    while (ret >= 0) {
        ret = avcodec_receive_packet(codecContext_, packet_);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
            break;
        else if (ret < 0) {
            std::cerr << "Error during encoding" << std::endl;
            return false;
        }

        outputData.insert(outputData.end(), packet_->data, packet_->data + packet_->size);
        if (packet_->flags & AV_PKT_FLAG_KEY) isKeyFrame = true;
        
        av_packet_unref(packet_);
    }

    return true;
}