hunting/utils/MediaBase/src/FfmpegOutputStream.cpp
2024-07-24 11:47:44 +08:00

207 lines
7.2 KiB
C++
Raw Blame History

/*
* 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 "FfmpegOutputStream.h"
#include "FfmpegDecoder.h"
#include "FfmpegEncoder.h"
#include "FfmpegThumbnail.h"
#include "ILog.h"
#include <stdio.h>
#ifdef __cplusplus
extern "C" {
#endif
#include <libavcodec/codec_id.h>
#include <libavcodec/packet.h>
#include <libavformat/avformat.h>
#include <libavutil/avutil.h>
#include <libavutil/frame.h>
#include <libavutil/mathematics.h>
#include <libavutil/mem.h>
#include <libavutil/pixfmt.h>
#ifdef __cplusplus
}
#endif
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <functional>
#include <memory>
#include <string.h>
#include <string>
#include <thread>
FfmpegOutputStream::FfmpegOutputStream(const AVCodecID &encodecId, const AVCodecID &dncodecId)
: mEncodecId(encodecId), mDeccodecId(dncodecId), mTmpPkt(nullptr), mStream(nullptr), mStreamHeaderWritten(false),
mH264Data2Jpeg(nullptr)
{
}
bool FfmpegOutputStream::Init(AVFormatContext *outputFormat)
{
mDecodeCallback = std::bind(&FfmpegOutputStream::GetDecodeDataCallback, this, std::placeholders::_1);
mTmpPkt = av_packet_alloc();
if (!mTmpPkt) {
LogError("Could not allocate AVPacket\n");
return false;
}
mStream = avformat_new_stream(outputFormat, nullptr);
if (!mStream) {
LogError("Could not allocate stream\n");
return false;
}
if (mDeccodecId != AV_CODEC_ID_NONE) {
mDecoder = std::make_shared<FfmpegDecoder>(mDeccodecId);
mDecoder->Init();
}
mStream->id = outputFormat->nb_streams - 1;
if (mEncodecId != AV_CODEC_ID_NONE) {
mEncoder = std::make_shared<FfmpegEncoder>(mEncodecId);
mEncoder->Init(outputFormat->flags);
mStream->time_base = mEncoder->GetTimeBase();
mEncoder->OpenEncoder(nullptr, mStream);
}
else {
/**
* @brief There is no need to set time_base here, time_base will be automatically corrected inside ffmpeg.
*
*/
mStream->time_base = (AVRational){1, 15};
mStream->codecpar->codec_id = AV_CODEC_ID_H264;
mStream->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
mStream->codecpar->width = 1920;
mStream->codecpar->height = 2160;
// mStream->codecpar->bit_rate = 2073;
mStream->codecpar->format = AV_PIX_FMT_YUV420P;
mStream->codecpar->codec_tag = 0;
mStream->codecpar->extradata = nullptr;
mStream->codecpar->extradata_size = 0;
}
return true;
}
bool FfmpegOutputStream::Init(AVFormatContext *outputFormat, const std::string &thumbnailFile)
{
mThumbnailFileName = thumbnailFile;
return Init(outputFormat);
}
void FfmpegOutputStream::UnInit(void)
{
if (mEncoder) {
mEncoder->UnInit();
mEncoder.reset();
}
if (mDecoder) {
mDecoder->UnInit();
mDecoder.reset();
}
av_packet_free(&mTmpPkt);
if (mThumbnailThread.joinable()) {
mThumbnailThread.join();
}
if (mH264Data2Jpeg) {
free(mH264Data2Jpeg);
mH264Data2Jpeg = nullptr;
}
}
void FfmpegOutputStream::WriteSourceData(const void *data, const size_t &size, const unsigned long long &pts)
{
if (mDecoder) {
mDecoder->DecodeData(data, size, pts, mDecodeCallback);
return;
}
AVPacket *tmpPkt = av_packet_alloc();
static unsigned long long u64Interval = 0;
// AVRational in_timebase = (AVRational){1, 15};
if (mEncodeCallback) {
tmpPkt->data = (uint8_t *)data;
tmpPkt->size = size;
tmpPkt->stream_index = mStream->index;
constexpr int64_t ZERO_MEANS_UNKNOWN = 0;
tmpPkt->duration = ZERO_MEANS_UNKNOWN;
// tmpPkt->pts = u64Interval * 1000; // ת<><D7AA><EFBFBD><EFBFBD> us
tmpPkt->pts = av_rescale_q(pts, (AVRational){1, 1000000}, mStream->time_base);
// LogInfo("vvvvvvvvvvvvvvvvvvvvvvvvv num:%d, den:%d\n", mStream->time_base.num, mStream->time_base.den);
// LogInfo("vvvvvvvvvvvvvvvvvvvvvvvvv pts:%llu, duration:%d\n", tmpPkt->pts, tmpPkt->duration);
// tmpPkt->pts = pts;
u64Interval++;
tmpPkt->dts = tmpPkt->pts;
/* copy packet */
// av_packet_rescale_ts(mTmpPkt, in_timebase, mStream->time_base);
tmpPkt->pos = -1;
mEncodeCallback(tmpPkt);
}
av_packet_unref(tmpPkt);
av_packet_free(&tmpPkt);
}
void FfmpegOutputStream::SetWriteSourceDataCallback(std::function<void(AVPacket *pkt)> callback)
{
mEncodeCallback = callback;
}
bool FfmpegOutputStream::CheckStreamHeader(const void *data, const size_t &size)
{
if (mStreamHeaderWritten || mEncodecId != AV_CODEC_ID_NONE) {
return true;
}
char *pData = (char *)data;
for (size_t i = 0; i < size; i++) {
if ((0x00 == pData[i]) && (0x00 == pData[i + 1]) && (0x00 == pData[i + 2]) && (0x01 == pData[i + 3]) &&
(0x5 == (pData[i + 4] & 0x1F))) {
uint8_t *extradata = (uint8_t *)av_mallocz(i + 1);
mH264Data2Jpeg = (char *)malloc(size + 1);
if (!extradata && !mH264Data2Jpeg) {
LogError("Could not allocate extradata\n");
return false;
}
LogInfo("Found extradata\n");
memcpy(mH264Data2Jpeg, data, size);
/**
* @brief Find the first I-frame and decode it ->encode it into a JPEG image file.
*/
FfmpegOutputStream::CreateThumbnailFile(mH264Data2Jpeg, size);
memcpy(extradata, pData, i);
mStream->codecpar->extradata = extradata;
mStream->codecpar->extradata_size = i;
mStreamHeaderWritten = true;
return mStreamHeaderWritten;
}
}
return false;
}
void FfmpegOutputStream::GetDecodeDataCallback(AVFrame *frame)
{
if (mEncoder) {
mEncoder->EncodeData(frame, mStream, mEncodeCallback);
return;
}
}
void FfmpegOutputStream::CreateThumbnailFile(const void *frame, const size_t &size)
{
if (mThumbnailFileName.empty()) {
LogError("mThumbnailFileName is empty.\n");
return;
}
auto thumbnailThread =
[](std::shared_ptr<FfmpegOutputStream> output, const void *frameData, const size_t dataSize) {
LogInfo("CreateThumbnailFile start.\n");
output->CreateThumbnailFileThread(frameData, dataSize);
};
std::shared_ptr<FfmpegOutputStream> impl = shared_from_this();
mThumbnailThread = std::thread(thumbnailThread, impl, frame, size);
}
void FfmpegOutputStream::CreateThumbnailFileThread(const void *frame, const size_t &size)
{
FfmpegThumbnail thumbnail(AV_CODEC_ID_MJPEG, AV_CODEC_ID_H264);
ThumbnailInfo info(1920, 2160, 192, 216); // TODO:
thumbnail.Init(info);
thumbnail.CreateThumbnail(mThumbnailFileName, frame, size);
thumbnail.UnInit();
LogInfo("CreateThumbnailFile end.\n");
}