mirror of
https://github.com/SatDump/SatDump
synced 2026-08-13 17:47:30 -04:00
668 lines
20 KiB
C++
668 lines
20 KiB
C++
/*
|
|
* This file is part of the SatDump
|
|
*
|
|
* Autors:
|
|
* Oleg Kutkov <contact@olegkutkov.me>
|
|
*
|
|
* This program is free software: you can redistribute it and/or modify
|
|
* it under the terms of the GNU General Public License as published by
|
|
* the Free Software Foundation, version 3.
|
|
*
|
|
* This program is distributed in the hope that it will be useful, but
|
|
* WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
|
|
* General Public License for more details.
|
|
*
|
|
* You should have received a copy of the GNU General Public License
|
|
* along with this program. If not, see <http://www.gnu.org/licenses/>.
|
|
*/
|
|
|
|
#ifdef USE_VIDEO_ENCODER
|
|
|
|
extern "C"
|
|
{
|
|
#include <libavformat/avformat.h>
|
|
#include <libavcodec/avcodec.h>
|
|
#include <libavfilter/avfilter.h>
|
|
#include <libavfilter/buffersink.h>
|
|
#include <libavfilter/buffersrc.h>
|
|
#include <libavutil/opt.h>
|
|
#include <stdio.h>
|
|
}
|
|
|
|
#include <algorithm>
|
|
#include <stdexcept>
|
|
|
|
#include "falcon_video_encoder.hpp"
|
|
#include "logger.h"
|
|
|
|
///
|
|
|
|
#define TS_SYNC_BYTE 0x47
|
|
#define TS_DATA_START 0x40
|
|
#define AVIO_BUF_SIZE 204800
|
|
|
|
static const std::string overlay_file = "telemetry_overlat.txt";
|
|
static const std::string overlay_params = ":reload=1:x=15:y=310:fontsize=15:fontcolor=white";
|
|
|
|
///
|
|
|
|
namespace spacex
|
|
{
|
|
FalconVideoEncoder::FalconVideoEncoder(std::string out_dir)
|
|
: io_ctx(NULL), in_ctx(NULL), filter_ctx(NULL), avio_buf_len(AVIO_BUF_SIZE), avio_buf(static_cast<uint8_t *>(av_malloc(AVIO_BUF_SIZE))), buffersrc(avfilter_get_by_name("buffer")), buffersink(avfilter_get_by_name("buffersink")), out_directory(out_dir), overlay_txt_file(out_dir + "/" + overlay_file), filter_desc("drawtext=textfile='" + overlay_txt_file + "'" + overlay_params), data_found(false), decoder_started(false)
|
|
{
|
|
if (!avio_buf)
|
|
{
|
|
throw std::runtime_error("Failed to allocate avio buffer\n");
|
|
}
|
|
|
|
stream_buffer.reserve(AVIO_BUF_SIZE);
|
|
|
|
io_ctx = avio_alloc_context(avio_buf, avio_buf_len, 0, &stream_buffer, &FalconVideoEncoder::avReader, NULL, NULL);
|
|
|
|
if (!io_ctx)
|
|
{
|
|
throw std::runtime_error("Failed to allocate avio context\n");
|
|
}
|
|
|
|
io_ctx->max_packet_size = avio_buf_len;
|
|
in_ctx = avformat_alloc_context();
|
|
in_ctx->pb = io_ctx;
|
|
|
|
memset(avio_buf, 0, AVIO_BUF_SIZE);
|
|
|
|
av_log_set_callback(FalconVideoEncoder::libraryLogger);
|
|
|
|
overlay.open(overlay_txt_file, std::ios::out | std::ios::trunc);
|
|
|
|
if (!overlay.is_open())
|
|
{
|
|
throw std::runtime_error("Failed to open overlay text file\n");
|
|
}
|
|
}
|
|
|
|
FalconVideoEncoder::~FalconVideoEncoder()
|
|
{
|
|
if (in_ctx)
|
|
{
|
|
avformat_close_input(&in_ctx);
|
|
avformat_free_context(in_ctx);
|
|
}
|
|
|
|
if (io_ctx)
|
|
{
|
|
av_free(io_ctx);
|
|
}
|
|
|
|
overlay.close();
|
|
remove(overlay_txt_file.c_str());
|
|
|
|
for (std::pair<uint32_t, AVCodecContext *> dec : video_decoders)
|
|
{
|
|
avcodec_free_context(&dec.second);
|
|
}
|
|
|
|
for (std::pair<uint32_t, AVCodecContext *> enc : video_encoders)
|
|
{
|
|
avcodec_free_context(&enc.second);
|
|
}
|
|
|
|
for (std::pair<uint32_t, VideoFilter> filter : video_filters)
|
|
{
|
|
avfilter_graph_free(&filter.second.filter_graph);
|
|
}
|
|
}
|
|
|
|
void FalconVideoEncoder::pushTelemetryText(std::string text)
|
|
{
|
|
overlay.seekp(0);
|
|
overlay << text << std::endl;
|
|
}
|
|
|
|
void FalconVideoEncoder::libraryLogger(void *ptr, int level, const char *fmt, va_list vl)
|
|
{
|
|
(void)ptr;
|
|
char str[1024] = {0};
|
|
|
|
vsnprintf(str, sizeof(str), fmt, vl);
|
|
str[strlen(str) - 1] = '\0';
|
|
|
|
switch (level)
|
|
{
|
|
case AV_LOG_PANIC:
|
|
case AV_LOG_FATAL:
|
|
logger->critical(str);
|
|
break;
|
|
|
|
case AV_LOG_ERROR:
|
|
logger->error(str);
|
|
break;
|
|
|
|
case AV_LOG_INFO:
|
|
logger->info(str);
|
|
break;
|
|
|
|
case AV_LOG_WARNING:
|
|
logger->warn(str);
|
|
break;
|
|
|
|
case AV_LOG_DEBUG:
|
|
logger->debug(str);
|
|
break;
|
|
|
|
default:
|
|
logger->trace(std::string("FalconVideoEncoder: ") + str);
|
|
};
|
|
}
|
|
|
|
bool FalconVideoEncoder::feedData(std::vector<uint8_t>::const_iterator &data_chunk_start,
|
|
std::vector<uint8_t>::const_iterator &data_chunk_end)
|
|
{
|
|
if ((stream_buffer.size() + std::distance(data_chunk_start, data_chunk_end)) < AVIO_BUF_SIZE)
|
|
{
|
|
stream_buffer.insert(stream_buffer.begin() + stream_buffer.size(), data_chunk_start, data_chunk_end);
|
|
return true;
|
|
}
|
|
|
|
if (!decoder_started)
|
|
{
|
|
decoder_started = initDecoder();
|
|
|
|
if (!decoder_started)
|
|
{
|
|
stream_buffer.clear();
|
|
return false;
|
|
}
|
|
}
|
|
|
|
frameWorker();
|
|
|
|
stream_buffer.clear();
|
|
|
|
return true;
|
|
}
|
|
|
|
bool FalconVideoEncoder::streamFinished()
|
|
{
|
|
if (!decoder_started && stream_buffer.size())
|
|
{
|
|
decoder_started = initDecoder();
|
|
|
|
if (decoder_started)
|
|
{
|
|
frameWorker();
|
|
stream_buffer.clear();
|
|
}
|
|
}
|
|
|
|
if (decoder_started)
|
|
{
|
|
for (std::pair<uint32_t, AVFormatContext *> out_fmt_ctx : video_out_ctx)
|
|
{
|
|
av_write_trailer(out_fmt_ctx.second);
|
|
avio_closep(&(out_fmt_ctx.second->pb));
|
|
avformat_free_context(out_fmt_ctx.second);
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
////
|
|
|
|
int FalconVideoEncoder::avReader(void *opaque, uint8_t *buf, int buf_size)
|
|
{
|
|
std::vector<uint8_t> *sb = (std::vector<uint8_t> *)opaque;
|
|
|
|
buf_size = sb->size();
|
|
|
|
if (sb->size())
|
|
{
|
|
std::copy(sb->begin(), sb->end(), buf);
|
|
sb->clear();
|
|
}
|
|
|
|
return buf_size;
|
|
}
|
|
|
|
bool FalconVideoEncoder::initDecoder()
|
|
{
|
|
int ret = avformat_open_input(&in_ctx, "VIDEO", NULL, NULL);
|
|
|
|
if (ret < 0)
|
|
{
|
|
logger->error("avformat_open_input() failed!");
|
|
return false;
|
|
}
|
|
|
|
if (avformat_find_stream_info(in_ctx, NULL) < 0)
|
|
{
|
|
logger->error("Failed to find stream info");
|
|
return false;
|
|
}
|
|
|
|
for (uint32_t i = 0; i < in_ctx->nb_streams; i++)
|
|
{
|
|
AVStream *stream = in_ctx->streams[i];
|
|
|
|
if (stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
|
|
{
|
|
AVCodec *dec = avcodec_find_decoder(stream->codecpar->codec_id);
|
|
|
|
if (!dec)
|
|
{
|
|
logger->error("Couldn't find decoder for stream #" + std::to_string(i));
|
|
return false;
|
|
}
|
|
|
|
AVCodecContext *codec_ctx = avcodec_alloc_context3(dec);
|
|
|
|
if (!codec_ctx)
|
|
{
|
|
logger->error("AVCODEC context allocation failed");
|
|
return false;
|
|
}
|
|
|
|
if (avcodec_parameters_to_context(codec_ctx, stream->codecpar) < 0)
|
|
{
|
|
logger->error("Failed to copy decoder parameters to input decoder context");
|
|
return false;
|
|
}
|
|
|
|
logger->info("Found Video stream, id = " + std::to_string(i));
|
|
|
|
codec_ctx->framerate = av_guess_frame_rate(in_ctx, stream, NULL);
|
|
|
|
if (avcodec_open2(codec_ctx, avcodec_find_decoder(codec_ctx->codec_id), NULL) < 0)
|
|
{
|
|
logger->error("Can't open decoder for the stream");
|
|
return false;
|
|
}
|
|
|
|
video_decoders[i] = codec_ctx;
|
|
}
|
|
else if (stream->codecpar->codec_type == AVMEDIA_TYPE_DATA)
|
|
{
|
|
if (stream->codecpar->codec_id == AV_CODEC_ID_SMPTE_KLV)
|
|
{
|
|
logger->info("Found KLV data stream, id = " + std::to_string(i));
|
|
klv_streams[i].open(buildOutFileName(i, "data", "klv").c_str());
|
|
}
|
|
}
|
|
}
|
|
|
|
av_dump_format(in_ctx, 0, "STREAM", 0);
|
|
|
|
if (!configureVideoOutput())
|
|
{
|
|
return false;
|
|
}
|
|
|
|
if (!configureVideoFilters())
|
|
{
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
std::string FalconVideoEncoder::buildOutFileName(uint32_t stream_id, const char *base, const char *ext)
|
|
{
|
|
std::string result = out_directory + "/" + base + "_" + std::to_string(stream_id) + "." + ext;
|
|
return result;
|
|
}
|
|
|
|
bool FalconVideoEncoder::configureVideoOutput()
|
|
{
|
|
AVStream *out_stream;
|
|
AVFormatContext *out_ctx;
|
|
AVCodecContext *enc_ctx;
|
|
AVCodec *encoder;
|
|
|
|
for (std::pair<uint32_t, AVCodecContext *> dec_ctx : video_decoders)
|
|
{
|
|
std::string out_file = buildOutFileName(dec_ctx.first, "video", "mp4");
|
|
|
|
avformat_alloc_output_context2(&out_ctx, NULL, NULL, out_file.c_str());
|
|
|
|
if (!out_ctx)
|
|
{
|
|
logger->error("Failed to create output context file " + out_file);
|
|
continue;
|
|
}
|
|
|
|
encoder = avcodec_find_encoder(dec_ctx.second->codec_id);
|
|
|
|
if (!encoder)
|
|
{
|
|
logger->error("Failed to find encoder for stream #" + std::to_string(dec_ctx.first));
|
|
continue;
|
|
}
|
|
|
|
out_stream = avformat_new_stream(out_ctx, encoder);
|
|
|
|
if (!out_stream)
|
|
{
|
|
logger->error("Failed allocating output stream " + out_file);
|
|
continue;
|
|
}
|
|
|
|
enc_ctx = avcodec_alloc_context3(encoder);
|
|
|
|
if (!enc_ctx)
|
|
{
|
|
logger->error("Failed to allocate the encoder context for " + out_file);
|
|
continue;
|
|
}
|
|
|
|
enc_ctx->width = dec_ctx.second->width;
|
|
enc_ctx->height = dec_ctx.second->height;
|
|
enc_ctx->sample_aspect_ratio = dec_ctx.second->sample_aspect_ratio;
|
|
|
|
if (encoder->pix_fmts)
|
|
{
|
|
enc_ctx->pix_fmt = encoder->pix_fmts[0];
|
|
}
|
|
else
|
|
{
|
|
enc_ctx->pix_fmt = dec_ctx.second->pix_fmt;
|
|
}
|
|
|
|
enc_ctx->time_base = av_inv_q(dec_ctx.second->framerate);
|
|
|
|
if (out_ctx->oformat->flags & AVFMT_GLOBALHEADER)
|
|
{
|
|
enc_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
|
|
}
|
|
|
|
if (avcodec_open2(enc_ctx, encoder, NULL) < 0)
|
|
{
|
|
logger->error("Couldn't open video encoder for stream #" + std::to_string(dec_ctx.first));
|
|
continue;
|
|
}
|
|
|
|
avcodec_parameters_from_context(out_stream->codecpar, enc_ctx);
|
|
|
|
out_stream->time_base = enc_ctx->time_base;
|
|
|
|
if (avio_open(&out_ctx->pb, out_file.c_str(), AVIO_FLAG_WRITE) < 0)
|
|
{
|
|
logger->error("Couldn't create output file " + out_file);
|
|
continue;
|
|
}
|
|
|
|
if (avformat_write_header(out_ctx, NULL) < 0)
|
|
{
|
|
logger->error("Unable to write AV header to " + out_file);
|
|
continue;
|
|
}
|
|
|
|
logger->info("Writing Video #" + std::to_string(dec_ctx.first) + " to " + out_file);
|
|
|
|
video_encoders[dec_ctx.first] = enc_ctx;
|
|
video_out_ctx[dec_ctx.first] = out_ctx;
|
|
|
|
av_dump_format(out_ctx, 0, out_file.c_str(), 1);
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
bool FalconVideoEncoder::configureVideoFilters()
|
|
{
|
|
char args[512];
|
|
|
|
for (std::pair<uint32_t, AVCodecContext *> dec_ctx : video_decoders)
|
|
{
|
|
AVFilterGraph *fgraph = avfilter_graph_alloc();
|
|
|
|
if (!fgraph)
|
|
{
|
|
logger->error("Failed to allocate filter graph");
|
|
return false;
|
|
}
|
|
|
|
AVStream *stream = in_ctx->streams[dec_ctx.first];
|
|
AVFilterInOut *outputs = avfilter_inout_alloc();
|
|
AVFilterInOut *inputs = avfilter_inout_alloc();
|
|
|
|
if (!outputs || !inputs)
|
|
{
|
|
logger->error("Failed to allocate filter inputs and outputs");
|
|
return false;
|
|
}
|
|
|
|
AVRational time_base = stream->time_base;
|
|
|
|
snprintf(args, sizeof(args),
|
|
"video_size=%dx%d:pix_fmt=%d:time_base=%d/%d:pixel_aspect=%d/%d",
|
|
dec_ctx.second->width, dec_ctx.second->height, dec_ctx.second->pix_fmt,
|
|
time_base.num, time_base.den,
|
|
dec_ctx.second->sample_aspect_ratio.num, dec_ctx.second->sample_aspect_ratio.den);
|
|
|
|
AVFilterContext *buffersrc_ctx;
|
|
AVFilterContext *buffersink_ctx;
|
|
|
|
if (avfilter_graph_create_filter(&buffersrc_ctx, buffersrc, "in", args, NULL, fgraph) < 0)
|
|
{
|
|
logger->error("Cannot create filter buffer source");
|
|
return false;
|
|
}
|
|
|
|
if (avfilter_graph_create_filter(&buffersink_ctx, buffersink, "out", NULL, NULL, fgraph) < 0)
|
|
{
|
|
logger->error("Cannot create filter buffer sink");
|
|
return false;
|
|
}
|
|
|
|
outputs->name = av_strdup("in");
|
|
outputs->filter_ctx = buffersrc_ctx;
|
|
outputs->pad_idx = 0;
|
|
outputs->next = NULL;
|
|
|
|
inputs->name = av_strdup("out");
|
|
inputs->filter_ctx = buffersink_ctx;
|
|
inputs->pad_idx = 0;
|
|
inputs->next = NULL;
|
|
|
|
if (avfilter_graph_parse_ptr(fgraph, filter_desc.c_str(), &inputs, &outputs, NULL) < 0)
|
|
{
|
|
logger->error("avfilter_graph_parse_ptr fail");
|
|
return false;
|
|
}
|
|
|
|
if (avfilter_graph_config(fgraph, NULL) < 0)
|
|
{
|
|
logger->error("avfilter_graph_config failed");
|
|
return false;
|
|
}
|
|
|
|
video_filters[dec_ctx.first] = {buffersrc_ctx, buffersink_ctx, fgraph};
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
bool FalconVideoEncoder::frameDecoder(AVCodecContext *avctx, AVFrame *frame, int *got_frame, AVPacket *pkt)
|
|
{
|
|
int ret;
|
|
*got_frame = 0;
|
|
|
|
if (pkt)
|
|
{
|
|
ret = avcodec_send_packet(avctx, pkt);
|
|
|
|
if (ret < 0 && ret != AVERROR_EOF)
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
|
|
ret = avcodec_receive_frame(avctx, frame);
|
|
|
|
if (ret < 0 && ret != AVERROR(EAGAIN))
|
|
{
|
|
return false;
|
|
}
|
|
|
|
if (ret >= 0)
|
|
{
|
|
*got_frame = 1;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
bool FalconVideoEncoder::frameWorker()
|
|
{
|
|
AVPacket packet;
|
|
AVPacket enc_pkt;
|
|
|
|
std::unordered_map<uint32_t, AVCodecContext *>::iterator dec_it;
|
|
std::unordered_map<uint32_t, AVCodecContext *>::iterator enc_it;
|
|
std::unordered_map<uint32_t, AVFormatContext *>::iterator out_ctx_it;
|
|
std::unordered_map<uint32_t, VideoFilter>::iterator vfilter_it;
|
|
std::unordered_map<uint32_t, std::ofstream>::iterator klv_it;
|
|
|
|
AVCodecContext *dec;
|
|
AVCodecContext *enc;
|
|
AVFormatContext *out_ctx;
|
|
|
|
int got_frame = 0, ret = 0;
|
|
|
|
AVFrame *frame = av_frame_alloc();
|
|
AVFrame *filt_frame = av_frame_alloc();
|
|
|
|
if (!frame || !filt_frame)
|
|
{
|
|
logger->error("Couldn't allocate frame");
|
|
return false;
|
|
}
|
|
|
|
av_init_packet(&packet);
|
|
|
|
while (av_read_frame(in_ctx, &packet) >= 0)
|
|
{
|
|
uint32_t stream_id = packet.stream_index;
|
|
dec_it = video_decoders.find(stream_id);
|
|
|
|
/* Found decoder for the stream */
|
|
if (dec_it != video_decoders.end())
|
|
{
|
|
|
|
enc_it = video_encoders.find(stream_id);
|
|
|
|
if (enc_it == video_encoders.end())
|
|
{
|
|
logger->error("Can't find video encoder for stream #" + std::to_string(stream_id));
|
|
return false;
|
|
}
|
|
|
|
out_ctx_it = video_out_ctx.find(stream_id);
|
|
|
|
if (out_ctx_it == video_out_ctx.end())
|
|
{
|
|
logger->error("Can't find output context for stream #" + std::to_string(stream_id));
|
|
return false;
|
|
}
|
|
|
|
vfilter_it = video_filters.find(stream_id);
|
|
|
|
if (vfilter_it == video_filters.end())
|
|
{
|
|
logger->error("Can't find video filter for stream #" + std::to_string(stream_id));
|
|
return false;
|
|
}
|
|
|
|
dec = dec_it->second;
|
|
enc = enc_it->second;
|
|
out_ctx = out_ctx_it->second;
|
|
|
|
av_packet_rescale_ts(&packet, in_ctx->streams[stream_id]->time_base, dec->time_base);
|
|
|
|
if (!frameDecoder(dec, frame, &got_frame, &packet))
|
|
{
|
|
logger->error("Stream #" + std::to_string(stream_id) + " decoder failed");
|
|
return false;
|
|
}
|
|
|
|
if (frame_counters.find(stream_id) == frame_counters.end())
|
|
{
|
|
frame_counters[stream_id] = 0;
|
|
}
|
|
|
|
frame->pts = frame_counters[stream_id]++;
|
|
|
|
if (got_frame)
|
|
{
|
|
|
|
if (av_buffersrc_add_frame_flags(vfilter_it->second.buf_src_ctx, frame, AV_BUFFERSRC_FLAG_KEEP_REF) < 0)
|
|
{
|
|
logger->error("Error while feeding the filtergraph");
|
|
continue;
|
|
}
|
|
|
|
while (1)
|
|
{
|
|
ret = av_buffersink_get_frame(vfilter_it->second.buf_sink_ctx, filt_frame);
|
|
|
|
if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
|
|
{
|
|
break;
|
|
}
|
|
else if (ret < 0)
|
|
{
|
|
return -1;
|
|
}
|
|
|
|
if (avcodec_send_frame(enc, filt_frame) < 0)
|
|
{
|
|
break;
|
|
}
|
|
|
|
av_init_packet(&enc_pkt);
|
|
|
|
while (ret >= 0)
|
|
{
|
|
ret = avcodec_receive_packet(enc, &enc_pkt);
|
|
|
|
if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
|
|
{
|
|
break;
|
|
}
|
|
else if (ret < 0)
|
|
{
|
|
logger->error("Encoder failure! Stream #" + std::to_string(stream_id));
|
|
return false;
|
|
}
|
|
|
|
enc_pkt.stream_index = 0;
|
|
av_packet_rescale_ts(&enc_pkt, enc->time_base, out_ctx->streams[0]->time_base);
|
|
|
|
av_interleaved_write_frame(out_ctx, &enc_pkt);
|
|
} /* while */
|
|
} /* while(1) */
|
|
} /* got_frame */
|
|
}
|
|
else
|
|
{ /* Data */
|
|
klv_it = klv_streams.find(stream_id);
|
|
|
|
if (klv_it != klv_streams.end())
|
|
{
|
|
klv_it->second.write(reinterpret_cast<const char *>(packet.data), packet.size);
|
|
}
|
|
}
|
|
|
|
av_packet_unref(&packet);
|
|
}
|
|
|
|
av_frame_unref(frame);
|
|
av_frame_unref(filt_frame);
|
|
|
|
return true;
|
|
}
|
|
|
|
} // namespace spacex
|
|
|
|
#endif // USE_VIDEO_ENCODER
|