hunting/utils/MediaBase/src/FfmpegThumbnail.cpp

/*
 * Copyright (c) 2023 Fancy Code.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "FfmpegThumbnail.h"
#include "FfmpegDecoder.h"
#include "FfmpegEncoder.h"
#include "ILog.h"
#include "LinuxApi.h"
#include <stdio.h>
#ifdef __cplusplus
extern "C" {
#endif
#include <libavcodec/codec_id.h>
#include <libavcodec/packet.h>
#include <libavformat/avformat.h>
#include <libavformat/avio.h>
#include <libavutil/avutil.h>
#include <libavutil/dict.h>
#include <libavutil/error.h>
#include <libavutil/frame.h>
#include <libavutil/imgutils.h>
#include <libavutil/mem.h>
#include <libavutil/pixfmt.h>
#include <libswscale/swscale.h>
#ifdef __cplusplus
}
#endif
#include <cstdlib>
#include <functional>
#include <memory>
#include <stdint.h>
#include <string>
thumbnail_info::thumbnail_info(const int &srouceWidth, const int &srouceHeight, const int &targetWidth,
                               const int &targetHeight)
    : mSrouceWidth(srouceWidth), mSrouceHeight(srouceHeight), mTargetWidth(targetWidth), mTargetHeight(targetHeight)
{
}
FfmpegThumbnail::FfmpegThumbnail(const AVCodecID &encodecId, const AVCodecID &decodecId)
    : mOutputFormat(nullptr), mStream(nullptr), mSwsCtx(nullptr), mEncodecId(encodecId), mDecodecId(decodecId)
{
}
void FfmpegThumbnail::Init(const ThumbnailInfo &thumbnailInfo)
{
    LogInfo("FfmpegThumbnail Init\n");
    mSrouceWidth = thumbnailInfo.mSrouceWidth;
    mSrouceHeight = thumbnailInfo.mSrouceHeight;
    mTargetWidth = thumbnailInfo.mTargetWidth;
    mTargetHeight = thumbnailInfo.mTargetHeight;
    mDecoder = std::make_shared<FfmpegDecoder>(mDecodecId, mSrouceWidth, mSrouceHeight);
    if (!mDecoder) {
        LogError("mDecoder = nullptr.\n");
    }
    mEncoder = std::make_shared<FfmpegEncoder>(mEncodecId, mTargetWidth, mTargetHeight);
    if (!mEncoder) {
        LogError("mEncoder = nullptr.\n");
    }
}
void FfmpegThumbnail::UnInit(void)
{
    if (mOutputFormat && mOutputFormat->pb) {
        av_write_trailer(mOutputFormat);
    }
    if (mEncoder) {
        mEncoder->UnInit();
        mEncoder.reset();
    }
    if (mDecoder) {
        mDecoder->UnInit();
        mDecoder.reset();
    }
    if (nullptr == mOutputFormat) {
        return;
    }
    if (!(mOutputFormat->oformat->flags & AVFMT_NOFILE)) {
        /* Close the output file. */
        avio_closep(&mOutputFormat->pb);
    }
    avformat_free_context(mOutputFormat);
    fx_system_v2("sync");
    if (mSwsCtx) {
        sws_freeContext(mSwsCtx);
        mSwsCtx = nullptr;
    }
}
bool FfmpegThumbnail::CreateThumbnail(const std::string &outputFile, const void *data, const size_t &size)
{
    if (!mDecoder) {
        LogError("CreateThumbnail mDecoder && mDecodeCallback\n");
        return true;
    }
    mDecodeCallback = std::bind(&FfmpegThumbnail::GetDecodeDataCallback, this, std::placeholders::_1);
    mEncodeCallback = std::bind(&FfmpegThumbnail::GetEncodeDataCallback, this, std::placeholders::_1, outputFile);
    /* allocate the output media context */
    avformat_alloc_output_context2(&mOutputFormat, nullptr, "image2", outputFile.c_str());
    if (!mOutputFormat) {
        LogError("Could not deduce output format from file.\n");
        return false;
    }
    /* Add the audio and video streams using the default format codecs
     * and initialize the codecs. */
    if (mOutputFormat->oformat->video_codec != AV_CODEC_ID_NONE) {
        /**
         * @brief Maybe there is no need to use avformat_alloc_output_context2 function to create ffmpeg container.
         * TODO: if mOutputFormat can be deleted here?
         */
        mStream = avformat_new_stream(mOutputFormat, nullptr);
        if (!mStream) {
            LogError("Could not allocate stream\n");
            return false;
        }
        mStream->id = mOutputFormat->nb_streams - 1;
        LogInfo("Create video stream\n");
    }
    mDecoder->Init();
    constexpr int NO_FLAGS = 0;
    mEncoder->Init(NO_FLAGS);
    mStream->time_base = mEncoder->GetTimeBase();
    mEncoder->OpenEncoder(nullptr, mStream);
    LogInfo("Start to decode data\n");
    mDecoder->DecodeData(data, size, AV_NOPTS_VALUE, mDecodeCallback);
    LogInfo("Decode data end\n");
    return false;
}
void FfmpegThumbnail::GetDecodeDataCallback(AVFrame *frame)
{
    EncodeDataToPicture(frame);
}
void FfmpegThumbnail::GetEncodeDataCallback(AVPacket *pkt, const std::string &fileName)
{
    SaveThumbnailFile(fileName, pkt->data, pkt->size);
}
void FfmpegThumbnail::EncodeDataToPicture(AVFrame *frame)
{
    LogInfo("Decode frame->width = %d, frame->height=%d\n", frame->width, frame->height);
    AVFrame *thumbnailFrame = av_frame_alloc();
    if (!thumbnailFrame) {
        LogError("thumbnailFrame = nullptr.\n");
        return;
    }
    thumbnailFrame->format = AV_PIX_FMT_YUV420P;
    thumbnailFrame->width = mTargetWidth;
    thumbnailFrame->height = mTargetHeight;

    int jpegBufSize = av_image_get_buffer_size(AV_PIX_FMT_YUV420P, mSrouceWidth, mSrouceHeight, 1);
    LogInfo("jpegBufSize: %d\n", jpegBufSize);
    uint8_t *jpegBuf = (uint8_t *)av_malloc(jpegBufSize);
    if (!jpegBuf) {
        LogError("jpegBuf = nullptr.\n");
        goto END;
    }
    av_image_fill_arrays(
        thumbnailFrame->data, thumbnailFrame->linesize, jpegBuf, AV_PIX_FMT_YUV420P, frame->width, frame->height, 1);

    mSwsCtx = sws_getContext(mSrouceWidth,
                             mSrouceHeight,
                             static_cast<AVPixelFormat>(frame->format),
                             thumbnailFrame->width,
                             thumbnailFrame->height,
                             AV_PIX_FMT_YUV420P,
                             SWS_BILINEAR,
                             nullptr,
                             nullptr,
                             nullptr);
    if (!mSwsCtx) {
        LogError("mSwsCtx = nullptr.\n");
        goto END;
    }
    fx_system_v2("echo 1 > /proc/sys/vm/drop_caches");
    /**
     * @brief Perform pixel format conversion.
     * NOTE: This function will crash when the chip does not have enough memory.
     */
    sws_scale(mSwsCtx, frame->data, frame->linesize, 0, frame->height, thumbnailFrame->data, thumbnailFrame->linesize);

    if (mEncoder) {
        mEncoder->EncodeData(thumbnailFrame, mStream, mEncodeCallback);
    }
END:
    if (thumbnailFrame) {
        av_frame_free(&thumbnailFrame);
    }
    if (jpegBuf) {
        av_free(jpegBuf);
    }
    return;
}
bool FfmpegThumbnail::SaveThumbnailFile(const std::string &fileName, const void *data, const size_t &size)
{
    FILE *file = nullptr;
    if (!data) {
        LogError("SaveThumbnailFile:%s failed, data is nullptr.\n", fileName.c_str());
        return false;
    }
    LogInfo("SaveThumbnailFile:%s, size = %u\n", fileName.c_str(), size);
    file = fopen(fileName.c_str(), "a+");
    if (!file) {
        LogError("fopen failed.\n");
        return false;
    }
    fwrite(data, 1, size, file);
    fflush(file);
    fclose(file);
    // system("sync");
    return true;
}