add transmuxer

1 files changed, 168 insertions, 0 deletions

diff --git a/transmuxer.c b/transmuxer.c
new file mode 100644
index 0000000..05074d6
--- /dev/null
+++ b/transmuxer.c
@@ -0,0 +1,168 @@
+ #include "transmuxer.h"
+#include "ringbuf.h"
+
+#include <bits/pthreadtypes.h>
+#include <libavformat/avformat.h>
+#include <libavutil/timestamp.h>
+#include <pthread.h>
+#include <stdlib.h>
+
+
+static int wait_for_new_stream(Transmuxer *self) {
+    // if no new stream, return 0
+    return 0;
+}
+
+typedef struct {
+    AVStream *audio_stream;
+    AVStream *video_stream;
+} StreamPair;
+
+static StreamPair start_new_output_file(
+        AVFormatContext *in_fmt_ctx,
+        AVFormatContext *out_fmt_ctx,
+        const char *out_filename,
+        int aidx, int vidx) {
+    avformat_alloc_output_context2(&out_fmt_ctx, NULL, "mpegts", out_filename);
+
+    AVStream *out_video_stream = avformat_new_stream(out_fmt_ctx, NULL);
+    AVStream *out_audio_stream = avformat_new_stream(out_fmt_ctx, NULL);
+
+    avcodec_parameters_copy(out_video_stream->codecpar, in_fmt_ctx->streams[vidx]->codecpar);
+    avcodec_parameters_copy(out_audio_stream->codecpar, in_fmt_ctx->streams[aidx]->codecpar);
+
+    if (!(out_fmt_ctx->oformat->flags & AVFMT_NOFILE)) {
+        if (avio_open(&out_fmt_ctx->pb, out_filename, AVIO_FLAG_WRITE) < 0) {
+            fprintf(stderr, "Could not open output file '%s'\n", out_filename);
+            exit(-1);
+        }
+    }
+    
+    if (avformat_write_header(out_fmt_ctx, NULL) < 0) {
+        fprintf(stderr, "avformat_write_header failed.\n");
+        abort();
+    }
+    return (StreamPair) {
+        .audio_stream = out_audio_stream,
+        .video_stream = out_video_stream,
+    };
+}
+
+static void finalize_output_file(AVFormatContext *out_fmt_ctx) {
+    av_write_trailer(out_fmt_ctx);
+    if (!(out_fmt_ctx->oformat->flags & AVFMT_NOFILE)) {
+        avio_closep(&out_fmt_ctx->pb);
+    }
+    avformat_free_context(out_fmt_ctx);
+    // TODO: update m3u8
+}
+
+#define SEGMENT_DURATION 5
+
+void* Transmuxer_main (void *vself) {
+    Transmuxer *self = vself;
+    pthread_mutex_lock(&self->lock);
+    while (wait_for_new_stream(self)) {
+        AVFormatContext *in_fmt_ctx = NULL;
+        if (avformat_open_input(&in_fmt_ctx, NULL, NULL, NULL) < 0) {
+            fprintf(stderr, "Could not open input file\n");
+            abort();
+        }
+
+        if (avformat_find_stream_info(in_fmt_ctx, NULL) < 0) {
+            fprintf(stderr, "Could not find stream info\n");
+            abort();
+        }
+
+        int video_stream_index = -1, audio_stream_index = -1;
+        for (int i = 0; i < in_fmt_ctx->nb_streams; i++) {
+            if (in_fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
+                video_stream_index = i;
+            else if (in_fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO)
+                audio_stream_index = i;
+        }
+
+        if (video_stream_index < 0) {
+            fprintf(stderr, "No video stream found\n");
+            abort();
+        }
+        if (audio_stream_index < 0) {
+            fprintf(stderr, "No audio stream found\n");
+            abort();
+        }
+
+        AVFormatContext *out_fmt_ctx = NULL;
+        int segment_index = 0;
+        int64_t segment_start_pts = 0;
+
+        char out_filename[256];
+        snprintf(out_filename, sizeof(out_filename), "segment_%03d.ts", segment_index);
+
+
+        int64_t pts_time;
+        AVPacket pkt;
+        StreamPair output_stream;
+
+        output_stream = start_new_output_file(in_fmt_ctx, out_fmt_ctx, out_filename, audio_stream_index, video_stream_index);
+        while (av_read_frame(in_fmt_ctx, &pkt) >= 0) {
+            AVStream *in_stream = in_fmt_ctx->streams[pkt.stream_index];
+            AVStream *out_stream = NULL;
+            if (pkt.stream_index == video_stream_index)
+                out_stream = output_stream.video_stream;
+            else if (pkt.stream_index == audio_stream_index)
+                out_stream = output_stream.video_stream;
+            else {
+                av_packet_unref(&pkt);
+                continue;
+            }
+
+            pts_time = av_rescale_q(pkt.pts, in_stream->time_base, AV_TIME_BASE_Q);
+
+            // if need split
+            if (pts_time - segment_start_pts >= SEGMENT_DURATION * AV_TIME_BASE) {
+                if (pkt.stream_index == video_stream_index && (pkt.flags & AV_PKT_FLAG_KEY)) {
+                    // close current ts
+                    printf("new ts: %ld\n", pts_time - segment_start_pts);
+                    finalize_output_file(out_fmt_ctx);
+                    
+                    // open new ts
+                    segment_index++;
+                    segment_start_pts = pts_time;
+                    snprintf(out_filename, sizeof(out_filename), "segment_%03d.ts", segment_index);
+                    start_new_output_file(in_fmt_ctx, out_fmt_ctx, out_filename, audio_stream_index, video_stream_index);
+                }
+            }
+            pkt.pts = av_rescale_q_rnd(pkt.pts, in_stream->time_base, out_stream->time_base, AV_ROUND_NEAR_INF|AV_ROUND_PASS_MINMAX);
+            pkt.dts = av_rescale_q_rnd(pkt.dts, in_stream->time_base, out_stream->time_base, AV_ROUND_NEAR_INF|AV_ROUND_PASS_MINMAX);
+            pkt.duration = av_rescale_q(pkt.duration, in_stream->time_base, out_stream->time_base);
+            pkt.pos = -1;
+            pkt.stream_index = (pkt.stream_index == video_stream_index) ? output_stream.video_stream->index : output_stream.audio_stream->index;
+
+            av_interleaved_write_frame(out_fmt_ctx, &pkt);
+            av_packet_unref(&pkt);
+        }
+        printf("new ts: %ld\n", pts_time - segment_start_pts);
+        finalize_output_file(out_fmt_ctx);
+
+        avformat_close_input(&in_fmt_ctx);
+        RingBuffer_destroy(self->stream);
+        self->stream = NULL;
+        free(self->stream);
+    }
+    pthread_mutex_unlock(&self->lock);
+    return NULL;
+}
+
+void Transmuxer_init(Transmuxer *self) {
+    pthread_mutex_init(&self->lock, NULL);
+    pthread_cond_init(&self->streaming_cond, NULL);
+    self->stream = NULL;
+    self->quit = false;
+}
+
+void Transmuxer_new_stream(Transmuxer *self, RingBuffer *ringbuf) {
+    pthread_mutex_lock(&self->lock);
+    self->stream = ringbuf;
+    pthread_mutex_unlock(&self->lock);
+    pthread_cond_signal(&self->streaming_cond);
+}