Video Privacy Shield: Face Tracking, Masking & Re-encoding

#include <iostream>
#include <vector>
#include <string>

extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libswscale/swscale.h>
#include <libavutil/imgutils.h>
#include <libavutil/opt.h>
}

#include <opencv2/opencv.hpp>
#include <opencv2/objdetect.hpp>

// Helper function: Blends a 4-channel PNG sticker onto a 3-channel BGR base image
void overlaySticker(cv::Mat& baseImg, const cv::Mat& sticker, const cv::Rect& faceRect) {
    if (sticker.empty()) return;
    
    // Scale sticker to match the exact size of detected face boundary
    cv::Mat resizedSticker;
    cv::resize(sticker, resizedSticker, faceRect.size());

    for (int y = 0; y < faceRect.height; ++y) {
        for (int x = 0; x < faceRect.width; ++x) {
            int base_y = faceRect.y + y;
            int base_x = faceRect.x + x;

            // Boundary check for safety
            if (base_y >= baseImg.rows || base_x >= baseImg.cols || base_y < 0 || base_x < 0)
                continue;

            cv::Vec4b stickerPixel = resizedSticker.at<cv::Vec4b>(y, x);
            float alpha = stickerPixel[3] / 255.0f; // Extract Alpha Channel weight

            if (alpha > 0.1f) { 
                cv::Vec3b& basePixel = baseImg.at<cv::Vec3b>(base_y, base_x);
                // Alpha blending formula
                for (int c = 0; c < 3; ++c) {
                    basePixel[c] = static_cast<uchar>((1.0f - alpha) * basePixel[c] + alpha * stickerPixel[c]);
                }
            }
        }
    }
}

// Helper function: Performs the actual encoding frame pipeline block
bool encodeAndWriteFrame(AVCodecContext* encCtx, AVFormatContext* outFormatCtx, AVStream* outStream, AVFrame* frame, AVPacket* encPacket) {
    int ret = avcodec_send_frame(encCtx, frame);
    if (ret < 0) return false;

    while (ret >= 0) {
        ret = avcodec_receive_packet(encCtx, encPacket);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            return true;
        } else if (ret < 0) {
            return false;
        }

        // Re-scale timestamps from encoder timebase to muxer stream timebase
        av_packet_rescale_ts(encPacket, encCtx->time_base, outStream->time_base);
        encPacket->stream_index = outStream->index;

        av_interleaved_write_frame(outFormatCtx, encPacket);
        av_packet_unref(encPacket);
    }
    return true;
}

int main(int argc, char** argv) {
    if (argc < 5) {
        std::cerr << "Usage: ./VideoFaceMasker <input_video> <output_video> <cascade_xml> <sticker_png>" << std::endl;
        return -1;
    }

    std::string inputPath = argv[1];
    std::string outputPath = argv[2];
    std::string cascadePath = argv[3];
    std::string stickerPath = argv[4];

    // Load OpenCV cascade model and the sticker asset (with alpha channel)
    cv::CascadeClassifier faceCascade;
    if (!faceCascade.load(cascadePath)) {
        std::cerr << "Error: Failed to load cascade XML." << std::endl;
        return -1;
    }
    cv::Mat stickerImg = cv::imread(stickerPath, cv::IMREAD_UNCHANGED);
    if (stickerImg.empty() || stickerImg.channels() != 4) {
        std::cerr << "Error: Sticker must be a valid 4-channel PNG image." << std::endl;
        return -1;
    }

    // ─── 1. 初始化输入解码器管线 ───
    AVFormatContext* inFormatCtx = nullptr;
    if (avformat_open_input(&inFormatCtx, inputPath.c_str(), nullptr, nullptr) != 0) return -1;
    if (avformat_find_stream_info(inFormatCtx, nullptr) < 0) return -1;

    int videoStreamIdx = -1;
    for (unsigned int i = 0; i < inFormatCtx->nb_streams; i++) {
        if (inFormatCtx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            videoStreamIdx = i;
            break;
        }
    }
    if (videoStreamIdx == -1) return -1;

    const AVCodec* decoder = avcodec_find_decoder(inFormatCtx->streams[videoStreamIdx]->codecpar->codec_id);
    AVCodecContext* decCtx = avcodec_alloc_context3(decoder);
    avcodec_parameters_to_context(decCtx, inFormatCtx->streams[videoStreamIdx]->codecpar);
    if (avcodec_open2(decCtx, decoder, nullptr) < 0) return -1;

    // ─── 2. 初始化输出编码器管线 (FFmpeg 7 标准) ───
    AVFormatContext* outFormatCtx = nullptr;
    avformat_alloc_output_context2(&outFormatCtx, nullptr, nullptr, outputPath.c_str());
    if (!outFormatCtx) return -1;

    const AVCodec* encoder = avcodec_find_encoder(AV_CODEC_ID_H264);
    AVStream* outStream = avformat_new_stream(outFormatCtx, nullptr);
    AVCodecContext* encCtx = avcodec_alloc_context3(encoder);

    // Configure hardware compression constraints
    encCtx->height = decCtx->height;
    encCtx->width = decCtx->width;
    encCtx->pix_fmt = AV_PIX_FMT_YUV420P;
    // Setup framerate metadata metrics
    encCtx->time_base = av_inv_q(av_guess_frame_rate(inFormatCtx, inFormatCtx->streams[videoStreamIdx], nullptr));
    outStream->time_base = encCtx->time_base;

    if (outFormatCtx->oformat->flags & AVFMT_GLOBALHEADER) {
        encCtx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
    }
    if (avcodec_open2(encCtx, encoder, nullptr) < 0) return -1;
    avcodec_parameters_from_context(outStream->codecpar, encCtx);

    // Open output disk file IO handler
    if (!(outFormatCtx->oformat->flags & AVFMT_NOFILE)) {
        if (avio_open(&outFormatCtx->pb, outputPath.c_str(), AVIO_FLAG_WRITE) < 0) return -1;
    }
    if (avformat_write_header(outFormatCtx, nullptr) < 0) return -1;

    // ─── 3. 像素缓冲区与像素色彩变换转换器配置 ───
    AVFrame* inFrame = av_frame_alloc();
    AVFrame* bgrFrame = av_frame_alloc();
    AVFrame* outFrame = av_frame_alloc();
    AVPacket* inPacket = av_packet_alloc();
    AVPacket* outPacket = av_packet_alloc();

    outFrame->width = encCtx->width;
    outFrame->height = encCtx->height;
    outFrame->format = encCtx->pix_fmt;
    av_frame_get_buffer(outFrame, 0);

    // Converter A: YUV to BGR24
    int bgrBufferSize = av_image_get_buffer_size(AV_PIX_FMT_BGR24, decCtx->width, decCtx->height, 1);
    uint8_t* bgrBuffer = (uint8_t*)av_malloc(bgrBufferSize);
    av_image_fill_arrays(bgrFrame->data, bgrFrame->linesize, bgrBuffer, AV_PIX_FMT_BGR24, decCtx->width, decCtx->height, 1);
    SwsContext* decSwsCtx = sws_getContext(
        decCtx->width, decCtx->height, decCtx->pix_fmt,
        decCtx->width, decCtx->height, AV_PIX_FMT_BGR24,
        SWS_BILINEAR, nullptr, nullptr, nullptr
    );

    // Converter B: BGR24 back to YUV420P
    SwsContext* encSwsCtx = sws_getContext(
        decCtx->width, decCtx->height, AV_PIX_FMT_BGR24,
        encCtx->width, encCtx->height, encCtx->pix_fmt,
        SWS_BILINEAR, nullptr, nullptr, nullptr
    );

    int frameCounter = 0;

    // ─── 4. 主循环处理链路 ───
    while (av_read_frame(inFormatCtx, inPacket) >= 0) {
        if (inPacket->stream_index == videoStreamIdx) {
            if (avcodec_send_packet(decCtx, inPacket) == 0) {
                while (avcodec_receive_frame(decCtx, inFrame) == 0) {
                    
                    // Execute Step A conversion
                    sws_scale(decSwsCtx, inFrame->data, inFrame->linesize, 0, decCtx->height, bgrFrame->data, bgrFrame->linesize);

                    // Direct mapping to OpenCV container
                    cv::Mat frameMat(decCtx->height, decCtx->width, CV_8UC3, bgrFrame->data[0], bgrFrame->linesize[0]);

                    // Face identification cascade phase
                    cv::Mat gray;
                    cv::cvtColor(frameMat, gray, cv::COLOR_BGR2GRAY);
                    cv::equalizeHist(gray, gray);
                    std::vector<cv::Rect> faces;
                    faceCascade.detectMultiScale(gray, faces, 1.1, 3, 0, cv::Size(40, 40));

                    // Iterate and overlay stickers onto targeted face geometries
                    for (const auto& face : faces) {
                        overlaySticker(frameMat, stickerImg, face);
                    }

                    // Execute Step B conversion back to YUV
                    av_frame_make_writable(outFrame);
                    sws_scale(encSwsCtx, bgrFrame->data, bgrFrame->linesize, 0, decCtx->height, outFrame->data, outFrame->linesize);
                    
                    // Sync Presentation Timestamps (PTS) directly
                    outFrame->pts = inFrame->pts;

                    // Send the modified frame to encoding module
                    encodeAndWriteFrame(encCtx, outFormatCtx, outStream, outFrame, outPacket);
                    frameCounter++;
                }
            }
        }
        av_packet_unref(inPacket);
    }

    // ─── 5. 解码与编码延迟包冲刷清理段 (Flush) ───
    // Flush encoder cached queues
    encodeAndWriteFrame(encCtx, outFormatCtx, outStream, nullptr, outPacket);
    av_write_trailer(outFormatCtx);

    std::cout << "\n[Process Completed] " << frameCounter << " Frames written successfully." << std::endl;

    // ─── 6. 销毁环境物理内存释放 ───
    av_free(bgrBuffer);
    av_frame_free(&inFrame);
    av_frame_free(&bgrFrame);
    av_frame_free(&outFrame);
    av_packet_free(&inPacket);
    av_packet_free(&outPacket);
    avcodec_free_context(&decCtx);
    avcodec_free_context(&encCtx);
    avformat_close_input(&inFormatCtx);
    if (!(outFormatCtx->oformat->flags & AVFMT_NOFILE)) {
        avio_closep(&outFormatCtx->pb);
    }
    avformat_free_context(outFormatCtx);
    sws_freeContext(decSwsCtx);
    sws_freeContext(encSwsCtx);

    return 0;
}