244 lines
10 KiB
C++
244 lines
10 KiB
C++
#include <atomic>
|
|
#include <cerrno>
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <iostream>
|
|
#include <optional>
|
|
#include <string>
|
|
#include <string_view>
|
|
#include <thread>
|
|
|
|
#ifdef _WIN32
|
|
#include <fcntl.h>
|
|
#include <io.h>
|
|
#endif
|
|
|
|
extern "C" {
|
|
#include <libavcodec/avcodec.h>
|
|
#include <libavdevice/avdevice.h>
|
|
#include <libavformat/avformat.h>
|
|
#include <libavutil/avutil.h>
|
|
#include <libavutil/error.h>
|
|
#include <libavutil/opt.h>
|
|
}
|
|
|
|
namespace {
|
|
|
|
struct Options {
|
|
std::string input = "ddagrab=framerate=60:output_idx=0:draw_mouse=1";
|
|
int framerate = 60;
|
|
int bitrate_kbps = 12000;
|
|
int gop_frames = 15;
|
|
std::optional<std::uint64_t> max_frames;
|
|
std::optional<std::uint64_t> force_idr_frame;
|
|
};
|
|
|
|
std::string ff_error(int code)
|
|
{
|
|
char text[AV_ERROR_MAX_STRING_SIZE]{};
|
|
av_strerror(code, text, sizeof(text));
|
|
return text;
|
|
}
|
|
|
|
[[noreturn]] void fail(const std::string &message)
|
|
{
|
|
std::cerr << "muninn-video-encoder: " << message << '\n';
|
|
std::exit(1);
|
|
}
|
|
|
|
int parse_positive(const char *value, const char *name)
|
|
{
|
|
char *end = nullptr;
|
|
errno = 0;
|
|
const long parsed = std::strtol(value, &end, 10);
|
|
if (errno != 0 || end == value || *end != '\0' || parsed <= 0 || parsed > INT32_MAX)
|
|
fail(std::string("invalid ") + name);
|
|
return static_cast<int>(parsed);
|
|
}
|
|
|
|
Options parse_options(int argc, char **argv)
|
|
{
|
|
Options result;
|
|
for (int index = 1; index < argc; ++index) {
|
|
const std::string_view arg(argv[index]);
|
|
const auto take = [&]() -> const char * {
|
|
if (++index >= argc) fail(std::string("missing value for ") + std::string(arg));
|
|
return argv[index];
|
|
};
|
|
if (arg == "--input") result.input = take();
|
|
else if (arg == "--framerate") result.framerate = parse_positive(take(), "framerate");
|
|
else if (arg == "--bitrate-kbps") result.bitrate_kbps = parse_positive(take(), "bitrate");
|
|
else if (arg == "--gop-frames") result.gop_frames = parse_positive(take(), "gop");
|
|
else if (arg == "--force-idr-frame")
|
|
result.force_idr_frame = static_cast<std::uint64_t>(parse_positive(take(), "force frame"));
|
|
else if (arg == "--frames")
|
|
result.max_frames = static_cast<std::uint64_t>(parse_positive(take(), "frames"));
|
|
else fail("unknown argument: " + std::string(arg));
|
|
}
|
|
return result;
|
|
}
|
|
|
|
struct FormatCloser { void operator()(AVFormatContext *value) const { avformat_close_input(&value); } };
|
|
struct CodecCloser { void operator()(AVCodecContext *value) const { avcodec_free_context(&value); } };
|
|
struct FrameCloser { void operator()(AVFrame *value) const { av_frame_free(&value); } };
|
|
struct PacketCloser { void operator()(AVPacket *value) const { av_packet_free(&value); } };
|
|
|
|
using FormatPtr = std::unique_ptr<AVFormatContext, FormatCloser>;
|
|
using CodecPtr = std::unique_ptr<AVCodecContext, CodecCloser>;
|
|
using FramePtr = std::unique_ptr<AVFrame, FrameCloser>;
|
|
using PacketPtr = std::unique_ptr<AVPacket, PacketCloser>;
|
|
|
|
void write_packet(const AVPacket *packet)
|
|
{
|
|
if (packet->size == 0) return;
|
|
if (std::fwrite(packet->data, 1, static_cast<size_t>(packet->size), stdout) !=
|
|
static_cast<size_t>(packet->size))
|
|
fail("writing Annex-B stdout failed");
|
|
std::fflush(stdout);
|
|
}
|
|
|
|
void drain_encoder(AVCodecContext *encoder, AVPacket *packet)
|
|
{
|
|
while (true) {
|
|
const int result = avcodec_receive_packet(encoder, packet);
|
|
if (result == AVERROR(EAGAIN) || result == AVERROR_EOF) return;
|
|
if (result < 0) fail("receiving encoded packet: " + ff_error(result));
|
|
write_packet(packet);
|
|
av_packet_unref(packet);
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main(int argc, char **argv)
|
|
{
|
|
#ifdef _WIN32
|
|
if (_setmode(_fileno(stdout), _O_BINARY) == -1)
|
|
fail("setting stdout binary mode failed");
|
|
#endif
|
|
const Options options = parse_options(argc, argv);
|
|
avdevice_register_all();
|
|
|
|
const AVInputFormat *lavfi = av_find_input_format("lavfi");
|
|
if (!lavfi) fail("this FFmpeg build has no lavfi input device");
|
|
AVFormatContext *raw_input = nullptr;
|
|
int result = avformat_open_input(&raw_input, options.input.c_str(), lavfi, nullptr);
|
|
if (result < 0) fail("opening capture input: " + ff_error(result));
|
|
FormatPtr input(raw_input);
|
|
result = avformat_find_stream_info(input.get(), nullptr);
|
|
if (result < 0) fail("reading capture stream information: " + ff_error(result));
|
|
const int video_stream = av_find_best_stream(input.get(), AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
|
|
if (video_stream < 0) fail("capture input has no video stream");
|
|
|
|
const AVCodec *decoder_codec = avcodec_find_decoder(
|
|
input->streams[video_stream]->codecpar->codec_id);
|
|
if (!decoder_codec) fail("capture frame decoder is unavailable");
|
|
CodecPtr decoder(avcodec_alloc_context3(decoder_codec));
|
|
if (!decoder) fail("allocating capture decoder");
|
|
result = avcodec_parameters_to_context(decoder.get(), input->streams[video_stream]->codecpar);
|
|
if (result < 0) fail("copying capture parameters: " + ff_error(result));
|
|
result = avcodec_open2(decoder.get(), decoder_codec, nullptr);
|
|
if (result < 0) fail("opening capture decoder: " + ff_error(result));
|
|
|
|
const AVCodec *encoder_codec = avcodec_find_encoder_by_name("h264_nvenc");
|
|
if (!encoder_codec) fail("h264_nvenc is unavailable");
|
|
CodecPtr encoder(avcodec_alloc_context3(encoder_codec));
|
|
if (!encoder) fail("allocating h264_nvenc context");
|
|
|
|
PacketPtr input_packet(av_packet_alloc());
|
|
PacketPtr output_packet(av_packet_alloc());
|
|
FramePtr frame(av_frame_alloc());
|
|
if (!input_packet || !output_packet || !frame) fail("allocating frame pipeline");
|
|
|
|
std::atomic_bool force_next_idr{false};
|
|
std::atomic_bool quit{false};
|
|
std::thread command_reader([&] {
|
|
std::string command;
|
|
while (std::getline(std::cin, command)) {
|
|
if (command == "IDR") force_next_idr.store(true, std::memory_order_release);
|
|
else if (command == "QUIT") { quit.store(true, std::memory_order_release); break; }
|
|
}
|
|
});
|
|
command_reader.detach();
|
|
|
|
bool encoder_open = false;
|
|
std::uint64_t frame_number = 0;
|
|
while (!quit.load(std::memory_order_acquire)) {
|
|
result = av_read_frame(input.get(), input_packet.get());
|
|
if (result == AVERROR_EOF) break;
|
|
if (result < 0) fail("reading capture frame: " + ff_error(result));
|
|
if (input_packet->stream_index != video_stream) {
|
|
av_packet_unref(input_packet.get());
|
|
continue;
|
|
}
|
|
result = avcodec_send_packet(decoder.get(), input_packet.get());
|
|
av_packet_unref(input_packet.get());
|
|
if (result < 0) fail("submitting capture packet: " + ff_error(result));
|
|
|
|
while ((result = avcodec_receive_frame(decoder.get(), frame.get())) >= 0) {
|
|
if (!encoder_open) {
|
|
encoder->width = frame->width;
|
|
encoder->height = frame->height;
|
|
encoder->pix_fmt = static_cast<AVPixelFormat>(frame->format);
|
|
encoder->time_base = AVRational{1, options.framerate};
|
|
encoder->framerate = AVRational{options.framerate, 1};
|
|
encoder->bit_rate = static_cast<int64_t>(options.bitrate_kbps) * 1000;
|
|
encoder->rc_max_rate = encoder->bit_rate;
|
|
encoder->rc_buffer_size = static_cast<int>(encoder->bit_rate / options.framerate * 2);
|
|
encoder->gop_size = options.gop_frames;
|
|
encoder->max_b_frames = 0;
|
|
if (frame->hw_frames_ctx) encoder->hw_frames_ctx = av_buffer_ref(frame->hw_frames_ctx);
|
|
AVDictionary *settings = nullptr;
|
|
av_dict_set(&settings, "preset", "p1", 0);
|
|
av_dict_set(&settings, "tune", "ull", 0);
|
|
av_dict_set(&settings, "zerolatency", "1", 0);
|
|
av_dict_set(&settings, "delay", "0", 0);
|
|
av_dict_set(&settings, "rc", "cbr", 0);
|
|
av_dict_set(&settings, "rc-lookahead", "0", 0);
|
|
av_dict_set(&settings, "multipass", "disabled", 0);
|
|
av_dict_set(&settings, "strict_gop", "1", 0);
|
|
av_dict_set(&settings, "forced-idr", "1", 0);
|
|
av_dict_set(&settings, "aud", "1", 0);
|
|
result = avcodec_open2(encoder.get(), encoder_codec, &settings);
|
|
av_dict_free(&settings);
|
|
if (result < 0) fail("opening h264_nvenc: " + ff_error(result));
|
|
if (encoder->extradata && encoder->extradata_size > 0) {
|
|
if (std::fwrite(encoder->extradata, 1,
|
|
static_cast<size_t>(encoder->extradata_size), stdout) !=
|
|
static_cast<size_t>(encoder->extradata_size))
|
|
fail("writing encoder headers failed");
|
|
}
|
|
std::cerr << "muninn-video-encoder: " << encoder->width << 'x'
|
|
<< encoder->height << " format=" << encoder->pix_fmt
|
|
<< " headers=" << encoder->extradata_size << '\n';
|
|
encoder_open = true;
|
|
}
|
|
|
|
const bool scheduled = options.force_idr_frame == frame_number;
|
|
if (scheduled || force_next_idr.exchange(false, std::memory_order_acq_rel))
|
|
frame->pict_type = AV_PICTURE_TYPE_I;
|
|
else
|
|
frame->pict_type = AV_PICTURE_TYPE_NONE;
|
|
frame->pts = static_cast<int64_t>(frame_number++);
|
|
result = avcodec_send_frame(encoder.get(), frame.get());
|
|
av_frame_unref(frame.get());
|
|
if (result < 0) fail("submitting frame to h264_nvenc: " + ff_error(result));
|
|
drain_encoder(encoder.get(), output_packet.get());
|
|
if (options.max_frames && frame_number >= *options.max_frames) {
|
|
quit.store(true, std::memory_order_release);
|
|
break;
|
|
}
|
|
}
|
|
if (quit.load(std::memory_order_acquire)) break;
|
|
if (result != AVERROR(EAGAIN) && result != AVERROR_EOF)
|
|
fail("decoding capture frame: " + ff_error(result));
|
|
}
|
|
|
|
if (encoder_open) {
|
|
result = avcodec_send_frame(encoder.get(), nullptr);
|
|
if (result >= 0) drain_encoder(encoder.get(), output_packet.get());
|
|
}
|
|
return 0;
|
|
}
|