From 0f942c806786bdf1fa608d5cf32fdf8a418236b0 Mon Sep 17 00:00:00 2001 From: Matthew Wozniak Date: Mon, 18 Nov 2024 20:48:48 -0500 Subject: switch to using FFmpeg for encoding, add audio --- render.c | 371 ++++++++++++++++++++++++++++++++++++++++----------------------- 1 file changed, 235 insertions(+), 136 deletions(-) (limited to 'render.c') diff --git a/render.c b/render.c index ecea85a..ec841e7 100644 --- a/render.c +++ b/render.c @@ -15,21 +15,31 @@ */ #include "api.h" +#include "libavutil/channel_layout.h" +#include "libavutil/mathematics.h" #include "log.h" #include "hook.h" #include "render.h" #include "os.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include extern struct { int width; @@ -38,67 +48,81 @@ extern struct { int fps; int quality; int bitrate; - int qvs; const char *out; const char **demo; int demo_count; - bool combine; + const char *encoder; + const char *preset; } args; struct renderctx { - GUID enc_format; - GUID in_fmt; - IMFSinkWriter *sink_writer; - ulong stream_index; - ulong audio_index; - int cur_demo; + AVFormatContext *format; + AVStream *vstream; + const AVCodec *vcodec; + AVCodecContext *vcodec_ctx; + AVFrame *yuv_frame; + struct SwsContext *sws; + + AVStream *astream; + const AVCodec *acodec; + AVCodecContext *acodec_ctx; + AVFrame *aframe; + AVFrame *aframe_tmp; + struct SwrContext *swr; + u64 nextpts; }; static struct renderctx ctx = { 0 }; static bool is_rendering = false; -#define HR(expr) { \ - HRESULT hr = expr; \ - if (!SUCCEEDED(hr)) \ - bail(#expr " failed!!! (code %lu)", hr); \ +bool encode(AVCodecContext *codec, AVFrame *frame, AVStream *stream) { + int ret; + AVPacket *pkt = av_packet_alloc(); + + if ((ret = avcodec_send_frame(codec, frame)) < 0) + bail("error sending a frame for encoding"); + + while (ret >= 0) { + ret = avcodec_receive_packet(codec, pkt); + if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) + break; + else if (ret < 0) + bail("encoding error!"); + + pkt->stream_index = stream->index; + + av_packet_rescale_ts(pkt, codec->time_base, stream->time_base); + + if (av_interleaved_write_frame(ctx.format, pkt) < 0) + bail("failed to write frame to file!"); + av_packet_unref(pkt); + } + av_packet_free(&pkt); + return true; } bool do_frame(struct videomode *this, struct movieinfo *info) { is_rendering = true; - const int szpx = 3; - - static struct { u8 bgr[3]; } *raw = 0; - if (!raw) raw = malloc(args.width * args.height * szpx); + static struct { u8 bgr[3]; } *pixels = 0; + if (!pixels) + pixels = malloc(args.width * args.height * sizeof(*pixels)); - IMFMediaBuffer *imf_buffer = 0; - HR(MFCreateMemoryBuffer(args.width * args.height * szpx, &imf_buffer)); - struct { u8 bgr[3]; } *data = 0; - HR(imf_buffer->lpVtbl->Lock(imf_buffer, (u8 **)&data, NULL, NULL)); - // THIS IS SLOW!! if (steampipe) this->vt->steampipe_read_screen_pixels(this, 0, 0, args.width, - args.height, raw, IMAGE_FORMAT_BGR888); + args.height, pixels, IMAGE_FORMAT_BGR888); else - this->vt->read_screen_pixels(this, 0, 0, args.width, args.height, raw, + this->vt->read_screen_pixels(this, 0, 0, args.width, args.height, pixels, IMAGE_FORMAT_BGR888); - for (int y = args.height - 1; y > 0; y--) { - memcpy(data + ((args.height - y) * args.width), raw + y * args.width, - szpx * args.width); - } - HR(imf_buffer->lpVtbl->Unlock(imf_buffer)); - HR(imf_buffer->lpVtbl->SetCurrentLength(imf_buffer, args.width * args.height * szpx)); - - IMFSample *sample = 0; - HR(MFCreateSample(&sample)); - HR(sample->lpVtbl->AddBuffer(sample, imf_buffer)); - u64 sample_duration = (u64)((1.0 / (double)args.fps) * 10000000.0); - HR(sample->lpVtbl->SetSampleDuration(sample, sample_duration)); - u64 sample_time = sample_duration * info->curframe; - HR(sample->lpVtbl->SetSampleTime(sample, sample_time)); - HR(ctx.sink_writer->lpVtbl->WriteSample(ctx.sink_writer, ctx.stream_index, sample)); - HR(sample->lpVtbl->Release(sample)); - HR(imf_buffer->lpVtbl->Release(imf_buffer)); + int stride = args.width * sizeof(*pixels); + sws_scale(ctx.sws, (const u8 **)&pixels, &stride, 0, args.height, + ctx.yuv_frame->data, ctx.yuv_frame->linesize); + + ctx.yuv_frame->pts = info->curframe; + ctx.yuv_frame->time_base = ctx.vcodec_ctx->time_base; + ctx.yuv_frame->duration = 1; + + encode(ctx.vcodec_ctx, ctx.yuv_frame, ctx.vstream); return true; } @@ -106,13 +130,34 @@ bool do_frame(struct videomode *this, struct movieinfo *info) { typeof((*videomode)->vt->write_movie_frame) orig_write_movie_frame; void VIRTUAL hook_write_movie_frame(struct videomode *this, struct movieinfo *info) { - if (!do_frame(this, info)) - die("oopsie!"); + if (!do_frame(this, info)) warn("do_frame failed!"); } bool do_audio_frame(void) { - // TODO: audio + if (!is_rendering) return true; + + static int frame_idx = 0; #define CLIP16(n) (i16)(n < -32768 ? -32768 : (n > 32767 ? 32767 : n)) + + for (int i = 0; i < *snd->snd_linear_count; i += 2) { + for (int c = 0; c < 2; ++c) { + ((i16 *)ctx.aframe_tmp->data[c])[frame_idx] = + CLIP16(((*snd->snd_p)[i + c] * *snd->snd_vol) >> 8); + } + + if (++frame_idx == ctx.aframe->nb_samples) { + if (swr_convert(ctx.swr, ctx.aframe->data, ctx.aframe->nb_samples, + (const u8 *const *)ctx.aframe_tmp->data, + ctx.aframe_tmp->nb_samples) < 0) + bail("failed to resample audio frame"); + + ctx.aframe->pts = ctx.nextpts; + ctx.nextpts += ctx.aframe->nb_samples; + encode(ctx.acodec_ctx, ctx.aframe, ctx.astream); + frame_idx = 0; + } + } + #undef CLIP16 return true; } @@ -127,30 +172,40 @@ void hook_snd_recordbuffer(void) { } bool do_stop() { - //HR(ctx.sink_writer->lpVtbl->Flush(ctx.sink_writer, ctx.stream_index)); - HR(ctx.sink_writer->lpVtbl->Finalize(ctx.sink_writer)); - HR(MFShutdown()); + if (!encode(ctx.vcodec_ctx, 0, ctx.vstream)) + warn("video flush failed"); + if (!encode(ctx.acodec_ctx, 0, ctx.astream)) + warn("audio flush failed!"); + + if (av_write_trailer(ctx.format)) + warn("couldn't write trailer!"); + + avio_flush(ctx.format->pb); + avio_close(ctx.format->pb); + + av_frame_free(&ctx.yuv_frame); + avcodec_free_context(&ctx.vcodec_ctx); + av_frame_free(&ctx.aframe); + avcodec_free_context(&ctx.acodec_ctx); + avformat_free_context(ctx.format); + + sws_freeContext(ctx.sws); + return true; } static void (VIRTUAL *orig_stop_playback)(struct demoplayer *); void VIRTUAL hook_stop_playback(struct demoplayer *this) { orig_stop_playback(this); - if (!is_rendering) return; is_rendering = false; - // do we have any more demos to play? if (args.demo_count--) { demoplayer->vt->start_playback(demoplayer, *++args.demo, false); return; } - info("finished!"); - - if (!do_stop()) - die("oopsie!"); - + do_stop(); exit(0); } @@ -182,79 +237,123 @@ bool render_init(void) { sprintf(framerate_cmd, "host_framerate %d;", args.fps); cbuf_addtext(framerate_cmd); - if (CoInitializeEx(NULL, COINIT_MULTITHREADED) == RPC_E_CHANGED_MODE) - warn("changed COM concurrency mode!"); - - HR(MFStartup(MF_VERSION, 0)); - - // init sinkwriter - { - IMFMediaType *mt_out = NULL; - IMFMediaType *mt_in = NULL; - - u16 path[MAX_PATH] = {0}; - mbstowcs(path, args.out, strlen(args.out)); - - info("rendering to %S", path); - - IMFAttributes *sinkwriter_attr; - HR(MFCreateAttributes(&sinkwriter_attr, 1)); - HR(sinkwriter_attr->lpVtbl->SetUINT32(sinkwriter_attr, &MF_READWRITE_ENABLE_HARDWARE_TRANSFORMS, TRUE)); - HR(MFCreateSinkWriterFromURL(path, NULL, sinkwriter_attr, &ctx.sink_writer)); - - HR(MFCreateMediaType(&mt_out)); - HR(mt_out->lpVtbl->SetGUID(mt_out, &MF_MT_MAJOR_TYPE, &MFMediaType_Video)); - HR(mt_out->lpVtbl->SetGUID(mt_out, &MF_MT_SUBTYPE, &MFVideoFormat_H264)); - HR(mt_out->lpVtbl->SetUINT32(mt_out, &MF_MT_AVG_BITRATE, args.bitrate ? args.bitrate : 9999999)); - HR(mt_out->lpVtbl->SetUINT32(mt_out, &MF_MT_MPEG2_PROFILE , eAVEncH264VProfile_Main)); - HR(mt_out->lpVtbl->SetUINT32(mt_out, &MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive)); - HR(mt_out->lpVtbl->SetUINT64(mt_out, &MF_MT_FRAME_SIZE, ((u64)args.width << 32) | args.height)); - HR(mt_out->lpVtbl->SetUINT64(mt_out, &MF_MT_FRAME_RATE, ((u64)args.fps << 32) | 1)); - HR(mt_out->lpVtbl->SetUINT64(mt_out, &MF_MT_PIXEL_ASPECT_RATIO, (((u64)1 << 32) | 1))); - HR(ctx.sink_writer->lpVtbl->AddStream(ctx.sink_writer, mt_out, &ctx.stream_index)); - - HR(MFCreateMediaType(&mt_in)); - HR(mt_in->lpVtbl->SetGUID(mt_in, &MF_MT_MAJOR_TYPE, &MFMediaType_Video)); - HR(mt_in->lpVtbl->SetGUID(mt_in, &MF_MT_SUBTYPE, &MFVideoFormat_RGB24)); - HR(mt_in->lpVtbl->SetUINT32(mt_in, &MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive)); - HR(mt_in->lpVtbl->SetUINT32(mt_in, &MF_MT_ALL_SAMPLES_INDEPENDENT, TRUE)); - HR(mt_in->lpVtbl->SetUINT64(mt_in, &MF_MT_FRAME_SIZE, ((u64)args.width << 32) | args.height)); - HR(mt_in->lpVtbl->SetUINT64(mt_in, &MF_MT_FRAME_RATE, ((u64)args.fps << 32) | 1)); - HR(mt_in->lpVtbl->SetUINT64(mt_in, &MF_MT_PIXEL_ASPECT_RATIO, (((u64)1 << 32) | 1))); - HR(ctx.sink_writer->lpVtbl->SetInputMediaType(ctx.sink_writer, ctx.stream_index, mt_in, NULL)); - } - - // video encoding parameters - { - ICodecAPI *codec; - HR(ctx.sink_writer->lpVtbl->GetServiceForStream(ctx.sink_writer, 0, &GUID_NULL, &IID_ICodecAPI, (void*)&codec)); - - int ratecontrol = !!args.bitrate ? eAVEncCommonRateControlMode_UnconstrainedVBR : eAVEncCommonRateControlMode_Quality; - VARIANT _ratecontrol = { .vt = VT_UI4, .ulVal = ratecontrol }; - HR(codec->lpVtbl->SetValue(codec, &CODECAPI_AVEncCommonRateControlMode, &_ratecontrol)); - - // set the quality - VARIANT _quality = { - .vt = VT_UI4, - .ulVal = !!args.bitrate ? args.bitrate : args.quality - }; - HR(codec->lpVtbl->SetValue(codec, &CODECAPI_AVEncCommonQuality, &_quality)); - - VARIANT _qvs = { .vt = VT_UI4, .ulVal = args.qvs }; - HR(codec->lpVtbl->SetValue(codec, &CODECAPI_AVEncCommonQualityVsSpeed, &_qvs)); - - // set the gop size to 2 seconds - VARIANT gop = { .vt = VT_UI4, .ulVal = 2 * args.fps }; - HR(codec->lpVtbl->SetValue(codec, &CODECAPI_AVEncMPVGOPSize, &gop)); - - // enable 2 b-frames - VARIANT bframes = { .vt = VT_UI4, .ulVal = 2 }; - HR(codec->lpVtbl->SetValue(codec, &CODECAPI_AVEncMPVDefaultBPictureCount, &bframes)); - - codec->lpVtbl->Release(codec); - } - - HR(ctx.sink_writer->lpVtbl->BeginWriting(ctx.sink_writer)); + info("started recording to %s (%dx%d @ %d)", args.out, args.width, + args.height, args.fps); + + avformat_alloc_output_context2(&ctx.format, 0, 0, args.out); + if (!ctx.format) + die("couldn't create output context (does the filename have an ext?)"); + + // video codec + ctx.vcodec = avcodec_find_encoder_by_name(args.encoder); + ctx.vcodec_ctx = avcodec_alloc_context3(ctx.vcodec); + ctx.vcodec_ctx->width = args.width; + ctx.vcodec_ctx->height = args.height; + ctx.vcodec_ctx->pix_fmt = AV_PIX_FMT_YUV420P; + ctx.vcodec_ctx->time_base = av_make_q(1, args.fps); + ctx.vcodec_ctx->framerate = av_make_q(args.fps, 1); + ctx.vcodec_ctx->gop_size = args.fps * 2; + ctx.vcodec_ctx->bit_rate = args.bitrate * 1000; + + if (!args.bitrate) + av_opt_set_int(ctx.vcodec_ctx->priv_data, "crf", args.quality, 0); + if (args.preset) + av_opt_set(ctx.vcodec_ctx->priv_data, "preset", args.preset, 0); + + if (avcodec_open2(ctx.vcodec_ctx, ctx.vcodec, NULL) < 0) + die("failed to open video encoder"); + + // video stream + ctx.vstream = avformat_new_stream(ctx.format, NULL); + if (!ctx.vstream) die("failed to create vide stream"); + avcodec_parameters_from_context(ctx.vstream->codecpar, ctx.vcodec_ctx); + ctx.vstream->time_base = ctx.vcodec_ctx->time_base; + ctx.vstream->avg_frame_rate = ctx.vcodec_ctx->framerate; + + // video frame + ctx.yuv_frame = av_frame_alloc(); + ctx.yuv_frame->width = args.width; + ctx.yuv_frame->height = args.height; + ctx.yuv_frame->color_range = AVCOL_RANGE_JPEG; + ctx.yuv_frame->time_base = ctx.vcodec_ctx->time_base; + ctx.yuv_frame->format = ctx.vcodec_ctx->pix_fmt; + + ctx.sws = sws_getContext(args.width, args.height, AV_PIX_FMT_BGR24, + args.width, args.height, ctx.yuv_frame->format, + SWS_FAST_BILINEAR | SWS_FULL_CHR_H_INT | SWS_ACCURATE_RND, 0, 0, 0); + + if (av_frame_get_buffer(ctx.yuv_frame, 0) < 0) + bail("failed to alloc frame buffer"); + + // audio codec + ctx.acodec = avcodec_find_encoder(AV_CODEC_ID_AAC); + if (!ctx.acodec) bail("failed to find aac codec"); + ctx.acodec_ctx = avcodec_alloc_context3(ctx.acodec); + ctx.acodec_ctx->bit_rate = 128000; + ctx.acodec_ctx->sample_fmt = AV_SAMPLE_FMT_FLTP; + ctx.acodec_ctx->sample_rate = 44100; + ctx.acodec_ctx->profile = FF_PROFILE_AAC_MAIN; + ctx.acodec_ctx->time_base = av_make_q(1, 44100); + + if (av_channel_layout_copy(&ctx.acodec_ctx->ch_layout, + &(AVChannelLayout)AV_CHANNEL_LAYOUT_STEREO)) + bail("failed to copy channel layout"); + + if (avcodec_open2(ctx.acodec_ctx, ctx.acodec, NULL) < 0) + die("failed to open audio encoder"); + + // audio resampler + ctx.swr = swr_alloc(); + av_opt_set_int(ctx.swr, "in_channel_count", 2, 0); + av_opt_set_int(ctx.swr, "in_sample_rate", 44100, 0); + av_opt_set_sample_fmt(ctx.swr, "in_sample_fmt", AV_SAMPLE_FMT_S16P, 0); + av_opt_set_chlayout(ctx.swr, "in_chlayout", &ctx.acodec_ctx->ch_layout, 0); + av_opt_set_int(ctx.swr, "out_channel_count", 2, 0); + av_opt_set_int(ctx.swr, "out_sample_rate", ctx.acodec_ctx->sample_rate, 0); + av_opt_set_sample_fmt(ctx.swr, "out_sample_fmt", ctx.acodec_ctx->sample_fmt, 0); + av_opt_set_chlayout(ctx.swr, "out_chlayout", &ctx.acodec_ctx->ch_layout, 0); + + if (swr_init(ctx.swr) < 0) + bail("failed to init resampler"); + + ctx.nextpts = 0; + + // audio frame + ctx.aframe = av_frame_alloc(); + if (!ctx.aframe) bail("failed to alloc audio frame"); + + ctx.aframe->nb_samples = ctx.acodec_ctx->frame_size; + ctx.aframe->format = ctx.acodec_ctx->sample_fmt; + if (av_channel_layout_copy(&ctx.aframe->ch_layout, &ctx.acodec_ctx->ch_layout)) + bail("failed to copy channel layout"); + if (av_frame_get_buffer(ctx.aframe, 0)) + bail("couldn't alloc audio frame buffer"); + + // audio tmp frame (pre resampling) + ctx.aframe_tmp = av_frame_alloc(); + if (!ctx.aframe_tmp) bail("failed to alloc tmp audio frame"); + + ctx.aframe_tmp->nb_samples = ctx.acodec_ctx->frame_size; + ctx.aframe_tmp->format = ctx.acodec_ctx->sample_fmt; + if (av_channel_layout_copy(&ctx.aframe_tmp->ch_layout, &ctx.acodec_ctx->ch_layout)) + bail("failed to copy channel layout"); + if (av_frame_get_buffer(ctx.aframe_tmp, 0)) + bail("couldn't alloc audio frame buffer"); + + // audio stream + ctx.astream = avformat_new_stream(ctx.format, NULL); + if (!ctx.astream) die("failed to create audio stream"); + avcodec_parameters_from_context(ctx.astream->codecpar, ctx.acodec_ctx); + ctx.astream->time_base = ctx.acodec_ctx->time_base; + + if (avcodec_open2(ctx.acodec_ctx, ctx.acodec, NULL) < 0) + bail("failed to open audio encoder"); + + if (avio_open(&ctx.format->pb, args.out, AVIO_FLAG_WRITE) < 0) + bail("failed to open file"); + + if (avformat_write_header(ctx.format, NULL) < 0) + bail("error when writing to output file"); return true; } -- cgit v1.2.3-54-g00ecf