Enable video output file, with pts set by server

Author: igorinov

app/src/decoder.c

@@ -1,6 +1,7 @@
 #include "decoder.h"
 
 #include <libavformat/avformat.h>
+#include <libavutil/time.h>
 #include <SDL2/SDL_events.h>
 #include <SDL2/SDL_mutex.h>
 #include <SDL2/SDL_thread.h>
@@ -14,9 +15,50 @@
 
 #define BUFSIZE 0x10000
 
+static AVRational us = {1, 1000000};
+
+static inline uint64_t from_be(uint8_t *b, int size)
+{
+    uint64_t x = 0;
+    int i;
+
+    for (i = 0; i < size; i += 1) {
+        x <<= 8;
+        x |= b[i];
+    }
+
+    return x;
+}
+
+#define HEADER_SIZE 16
+
 static int read_packet(void *opaque, uint8_t *buf, int buf_size) {
     struct decoder *decoder = opaque;
-    return net_recv(decoder->video_socket, buf, buf_size);
+    uint8_t header[HEADER_SIZE];
+    int remaining;
+    int ret;
+
+    remaining = decoder->remaining;
+    if (remaining == 0) {
+        ret = net_recv(decoder->video_socket, header, HEADER_SIZE);
+        if (ret <= 0)
+            return ret;
+
+        decoder->pts = from_be(header, 8);
+        remaining = from_be(header + 12, 4);
+    }
+
+    if (buf_size > remaining)
+        buf_size = remaining;
+
+    ret = net_recv(decoder->video_socket, buf, buf_size);
+    if (ret <= 0)
+        return ret;
+
+    remaining -= ret;
+    decoder->remaining = remaining;
+
+    return ret;
 }
 
 // set the decoded frame as ready for rendering, and notify
@@ -40,6 +82,7 @@ static void notify_stopped(void) {
 
 static int run_decoder(void *data) {
     struct decoder *decoder = data;
+    int ret;
 
     AVCodec *codec = avcodec_find_decoder(AV_CODEC_ID_H264);
     if (!codec) {
@@ -86,16 +129,55 @@ static int run_decoder(void *data) {
         goto run_finally_free_avio_ctx;
     }
 
+    AVStream *outstream = NULL;
+    AVFormatContext *output_ctx = NULL;
+    if (decoder->out_filename) {
+        avformat_alloc_output_context2(&output_ctx, NULL, NULL, decoder->out_filename);
+        if (!output_ctx) {
+            LOGE("Could not allocate output format context");
+            goto run_finally_free_avio_ctx;
+        } else {
+            outstream = avformat_new_stream(output_ctx, codec);
+            if (!outstream) {
+                LOGE("Could not allocate output stream");
+                goto run_finally_free_output_ctx;
+            }
+            outstream->codec = avcodec_alloc_context3(codec);
+            outstream->codec->pix_fmt = AV_PIX_FMT_YUV420P;
+            outstream->codec->width = decoder->frame_size.width;
+            outstream->codec->height = decoder->frame_size.height;
+            outstream->time_base = (AVRational) {1, 60};
+            outstream->codec->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
+            ret = avio_open(&output_ctx->pb, decoder->out_filename, AVIO_FLAG_WRITE);
+            if (ret < 0) {
+                LOGE("Failed to open output file");
+                goto run_finally_free_output_ctx;
+            }
+            ret = avformat_write_header(output_ctx, NULL);
+            if (ret < 0) {
+                LOGE("Error writing output header");
+                avio_closep(&output_ctx->pb);
+                goto run_finally_free_output_ctx;
+            }
+        }
+    }
+
     AVPacket packet;
     av_init_packet(&packet);
     packet.data = NULL;
     packet.size = 0;
 
     while (!av_read_frame(format_ctx, &packet)) {
+
+        if (output_ctx) {
+            packet.pts = decoder->pts;
+            av_packet_rescale_ts(&packet, us, outstream->time_base);
+            ret = av_write_frame(output_ctx, &packet);
+        }
+
 // the new decoding/encoding API has been introduced by:
 // <http://git.videolan.org/?p=ffmpeg.git;a=commitdiff;h=7fc329e2dd6226dfecaa4a1d7adf353bf2773726>
 #if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 37, 0)
-        int ret;
         if ((ret = avcodec_send_packet(codec_ctx, &packet)) < 0) {
             LOGE("Could not send video packet: %d", ret);
             goto run_quit;
@@ -125,6 +207,7 @@ static int run_decoder(void *data) {
             packet.data += len;
         }
 #endif
+
         av_packet_unref(&packet);
 
         if (avio_ctx->eof_reached) {
@@ -135,7 +218,14 @@ static int run_decoder(void *data) {
     LOGD("End of frames");
 
 run_quit:
+    if (output_ctx) {
+        ret = av_write_trailer(output_ctx);
+        avio_closep(&output_ctx->pb);
+    }
     avformat_close_input(&format_ctx);
+run_finally_free_output_ctx:
+    if (output_ctx)
+        avformat_free_context(output_ctx);
 run_finally_free_avio_ctx:
     av_freep(&avio_ctx);
 run_finally_free_format_ctx:
@@ -149,20 +239,21 @@ static int run_decoder(void *data) {
     return 0;
 }
 
-void decoder_init(struct decoder *decoder, struct frames *frames, socket_t video_socket) {
+void decoder_init(struct decoder *decoder, struct frames *frames, socket_t video_socket, struct size frame_size) {
     decoder->frames = frames;
     decoder->video_socket = video_socket;
+    decoder->frame_size = frame_size;
 }
 
-SDL_bool decoder_start(struct decoder *decoder) {
+SDL_bool decoder_start(struct decoder *decoder, const char *out_filename) {
     LOGD("Starting decoder thread");
 
+    decoder->out_filename = out_filename;
     decoder->thread = SDL_CreateThread(run_decoder, "video_decoder", decoder);
     if (!decoder->thread) {
         LOGC("Could not start decoder thread");
         return SDL_FALSE;
     }
-
     return SDL_TRUE;
 }
 

app/src/decoder.h

@@ -4,19 +4,24 @@
 #include <SDL2/SDL_stdinc.h>
 #include <SDL2/SDL_thread.h>
 
+#include "common.h"
 #include "net.h"
 
 struct frames;
 
 struct decoder {
+    uint64_t pts;
     struct frames *frames;
     socket_t video_socket;
     SDL_Thread *thread;
     SDL_mutex *mutex;
+    const char *out_filename;
+    struct size frame_size;
+    int remaining;
 };
 
-void decoder_init(struct decoder *decoder, struct frames *frames, socket_t video_socket);
-SDL_bool decoder_start(struct decoder *decoder);
+void decoder_init(struct decoder *decoder, struct frames *frames, socket_t video_socket, struct size frame_size);
+SDL_bool decoder_start(struct decoder *decoder, const char *out_filename);
 void decoder_stop(struct decoder *decoder);
 void decoder_join(struct decoder *decoder);
 

app/src/main.c

@@ -11,6 +11,7 @@
 struct args {
     const char *serial;
     const char *crop;
+    const char *out_filename;
     SDL_bool fullscreen;
     SDL_bool help;
     SDL_bool version;
@@ -49,6 +50,9 @@ static void usage(const char *arg0) {
         "        is preserved.\n"
         "        Default is %d%s.\n"
         "\n"
+        "    -o, --output-file\n"
+        "        Write video output to file.\n"
+        "\n"
         "    -p, --port port\n"
         "        Set the TCP port the client listens on.\n"
         "        Default is %d.\n"
@@ -207,14 +211,15 @@ static SDL_bool parse_args(struct args *args, int argc, char *argv[]) {
         {"fullscreen",   no_argument,       NULL, 'f'},
         {"help",         no_argument,       NULL, 'h'},
         {"max-size",     required_argument, NULL, 'm'},
+        {"output-file",  required_argument, NULL, 'o'},
         {"port",         required_argument, NULL, 'p'},
         {"serial",       required_argument, NULL, 's'},
         {"show-touches", no_argument,       NULL, 't'},
         {"version",      no_argument,       NULL, 'v'},
         {NULL,           0,                 NULL, 0  },
     };
     int c;
-    while ((c = getopt_long(argc, argv, "b:c:fhm:p:s:tv", long_options, NULL)) != -1) {
+    while ((c = getopt_long(argc, argv, "b:c:fhm:o:p:s:tv", long_options, NULL)) != -1) {
         switch (c) {
             case 'b':
                 if (!parse_bit_rate(optarg, &args->bit_rate)) {
@@ -235,6 +240,9 @@ static SDL_bool parse_args(struct args *args, int argc, char *argv[]) {
                     return SDL_FALSE;
                 }
                 break;
+            case 'o':
+                args->out_filename = optarg;
+                break;
             case 'p':
                 if (!parse_port(optarg, &args->port)) {
                     return SDL_FALSE;
@@ -310,6 +318,7 @@ int main(int argc, char *argv[]) {
         .serial = args.serial,
         .crop = args.crop,
         .port = args.port,
+        .out_filename = args.out_filename,
         .max_size = args.max_size,
         .bit_rate = args.bit_rate,
         .show_touches = args.show_touches,

app/src/scrcpy.c

@@ -193,11 +193,11 @@ SDL_bool scrcpy(const struct scrcpy_options *options) {
         goto finally_destroy_frames;
     }
 
-    decoder_init(&decoder, &frames, device_socket);
+    decoder_init(&decoder, &frames, device_socket, frame_size);
 
     // now we consumed the header values, the socket receives the video stream
     // start the decoder
-    if (!decoder_start(&decoder)) {
+    if (!decoder_start(&decoder, options->out_filename)) {
         ret = SDL_FALSE;
         server_stop(&server);
         goto finally_destroy_file_handler;

app/src/scrcpy.h

@@ -6,6 +6,7 @@
 struct scrcpy_options {
     const char *serial;
     const char *crop;
+    const char *out_filename;
     Uint16 port;
     Uint16 max_size;
     Uint32 bit_rate;

gradle.properties

@@ -15,3 +15,4 @@ org.gradle.jvmargs=-Xmx1536m
 # This option should only be used with decoupled projects. More details, visit
 # http://www.gradle.org/docs/current/userguide/multi_project_builds.html#sec:decoupled_projects
 # org.gradle.parallel=true
+

server/src/main/java/com/genymobile/scrcpy/ScreenEncoder.java

@@ -3,6 +3,7 @@
 import com.genymobile.scrcpy.wrappers.SurfaceControl;
 
 import android.graphics.Rect;
+import android.media.MediaMuxer;
 import android.media.MediaCodec;
 import android.media.MediaCodecInfo;
 import android.media.MediaFormat;
@@ -80,6 +81,8 @@ public void streamScreen(Device device, FileDescriptor fd) throws IOException {
     private boolean encode(MediaCodec codec, FileDescriptor fd) throws IOException {
         boolean eof = false;
         MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
+        ByteBuffer bBuffer = ByteBuffer.allocate(16);
+
         while (!consumeRotationChange() && !eof) {
             int outputBufferId = codec.dequeueOutputBuffer(bufferInfo, -1);
             eof = (bufferInfo.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0;
@@ -90,6 +93,12 @@ private boolean encode(MediaCodec codec, FileDescriptor fd) throws IOException {
                 }
                 if (outputBufferId >= 0) {
                     ByteBuffer codecBuffer = codec.getOutputBuffer(outputBufferId);
+                    bBuffer.position(0);
+                    bBuffer.putLong(bufferInfo.presentationTimeUs);
+                    bBuffer.putInt(bufferInfo.flags);
+                    bBuffer.putInt(codecBuffer.remaining());
+                    bBuffer.position(0);
+                    IO.writeFully(fd, bBuffer);
                     IO.writeFully(fd, codecBuffer);
                 }
             } finally {