switch to using FFmpeg for encoding, add audio

1 files changed, 235 insertions, 136 deletions

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 <combaseapi.h>
-#include <mfapi.h>
-#include <mfidl.h>
-#include <mfreadwrite.h>
-#include <codecapi.h>
-#include <mftransform.h>
-#include <stdlib.h>
-#include <string.h>
-#include <strmif.h>
-#include <winerror.h>
+#include <libavformat/avio.h>
+#include <libavformat/avformat.h>
+#include <libavcodec/avcodec.h>
+#include <libavcodec/codec_id.h>
+#include <libavutil/error.h>
+#include <libavutil/opt.h>
+#include <libavutil/log.h>
+#include <libavutil/rational.h>
+#include <libavutil/imgutils.h>
+#include <libavcodec/codec.h>
+#include <libavcodec/packet.h>
+#include <libavutil/frame.h>
+#include <libavutil/avutil.h>
+#include <libavutil/samplefmt.h>
+#include <libavcodec/codec_par.h>
+#include <libavutil/log.h>
+#include <libswscale/swscale.h> 
+#include <libswresample/swresample.h>
 
 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;
 }